make eviction work in Cython

cpp/src/plasma/plasma.pyx

@@ -43,6 +43,8 @@ cdef extern from "plasma/client.h" nogil:
 
     CStatus Seal(const CUniqueID& object_id)
 
+    CStatus Release(const CUniqueID& object_id)
+
     CStatus Disconnect()
 
 cdef extern from "plasma/client.h" nogil:
@@ -64,6 +66,19 @@ cdef class ObjectID:
   def __repr__(self):
     return "ObjectID(" + self.data.hex().decode() + ")"
 
+cdef class PlasmaBuffer(Buffer):
+
+  cdef:
+    ObjectID object_id
+    PlasmaClient client
+
+  def __cinit__(self, ObjectID object_id, PlasmaClient client):
+    self.object_id = object_id
+    self.client = client
+
+  def __dealloc__(self):
+    self.client.release(self.object_id)
+
 cdef class PlasmaClient:
 
   cdef:
@@ -80,10 +95,10 @@ cdef class PlasmaClient:
     result[0].resize(ids.size())
     check_status(self.client.get().Get(ids.data(), ids.size(), timeout_ms, result[0].data()))
 
-  cdef _make_buffer(self, uint8_t* data, int64_t size):
+  cdef _make_plasma_buffer(self, ObjectID object_id, uint8_t* data, int64_t size):
     cdef shared_ptr[MutableBuffer] buffer
     buffer.reset(new MutableBuffer(data, size))
-    result = Buffer()
+    result = PlasmaBuffer(object_id, self)
     result.init(<shared_ptr[CBuffer]>(buffer))
     return result
 
@@ -93,15 +108,15 @@ cdef class PlasmaClient:
   def create(self, ObjectID object_id, data_size, c_string metadata=b""):
     cdef uint8_t* data
     check_status(self.client.get().Create(object_id.data, data_size, <uint8_t*>(metadata.data()), metadata.size(), &data))
-    return self._make_buffer(data, data_size)
+    return self._make_plasma_buffer(object_id, data, data_size)
 
   def get(self, object_ids, timeout_ms=-1):
     cdef c_vector[CObjectBuffer] object_buffers
     self._get_object_buffers(object_ids, timeout_ms, &object_buffers)
     result = []
-    for object_buffer in object_buffers:
-      if object_buffer.data_size != -1:
-        result.append(self._make_buffer(object_buffer.data, object_buffer.data_size))
+    for i in range(object_buffers.size()):
+      if object_buffers[i].data_size != -1:
+        result.append(self._make_plasma_buffer(object_ids[i], object_buffers[i].data, object_buffers[i].data_size))
       else:
         result.append(None)
     return result
@@ -109,10 +124,16 @@ cdef class PlasmaClient:
   def get_metadata(self, object_ids, timeout_ms=-1):
     cdef c_vector[CObjectBuffer] object_buffers
     self._get_object_buffers(object_ids, timeout_ms, &object_buffers)
-    return [self._make_buffer(b.metadata, b.metadata_size) for b in object_buffers]
+    result = []
+    for i in range(object_buffers.size()):
+      result.append(self._make_plasma_buffer(object_ids[i], object_buffers[i].metadata, object_buffers[i].metadata_size))
+    return result
 
   def seal(self, ObjectID object_id):
     check_status(self.client.get().Seal(object_id.data))
 
+  def release(self, ObjectID object_id):
+    check_status(self.client.get().Release(object_id.data))
+
   def disconnect(self):
     check_status(self.client.get().Disconnect())

cpp/src/plasma/test/test.py

@@ -51,6 +51,12 @@ def create_object_with_id(client, object_id, data_size, metadata_size,
     client.seal(object_id)
   return memory_buffer, metadata
 
+def create_object(client, data_size, metadata_size, seal=True):
+  object_id = random_object_id()
+  memory_buffer, metadata = create_object_with_id(client, object_id, data_size,
+                                                  metadata_size, seal=seal)
+  return object_id, memory_buffer, metadata
+
 def assert_get_object_equal(unit_test, client1, client2, object_id,
                             memory_buffer=None, metadata=None):
   client1_buff = client1.get([object_id])[0]
@@ -225,6 +231,45 @@ class TestPlasmaClient(unittest.TestCase):
         else:
           self.assertIsNone(results[i])
 
+  def test_store_full(self):
+    # The store is started with 1GB, so make sure that create throws an
+    # exception when it is full.
+    def assert_create_raises_plasma_full(unit_test, size):
+      partial_size = np.random.randint(size)
+      try:
+        _, memory_buffer, _ = create_object(unit_test.plasma_client,
+                                            partial_size,
+                                            size - partial_size)
+      # TODO(pcm): More specific error here.
+      except pa.lib.ArrowException as e:
+        pass
+      else:
+        # For some reason the above didn't throw an exception, so fail.
+        unit_test.assertTrue(False)
+
+    # Create a list to keep some of the buffers in scope.
+    memory_buffers = []
+    _, memory_buffer, _ = create_object(self.plasma_client, 5 * 10 ** 8, 0)
+    memory_buffers.append(memory_buffer)
+    # Remaining space is 5 * 10 ** 8. Make sure that we can't create an object
+    # of size 5 * 10 ** 8 + 1, but we can create one of size 2 * 10 ** 8.
+    assert_create_raises_plasma_full(self, 5 * 10 ** 8 + 1)
+    _, memory_buffer, _ = create_object(self.plasma_client, 2 * 10 ** 8, 0)
+    del memory_buffer
+    _, memory_buffer, _ = create_object(self.plasma_client, 2 * 10 ** 8, 0)
+    del memory_buffer
+    assert_create_raises_plasma_full(self, 5 * 10 ** 8 + 1)
+
+    _, memory_buffer, _ = create_object(self.plasma_client, 2 * 10 ** 8, 0)
+    memory_buffers.append(memory_buffer)
+    # Remaining space is 3 * 10 ** 8.
+    assert_create_raises_plasma_full(self, 3 * 10 ** 8 + 1)
+
+    _, memory_buffer, _ = create_object(self.plasma_client, 10 ** 8, 0)
+    memory_buffers.append(memory_buffer)
+    # Remaining space is 2 * 10 ** 8.
+    assert_create_raises_plasma_full(self, 2 * 10 ** 8 + 1)
+
 if __name__ == "__main__":
   if len(sys.argv) > 1:
     # Pop the argument so we don't mess with unittest's own argument parser.