implement reference counting and much more (#43)

src/ipc.cc

@@ -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());
   }
 }

src/ipc.h

@@ -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

src/objstore.cc

@@ -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) {

src/objstore.h

@@ -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_;

src/orchpylib.cc

@@ -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 }
 };
 

src/scheduler.cc

@@ -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;

src/scheduler.h

@@ -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

src/worker.cc

@@ -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

src/worker.h

@@ -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_;