Unmap segments properly to avoid memory leaks

src/ipc.cc

@@ -56,6 +56,11 @@ void MemorySegmentPool::open_segment(SegmentId segmentid, size_t size) {
   }
 }
 
+void MemorySegmentPool::unmap_segment(SegmentId segmentid) {
+  segments_[segmentid].first.reset();
+  segments_[segmentid].second = SegmentStatusType::UNOPENED;
+}
+
 void MemorySegmentPool::close_segment(SegmentId segmentid) {
   RAY_LOG(RAY_DEBUG, "closing segmentid " << segmentid);
   std::string segment_name = get_segment_name(segmentid);

src/ipc.h

@@ -152,6 +152,7 @@ public:
   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
   std::string get_segment_name(SegmentId segmentid); // get the name of a segment
+  void unmap_segment(SegmentId segmentid); // unmap a memory segment from a client (only to be called by clients)
 private:
   void open_segment(SegmentId segmentid, size_t size = 0); // create a segment or map an existing one into memory
   void close_segment(SegmentId segmentid); // close a segment

src/raylib.cc

@@ -460,7 +460,11 @@ PyObject* put_arrow(PyObject* self, PyObject* args) {
   if (!PyArg_ParseTuple(args, "O&O&O", &PyObjectToWorker, &worker, &PyObjectToObjRef, &objref, &value)) {
     return NULL;
   }
-  worker->put_arrow(objref, value);
+  // The following is reqired, because numbuf expects contiguous arrays at the moment.
+  // This is to make sure that we do not have to do reference counting inside numbuf, it is expected to change.
+  PyArrayObject* array = PyArray_GETCONTIGUOUS((PyArrayObject*) value); // TODO(pcm): put that into numbuf
+  worker->put_arrow(objref, (PyObject*) array);
+  Py_XDECREF(array); // GETCONTIGUOUS from above returned a new reference
   Py_RETURN_NONE;
 }
 
@@ -490,6 +494,16 @@ PyObject* is_arrow(PyObject* self, PyObject* args) {
     Py_RETURN_FALSE;
 }
 
+PyObject* unmap_object(PyObject* self, PyObject* args) {
+  Worker* worker;
+  int segmentid;
+  if (!PyArg_ParseTuple(args, "O&i", &PyObjectToWorker, &worker, &segmentid)) {
+    return NULL;
+  }
+  worker->unmap_object(segmentid);
+  Py_RETURN_NONE;
+}
+
 PyObject* deserialize_object(PyObject* self, PyObject* args) {
   PyObject* worker_capsule;
   Obj* obj;
@@ -827,6 +841,7 @@ static PyMethodDef RayLibMethods[] = {
  { "put_arrow", put_arrow, METH_VARARGS, "put an arrow array on the local object store"},
  { "get_arrow", get_arrow, METH_VARARGS, "get an arrow array from the local object store"},
  { "is_arrow", is_arrow, METH_VARARGS, "is the object in the local object store an arrow object?"},
+ { "unmap_object", unmap_object, METH_VARARGS, "unmap the object from the client's shared memory pool"},
  { "serialize_task", serialize_task, METH_VARARGS, "serialize a task to protocol buffers" },
  { "deserialize_task", deserialize_task, METH_VARARGS, "deserialize a task from protocol buffers" },
  { "create_worker", create_worker, METH_VARARGS, "connect to the scheduler and the object store" },

src/worker.cc

@@ -117,6 +117,9 @@ void Worker::put_object(ObjRef objref, const Obj* obj, std::vector<ObjRef> &cont
   receive_obj_queue_.receive(&result);
   uint8_t* target = segmentpool_->get_address(result);
   std::memcpy(target, &data[0], data.size());
+  // We immediately unmap here; if the object is going to be accessed again, it will be mapped again;
+  // This is reqired because we do not have a mechanism to unmap the object later.
+  segmentpool_->unmap_segment(result.segmentid());
   request.type = ObjRequestType::WORKER_DONE;
   request.metadata_offset = 0;
   request_obj_queue_.send(&request);
@@ -160,6 +163,9 @@ PyObject* Worker::put_arrow(ObjRef objref, PyObject* value) {
   uint8_t* address = segmentpool_->get_address(result);
   auto source = std::make_shared<BufferMemorySource>(address, size);
   CHECK_ARROW_STATUS(writer.Write(source.get(), &metadata_offset), "error during Write: ");
+  // We immediately unmap here; if the object is going to be accessed again, it will be mapped again;
+  // This is reqired because we do not have a mechanism to unmap the object later.
+  segmentpool_->unmap_segment(result.segmentid());
   request.type = ObjRequestType::WORKER_DONE;
   request.metadata_offset = metadata_offset;
   request_obj_queue_.send(&request);
@@ -197,6 +203,14 @@ bool Worker::is_arrow(ObjRef objref) {
   return result.metadata_offset() != 0;
 }
 
+void Worker::unmap_object(ObjRef objref) {
+  if (!connected_) {
+    RAY_LOG(RAY_DEBUG, "Attempted to perform unmap_object but failed.");
+    return;
+  }
+  segmentpool_->unmap_segment(objref);
+}
+
 void Worker::alias_objrefs(ObjRef alias_objref, ObjRef target_objref) {
   RAY_CHECK(connected_, "Attempted to perform alias_objrefs but failed.");
   ClientContext context;

src/worker.h

@@ -60,6 +60,8 @@ class Worker {
   PyObject* get_arrow(ObjRef objref, SegmentId& segmentid);
   // determine if the object stored in objref is an arrow object // TODO(pcm): more general mechanism for this?
   bool is_arrow(ObjRef objref);
+  // unmap the segment containing an object from the local address space
+  void unmap_object(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