implement reference counting and much more (#43)

2026-07-17 11:32:33 +08:00 · 2016-04-18 13:05:36 -07:00
parent a6a77bc416
commit 1548a1a523
22 changed files with 1063 additions and 258 deletions
@@ -2,6 +2,13 @@

 Orchestra is a distributed execution framework with a Python-like programming model.

+## Design Decisions
+
+For a description of our design decisions, see
+
+- [Reference Counting](doc/reference-counting.md)
+- [Aliasing](doc/aliasing.md)
+
 ## Setup

 **Install Arrow**
@@ -0,0 +1,95 @@
+# Aliasing
+
+An important feature of Photon is that a remote call sent to the scheduler
+immediately returns object references to the outputs of the task, and the actual
+outputs of the task are only associated with the relevant object references
+after the task has been executed and the outputs have been computed. This allows
+the worker to continue without blocking.
+
+However, to provide a more flexible API, we allow tasks to not only return
+values, but to also return object references to values. As an examples, consider
+the following code.
+```python
+@op.distributed([], [np.ndarray])
+def f()
+  return np.zeros(5)
+
+@op.distributed([], [np.ndarray])
+def g()
+  return f()
+
+@op.distributed([], [np.ndarray])
+def h()
+  return g()
+```
+A call to `h` will immediate return an object reference `ref_h` for the return
+value of `h`. The task of executing `h` (call it `task_h`) will then be
+scheduled for execution. When `task_h` is executed, it will call `g`, which will
+immediately return an object reference `ref_g` for the output of `g`. Then two
+things will happen and can happen in any order: `AliasObjRefs(ref_h, ref_g)`
+will be called and `task_g` will be scheduled and executed. When `task_g` is
+executed, it will call `f`, and immediately obtain an object reference `ref_f`
+for the output of `f`. Then two things will happen and can happen in either
+order, `AliasObjRefs(ref_g, ref_f)` will be called, and `f` will be executed.
+When `f` is executed, it will create an actual array and `put_object` will be
+called, which will store the array in the object store (it will also call
+`SchedulerService::AddCanonicalObjRef(ref_f)`).
+
+From the scheduler's perspective, there are three important calls,
+`AliasObjRefs(ref_h, ref_g)`, `AliasObjRefs(ref_g, ref_f)`, and
+`AddCanonicalObjRef(ref_f)`. These three calls can happen in any order.
+
+The scheduler maintains a data structure called `target_objrefs_`, which keeps
+track of which object references have been aliased together (`target_objrefs_`
+is a vector, but we can think of it as a graph). The call
+`AliasObjRefs(ref_h, ref_g)` updates `target_objrefs_` with `ref_h -> ref_g`.
+The call `AliasObjRefs(ref_g, ref_f)` updates it with `ref_g -> ref_f`, and the
+call `AddCanonicalObjRef(ref_f)` updates it with `ref_f -> ref_f`. The data
+structure is initialized with `ref -> UNINITIALIZED_ALIAS` for each object
+reference `ref`.
+
+We refer to `ref_f` as a "canonical" object reference. And in a pair such as
+`ref_h -> ref_g`, we refer to `ref_h` as the "alias" object reference and to
+`ref_g` as the "target" object reference. These details are available to the
+scheduler, but a worker process just has an object reference and doesn't know if
+it is canonical or not.
+
+We also maintain a data structure `reverse_target_objrefs_`, which maps in the
+reverse direction (in the above example, we would have `ref_g -> ref_h`,
+`ref_f -> ref_g`, and `ref_h -> UNINITIALIZED_ALIAS`). This data structure is
+not particuarly important for the task of aliasing, but when we do reference
+counting and attempt to deallocate an object, we need to be able to determine
+all of the object references that refer to the same object, and this data
+structure comes in handy for that purpose.
+
+## Pulls and Remote Calls
+
+When a worker calls `pull(ref)`, it first sends a message to the scheduler
+asking the scheduler to ship the object referred to by `ref` to the worker's
+local object store. Then the worker asks its local object store for the object
+referred to by `ref`. If `ref` is a canonical object reference, then that's all
+there is too it. However, if `ref` is not a canonical object reference but
+rather is an alias for the canonical object reference `c_ref`, then the
+scheduler also notifies the worker's local object store that `ref` is an
+alias for `c_ref`. This is important because the object store does not keep
+track of aliasing on its own (it only knows the bits about aliasing that the
+scheduler tells it). Lastly, if the scheduler does not yet have enough
+information to determine if `ref` is canonical, or if the scheduler cannot
+yet determine what the canonical object reference for `ref` is, then the
+scheduler will wait until it has the relevant information.
+
+Similar things happen when a worker performs a remote call. If an object
+reference is passed to a remote call, the object referred to by that object
+reference will be shipped to the local object store of the worker that executes
+the task. The scheduler will notify that object store about any aliasing that it
+needs to be aware of.
+
+## Passing Object References by Value
+Currently, the graph of aliasing looks like a collection of chains, as in the
+above example with `ref_h -> ref_g -> ref_f -> ref_f`. In the future, we will
+allow object references to be passed by value to remote calls (so the worker
+has access to the object reference object and not the object that the object
+reference refers to). If an object reference that is passed by value is then
+returned by the task, it is possible that a given object reference could be
+the target of multiple alias object references. In this case, the graph of
+aliasing will be a tree.
@@ -0,0 +1,91 @@
+# Reference Counting
+
+In Photon, each object is assigned a globally unique object reference by the
+scheduler (starting with 0 and incrementing upward). The objects are stored in
+object stores. In order to avoid running out of memory, the object stores must
+know when it is ok to deallocate an object. Since a worker on one node may have
+an object reference for an object that lives in an object store on a different
+node, knowing when we can safely deallocate an object requires cluster-wide
+information.
+
+## Reference Counting
+
+Two approaches to reclaiming memory are garbage collection and reference
+counting. We choose to use a reference counting approach in Photon. There are a
+couple of reasons for this. Reference counting allows us to reclaim memory as
+early as possible. It also avoids pausing the system for garbage collection. We
+also note that implementing reference counting at the cluster level plays nicely
+with worker processes that use reference counting internally (currently our
+worker processes are Python processes). However, this could be made to work with
+worker processes that use garbage collection, for example, if each worker
+process is a Java Virtual Machine.
+
+At a high level, the scheduler keeps track of the number of object references
+that exist on the cluster for each object reference. When the number of object
+references reaches 0 for a particular object, the scheduler notifies all of the
+object stores that contain that object to deallocate it.
+
+Object references can exist in several places.
+
+1. They can be Python objects on a worker.
+2. They can be serialized within an object in an object store.
+3. They can be in a message being sent between processes (e.g., as an argument
+to a remote procedure call).
+
+## When to Increment and Decrement the Reference Count
+
+We handle these three cases by calling the SchedulerService methods
+`IncrementRefCount` and `DecrementRefCount` as follows:
+
+1. To handle the first case, we increment in the ObjRef constructor and
+decrement in the ObjRef destructor.
+2. To handle the second case, when an object is written to an object store with
+a call to `put_object`, we call `IncrementRefCount` for each object reference
+that is contained internally in the serialized object (for example, if we
+serialize a `DistArray`, we increment the reference counts for its blocks). This
+will notify the scheduler that those object references are in the object store.
+Then when the scheduler deallocates the object, we call `DecrementRefCount` for
+the object references that it holds internally (the scheduler keeps track of
+these internal object references in the `contained_objrefs_` data structure).
+3. To handle the third case, we increment in the `serialize_call` method and
+decrement in the `deserialize_call` method.
+
+## How to Handle Aliasing
+Reference counting interacts with aliasing. Since multiple object references
+may refer to the same object, we cannot deallocate that object until all of the
+object references that refer to it have reference counts of 0. We keep track of
+the number of separate aliases separately. If two object references refer to the
+same object, the scheduler keeps track the number of occurrences of each of
+those object references separately. This simplifies the scheduler's job because
+it may not always know if two object references refer to the same object or not
+(since it assigns them before hearing back about what they refer to).
+
+When we decrement the count for an object reference, if the count reaches 0,
+we compute all object references that the scheduler knows to reference the same
+object. If these object references all have count 0, then we deallocate the
+object. Otherwise, we do not deallocate the object.
+
+You may ask, what if there is some object reference with a nonzero count which
+refers to the same object, but the scheduler does not know it? This cannot
+happen because the following invariant holds. If `a` and `b` are object
+references that will be aliased together (through a call to
+`AliasObjRefs(a, b)`), then either the call has already happened, or both `a`
+and `b` have positive reference counts (they must have positive reference counts
+because they must be passed into `AliasObjRefs` at some point).
+
+## Complications
+The following problem has not yet been resolved. In the following code, the
+result `x` will be garbage.
+```python
+x = orchpy.pull(single.zeros([10, 10], "float"))
+```
+When `single.zeros` is called, a worker will create an array of zeros and store
+it in an object store. An object reference to the output is returned. The call
+to `orchpy.pull` will not copy data from the object store process to the worker
+process, but will instead give the worker process a pointer to shared memory.
+After the `orchpy.pull` call completes, the object reference returned by
+`single.zeros` will go out of scope, and the object it refers to will be
+deallocated from the object store. This will cause the memory that `x` points to
+to be garbage.
+
+This problem is currently unresolved.
@@ -1,3 +1,3 @@
 import liborchpylib as lib
 import serialization
-from worker import register_module, connect, pull, push, distributed
+from worker import scheduler_info, register_module, connect, disconnect, pull, push, distributed
@@ -20,19 +20,20 @@ def from_primitive(primitive_obj):
    obj.deserialize(primitive_obj[1])
  return obj

-def serialize(obj):
+def serialize(worker_capsule, obj):
  primitive_obj = to_primitive(obj)
-  return orchpy.lib.serialize_object(primitive_obj)
+  obj_capsule, contained_objrefs = orchpy.lib.serialize_object(worker_capsule, primitive_obj) # contained_objrefs is a list of the objrefs contained in obj
+  return obj_capsule, contained_objrefs

-def deserialize(capsule):
-  primitive_obj = orchpy.lib.deserialize_object(capsule)
+def deserialize(worker_capsule, capsule):
+  primitive_obj = orchpy.lib.deserialize_object(worker_capsule, capsule)
  return from_primitive(primitive_obj)

-def serialize_call(func_name, args):
+def serialize_call(worker_capsule, func_name, args):
  primitive_args = [(arg if isinstance(arg, orchpy.lib.ObjRef) else to_primitive(arg)) for arg in args]
-  return orchpy.lib.serialize_call(func_name, primitive_args)
+  return orchpy.lib.serialize_call(worker_capsule, func_name, primitive_args)

-def deserialize_call(call):
-  func_name, primitive_args, return_objrefs = orchpy.lib.deserialize_call(call)
+def deserialize_call(worker_capsule, call):
+  func_name, primitive_args, return_objrefs = orchpy.lib.deserialize_call(worker_capsule, call)
  args = [(arg if isinstance(arg, orchpy.lib.ObjRef) else from_primitive(arg)) for arg in primitive_args]
  return func_name, args, return_objrefs
@@ -9,6 +9,7 @@ import orchpy.worker as worker
 _services_path = os.path.dirname(os.path.abspath(__file__))

 all_processes = []
+driver = None

 IP_ADDRESS = "127.0.0.1"
 TIMEOUT_SECONDS = 5
@@ -55,7 +56,14 @@ def cleanup():
    print "Termination attempt failed, giving up."
  all_processes = []

-atexit.register(cleanup)
+  global driver
+  if driver is not None:
+    orchpy.disconnect(driver)
+  else:
+    orchpy.disconnect()
+  driver = None
+
+# atexit.register(cleanup)

 def start_scheduler(scheduler_address):
  p = subprocess.Popen([os.path.join(_services_path, "scheduler"), scheduler_address])
@@ -74,6 +82,7 @@ def start_worker(test_path, scheduler_address, objstore_address, worker_address)
  all_processes.append((p, worker_address))

 def start_cluster(driver_worker=None, num_workers=0, worker_path=None):
+  global driver
  if num_workers > 0 and worker_path is None:
    raise Exception("Attempting to start a cluster with some workers, but `worker_path` is None.")
  scheduler_address = address(IP_ADDRESS, new_scheduler_port())
@@ -84,6 +93,7 @@ def start_cluster(driver_worker=None, num_workers=0, worker_path=None):
  time.sleep(0.2)
  if driver_worker is not None:
    orchpy.connect(scheduler_address, objstore_address, address(IP_ADDRESS, new_worker_port()), driver_worker)
+    driver = driver_worker
  else:
    orchpy.connect(scheduler_address, objstore_address, address(IP_ADDRESS, new_worker_port()))
  for _ in range(num_workers):
@@ -10,7 +10,6 @@ class Worker(object):

  def __init__(self):
    self.functions = {}
-    self.connected = False
    self.handle = None

  def put_object(self, objref, value):
@@ -18,8 +17,8 @@ class Worker(object):
    if type(value) == np.ndarray:
      orchpy.lib.put_arrow(self.handle, objref, value)
    else:
-      object_capsule = serialization.serialize(value)
-      orchpy.lib.put_object(self.handle, objref, object_capsule)
+      object_capsule, contained_objrefs = serialization.serialize(self.handle, value) # contained_objrefs is a list of the objrefs contained in object_capsule
+      orchpy.lib.put_object(self.handle, objref, object_capsule, contained_objrefs)

  def get_object(self, objref):
    """
@@ -32,7 +31,7 @@ class Worker(object):
      return orchpy.lib.get_arrow(self.handle, objref)
    else:
      object_capsule = orchpy.lib.get_object(self.handle, objref)
-      return serialization.deserialize(object_capsule)
+      return serialization.deserialize(self.handle, object_capsule)

  def alias_objrefs(self, alias_objref, target_objref):
    """Make `alias_objref` refer to the same object that `target_objref` refers to."""
@@ -45,13 +44,16 @@ class Worker(object):

  def remote_call(self, func_name, args):
    """Tell the scheduler to schedule the execution of the function with name `func_name` with arguments `args`. Retrieve object references for the outputs of the function from the scheduler and immediately return them."""
-    call_capsule = serialization.serialize_call(func_name, args)
+    call_capsule = serialization.serialize_call(self.handle, func_name, args)
    objrefs = orchpy.lib.remote_call(self.handle, call_capsule)
    return objrefs

 # We make `global_worker` a global variable so that there is one worker per worker process.
 global_worker = Worker()

+def scheduler_info(worker=global_worker):
+  return orchpy.lib.scheduler_info(worker.handle);
+
 def register_module(module, recursive=False, worker=global_worker):
  print "registering functions in module {}.".format(module.__name__)
  for name in dir(module):
@@ -63,10 +65,12 @@ def register_module(module, recursive=False, worker=global_worker):
    #   register_module(val, recursive, worker)

 def connect(scheduler_addr, objstore_addr, worker_addr, worker=global_worker):
-  if worker.connected:
-    del worker.handle # TODO(rkn): Make sure this actually deallocates (need a destructor for the capsule)
+  if hasattr(worker, "handle"):
+    del worker.handle
  worker.handle = orchpy.lib.create_worker(scheduler_addr, objstore_addr, worker_addr)
-  worker.connected = True
+
+def disconnect(worker=global_worker):
+  orchpy.lib.disconnect(worker.handle)

 def pull(objref, worker=global_worker):
  orchpy.lib.request_object(worker.handle, objref)
@@ -78,16 +82,18 @@ def push(value, worker=global_worker):
  return objref

 def main_loop(worker=global_worker):
-  if not worker.connected:
+  if not orchpy.lib.connected(worker.handle):
    raise Exception("Worker is attempting to enter main_loop but has not been connected yet.")
  orchpy.lib.start_worker_service(worker.handle)
-  while True:
-    call = orchpy.lib.wait_for_next_task(worker.handle)
-    func_name, args, return_objrefs = serialization.deserialize_call(call)
+  def process_call(call): # wrapping these calls in a function should cause the local variables to go out of scope more quickly, which is useful for inspecting reference counts
+    func_name, args, return_objrefs = serialization.deserialize_call(worker.handle, call)
    arguments = get_arguments_for_execution(worker.functions[func_name], args, worker) # get args from objstore
    outputs = worker.functions[func_name].executor(arguments) # execute the function
    store_outputs_in_objstore(return_objrefs, outputs, worker) # store output in local object store
    orchpy.lib.notify_task_completed(worker.handle) # notify the scheduler that the task has completed
+  while True:
+    call = orchpy.lib.wait_for_next_task(worker.handle)
+    process_call(call)

 def distributed(arg_types, return_types, worker=global_worker):
  def distributed_decorator(func):
@@ -37,10 +37,16 @@ service Scheduler {
  rpc AliasObjRefs(AliasObjRefsRequest) returns (AckReply);
  // Used by an object store to tell the scheduler that an object is ready (i.e. has been finalized and can be shared)
  rpc ObjReady(ObjReadyRequest) returns (AckReply);
+  // Increments the reference count of a particular object reference
+  rpc IncrementRefCount(IncrementRefCountRequest) returns (AckReply);
+  // Decrements the reference count of a particular object reference
+  rpc DecrementRefCount(DecrementRefCountRequest) returns (AckReply);
+  // Used by the worker to notify the scheduler about which objrefs a particular object contains
+  rpc AddContainedObjRefs(AddContainedObjRefsRequest) returns (AckReply);
  // Used by the worker to report back and ask for more work
  rpc WorkerReady(WorkerReadyRequest) returns (AckReply);
-  // Get debugging information from the scheduler
-  rpc SchedulerDebugInfo(SchedulerDebugInfoRequest) returns (SchedulerDebugInfoReply);
+  // Get information about the scheduler state
+  rpc SchedulerInfo(SchedulerInfoRequest) returns (SchedulerInfoReply);
 }

 message AckReply {
@@ -100,6 +106,19 @@ message ObjReadyRequest {
  uint64 objstoreid = 2; // ID of the object store the object lives on
 }

+message IncrementRefCountRequest {
+  repeated uint64 objref = 1; // Object references whose reference count should be incremented. Duplicates will be incremented multiple times.
+}
+
+message AddContainedObjRefsRequest {
+  uint64 objref = 1; // The objref of the object in question
+  repeated uint64 contained_objref = 2; // Object references contained in the object
+}
+
+message DecrementRefCountRequest {
+  repeated uint64 objref = 1; // Object references whose reference count should be decremented. Duplicates will be decremented multiple times.
+}
+
 message WorkerReadyRequest {
  uint64 workerid = 1; // ID of the worker which is ready
 }
@@ -108,9 +127,9 @@ message ChangeCountRequest {
  uint64 objref = 1; // Object reference of the object whose reference count is increased or decreased
 }

-// The following messages are used for debugging purposes:
+// The following messages are used to get information about the scheduler state

-message SchedulerDebugInfoRequest {
+message SchedulerInfoRequest {
 }

 message FnTableEntry {
@@ -118,10 +137,12 @@ message FnTableEntry {
  uint64 num_return_vals = 2; // Number of return values of the function
 }

-message SchedulerDebugInfoReply {
+message SchedulerInfoReply {
  repeated Call task = 1; // Tasks on the task queue
  repeated uint64 avail_worker = 3; // List of workers waiting to get a task assigned
  map<string, FnTableEntry> function_table = 2; // Table of all available remote function
+  repeated uint64 target_objref = 4; // The target_objrefs_ data structure
+  repeated uint64 reference_count = 5; // The reference_counts_ data structure
 }

 // Object stores
@@ -133,8 +154,10 @@ service ObjStore {
  rpc StreamObj(stream ObjChunk) returns (AckReply);
  // Notify the object store about objref aliasing. This is called by the scheduler
  rpc NotifyAlias(NotifyAliasRequest) returns (AckReply);
-  // Get debug info from the object store
-  rpc ObjStoreDebugInfo(ObjStoreDebugInfoRequest) returns (ObjStoreDebugInfoReply);
+  // Tell the object store to deallocate an object held by the object store. This is called by the scheduler.
+  rpc DeallocateObject(DeallocateObjectRequest) returns (AckReply);
+  // Get info about the object store state
+  rpc ObjStoreInfo(ObjStoreInfoRequest) returns (ObjStoreInfoReply);
 }

 message DeliverObjRequest {
@@ -161,6 +184,10 @@ message NotifyAliasRequest {
  uint64 canonical_objref = 2; // The canonical objref that points to the actual object
 }

+message DeallocateObjectRequest {
+  uint64 canonical_objref = 1; // The canonical objref of the object to deallocate
+}
+
 message GetObjRequest {
  uint64 objref = 1; // Object reference of the object being requested by the worker
 }
@@ -171,13 +198,13 @@ message GetObjReply {
  uint64 size = 3; // Total size of the object in bytes
 }

-// These messages are for debugging purposes:
+// These messages are for getting information about the object store state

-message ObjStoreDebugInfoRequest {
+message ObjStoreInfoRequest {
  repeated uint64 objref = 1; // Object references we want to retrieve from the store for inspection
 }

-message ObjStoreDebugInfoReply {
+message ObjStoreInfoReply {
  repeated uint64 objref = 1; // List of object references in the store
  repeated Obj obj = 2; // Protocol buffer objects that were requested
 }
@@ -31,48 +31,74 @@ MemorySegmentPool::MemorySegmentPool(bool create) : create_mode_(create) { }
 // creates a memory segment if it is not already there; if the pool is in create mode,
 // space is allocated, if it is in open mode, the shared memory is mapped into the process
 void MemorySegmentPool::open_segment(SegmentId segmentid, size_t size) {
-  if (segmentid < segments_.size()) {
+  ORCH_LOG(ORCH_DEBUG, "OPENING segmentid " << segmentid);
+  if (segmentid != segments_.size() && create_mode_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to open segmentid " << segmentid << " on the object store, but segments_.size() = " << segments_.size());
+  }
+  if (segmentid >= segments_.size()) { // resize and initialize segments_
+    int current_size = segments_.size();
+    segments_.resize(segmentid + 1);
+    for (int i = current_size; i < segments_.size(); ++i) {
+      segments_[i].first = nullptr;
+      segments_[i].second = SegmentStatusType::UNOPENED;
+    }
+  }
+  if (segments_[segmentid].second == SegmentStatusType::OPENED) {
    return;
  }
-  segment_names_.resize(segmentid + 1);
-  segments_.resize(segmentid + 1);
+  if (segments_[segmentid].second == SegmentStatusType::CLOSED) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to open segmentid " << segmentid << ", but segments_[segmentid].second == SegmentStatusType::CLOSED.");
+  }
  std::string segment_name = std::string("segment:") + std::to_string(segmentid);
  if (create_mode_) {
    assert(size > 0);
    shared_memory_object::remove(segment_name.c_str()); // remove segment if it has not been properly removed from last run
    size_t new_size = (size / page_size_ + 2) * page_size_; // additional room for boost's bookkeeping
-    segments_[segmentid] = std::unique_ptr<managed_shared_memory>(new managed_shared_memory(create_only, segment_name.c_str(), new_size));
+    segments_[segmentid] = std::make_pair(std::unique_ptr<managed_shared_memory>(new managed_shared_memory(create_only, segment_name.c_str(), new_size)), SegmentStatusType::OPENED);
  } else {
-    segments_[segmentid] = std::unique_ptr<managed_shared_memory>(new managed_shared_memory(open_only, segment_name.c_str()));
+    segments_[segmentid] = std::make_pair(std::unique_ptr<managed_shared_memory>(new managed_shared_memory(open_only, segment_name.c_str())), SegmentStatusType::OPENED);
  }
-  segment_names_[segmentid] = segment_name;
+}
+
+void MemorySegmentPool::close_segment(SegmentId segmentid) {
+  ORCH_LOG(ORCH_DEBUG, "CLOSING segmentid " << segmentid);
+  std::string segment_name = std::string("segment:") + std::to_string(segmentid);
+  shared_memory_object::remove(segment_name.c_str());
+  segments_[segmentid].first.reset();
+  segments_[segmentid].second = SegmentStatusType::CLOSED;
 }

 ObjHandle MemorySegmentPool::allocate(size_t size) {
+  if (!create_mode_) { // allocate is called only by the object store
+    ORCH_LOG(ORCH_FATAL, "Attempting to call allocate, but create_mode_ is false");
+  }
  // TODO(pcm): at the moment, this always creates a new segment, this will be changed
-  SegmentId segmentid = segment_names_.size();
+  SegmentId segmentid = segments_.size();
  open_segment(segmentid, size);
-  void* ptr = segments_[segmentid]->allocate(size);
-  auto handle = segments_[segmentid]->get_handle_from_address(ptr);
+  void* ptr = segments_[segmentid].first->allocate(size);
+  auto handle = segments_[segmentid].first->get_handle_from_address(ptr);
  return ObjHandle(segmentid, size, handle);
 }

+void MemorySegmentPool::deallocate(ObjHandle pointer) {
+  SegmentId segmentid = pointer.segmentid();
+  void* ptr = segments_[segmentid].first->get_address_from_handle(pointer.ipcpointer());
+  segments_[segmentid].first->deallocate(ptr);
+  close_segment(segmentid);
+}
+
 // returns address of the object refered to by the handle, needs to be called on
 // the process that will use the address
 uint8_t* MemorySegmentPool::get_address(ObjHandle pointer) {
-  if (pointer.segmentid() >= segments_.size()) {
-    for (int i = segments_.size(); i <= pointer.segmentid(); ++i) {
-      open_segment(i);
-    }
-  }
-  managed_shared_memory* segment = segments_[pointer.segmentid()].get();
+  open_segment(pointer.segmentid());
+  managed_shared_memory* segment = segments_[pointer.segmentid()].first.get();
  return static_cast<uint8_t*>(segment->get_address_from_handle(pointer.ipcpointer()));
 }

 MemorySegmentPool::~MemorySegmentPool() {
-  assert(segment_names_.size() == segments_.size());
-  for (size_t i = 0; i < segment_names_.size(); ++i) {
-    segments_[i].reset();
-    shared_memory_object::remove(segment_names_[i].c_str());
+  for (size_t segmentid = 0; segmentid < segments_.size(); ++segmentid) {
+    std::string segment_name = std::string("segment:") + std::to_string(segmentid);
+    segments_[segmentid].first.reset();
+    shared_memory_object::remove(segment_name.c_str());
  }
 }
@@ -89,7 +89,7 @@ private:
 // ALIAS_DONE: objref -> ():
 // objstore tells itself that it has finalized something (perhaps an alias)

-enum ObjRequestType {ALLOC = 0, GET = 1, WORKER_DONE = 2, ALIAS_DONE};
+enum ObjRequestType {ALLOC = 0, GET = 1, WORKER_DONE = 2, ALIAS_DONE = 3};

 struct ObjRequest {
  WorkerId workerid; // worker that sends the request
@@ -141,18 +141,21 @@ public:
 // the segments, which have been created by the object store, are just mapped
 // into memory

+enum SegmentStatusType {UNOPENED = 0, OPENED = 1, CLOSED = 2};
+
 class MemorySegmentPool {
 public:
  MemorySegmentPool(bool create = false); // can be used in two modes: create mode and open mode (see above)
  ~MemorySegmentPool();
  ObjHandle allocate(size_t nbytes); // allocate memory, potentially creating a new segment (only run on object store)
+  void deallocate(ObjHandle pointer); // deallocate object, potentially deallocating a new segment (only run on object store)
  uint8_t* get_address(ObjHandle pointer); // get address of shared object
 private:
  void open_segment(SegmentId segmentid, size_t size = 0); // create a segment or map an existing one into memory
-  bool create_mode_;
+  void close_segment(SegmentId segmentid); // close a segment
+  bool create_mode_; // true in the object stores, false on the workers
  size_t page_size_ = mapped_region::get_page_size();
-  std::vector<std::string> segment_names_;
-  std::vector<std::unique_ptr<managed_shared_memory> > segments_;
+  std::vector<std::pair<std::unique_ptr<managed_shared_memory>, SegmentStatusType> > segments_;
 };

 #endif
@@ -53,10 +53,10 @@ Status ObjStoreService::DeliverObj(ServerContext* context, const DeliverObjReque
 }
 */

-Status ObjStoreService::ObjStoreDebugInfo(ServerContext* context, const ObjStoreDebugInfoRequest* request, ObjStoreDebugInfoReply* reply) {
+Status ObjStoreService::ObjStoreInfo(ServerContext* context, const ObjStoreInfoRequest* request, ObjStoreInfoReply* reply) {
  std::lock_guard<std::mutex> memory_lock(memory_lock_);
  for (size_t i = 0; i < memory_.size(); ++i) {
-    if (memory_[i].second) { // is the object available?
+    if (memory_[i].second == MemoryStatusType::READY) { // is the object available?
      reply->add_objref(i);
    }
  }
@@ -116,18 +116,25 @@ Status ObjStoreService::NotifyAlias(ServerContext* context, const NotifyAliasReq
  // NotifyAlias assumes that the objstore already holds canonical_objref
  ObjRef alias_objref = request->alias_objref();
  ObjRef canonical_objref = request->canonical_objref();
+  ORCH_LOG(ORCH_DEBUG, "Aliasing objref " << alias_objref << " with objref " << canonical_objref);
  std::lock_guard<std::mutex> memory_lock(memory_lock_);
  if (canonical_objref >= memory_.size()) {
    ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not in the objstore.")
  }
-  if (!memory_[canonical_objref].second) {
+  if (memory_[canonical_objref].second == MemoryStatusType::NOT_READY) {
    ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not ready yet in the objstore.")
  }
+  if (memory_[canonical_objref].second == MemoryStatusType::NOT_PRESENT) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not present in the objstore.")
+  }
+  if (memory_[canonical_objref].second == MemoryStatusType::DEALLOCATED) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " has already been deallocated.")
+  }
  if (alias_objref >= memory_.size()) {
-    memory_.resize(alias_objref + 1);
+    memory_.resize(alias_objref + 1, std::make_pair(ObjHandle(), MemoryStatusType::NOT_PRESENT));
  }
  memory_[alias_objref].first = memory_[canonical_objref].first;
-  memory_[alias_objref].second = true;
+  memory_[alias_objref].second = MemoryStatusType::READY;

  ObjRequest done_request;
  done_request.type = ObjRequestType::ALIAS_DONE;
@@ -136,6 +143,31 @@ Status ObjStoreService::NotifyAlias(ServerContext* context, const NotifyAliasReq
  return Status::OK;
 }

+Status ObjStoreService::DeallocateObject(ServerContext* context, const DeallocateObjectRequest* request, AckReply* reply) {
+  ObjRef canonical_objref = request->canonical_objref();
+  ORCH_LOG(ORCH_DEBUG, "Deallocating canonical_objref " << canonical_objref);
+  std::lock_guard<std::mutex> memory_lock(memory_lock_);
+  if (memory_[canonical_objref].second != MemoryStatusType::READY) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to deallocate canonical_objref " << canonical_objref << ", but memory_[canonical_objref].second = " << memory_[canonical_objref].second);
+  }
+  if (canonical_objref >= memory_.size()) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to deallocate canonical_objref " << canonical_objref << ", but it is not in the objstore.");
+  }
+  segmentpool_.deallocate(memory_[canonical_objref].first);
+  memory_[canonical_objref].second = MemoryStatusType::DEALLOCATED;
+  return Status::OK;
+}
+
+// This table describes how the memory status changes in response to requests.
+//
+// MemoryStatus | ObjRequest  | New MemoryStatus | action performed
+// -------------+-------------+------------------+----------------------------
+// NOT_PRESENT  | ALLOC       | NOT_READY        | allocate object
+// NOT_READY    | WORKER_DONE | READY            | send ObjReady to scheduler
+// NOT_READY    | GET         | NOT_READY        | add to pull queue
+// READY        | GET         | READY            | return handle
+// READY        | DEALLOC     | DEALLOCATED      | deallocate
+// -------------+-------------+------------------+----------------------------
 void ObjStoreService::process_objstore_request(const ObjRequest request) {
  switch (request.type) {
    case ObjRequestType::ALIAS_DONE: {
@@ -156,35 +188,49 @@ void ObjStoreService::process_worker_request(const ObjRequest request) {
    std::string queue_name = std::string("queue:") + objstore_address_ + std::string(":worker:") + std::to_string(request.workerid) + std::string(":obj");
    send_queues_[request.workerid].connect(queue_name, false);
  }
-  if (request.objref >= memory_.size()) {
-    memory_.resize(request.objref + 1);
-    memory_[request.objref].second = false;
+  {
+    std::lock_guard<std::mutex> memory_lock(memory_lock_);
+    if (request.objref >= memory_.size()) {
+      memory_.resize(request.objref + 1, std::make_pair(ObjHandle(), MemoryStatusType::NOT_PRESENT));
+    }
  }
  switch (request.type) {
    case ObjRequestType::ALLOC: {
+        // TODO(rkn): Does segmentpool_ need a lock around it?
        ObjHandle reply = segmentpool_.allocate(request.size);
        send_queues_[request.workerid].send(&reply);
-        if (request.objref >= memory_.size()) {
-          memory_.resize(request.objref + 1);
+        std::lock_guard<std::mutex> memory_lock(memory_lock_);
+        if (memory_[request.objref].second != MemoryStatusType::NOT_PRESENT) {
+          ORCH_LOG(ORCH_FATAL, "Attempting to allocate space for objref " << request.objref << ", but memory_[objref].second != MemoryStatusType::NOT_PRESENT, it equals " << memory_[request.objref].second);
        }
        memory_[request.objref].first = reply;
-        memory_[request.objref].second = false;
+        memory_[request.objref].second = MemoryStatusType::NOT_READY;
      }
      break;
    case ObjRequestType::GET: {
-        std::pair<ObjHandle, bool>& item = memory_[request.objref];
-        if (item.second) {
+        std::lock_guard<std::mutex> memory_lock(memory_lock_);
+        std::pair<ObjHandle, MemoryStatusType>& item = memory_[request.objref];
+        if (item.second == MemoryStatusType::READY) {
+          ORCH_LOG(ORCH_DEBUG, "Responding to GET request: returning objref " << request.objref);
          send_queues_[request.workerid].send(&item.first);
-        } else {
+        } else if (item.second == MemoryStatusType::NOT_READY || item.second == MemoryStatusType::NOT_PRESENT) {
          std::lock_guard<std::mutex> lock(pull_queue_lock_);
          pull_queue_.push_back(std::make_pair(request.workerid, request.objref));
+        } else {
+          ORCH_LOG(ORCH_FATAL, "A worker requested objref " << request.objref << ", but memory_[objref].second = " << memory_[request.objref].second);
        }
      }
      break;
    case ObjRequestType::WORKER_DONE: {
-        std::pair<ObjHandle, bool>& item = memory_[request.objref];
-        item.first.set_metadata_offset(request.metadata_offset);
-        item.second = true;
+        {
+          std::lock_guard<std::mutex> memory_lock(memory_lock_);
+          std::pair<ObjHandle, MemoryStatusType>& item = memory_[request.objref];
+          if (item.second != MemoryStatusType::NOT_READY) {
+            ORCH_LOG(ORCH_FATAL, "A worker notified the object store that objref " << request.objref << " has been written to the object store, but memory_[objref].second != NOT_READY.");
+          }
+          item.first.set_metadata_offset(request.metadata_offset);
+          item.second = MemoryStatusType::READY;
+        }
        process_pulls_for_objref(request.objref);
        // Tell the scheduler that the object arrived
        // TODO(pcm): put this in a separate thread so we don't have to pay the latency here
@@ -232,7 +278,7 @@ void ObjStoreService::process_requests() {
 }

 void ObjStoreService::process_pulls_for_objref(ObjRef objref) {
-  std::pair<ObjHandle, bool>& item = memory_[objref];
+  std::pair<ObjHandle, MemoryStatusType>& item = memory_[objref];
  std::lock_guard<std::mutex> pull_queue_lock(pull_queue_lock_);
  for (size_t i = 0; i < pull_queue_.size(); ++i) {
    if (pull_queue_[i].second == objref) {
@@ -28,6 +28,8 @@ public:
  static Status upload_data_to(slice data, ObjRef objref, ObjStore::Stub& stub);
 };

+enum MemoryStatusType {READY = 0, NOT_READY = 1, DEALLOCATED = 2, NOT_PRESENT = 3};
+
 class ObjStoreService final : public ObjStore::Service {
 public:
  ObjStoreService(const std::string& objstore_address, std::shared_ptr<Channel> scheduler_channel);
@@ -35,7 +37,8 @@ public:
  // Status DeliverObj(ServerContext* context, const DeliverObjRequest* request, AckReply* reply) override;
  // Status StreamObj(ServerContext* context, ServerReader<ObjChunk>* reader, AckReply* reply) override;
  Status NotifyAlias(ServerContext* context, const NotifyAliasRequest* request, AckReply* reply) override;
-  Status ObjStoreDebugInfo(ServerContext* context, const ObjStoreDebugInfoRequest* request, ObjStoreDebugInfoReply* reply) override;
+  Status DeallocateObject(ServerContext* context, const DeallocateObjectRequest* request, AckReply* reply) override;
+  Status ObjStoreInfo(ServerContext* context, const ObjStoreInfoRequest* request, ObjStoreInfoReply* reply) override;
  void start_objstore_service();
 private:
  // check if we already connected to the other objstore, if yes, return reference to connection, otherwise connect
@@ -48,7 +51,7 @@ private:
  std::string objstore_address_;
  ObjStoreId objstoreid_; // id of this objectstore in the scheduler object store table
  MemorySegmentPool segmentpool_;
-  std::vector<std::pair<ObjHandle, bool> > memory_; // object reference -> (memory address, memory finalized?)
+  std::vector<std::pair<ObjHandle, MemoryStatusType> > memory_; // object reference -> (memory address, memory status)
  std::mutex memory_lock_;
  std::unordered_map<std::string, std::unique_ptr<ObjStore::Stub>> objstores_;
  std::mutex objstores_lock_;
@@ -16,14 +16,20 @@

 extern "C" {

+int PyObjectToWorker(PyObject* object, Worker **worker);
+
 // Object references

 typedef struct {
    PyObject_HEAD
    ObjRef val;
+    Worker* worker;
 } PyObjRef;

 static void PyObjRef_dealloc(PyObjRef *self) {
+  std::vector<ObjRef> objrefs;
+  objrefs.push_back(self->val);
+  self->worker->decrement_reference_count(objrefs);
  self->ob_type->tp_free((PyObject*) self);
 }

@@ -36,9 +42,13 @@ static PyObject* PyObjRef_new(PyTypeObject *type, PyObject *args, PyObject *kwds
 }

 static int PyObjRef_init(PyObjRef *self, PyObject *args, PyObject *kwds) {
-  if (!PyArg_ParseTuple(args, "i", &self->val)) {
+  if (!PyArg_ParseTuple(args, "iO&", &self->val, &PyObjectToWorker, &self->worker)) {
    return -1;
  }
+  std::vector<ObjRef> objrefs;
+  objrefs.push_back(self->val);
+  ORCH_LOG(ORCH_DEBUG, "In PyObjRef_init, calling increment_reference_count for objref " << objrefs[0]);
+  self->worker->increment_reference_count(objrefs);
  return 0;
 };

@@ -65,7 +75,7 @@ static PyTypeObject PyObjRefType = {
  "orchpy.ObjRef",           /* tp_name */
  sizeof(PyObjRef),          /* tp_basicsize */
  0,                         /* tp_itemsize */
-  0,                         /* tp_dealloc */
+  (destructor)PyObjRef_dealloc,          /* tp_dealloc */
  0,                         /* tp_print */
  0,                         /* tp_getattr */
  0,                         /* tp_setattr */
@@ -102,8 +112,8 @@ static PyTypeObject PyObjRefType = {
 };

 // create PyObjRef from C++ (could be made more efficient if neccessary)
-PyObject* make_pyobjref(ObjRef objref) {
-  PyObject* arglist = Py_BuildValue("(i)", objref);
+PyObject* make_pyobjref(PyObject* worker_capsule, ObjRef objref) {
+  PyObject* arglist = Py_BuildValue("(iO)", objref, worker_capsule);
  PyObject* result = PyObject_CallObject((PyObject*) &PyObjRefType, arglist);
  Py_DECREF(arglist);
  return result;
@@ -113,6 +123,8 @@ PyObject* make_pyobjref(ObjRef objref) {

 static PyObject *OrchPyError;

+// Pass arguments from Python to C++
+
 int PyObjectToCall(PyObject* object, Call **call) {
  if (PyCapsule_IsValid(object, "call")) {
    *call = static_cast<Call*>(PyCapsule_GetPointer(object, "call"));
@@ -153,12 +165,30 @@ int PyObjectToObjRef(PyObject* object, ObjRef *objref) {
  }
 }

+// Destructors
+
+void ObjCapsule_Destructor(PyObject* capsule) {
+  Obj* obj = static_cast<Obj*>(PyCapsule_GetPointer(capsule, "obj"));
+  delete obj;
+}
+
+void WorkerCapsule_Destructor(PyObject* capsule) {
+  Worker* obj = static_cast<Worker*>(PyCapsule_GetPointer(capsule, "worker"));
+  delete obj;
+}
+
+void CallCapsule_Destructor(PyObject* capsule) {
+  Call* obj = static_cast<Call*>(PyCapsule_GetPointer(capsule, "call"));
+  delete obj;
+}
+
 // Serialization

 // serialize will serialize the python object val into the protocol buffer
 // object obj, returns 0 if successful and something else if not
 // FIXME(pcm): This currently only works for contiguous arrays
-int serialize(PyObject* val, Obj* obj) {
+// This method will push all of the object references contained in `obj` to the `objrefs` vector.
+int serialize(PyObject* worker_capsule, PyObject* val, Obj* obj, std::vector<ObjRef> &objrefs) {
  if (PyInt_Check(val)) {
    Int* data = obj->mutable_int_data();
    long d = PyInt_AsLong(val);
@@ -171,7 +201,7 @@ int serialize(PyObject* val, Obj* obj) {
    Tuple* data = obj->mutable_tuple_data();
    for (size_t i = 0, size = PyTuple_Size(val); i < size; ++i) {
      Obj* elem = data->add_elem();
-      if (serialize(PyTuple_GetItem(val, i), elem) != 0) {
+      if (serialize(worker_capsule, PyTuple_GetItem(val, i), elem, objrefs) != 0) {
        return -1;
      }
    }
@@ -179,7 +209,7 @@ int serialize(PyObject* val, Obj* obj) {
    List* data = obj->mutable_list_data();
    for (size_t i = 0, size = PyList_Size(val); i < size; ++i) {
      Obj* elem = data->add_elem();
-      if (serialize(PyList_GetItem(val, i), elem) != 0) {
+      if (serialize(worker_capsule, PyList_GetItem(val, i), elem, objrefs) != 0) {
        return -1;
      }
    }
@@ -190,11 +220,11 @@ int serialize(PyObject* val, Obj* obj) {
    while (PyDict_Next(val, &pos, &pykey, &pyvalue)) {
      DictEntry* elem = data->add_elem();
      Obj* key = elem->mutable_key();
-      if (serialize(pykey, key) != 0) {
+      if (serialize(worker_capsule, pykey, key, objrefs) != 0) {
        return -1;
      }
      Obj* value = elem->mutable_value();
-      if (serialize(pyvalue, value) != 0) {
+      if (serialize(worker_capsule, pyvalue, value, objrefs) != 0) {
        return -1;
      }
    }
@@ -267,6 +297,7 @@ int serialize(PyObject* val, Obj* obj) {
              return -1;
            }
            data->add_objref_data(objref);
+            objrefs.push_back(objref);
            PyArray_ITER_NEXT(iter);
          }
          Py_XDECREF(iter);
@@ -276,6 +307,7 @@ int serialize(PyObject* val, Obj* obj) {
        PyErr_SetString(OrchPyError, "serialization: numpy datatype not know");
        return -1;
    }
+    Py_DECREF(array); // TODO(rkn): is this right?
  } else {
    PyErr_SetString(OrchPyError, "serialization: type not know");
    return -1;
@@ -283,7 +315,8 @@ int serialize(PyObject* val, Obj* obj) {
  return 0;
 }

-PyObject* deserialize(const Obj& obj) {
+// This method will push all of the object references contained in `obj` to the `objrefs` vector.
+PyObject* deserialize(PyObject* worker_capsule, const Obj& obj, std::vector<ObjRef> &objrefs) {
  if (obj.has_int_data()) {
    return PyInt_FromLong(obj.int_data().data());
  } else if (obj.has_double_data()) {
@@ -293,7 +326,7 @@ PyObject* deserialize(const Obj& obj) {
    size_t size = data.elem_size();
    PyObject* tuple = PyTuple_New(size);
    for (size_t i = 0; i < size; ++i) {
-      PyTuple_SetItem(tuple, i, deserialize(data.elem(i)));
+      PyTuple_SetItem(tuple, i, deserialize(worker_capsule, data.elem(i), objrefs));
    }
    return tuple;
  } else if (obj.has_list_data()) {
@@ -301,7 +334,7 @@ PyObject* deserialize(const Obj& obj) {
    size_t size = data.elem_size();
    PyObject* list = PyList_New(size);
    for (size_t i = 0; i < size; ++i) {
-      PyList_SetItem(list, i, deserialize(data.elem(i)));
+      PyList_SetItem(list, i, deserialize(worker_capsule, data.elem(i), objrefs));
    }
    return list;
  } else if (obj.has_dict_data()) {
@@ -309,7 +342,7 @@ PyObject* deserialize(const Obj& obj) {
    PyObject* dict = PyDict_New();
    size_t size = data.elem_size();
    for (size_t i = 0; i < size; ++i) {
-      PyDict_SetItem(dict, deserialize(data.elem(i).key()), deserialize(data.elem(i).value()));
+      PyDict_SetItem(dict, deserialize(worker_capsule, data.elem(i).key(), objrefs), deserialize(worker_capsule, data.elem(i).value(), objrefs));
    }
    return dict;
  } else if (obj.has_string_data()) {
@@ -381,7 +414,8 @@ PyObject* deserialize(const Obj& obj) {
      npy_intp size = array.objref_data_size();
      PyObject** buffer = (PyObject**) PyArray_DATA(pyarray);
      for (npy_intp i = 0; i < size; ++i) {
-        buffer[i] = make_pyobjref(array.objref_data(i));
+        buffer[i] = make_pyobjref(worker_capsule, array.objref_data(i));
+        objrefs.push_back(array.objref_data(i));
      }
    } else {
      PyErr_SetString(OrchPyError, "deserialization: internal error (array type not implemented)");
@@ -394,16 +428,28 @@ PyObject* deserialize(const Obj& obj) {
  }
 }

+// This returns the serialized object and a list of the object references contained in that object.
 PyObject* serialize_object(PyObject* self, PyObject* args) {
  Obj* obj = new Obj(); // TODO: to be freed in capsul destructor
+  PyObject* worker_capsule;
  PyObject* pyval;
-  if (!PyArg_ParseTuple(args, "O", &pyval)) {
+  if (!PyArg_ParseTuple(args, "OO", &worker_capsule, &pyval)) {
    return NULL;
  }
-  if (serialize(pyval, obj) != 0) {
+  std::vector<ObjRef> objrefs;
+  if (serialize(worker_capsule, pyval, obj, objrefs) != 0) {
    return NULL;
  }
-  return PyCapsule_New(static_cast<void*>(obj), "obj", NULL);
+  Worker* worker;
+  PyObjectToWorker(worker_capsule, &worker);
+  PyObject* contained_objrefs = PyList_New(objrefs.size());
+  for (int i = 0; i < objrefs.size(); ++i) {
+    PyList_SetItem(contained_objrefs, i, make_pyobjref(worker_capsule, objrefs[i]));
+  }
+  PyObject* t = PyTuple_New(2); // We set the items of the tuple using PyTuple_SetItem, because that transfers ownership to the tuple.
+  PyTuple_SetItem(t, 0, PyCapsule_New(static_cast<void*>(obj), "obj", &ObjCapsule_Destructor));
+  PyTuple_SetItem(t, 1, contained_objrefs);
+  return t;
 }

 PyObject* put_arrow(PyObject* self, PyObject* args) {
@@ -444,62 +490,89 @@ PyObject* is_arrow(PyObject* self, PyObject* args) {
 }

 PyObject* deserialize_object(PyObject* self, PyObject* args) {
+  PyObject* worker_capsule;
  Obj* obj;
-  if (!PyArg_ParseTuple(args, "O&", &PyObjectToObj, &obj)) {
+  if (!PyArg_ParseTuple(args, "OO&", &worker_capsule, &PyObjectToObj, &obj)) {
    return NULL;
  }
-  return deserialize(*obj);
+  std::vector<ObjRef> objrefs; // This is a vector of all the objrefs that are serialized in this call, including objrefs that are contained in Python objects that are passed by value.
+  return deserialize(worker_capsule, *obj, objrefs);
+  // TODO(rkn): Should we do anything with objrefs?
 }

 PyObject* serialize_call(PyObject* self, PyObject* args) {
+  PyObject* worker_capsule;
  Call* call = new Call(); // TODO: to be freed in capsul destructor
  char* name;
  int len;
  PyObject* arguments;
-  if (!PyArg_ParseTuple(args, "s#O", &name, &len, &arguments)) {
+  if (!PyArg_ParseTuple(args, "Os#O", &worker_capsule, &name, &len, &arguments)) {
    return NULL;
  }
  call->set_name(name, len);
+  std::vector<ObjRef> objrefs; // This is a vector of all the objrefs that are serialized in this call, including objrefs that are contained in Python objects that are passed by value.
  if (PyList_Check(arguments)) {
    for (size_t i = 0, size = PyList_Size(arguments); i < size; ++i) {
      PyObject* element = PyList_GetItem(arguments, i);
      if (PyObject_IsInstance(element, (PyObject*)&PyObjRefType)) {
        ObjRef objref = ((PyObjRef*) element)->val;
        call->add_arg()->set_ref(objref);
+        objrefs.push_back(objref);
      } else {
        Obj* arg = call->add_arg()->mutable_obj();
-        serialize(PyList_GetItem(arguments, i), arg);
+        serialize(worker_capsule, PyList_GetItem(arguments, i), arg, objrefs);
      }
    }
  } else {
    PyErr_SetString(OrchPyError, "serialize_call: second argument needs to be a list");
    return NULL;
  }
-  return PyCapsule_New(static_cast<void*>(call), "call", NULL);
+  Worker* worker;
+  PyObjectToWorker(worker_capsule, &worker);
+  if (objrefs.size() > 0) {
+    ORCH_LOG(ORCH_DEBUG, "In serialize_call, calling increment_reference_count for objrefs:");
+    for (int i = 0; i < objrefs.size(); ++i) {
+      ORCH_LOG(ORCH_DEBUG, "----" << objrefs[i]);
+    }
+    worker->increment_reference_count(objrefs);
+  }
+  return PyCapsule_New(static_cast<void*>(call), "call", &CallCapsule_Destructor);
 }

 PyObject* deserialize_call(PyObject* self, PyObject* args) {
+  PyObject* worker_capsule;
  Call* call;
-  if (!PyArg_ParseTuple(args, "O&", &PyObjectToCall, &call)) {
+  if (!PyArg_ParseTuple(args, "OO&", &worker_capsule, &PyObjectToCall, &call)) {
    return NULL;
  }
+  std::vector<ObjRef> objrefs; // This is a vector of all the objrefs that were serialized in this call, including objrefs that are contained in Python objects that are passed by value.
  PyObject* string = PyString_FromStringAndSize(call->name().c_str(), call->name().size());
  int argsize = call->arg_size();
  PyObject* arglist = PyList_New(argsize);
  for (int i = 0; i < argsize; ++i) {
    const Value& val = call->arg(i);
    if (!val.has_obj()) {
-      PyList_SetItem(arglist, i, make_pyobjref(val.ref()));
+      PyList_SetItem(arglist, i, make_pyobjref(worker_capsule, val.ref()));
+      objrefs.push_back(val.ref());
    } else {
-      PyList_SetItem(arglist, i, deserialize(val.obj()));
+      PyList_SetItem(arglist, i, deserialize(worker_capsule, val.obj(), objrefs));
    }
  }
+  Worker* worker;
+  PyObjectToWorker(worker_capsule, &worker);
+  if (objrefs.size() > 0) {
+    worker->decrement_reference_count(objrefs);
+  }
  int resultsize = call->result_size();
  PyObject* resultlist = PyList_New(resultsize);
  for (int i = 0; i < resultsize; ++i) {
-    PyList_SetItem(resultlist, i, make_pyobjref(call->result(i)));
+    PyList_SetItem(resultlist, i, make_pyobjref(worker_capsule, call->result(i)));
  }
-  return PyTuple_Pack(3, string, arglist, resultlist);
+  PyObject* t = PyTuple_New(3); // We set the items of the tuple using PyTuple_SetItem, because that transfers ownership to the tuple.
+  PyTuple_SetItem(t, 0, string);
+  PyTuple_SetItem(t, 1, arglist);
+  PyTuple_SetItem(t, 2, resultlist);
+  return t;
 }

 // Orchestra Python API
@@ -515,7 +588,27 @@ PyObject* create_worker(PyObject* self, PyObject* args) {
  auto objstore_channel = grpc::CreateChannel(objstore_addr, grpc::InsecureChannelCredentials());
  Worker* worker = new Worker(std::string(worker_addr), scheduler_channel, objstore_channel);
  worker->register_worker(std::string(worker_addr), std::string(objstore_addr));
-  return PyCapsule_New(static_cast<void*>(worker), "worker", NULL); // TODO: add destructor the deallocates worker
+  return PyCapsule_New(static_cast<void*>(worker), "worker", &WorkerCapsule_Destructor);
+}
+
+PyObject* disconnect(PyObject* self, PyObject* args) {
+  Worker* worker;
+  if (!PyArg_ParseTuple(args, "O&", &PyObjectToWorker, &worker)) {
+    return NULL;
+  }
+  worker->disconnect();
+  Py_RETURN_NONE;
+}
+
+PyObject* connected(PyObject* self, PyObject* args) {
+  Worker* worker;
+  if (!PyArg_ParseTuple(args, "O&", &PyObjectToWorker, &worker)) {
+    return NULL;
+  }
+  if (worker->connected()) {
+    Py_RETURN_TRUE;
+  }
+  Py_RETURN_FALSE;
 }

 PyObject* wait_for_next_task(PyObject* self, PyObject* args) {
@@ -524,15 +617,17 @@ PyObject* wait_for_next_task(PyObject* self, PyObject* args) {
    return NULL;
  }
  Call* call = worker->receive_next_task();
-  return PyCapsule_New(static_cast<void*>(call), "call", NULL); // TODO: how is destruction going to be handled here?
+  return PyCapsule_New(static_cast<void*>(call), "call", NULL); // This call is owned by the C++ worker class, so we do not deallocate it.
 }

 PyObject* remote_call(PyObject* self, PyObject* args) {
-  Worker* worker;
+  PyObject* worker_capsule;
  Call* call;
-  if (!PyArg_ParseTuple(args, "O&O&", &PyObjectToWorker, &worker, &PyObjectToCall, &call)) {
+  if (!PyArg_ParseTuple(args, "OO&", &worker_capsule, &PyObjectToCall, &call)) {
    return NULL;
  }
+  Worker* worker;
+  PyObjectToWorker(worker_capsule, &worker);
  RemoteCallRequest request;
  request.set_allocated_call(call);
  RemoteCallReply reply = worker->remote_call(&request);
@@ -540,7 +635,7 @@ PyObject* remote_call(PyObject* self, PyObject* args) {
  int size = reply.result_size();
  PyObject* list = PyList_New(size);
  for (int i = 0; i < size; ++i) {
-    PyList_SetItem(list, i, make_pyobjref(reply.result(i)));
+    PyList_SetItem(list, i, make_pyobjref(worker_capsule, reply.result(i)));
  }
  return list;
 }
@@ -566,22 +661,35 @@ PyObject* register_function(PyObject* self, PyObject* args) {
 }

 PyObject* get_objref(PyObject* self, PyObject* args) {
-  Worker* worker;
-  if (!PyArg_ParseTuple(args, "O&", &PyObjectToWorker, &worker)) {
+  PyObject* worker_capsule;
+  if (!PyArg_ParseTuple(args, "O", &worker_capsule)) {
    return NULL;
  }
+  Worker* worker;
+  PyObjectToWorker(worker_capsule, &worker);
  ObjRef objref = worker->get_objref();
-  return make_pyobjref(objref);
+  return make_pyobjref(worker_capsule, objref);
 }

 PyObject* put_object(PyObject* self, PyObject* args) {
  Worker* worker;
  ObjRef objref;
  Obj* obj;
-  if (!PyArg_ParseTuple(args, "O&O&O&", &PyObjectToWorker, &worker, &PyObjectToObjRef, &objref, &PyObjectToObj, &obj)) {
+  PyObject* contained_objrefs;
+  if (!PyArg_ParseTuple(args, "O&O&O&O", &PyObjectToWorker, &worker, &PyObjectToObjRef, &objref, &PyObjectToObj, &obj, &contained_objrefs)) {
    return NULL;
  }
-  worker->put_object(objref, obj);
+  if (!PyList_Check(contained_objrefs)) {
+    ORCH_LOG(ORCH_FATAL, "The contained_objrefs argument must be a list.")
+  }
+  std::vector<ObjRef> vec_contained_objrefs;
+  size_t size = PyList_Size(contained_objrefs);
+  for (size_t i = 0; i < size; ++i) {
+    ObjRef contained_objref;
+    PyObjectToObjRef(PyList_GetItem(contained_objrefs, i), &contained_objref);
+    vec_contained_objrefs.push_back(contained_objref);
+  }
+  worker->put_object(objref, obj, vec_contained_objrefs);
  Py_RETURN_NONE;
 }

@@ -595,7 +703,7 @@ PyObject* get_object(PyObject* self, PyObject* args) {
  slice s = worker->get_object(objref);
  Obj* obj = new Obj(); // TODO: Make sure this will get deleted
  obj->ParseFromString(std::string(reinterpret_cast<char*>(s.data), s.len));
-  return PyCapsule_New(static_cast<void*>(obj), "obj", NULL);
+  return PyCapsule_New(static_cast<void*>(obj), "obj", &ObjCapsule_Destructor);
 }

 PyObject* request_object(PyObject* self, PyObject* args) {
@@ -628,6 +736,31 @@ PyObject* start_worker_service(PyObject* self, PyObject* args) {
  Py_RETURN_NONE;
 }

+PyObject* scheduler_info(PyObject* self, PyObject* args) {
+  Worker* worker;
+  if (!PyArg_ParseTuple(args, "O&", &PyObjectToWorker, &worker)) {
+    return NULL;
+  }
+  ClientContext context;
+  SchedulerInfoRequest request;
+  SchedulerInfoReply reply;
+  worker->scheduler_info(context, request, reply);
+
+  PyObject* target_objref_list = PyList_New(reply.target_objref_size());
+  for (size_t i = 0; i < reply.target_objref_size(); ++i) {
+    PyList_SetItem(target_objref_list, i, PyInt_FromLong(reply.target_objref(i)));
+  }
+  PyObject* reference_count_list = PyList_New(reply.reference_count_size());
+  for (size_t i = 0; i < reply.reference_count_size(); ++i) {
+    PyList_SetItem(reference_count_list, i, PyInt_FromLong(reply.reference_count(i)));
+  }
+
+  PyObject* dict = PyDict_New();
+  PyDict_SetItem(dict, PyString_FromString("target_objrefs"), target_objref_list);
+  PyDict_SetItem(dict, PyString_FromString("reference_counts"), reference_count_list);
+  return dict;
+}
+
 static PyMethodDef OrchPyLibMethods[] = {
 { "serialize_object", serialize_object, METH_VARARGS, "serialize an object to protocol buffers" },
 { "deserialize_object", deserialize_object, METH_VARARGS, "deserialize an object from protocol buffers" },
@@ -637,6 +770,8 @@ static PyMethodDef OrchPyLibMethods[] = {
 { "serialize_call", serialize_call, METH_VARARGS, "serialize a call to protocol buffers" },
 { "deserialize_call", deserialize_call, METH_VARARGS, "deserialize a call from protocol buffers" },
 { "create_worker", create_worker, METH_VARARGS, "connect to the scheduler and the object store" },
+ { "disconnect", disconnect, METH_VARARGS, "disconnect the worker from the scheduler and the object store" },
+ { "connected", connected, METH_VARARGS, "check if the worker is connected to the scheduler and the object store" },
 { "register_function", register_function, METH_VARARGS, "register a function with the scheduler" },
 { "put_object", put_object, METH_VARARGS, "put a protocol buffer object (given as a capsule) on the local object store" },
 { "get_object", get_object, METH_VARARGS, "get protocol buffer object from the local object store" },
@@ -647,6 +782,7 @@ static PyMethodDef OrchPyLibMethods[] = {
 { "remote_call", remote_call, METH_VARARGS, "call a remote function" },
 { "notify_task_completed", notify_task_completed, METH_VARARGS, "notify the scheduler that a task has been completed" },
 { "start_worker_service", start_worker_service, METH_VARARGS, "start the worker service" },
+ { "scheduler_info", scheduler_info, METH_VARARGS, "get info about scheduler state" },
 { NULL, NULL, 0, NULL }
 };

@@ -9,7 +9,7 @@ Status SchedulerService::RemoteCall(ServerContext* context, const RemoteCallRequ
  fntable_lock_.lock();

  if (fntable_.find(task->name()) == fntable_.end()) {
-    // TODO(rkn): In the future, this should probably not be fatal.
+    // TODO(rkn): In the future, this should probably not be fatal. Instead, propagate the error back to the worker.
    ORCH_LOG(ORCH_FATAL, "The function " << task->name() << " has not been registered by any worker.");
  }

@@ -51,7 +51,6 @@ Status SchedulerService::RequestObj(ServerContext* context, const RequestObjRequ
  pull_queue_lock_.lock();
  pull_queue_.push_back(std::make_pair(request->workerid(), objref));
  pull_queue_lock_.unlock();
-
  schedule();
  return Status::OK;
 }
@@ -59,26 +58,30 @@ Status SchedulerService::RequestObj(ServerContext* context, const RequestObjRequ
 Status SchedulerService::AliasObjRefs(ServerContext* context, const AliasObjRefsRequest* request, AckReply* reply) {
  ObjRef alias_objref = request->alias_objref();
  ObjRef target_objref = request->target_objref();
-
+  ORCH_LOG(ORCH_DEBUG, "Aliasing objref " << alias_objref << " with objref " << target_objref);
+  if (alias_objref == target_objref) {
+    ORCH_LOG(ORCH_FATAL, "internal error: attempting to alias objref " << alias_objref << " with itself.");
+  }
  objtable_lock_.lock();
  size_t size = objtable_.size();
  objtable_lock_.unlock();
-
  if (alias_objref >= size) {
    ORCH_LOG(ORCH_FATAL, "internal error: no object with objref " << alias_objref << " exists");
  }
  if (target_objref >= size) {
    ORCH_LOG(ORCH_FATAL, "internal error: no object with objref " << target_objref << " exists");
  }
-
  {
-    std::lock_guard<std::mutex> objstore_lock(target_objrefs_lock_);
+    std::lock_guard<std::mutex> target_objrefs_lock(target_objrefs_lock_);
    if (target_objrefs_[alias_objref] != UNITIALIZED_ALIAS) {
      ORCH_LOG(ORCH_FATAL, "internal error: attempting to alias objref " << alias_objref << " with objref " << target_objref << ", but objref " << alias_objref << " has already been aliased with objref " << target_objrefs_[alias_objref]);
    }
    target_objrefs_[alias_objref] = target_objref;
  }
-
+  {
+    std::lock_guard<std::mutex> reverse_target_objrefs_lock(reverse_target_objrefs_lock_);
+    reverse_target_objrefs_[target_objref].push_back(alias_objref);
+  }
  schedule();
  return Status::OK;
 }
@@ -112,7 +115,7 @@ Status SchedulerService::RegisterFunction(ServerContext* context, const Register

 Status SchedulerService::ObjReady(ServerContext* context, const ObjReadyRequest* request, AckReply* reply) {
  ObjRef objref = request->objref();
-  ORCH_LOG(ORCH_VERBOSE, "object " << objref << " ready on store " << request->objstoreid());
+  ORCH_LOG(ORCH_DEBUG, "object " << objref << " ready on store " << request->objstoreid());
  add_canonical_objref(objref);
  add_location(objref, request->objstoreid());
  schedule();
@@ -129,8 +132,52 @@ Status SchedulerService::WorkerReady(ServerContext* context, const WorkerReadyRe
  return Status::OK;
 }

-Status SchedulerService::SchedulerDebugInfo(ServerContext* context, const SchedulerDebugInfoRequest* request, SchedulerDebugInfoReply* reply) {
-  debug_info(*request, reply);
+Status SchedulerService::IncrementRefCount(ServerContext* context, const IncrementRefCountRequest* request, AckReply* reply) {
+  int num_objrefs = request->objref_size();
+  if (num_objrefs == 0) {
+    ORCH_LOG(ORCH_FATAL, "Scheduler received IncrementRefCountRequest with 0 objrefs.");
+  }
+  std::vector<ObjRef> objrefs;
+  for (int i = 0; i < num_objrefs; ++i) {
+    objrefs.push_back(request->objref(i));
+  }
+  std::lock_guard<std::mutex> reference_counts_lock(reference_counts_lock_); // we grab this lock because increment_ref_count assumes it has been acquired
+  increment_ref_count(objrefs);
+  return Status::OK;
+}
+
+Status SchedulerService::DecrementRefCount(ServerContext* context, const DecrementRefCountRequest* request, AckReply* reply) {
+  int num_objrefs = request->objref_size();
+  if (num_objrefs == 0) {
+    ORCH_LOG(ORCH_FATAL, "Scheduler received DecrementRefCountRequest with 0 objrefs.");
+  }
+  std::vector<ObjRef> objrefs;
+  for (int i = 0; i < num_objrefs; ++i) {
+    objrefs.push_back(request->objref(i));
+  }
+  std::lock_guard<std::mutex> reference_counts_lock(reference_counts_lock_); // we grab this lock, because decrement_ref_count assumes it has been acquired
+  decrement_ref_count(objrefs);
+  return Status::OK;
+}
+
+Status SchedulerService::AddContainedObjRefs(ServerContext* context, const AddContainedObjRefsRequest* request, AckReply* reply) {
+  ObjRef objref = request->objref();
+  // if (!is_canonical(objref)) {
+    // TODO(rkn): Perhaps we don't need this check. It won't work because the objstore may not have called ObjReady yet.
+    // ORCH_LOG(ORCH_FATAL, "Attempting to add contained objrefs for non-canonical objref " << objref);
+  // }
+  std::lock_guard<std::mutex> contained_objrefs_lock(contained_objrefs_lock_);
+  if (contained_objrefs_[objref].size() != 0) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to add contained objrefs for objref " << objref << ", but contained_objrefs_[objref].size() != 0.");
+  }
+  for (int i = 0; i < request->contained_objref_size(); ++i) {
+    contained_objrefs_[objref].push_back(request->contained_objref(i));
+  }
+  return Status::OK;
+}
+
+Status SchedulerService::SchedulerInfo(ServerContext* context, const SchedulerInfoRequest* request, SchedulerInfoReply* reply) {
+  get_info(*request, reply);
  return Status::OK;
 }

@@ -239,15 +286,34 @@ WorkerId SchedulerService::register_worker(const std::string& worker_address, co

 ObjRef SchedulerService::register_new_object() {
  // If we don't simultaneously lock objtable_ and target_objrefs_, we will probably get errors.
+  // TODO(rkn): increment/decrement_reference_count also acquire reference_counts_lock_ and target_objrefs_lock_ (through has_canonical_objref()), which caused deadlock in the past
+  std::lock_guard<std::mutex> reference_counts_lock(reference_counts_lock_);
+  std::lock_guard<std::mutex> contained_objrefs_lock(contained_objrefs_lock_);
  std::lock_guard<std::mutex> objtable_lock(objtable_lock_);
  std::lock_guard<std::mutex> target_objrefs_lock(target_objrefs_lock_);
+  std::lock_guard<std::mutex> reverse_target_objrefs_lock(reverse_target_objrefs_lock_);
  ObjRef objtable_size = objtable_.size();
  ObjRef target_objrefs_size = target_objrefs_.size();
+  ObjRef reverse_target_objrefs_size = reverse_target_objrefs_.size();
+  ObjRef reference_counts_size = reference_counts_.size();
+  ObjRef contained_objrefs_size = contained_objrefs_.size();
  if (objtable_size != target_objrefs_size) {
    ORCH_LOG(ORCH_FATAL, "objtable_ and target_objrefs_ should have the same size, but objtable_.size() = " << objtable_size << " and target_objrefs_.size() = " << target_objrefs_size);
  }
+  if (objtable_size != reverse_target_objrefs_size) {
+    ORCH_LOG(ORCH_FATAL, "objtable_ and reverse_target_objrefs_ should have the same size, but objtable_.size() = " << objtable_size << " and reverse_target_objrefs_.size() = " << reverse_target_objrefs_size);
+  }
+  if (objtable_size != reference_counts_size) {
+    ORCH_LOG(ORCH_FATAL, "objtable_ and reference_counts_ should have the same size, but objtable_.size() = " << objtable_size << " and reference_counts_.size() = " << reference_counts_size);
+  }
+  if (objtable_size != contained_objrefs_size) {
+    ORCH_LOG(ORCH_FATAL, "objtable_ and contained_objrefs_ should have the same size, but objtable_.size() = " << objtable_size << " and contained_objrefs_.size() = " << contained_objrefs_size);
+  }
  objtable_.push_back(std::vector<ObjStoreId>());
  target_objrefs_.push_back(UNITIALIZED_ALIAS);
+  reverse_target_objrefs_.push_back(std::vector<ObjRef>());
+  reference_counts_.push_back(0);
+  contained_objrefs_.push_back(std::vector<ObjRef>());
  return objtable_size;
 }

@@ -272,7 +338,7 @@ void SchedulerService::add_canonical_objref(ObjRef objref) {
  if (objref >= target_objrefs_.size()) {
    ORCH_LOG(ORCH_FATAL, "internal error: attempting to insert objref " << objref << " in target_objrefs_, but target_objrefs_.size() is " << target_objrefs_.size());
  }
-  if (target_objrefs_[objref] != UNITIALIZED_ALIAS) {
+  if (target_objrefs_[objref] != UNITIALIZED_ALIAS && target_objrefs_[objref] != objref) {
    ORCH_LOG(ORCH_FATAL, "internal error: attempting to declare objref " << objref << " as a canonical objref, but target_objrefs_[objref] is already aliased with objref " << target_objrefs_[objref]);
  }
  target_objrefs_[objref] = objref;
@@ -291,8 +357,25 @@ void SchedulerService::register_function(const std::string& name, WorkerId worke
  info.add_worker(workerid);
 }

-void SchedulerService::debug_info(const SchedulerDebugInfoRequest& request, SchedulerDebugInfoReply* reply) {
-  fntable_lock_.lock();
+void SchedulerService::get_info(const SchedulerInfoRequest& request, SchedulerInfoReply* reply) {
+  // TODO(rkn): Also grab the objstores_lock_
+  // alias_notification_queue_lock_ may need to come before objtable_lock_
+  std::lock_guard<std::mutex> reference_counts_lock(reference_counts_lock_);
+  std::lock_guard<std::mutex> contained_objrefs_lock(contained_objrefs_lock_);
+  std::lock_guard<std::mutex> objtable_lock(objtable_lock_);
+  std::lock_guard<std::mutex> pull_queue_lock(pull_queue_lock_);
+  std::lock_guard<std::mutex> target_objrefs_lock(target_objrefs_lock_);
+  std::lock_guard<std::mutex> reverse_target_objrefs_lock(reverse_target_objrefs_lock_);
+  std::lock_guard<std::mutex> fntable_lock(fntable_lock_);
+  std::lock_guard<std::mutex> avail_workers_lock(avail_workers_lock_);
+  std::lock_guard<std::mutex> task_queue_lock(task_queue_lock_);
+  std::lock_guard<std::mutex> alias_notification_queue_lock(alias_notification_queue_lock_);
+  for (int i = 0; i < reference_counts_.size(); ++i) {
+    reply->add_reference_count(reference_counts_[i]);
+  }
+  for (int i = 0; i < target_objrefs_.size(); ++i) {
+    reply->add_target_objref(target_objrefs_[i]);
+  }
  auto function_table = reply->mutable_function_table();
  for (const auto& entry : fntable_) {
    (*function_table)[entry.first].set_num_return_vals(entry.second.num_return_vals());
@@ -300,18 +383,14 @@ void SchedulerService::debug_info(const SchedulerDebugInfoRequest& request, Sche
      (*function_table)[entry.first].add_workerid(worker);
    }
  }
-  fntable_lock_.unlock();
-  task_queue_lock_.lock();
  for (const auto& entry : task_queue_) {
    Call* call = reply->add_task();
    call->CopyFrom(*entry);
  }
-  task_queue_lock_.unlock();
-  avail_workers_lock_.lock();
  for (const WorkerId& entry : avail_workers_) {
    reply->add_avail_worker(entry);
  }
-  avail_workers_lock_.unlock();
+
 }

 ObjStoreId SchedulerService::pick_objstore(ObjRef canonical_objref) {
@@ -333,7 +412,6 @@ bool SchedulerService::is_canonical(ObjRef objref) {
 }

 void SchedulerService::perform_pulls() {
-  std::lock_guard<std::mutex> objtable_lock(objtable_lock_);
  std::lock_guard<std::mutex> pull_queue_lock(pull_queue_lock_);
  // Complete all pull tasks that can be completed.
  for (int i = 0; i < pull_queue_.size(); ++i) {
@@ -346,12 +424,20 @@ void SchedulerService::perform_pulls() {
    }
    ObjRef canonical_objref = get_canonical_objref(objref);
    ORCH_LOG(ORCH_DEBUG, "attempting to pull objref " << pull.second << " with canonical objref " << canonical_objref);
-    if (objtable_[canonical_objref].size() > 0) {
-      if (!std::binary_search(objtable_[canonical_objref].begin(), objtable_[canonical_objref].end(), get_store(workerid))) {
-        // The worker's local object store does not already contain objref, so ship
-        // it there from an object store that does have it.
-        ObjStoreId objstoreid = pick_objstore(canonical_objref);
-        deliver_object(canonical_objref, objstoreid, get_store(workerid));
+
+    objtable_lock_.lock();
+    int num_stores = objtable_[canonical_objref].size();
+    objtable_lock_.unlock();
+
+    if (num_stores > 0) {
+      {
+        std::lock_guard<std::mutex> objtable_lock(objtable_lock_);
+        if (!std::binary_search(objtable_[canonical_objref].begin(), objtable_[canonical_objref].end(), get_store(workerid))) {
+          // The worker's local object store does not already contain objref, so ship
+          // it there from an object store that does have it.
+          ObjStoreId objstoreid = pick_objstore(canonical_objref);
+          deliver_object(canonical_objref, objstoreid, get_store(workerid));
+        }
      }
      {
        // Notify the relevant objstore about potential aliasing when it's ready
@@ -467,6 +553,96 @@ bool SchedulerService::attempt_notify_alias(ObjStoreId objstoreid, ObjRef alias_
  return true;
 }

+void SchedulerService::deallocate_object(ObjRef canonical_objref) {
+  // deallocate_object should only be called from decrement_ref_count (note that
+  // deallocate_object also recursively calls decrement_ref_count). Both of
+  // these methods require reference_counts_lock_ to have been acquired, and
+  // so the lock must before outside of these methods (it is acquired in
+  // DecrementRefCount).
+  ORCH_LOG(ORCH_DEBUG, "Deallocating canonical_objref " << canonical_objref << ".");
+  ClientContext context;
+  AckReply reply;
+  DeallocateObjectRequest request;
+  request.set_canonical_objref(canonical_objref);
+  {
+    std::lock_guard<std::mutex> objtable_lock(objtable_lock_);
+    auto &objstores = objtable_[canonical_objref];
+    std::lock_guard<std::mutex> objstores_lock(objstores_lock_); // TODO(rkn): Should this be inside the for loop instead?
+    for (int i = 0; i < objstores.size(); ++i) {
+      ObjStoreId objstoreid = objstores[i];
+      objstores_[objstoreid].objstore_stub->DeallocateObject(&context, request, &reply);
+    }
+  }
+  decrement_ref_count(contained_objrefs_[canonical_objref]);
+}
+
+void SchedulerService::increment_ref_count(std::vector<ObjRef> &objrefs) {
+  // increment_ref_count assumes that reference_counts_lock_ has been acquired already
+  for (int i = 0; i < objrefs.size(); ++i) {
+    ObjRef objref = objrefs[i];
+    if (reference_counts_[objref] == DEALLOCATED) {
+      ORCH_LOG(ORCH_FATAL, "Attempting to increment the reference count for objref " << objref << ", but this object appears to have been deallocated already.");
+    }
+    reference_counts_[objref] += 1;
+    ORCH_LOG(ORCH_DEBUG, "Incremented ref count for objref " << objref <<". New reference count is " << reference_counts_[objref]);
+  }
+}
+
+void SchedulerService::decrement_ref_count(std::vector<ObjRef> &objrefs) {
+  // decrement_ref_count assumes that reference_counts_lock_ has been acquired already
+  for (int i = 0; i < objrefs.size(); ++i) {
+    ObjRef objref = objrefs[i];
+    if (reference_counts_[objref] == DEALLOCATED) {
+      ORCH_LOG(ORCH_FATAL, "Attempting to decrement the reference count for objref " << objref << ", but this object appears to have been deallocated already.");
+    }
+    if (reference_counts_[objref] == 0) {
+      ORCH_LOG(ORCH_FATAL, "Attempting to decrement the reference count for objref " << objref << ", but the reference count for this object is already 0.");
+    }
+    reference_counts_[objref] -= 1;
+    ORCH_LOG(ORCH_DEBUG, "Decremented ref count for objref " << objref << ". New reference count is " << reference_counts_[objref]);
+    // See if we can deallocate the object
+    std::vector<ObjRef> equivalent_objrefs;
+    get_equivalent_objrefs(objref, equivalent_objrefs);
+    bool can_deallocate = true;
+    for (int j = 0; j < equivalent_objrefs.size(); ++j) {
+      if (reference_counts_[equivalent_objrefs[j]] != 0) {
+        can_deallocate = false;
+        break;
+      }
+    }
+    if (can_deallocate) {
+      ObjRef canonical_objref = equivalent_objrefs[0];
+      if (!is_canonical(canonical_objref)) {
+        ORCH_LOG(ORCH_FATAL, "canonical_objref is not canonical.");
+      }
+      deallocate_object(canonical_objref);
+      for (int j = 0; j < equivalent_objrefs.size(); ++j) {
+        reference_counts_[equivalent_objrefs[j]] = DEALLOCATED;
+      }
+    }
+  }
+  ORCH_LOG(ORCH_DEBUG, "Exiting decrement_ref_count");
+}
+
+void SchedulerService::upstream_objrefs(ObjRef objref, std::vector<ObjRef> &objrefs) {
+  // upstream_objrefs assumes that the lock reverse_target_objrefs_lock_ has been acquired
+  objrefs.push_back(objref);
+  for (int i = 0; i < reverse_target_objrefs_[objref].size(); ++i) {
+    upstream_objrefs(reverse_target_objrefs_[objref][i], objrefs);
+  }
+}
+
+void SchedulerService::get_equivalent_objrefs(ObjRef objref, std::vector<ObjRef> &equivalent_objrefs) {
+  std::lock_guard<std::mutex> target_objrefs_lock(target_objrefs_lock_);
+  ObjRef downstream_objref = objref;
+  while (target_objrefs_[downstream_objref] != downstream_objref && target_objrefs_[downstream_objref] != UNITIALIZED_ALIAS) {
+    ORCH_LOG(ORCH_DEBUG, "Looping in get_equivalent_objrefs");
+    downstream_objref = target_objrefs_[downstream_objref];
+  }
+  std::lock_guard<std::mutex> reverse_target_objrefs_lock(reverse_target_objrefs_lock_);
+  upstream_objrefs(downstream_objref, equivalent_objrefs);
+}
+
 void start_scheduler_service(const char* server_address) {
  SchedulerService service;
  ServerBuilder builder;
@@ -24,7 +24,10 @@ using grpc::ClientContext;

 using grpc::Channel;

+typedef size_t RefCount;
+
 const ObjRef UNITIALIZED_ALIAS = std::numeric_limits<ObjRef>::max();
+const RefCount DEALLOCATED = std::numeric_limits<RefCount>::max();

 struct WorkerHandle {
  std::shared_ptr<Channel> channel;
@@ -49,7 +52,10 @@ public:
  Status RegisterFunction(ServerContext* context, const RegisterFunctionRequest* request, AckReply* reply) override;
  Status ObjReady(ServerContext* context, const ObjReadyRequest* request, AckReply* reply) override;
  Status WorkerReady(ServerContext* context, const WorkerReadyRequest* request, AckReply* reply) override;
-  Status SchedulerDebugInfo(ServerContext* context, const SchedulerDebugInfoRequest* request, SchedulerDebugInfoReply* reply) override;
+  Status IncrementRefCount(ServerContext* context, const IncrementRefCountRequest* request, AckReply* reply) override;
+  Status DecrementRefCount(ServerContext* context, const DecrementRefCountRequest* request, AckReply* reply) override;
+  Status AddContainedObjRefs(ServerContext* context, const AddContainedObjRefsRequest* request, AckReply* reply) override;
+  Status SchedulerInfo(ServerContext* context, const SchedulerInfoRequest* request, SchedulerInfoReply* reply) override;

  // ask an object store to send object to another objectstore
  void deliver_object(ObjRef objref, ObjStoreId from, ObjStoreId to);
@@ -71,8 +77,8 @@ public:
  ObjStoreId get_store(WorkerId workerid);
  // register a function with the scheduler
  void register_function(const std::string& name, WorkerId workerid, size_t num_return_vals);
-  // get debugging information for the scheduler
-  void debug_info(const SchedulerDebugInfoRequest& request, SchedulerDebugInfoReply* reply);
+  // get information about the scheduler state
+  void get_info(const SchedulerInfoRequest& request, SchedulerInfoReply* reply);
 private:
  // pick an objectstore that holds a given object (needs protection by objtable_lock_)
  ObjStoreId pick_objstore(ObjRef objref);
@@ -89,6 +95,16 @@ private:
  ObjRef get_canonical_objref(ObjRef objref);
  // attempt to notify the objstore about potential objref aliasing, returns true if successful, if false then retry later
  bool attempt_notify_alias(ObjStoreId objstoreid, ObjRef alias_objref, ObjRef canonical_objref);
+  // tell all of the objstores holding canonical_objref to deallocate it
+  void deallocate_object(ObjRef canonical_objref);
+  // increment the ref counts for the object references in objrefs
+  void increment_ref_count(std::vector<ObjRef> &objrefs);
+  // decrement the ref counts for the object references in objrefs
+  void decrement_ref_count(std::vector<ObjRef> &objrefs);
+  // Find all of the object references which are upstream of objref (including objref itself). That is, you can get from everything in objrefs to objref by repeatedly indexing in target_objrefs_.
+  void upstream_objrefs(ObjRef objref, std::vector<ObjRef> &objrefs);
+  // Find all of the object references that refer to the same object as objref (as best as we can determine at the moment). The information may be incomplete because not all of the aliases may be known.
+  void get_equivalent_objrefs(ObjRef objref, std::vector<ObjRef> &equivalent_objrefs);

  // Vector of all workers registered in the system. Their index in this vector
  // is the workerid.
@@ -101,7 +117,6 @@ private:
  // vector is the objstoreid.
  std::vector<ObjStoreHandle> objstores_;
  grpc::mutex objstores_lock_;
-
  // Mapping from an aliased objref to the objref it is aliased with. If an
  // objref is a canonical objref (meaning it is not aliased), then
  // target_objrefs_[objref] == objref. For each objref, target_objrefs_[objref]
@@ -109,7 +124,9 @@ private:
  // when it is known.
  std::vector<ObjRef> target_objrefs_;
  std::mutex target_objrefs_lock_;
-
+  // This data structure maps an objref to all of the objrefs that alias it (there could be multiple such objrefs).
+  std::vector<std::vector<ObjRef> > reverse_target_objrefs_;
+  std::mutex reverse_target_objrefs_lock_;
  // Mapping from canonical objref to list of object stores where the object is stored. Non-canonical (aliased) objrefs should not be used to index objtable_.
  ObjTable objtable_;
  std::mutex objtable_lock_;
@@ -125,6 +142,13 @@ private:
  // List of pending alias notifications. Each element consists of (objstoreid, (alias_objref, canonical_objref)).
  std::vector<std::pair<ObjStoreId, std::pair<ObjRef, ObjRef> > > alias_notification_queue_;
  std::mutex alias_notification_queue_lock_;
+  // Reference counts. Currently, reference_counts_[objref] is the number of existing references
+  // held to objref. This is done for all objrefs, not just canonical_objrefs. This data structure completely ignores aliasing.
+  std::vector<RefCount> reference_counts_;
+  std::mutex reference_counts_lock_;
+  // contained_objrefs_[objref] is a vector of all of the objrefs contained inside the object referred to by objref
+  std::vector<std::vector<ObjRef> > contained_objrefs_;
+  std::mutex contained_objrefs_lock_;
 };

 #endif
@@ -13,9 +13,13 @@ Worker::Worker(const std::string& worker_address, std::shared_ptr<Channel> sched
      scheduler_stub_(Scheduler::NewStub(scheduler_channel)),
      objstore_stub_(ObjStore::NewStub(objstore_channel)) {
  receive_queue_.connect(worker_address_, true);
+  connected_ = true;
 }

 RemoteCallReply Worker::remote_call(RemoteCallRequest* request) {
+  if (!connected_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to perform remote_call, but connected_ = " << connected_ << ".");
+  }
  RemoteCallReply reply;
  ClientContext context;
  Status status = scheduler_stub_->RemoteCall(&context, *request, &reply);
@@ -37,6 +41,9 @@ void Worker::register_worker(const std::string& worker_address, const std::strin
 }

 void Worker::request_object(ObjRef objref) {
+  if (!connected_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to perform request_object, but connected_ = " << connected_ << ".");
+  }
  RequestObjRequest request;
  request.set_workerid(workerid_);
  request.set_objref(objref);
@@ -48,6 +55,9 @@ void Worker::request_object(ObjRef objref) {

 ObjRef Worker::get_objref() {
  // first get objref for the new object
+  if (!connected_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to perform get_objref, but connected_ = " << connected_ << ".");
+  }
  PushObjRequest push_request;
  PushObjReply push_reply;
  ClientContext push_context;
@@ -57,6 +67,9 @@ ObjRef Worker::get_objref() {

 slice Worker::get_object(ObjRef objref) {
  // get_object assumes that objref is a canonical objref
+  if (!connected_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to perform get_object, but connected_ = " << connected_ << ".");
+  }
  ObjRequest request;
  request.workerid = workerid_;
  request.type = ObjRequestType::GET;
@@ -71,7 +84,11 @@ slice Worker::get_object(ObjRef objref) {
 }

 // TODO(pcm): More error handling
-void Worker::put_object(ObjRef objref, const Obj* obj) {
+// contained_objrefs is a vector of all the objrefs contained in obj
+void Worker::put_object(ObjRef objref, const Obj* obj, std::vector<ObjRef> &contained_objrefs) {
+  if (!connected_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to perform put_object, but connected_ = " << connected_ << ".");
+  }
  std::string data;
  obj->SerializeToString(&data); // TODO(pcm): get rid of this serialization
  ObjRequest request;
@@ -80,6 +97,13 @@ void Worker::put_object(ObjRef objref, const Obj* obj) {
  request.objref = objref;
  request.size = data.size();
  request_obj_queue_.send(&request);
+  if (contained_objrefs.size() > 0) {
+    ORCH_LOG(ORCH_DEBUG, "In put_object, calling increment_reference_count for objrefs:");
+    for (int i = 0; i < contained_objrefs.size(); ++i){
+       ORCH_LOG(ORCH_DEBUG, "----" << contained_objrefs[i]);
+    }
+    increment_reference_count(contained_objrefs); // Notify the scheduler that some object references are serialized in the objstore.
+  }
  ObjHandle result;
  receive_obj_queue_.receive(&result);
  uint8_t* target = segmentpool_.get_address(result);
@@ -87,9 +111,22 @@ void Worker::put_object(ObjRef objref, const Obj* obj) {
  request.type = ObjRequestType::WORKER_DONE;
  request.metadata_offset = 0;
  request_obj_queue_.send(&request);
+
+  // Notify the scheduler about the objrefs that we are serializing in the objstore.
+  AddContainedObjRefsRequest contained_objrefs_request;
+  contained_objrefs_request.set_objref(objref);
+  for (int i = 0; i < contained_objrefs.size(); ++i) {
+    contained_objrefs_request.add_contained_objref(contained_objrefs[i]); // TODO(rkn): The naming here is bad
+  }
+  AckReply reply;
+  ClientContext context;
+  scheduler_stub_->AddContainedObjRefs(&context, contained_objrefs_request, &reply);
 }

 void Worker::put_arrow(ObjRef objref, PyArrayObject* array) {
+  if (!connected_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to perform put_arrow, but connected_ = " << connected_ << ".");
+  }
  ObjRequest request;
  size_t size = arrow_size(array);
  request.workerid = workerid_;
@@ -106,6 +143,9 @@ void Worker::put_arrow(ObjRef objref, PyArrayObject* array) {
 }

 PyArrayObject* Worker::get_arrow(ObjRef objref) {
+  if (!connected_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to perform get_arrow, but connected_ = " << connected_ << ".");
+  }
  ObjRequest request;
  request.workerid = workerid_;
  request.type = ObjRequestType::GET;
@@ -117,6 +157,9 @@ PyArrayObject* Worker::get_arrow(ObjRef objref) {
 }

 bool Worker::is_arrow(ObjRef objref) {
+  if (!connected_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to perform is_arrow, but connected_ = " << connected_ << ".");
+  }
  ObjRequest request;
  request.workerid = workerid_;
  request.type = ObjRequestType::GET;
@@ -128,6 +171,9 @@ bool Worker::is_arrow(ObjRef objref) {
 }

 void Worker::alias_objrefs(ObjRef alias_objref, ObjRef target_objref) {
+  if (!connected_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to perform alias_objrefs, but connected_ = " << connected_ << ".");
+  }
  ClientContext context;
  AliasObjRefsRequest request;
  request.set_alias_objref(alias_objref);
@@ -136,7 +182,40 @@ void Worker::alias_objrefs(ObjRef alias_objref, ObjRef target_objref) {
  scheduler_stub_->AliasObjRefs(&context, request, &reply);
 }

+void Worker::increment_reference_count(std::vector<ObjRef> &objrefs) {
+  if (!connected_) {
+    ORCH_LOG(ORCH_DEBUG, "Attempting to increment_reference_count for objrefs, but connected_ = " << connected_ << " so returning instead.");
+    return;
+  }
+  ClientContext context;
+  IncrementRefCountRequest request;
+  for (int i = 0; i < objrefs.size(); ++i) {
+    ORCH_LOG(ORCH_DEBUG, "Incrementing reference count for objref " << objrefs[i]);
+    request.add_objref(objrefs[i]);
+  }
+  AckReply reply;
+  scheduler_stub_->IncrementRefCount(&context, request, &reply);
+}
+
+void Worker::decrement_reference_count(std::vector<ObjRef> &objrefs) {
+  if (!connected_) {
+    ORCH_LOG(ORCH_DEBUG, "Attempting to decrement_reference_count, but connected_ = " << connected_ << " so returning instead.");
+    return;
+  }
+  ClientContext context;
+  DecrementRefCountRequest request;
+  for (int i = 0; i < objrefs.size(); ++i) {
+    ORCH_LOG(ORCH_DEBUG, "Decrementing reference count for objref " << objrefs[i]);
+    request.add_objref(objrefs[i]);
+  }
+  AckReply reply;
+  scheduler_stub_->DecrementRefCount(&context, request, &reply);
+}
+
 void Worker::register_function(const std::string& name, size_t num_return_vals) {
+  if (!connected_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to perform register_function, but connected_ = " << connected_ << ".");
+  }
  ClientContext context;
  RegisterFunctionRequest request;
  request.set_fnname(name);
@@ -153,6 +232,9 @@ Call* Worker::receive_next_task() {
 }

 void Worker::notify_task_completed() {
+  if (!connected_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to perform notify_task_completed, but connected_ = " << connected_ << ".");
+  }
  ClientContext context;
  WorkerReadyRequest request;
  request.set_workerid(workerid_);
@@ -160,6 +242,22 @@ void Worker::notify_task_completed() {
  scheduler_stub_->WorkerReady(&context, request, &reply);
 }

+void Worker::disconnect() {
+  connected_ = false;
+}
+
+bool Worker::connected() {
+  return connected_;
+}
+
+// TODO(rkn): Should we be using pointers or references? And should they be const?
+void Worker::scheduler_info(ClientContext &context, SchedulerInfoRequest &request, SchedulerInfoReply &reply) {
+  if (!connected_) {
+    ORCH_LOG(ORCH_FATAL, "Attempting to get scheduler info, but connected_ = " << connected_ << ".");
+  }
+  scheduler_stub_->SchedulerInfo(&context, request, &reply);
+}
+
 // Communication between the WorkerServer and the Worker happens via a message
 // queue. This is because the Python interpreter needs to be single threaded
 // (in our case running in the main thread), whereas the WorkerService will
@@ -48,7 +48,7 @@ class Worker {
  // request an object to be delivered to the local object store
  void request_object(ObjRef objref);
  // stores an object to the local object store
-  void put_object(ObjRef objref, const Obj* obj);
+  void put_object(ObjRef objref, const Obj* obj, std::vector<ObjRef> &contained_objrefs);
  // retrieve serialized object from local object store
  slice get_object(ObjRef objref);
  // stores an arrow object to the local object store
@@ -60,6 +60,10 @@ class Worker {
  bool is_arrow(ObjRef objref);
  // make `alias_objref` refer to the same object that `target_objref` refers to
  void alias_objrefs(ObjRef alias_objref, ObjRef target_objref);
+  // increment the reference count for objref
+  void increment_reference_count(std::vector<ObjRef> &objref);
+  // decrement the reference count for objref
+  void decrement_reference_count(std::vector<ObjRef> &objref);
  // register function with scheduler
  void register_function(const std::string& name, size_t num_return_vals);
  // start the worker server which accepts tasks from the scheduler and stores
@@ -69,8 +73,15 @@ class Worker {
  Call* receive_next_task();
  // tell the scheduler that we are done with the current task and request the next one
  void notify_task_completed();
+  // disconnect the worker
+  void disconnect();
+  // return connected_
+  bool connected();
+  // get info about scheduler state
+  void scheduler_info(ClientContext &context, SchedulerInfoRequest &request, SchedulerInfoReply &reply);

 private:
+  bool connected_;
  const size_t CHUNK_SIZE = 8 * 1024;
  std::unique_ptr<Scheduler::Stub> scheduler_stub_;
  std::unique_ptr<ObjStore::Stub> objstore_stub_;
@@ -2,6 +2,7 @@ import unittest
 import orchpy
 import orchpy.serialization as serialization
 import orchpy.services as services
+import orchpy.worker as worker
 import numpy as np
 import time
 import subprocess32 as subprocess
@@ -53,13 +54,18 @@ class ArraysSingleTest(unittest.TestCase):
 class ArraysDistTest(unittest.TestCase):

  def testSerialization(self):
+    w = worker.Worker()
+    services.start_cluster(driver_worker=w)
+
    x = dist.DistArray()
-    x.construct([2, 3, 4], np.array([[[orchpy.lib.ObjRef(0)]]]))
-    capsule = serialization.serialize(x)
-    y = serialization.deserialize(capsule)
+    x.construct([2, 3, 4], np.array([[[orchpy.push(0, w)]]]))
+    capsule, _ = serialization.serialize(w.handle, x) # TODO(rkn): THIS REQUIRES A WORKER_HANDLE
+    y = serialization.deserialize(w.handle, capsule) # TODO(rkn): THIS REQUIRES A WORKER_HANDLE
    self.assertEqual(x.shape, y.shape)
    self.assertEqual(x.objrefs[0, 0, 0].val, y.objrefs[0, 0, 0].val)

+    services.cleanup()
+
  def testAssemble(self):
    test_dir = os.path.dirname(os.path.abspath(__file__))
    test_path = os.path.join(test_dir, "testrecv.py")
@@ -76,33 +82,46 @@ class ArraysDistTest(unittest.TestCase):
  def testMethods(self):
    test_dir = os.path.dirname(os.path.abspath(__file__))
    test_path = os.path.join(test_dir, "testrecv.py")
-    services.start_cluster(num_workers=3, worker_path=test_path)
+    services.start_cluster(num_workers=4, worker_path=test_path)

    x = dist.zeros([9, 25, 51], "float")
-    self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(x)) == np.zeros([9, 25, 51])))
+    y = dist.assemble(x)
+    self.assertTrue(np.alltrue(orchpy.pull(y) == np.zeros([9, 25, 51])))

    x = dist.ones([11, 25, 49], "float")
-    self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(x)) == np.ones([11, 25, 49])))
+    y = dist.assemble(x)
+    self.assertTrue(np.alltrue(orchpy.pull(y) == np.ones([11, 25, 49])))

    x = dist.random.normal([11, 25, 49])
    y = dist.copy(x)
-    self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(x)) == orchpy.pull(dist.assemble(y))))
+    z = dist.assemble(x)
+    w = dist.assemble(y)
+    self.assertTrue(np.alltrue(orchpy.pull(z) == orchpy.pull(w)))

    x = dist.eye(25, "float")
-    self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(x)) == np.eye(25)))
+    y = dist.assemble(x)
+    self.assertTrue(np.alltrue(orchpy.pull(y) == np.eye(25)))

    x = dist.random.normal([25, 49])
    y = dist.triu(x)
-    self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(y)) == np.triu(orchpy.pull(dist.assemble(x)))))
+    z = dist.assemble(y)
+    w = dist.assemble(x)
+    self.assertTrue(np.alltrue(orchpy.pull(z) == np.triu(orchpy.pull(w))))

    x = dist.random.normal([25, 49])
    y = dist.tril(x)
-    self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(y)) == np.tril(orchpy.pull(dist.assemble(x)))))
+    z = dist.assemble(y)
+    w = dist.assemble(x)
+    self.assertTrue(np.alltrue(orchpy.pull(z) == np.tril(orchpy.pull(w))))

    x = dist.random.normal([25, 49])
    y = dist.random.normal([49, 18])
    z = dist.dot(x, y)
-    self.assertTrue(np.allclose(orchpy.pull(dist.assemble(z)), np.dot(orchpy.pull(dist.assemble(x)), orchpy.pull(dist.assemble(y)))))
+    w = dist.assemble(z)
+    u = dist.assemble(x)
+    v = dist.assemble(y)
+    np.allclose(orchpy.pull(w), np.dot(orchpy.pull(u), orchpy.pull(v)))
+    self.assertTrue(np.allclose(orchpy.pull(w), np.dot(orchpy.pull(u), orchpy.pull(v))))

    services.cleanup()

@@ -13,58 +13,58 @@ from google.protobuf.text_format import *
 import orchestra_pb2
 import types_pb2

+import test_functions
+import arrays.single as single
+import arrays.dist as dist
+
 class SerializationTest(unittest.TestCase):

-  def roundTripTest(self, data):
-    serialized = serialization.serialize(data)
-    result = serialization.deserialize(serialized)
+  def roundTripTest(self, worker, data):
+    serialized, _ = serialization.serialize(worker.handle, data)
+    result = serialization.deserialize(worker.handle, serialized)
    self.assertEqual(data, result)

-  def numpyTypeTest(self, typ):
+  def numpyTypeTest(self, worker, typ):
    a = np.random.randint(0, 10, size=(100, 100)).astype(typ)
-    b = serialization.serialize(a)
-    c = serialization.deserialize(b)
+    b, _ = serialization.serialize(worker.handle, a)
+    c = serialization.deserialize(worker.handle, b)
    self.assertTrue((a == c).all())

  def testSerialize(self):
-    self.roundTripTest([1, "hello", 3.0])
-    self.roundTripTest(42)
-    self.roundTripTest("hello world")
-    self.roundTripTest(42.0)
-    self.roundTripTest((1.0, "hi"))
+    w = worker.Worker()
+    services.start_cluster(driver_worker=w)

-    self.roundTripTest({"hello" : "world", 1: 42, 1.0: 45})
-    self.roundTripTest({})
+    self.roundTripTest(w, [1, "hello", 3.0])
+    self.roundTripTest(w, 42)
+    self.roundTripTest(w, "hello world")
+    self.roundTripTest(w, 42.0)
+    self.roundTripTest(w, (1.0, "hi"))
+
+    self.roundTripTest(w, {"hello" : "world", 1: 42, 1.0: 45})
+    self.roundTripTest(w, {})

    a = np.zeros((100, 100))
-    res = serialization.serialize(a)
-    b = serialization.deserialize(res)
+    res, _ = serialization.serialize(w.handle, a)
+    b = serialization.deserialize(w.handle, res)
    self.assertTrue((a == b).all())

-    self.numpyTypeTest('int8')
-    self.numpyTypeTest('uint8')
+    self.numpyTypeTest(w, 'int8')
+    self.numpyTypeTest(w, 'uint8')
    # self.numpyTypeTest('int16') # TODO(pcm): implement this
    # self.numpyTypeTest('int32') # TODO(pcm): implement this
-    self.numpyTypeTest('float32')
-    self.numpyTypeTest('float64')
+    self.numpyTypeTest(w, 'float32')
+    self.numpyTypeTest(w, 'float64')

-    a = np.array([[orchpy.lib.ObjRef(0), orchpy.lib.ObjRef(1)], [orchpy.lib.ObjRef(41), orchpy.lib.ObjRef(42)]])
-    capsule = serialization.serialize(a)
-    result = serialization.deserialize(capsule)
+    ref0 = orchpy.push(0, w)
+    ref1 = orchpy.push(0, w)
+    ref2 = orchpy.push(0, w)
+    ref3 = orchpy.push(0, w)
+    a = np.array([[ref0, ref1], [ref2, ref3]])
+    capsule, _ = serialization.serialize(w.handle, a)
+    result = serialization.deserialize(w.handle, capsule)
    self.assertTrue((a == result).all())

-class OrchPyLibTest(unittest.TestCase):
-
-    def testOrchPyLib(self):
-      w = worker.Worker()
-      services.start_cluster(driver_worker=w)
-
-      w.put_object(orchpy.lib.ObjRef(0), 'hello world')
-      result = w.get_object(orchpy.lib.ObjRef(0))
-
-      self.assertEqual(result, 'hello world')
-
-      services.cleanup()
+    services.cleanup()

 class ObjStoreTest(unittest.TestCase):

@@ -97,7 +97,7 @@ class SchedulerTest(unittest.TestCase):
    services.start_cluster(driver_worker=w, num_workers=1, worker_path=test_path)

    value_before = "test_string"
-    objref = w.remote_call("__main__.print_string", [value_before])
+    objref = w.remote_call("test_functions.print_string", [value_before])

    time.sleep(0.2)

@@ -148,12 +148,67 @@ class APITest(unittest.TestCase):
    test_path = os.path.join(test_dir, "testrecv.py")
    services.start_cluster(num_workers=3, worker_path=test_path, driver_worker=w)

-    objref = w.remote_call("__main__.test_alias_f", [])
+    objref = w.remote_call("test_functions.test_alias_f", [])
    self.assertTrue(np.alltrue(orchpy.pull(objref[0], w) == np.ones([3, 4, 5])))
-    objref = w.remote_call("__main__.test_alias_g", [])
+    objref = w.remote_call("test_functions.test_alias_g", [])
    self.assertTrue(np.alltrue(orchpy.pull(objref[0], w) == np.ones([3, 4, 5])))
-    objref = w.remote_call("__main__.test_alias_h", [])
+    objref = w.remote_call("test_functions.test_alias_h", [])
    self.assertTrue(np.alltrue(orchpy.pull(objref[0], w) == np.ones([3, 4, 5])))

+    services.cleanup()
+
+class ReferenceCountingTest(unittest.TestCase):
+
+  def testDeallocation(self):
+    test_dir = os.path.dirname(os.path.abspath(__file__))
+    test_path = os.path.join(test_dir, "testrecv.py")
+    services.start_cluster(num_workers=3, worker_path=test_path)
+
+    x = test_functions.test_alias_f()
+    orchpy.pull(x)
+    time.sleep(0.1)
+    objref_val = x.val
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val] == 1)
+
+    del x
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val] == -1) # -1 indicates deallocated
+
+    y = test_functions.test_alias_h()
+    orchpy.pull(y)
+    time.sleep(0.1)
+    objref_val = y.val
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [1, 0, 0])
+
+    del y
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [-1, -1, -1])
+
+    z = dist.zeros([dist.BLOCK_SIZE, 2 * dist.BLOCK_SIZE], "float")
+    time.sleep(0.1)
+    objref_val = z.val
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [1, 1, 1])
+
+    del z
+    time.sleep(0.1)
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [-1, -1, -1])
+
+    x = single.zeros([10, 10], "float")
+    y = single.zeros([10, 10], "float")
+    z = single.dot(x, y)
+    objref_val = x.val
+    time.sleep(0.1)
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [1, 1, 1])
+
+    del x
+    time.sleep(0.1)
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [-1, 1, 1])
+    del y
+    time.sleep(0.1)
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [-1, -1, 1])
+    del z
+    time.sleep(0.1)
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [-1, -1, -1])
+
+    services.cleanup()
+
 if __name__ == '__main__':
    unittest.main()
@@ -5,6 +5,7 @@ import orchpy
 import orchpy.services as services
 import orchpy.worker as worker

+import test_functions
 import arrays.single as single
 import arrays.dist as dist

@@ -19,50 +20,9 @@ parser.add_argument("--scheduler-address", default="127.0.0.1:10001", type=str,
 parser.add_argument("--objstore-address", default="127.0.0.1:20001", type=str, help="the objstore's address")
 parser.add_argument("--worker-address", default="127.0.0.1:30001", type=str, help="the worker's address")

-@orchpy.distributed([], [np.ndarray])
-def test_alias_f():
-  return np.ones([3, 4, 5])
-
-@orchpy.distributed([], [np.ndarray])
-def test_alias_g():
-  return test_alias_f()
-
-@orchpy.distributed([], [np.ndarray])
-def test_alias_h():
-  return test_alias_g()
-
-@orchpy.distributed([str], [str])
-def print_string(string):
-  print "called print_string with", string
-  f = open("asdfasdf.txt", "w")
-  f.write("successfully called print_string with argument {}.".format(string))
-  return string
-
-@orchpy.distributed([int, int], [int, int])
-def handle_int(a, b):
-  return a + 1, b + 1
-
-def connect_to_scheduler(host, port):
-  channel = implementations.insecure_channel(host, port)
-  return orchestra_pb2.beta_create_Scheduler_stub(channel)
-
-def connect_to_objstore(host, port):
-  channel = implementations.insecure_channel(host, port)
-  return orchestra_pb2.beta_create_ObjStore_stub(channel)
-
 if __name__ == '__main__':
  args = parser.parse_args()
-  scheduler_ip_address, scheduler_port = args.scheduler_address.split(":")
-  scheduler_stub = connect_to_scheduler(scheduler_ip_address, int(scheduler_port))
-  objstore_ip_address, objstore_port = args.objstore_address.split(":")
-  objstore_stub = connect_to_objstore(objstore_ip_address, int(objstore_port))
  worker.connect(args.scheduler_address, args.objstore_address, args.worker_address)

-  def scheduler_debug_info():
-    return scheduler_stub.SchedulerDebugInfo(orchestra_pb2.SchedulerDebugInfoRequest(), TIMEOUT_SECONDS)
-
-  def objstore_debug_info():
-    return objstore_stub.ObjStoreDebugInfo(orchestra_pb2.ObjStoreDebugInfoRequest(), TIMEOUT_SECONDS)
-
  import IPython
  IPython.embed()
@@ -0,0 +1,32 @@
+import orchpy
+
+import numpy as np
+
+# Test simple functionality
+
+@orchpy.distributed([str], [str])
+def print_string(string):
+  print "called print_string with", string
+  f = open("asdfasdf.txt", "w")
+  f.write("successfully called print_string with argument {}.".format(string))
+  return string
+
+@orchpy.distributed([int, int], [int, int])
+def handle_int(a, b):
+  return a + 1, b + 1
+
+# Test aliasing
+
+@orchpy.distributed([], [np.ndarray])
+def test_alias_f():
+  return np.ones([3, 4, 5])
+
+@orchpy.distributed([], [np.ndarray])
+def test_alias_g():
+  return test_alias_f()
+
+@orchpy.distributed([], [np.ndarray])
+def test_alias_h():
+  return test_alias_g()
+
+# Test reference counting
@@ -2,6 +2,7 @@ import sys
 import argparse
 import numpy as np

+import test_functions
 import arrays.single as single
 import arrays.dist as dist

@@ -14,33 +15,11 @@ parser.add_argument("--scheduler-address", default="127.0.0.1:10001", type=str,
 parser.add_argument("--objstore-address", default="127.0.0.1:20001", type=str, help="the objstore's address")
 parser.add_argument("--worker-address", default="127.0.0.1:40001", type=str, help="the worker's address")

-@orchpy.distributed([], [np.ndarray])
-def test_alias_f():
-  return np.ones([3, 4, 5])
-
-@orchpy.distributed([], [np.ndarray])
-def test_alias_g():
-  return test_alias_f()
-
-@orchpy.distributed([], [np.ndarray])
-def test_alias_h():
-  return test_alias_g()
-
-@orchpy.distributed([str], [str])
-def print_string(string):
-  print "called print_string with", string
-  f = open("asdfasdf.txt", "w")
-  f.write("successfully called print_string with argument {}.".format(string))
-  return string
-
-@orchpy.distributed([int, int], [int, int])
-def handle_int(a, b):
-  return a + 1, b + 1
-
 if __name__ == '__main__':
  args = parser.parse_args()
  worker.connect(args.scheduler_address, args.objstore_address, args.worker_address)

+  orchpy.register_module(test_functions)
  orchpy.register_module(single)
  orchpy.register_module(single.random)
  orchpy.register_module(single.linalg)