get all python tests in place

cpp/src/plasma/client.cc

@@ -388,22 +388,6 @@ static uint64_t compute_object_hash(const ObjectBuffer& obj_buffer) {
   return XXH64_digest(&hash_state);
 }
 
-bool plasma_compute_object_hash(
-    PlasmaClient* conn, ObjectID object_id, unsigned char* digest) {
-  // Get the plasma object data. We pass in a timeout of 0 to indicate that
-  // the operation should timeout immediately.
-  ObjectBuffer object_buffer;
-  ARROW_CHECK_OK(conn->Get(&object_id, 1, 0, &object_buffer));
-  // If the object was not retrieved, return false.
-  if (object_buffer.data_size == -1) { return false; }
-  // Compute the hash.
-  uint64_t hash = compute_object_hash(object_buffer);
-  memcpy(digest, &hash, sizeof(hash));
-  // Release the plasma object.
-  ARROW_CHECK_OK(conn->Release(object_id));
-  return true;
-}
-
 Status PlasmaClient::Seal(const ObjectID& object_id) {
   // Make sure this client has a reference to the object before sending the
   // request to Plasma.
@@ -415,7 +399,7 @@ Status PlasmaClient::Seal(const ObjectID& object_id) {
   object_entry->second->is_sealed = true;
   /// Send the seal request to Plasma.
   static unsigned char digest[kDigestSize];
-  ARROW_CHECK(plasma_compute_object_hash(this, object_id, &digest[0]));
+  RETURN_NOT_OK(Hash(object_id, &digest[0]));
   RETURN_NOT_OK(SendSealRequest(store_conn_, object_id, &digest[0]));
   // We call PlasmaClient::Release to decrement the number of instances of this
   // object
@@ -441,6 +425,22 @@ Status PlasmaClient::Evict(int64_t num_bytes, int64_t& num_bytes_evicted) {
   return ReadEvictReply(buffer.data(), num_bytes_evicted);
 }
 
+Status PlasmaClient::Hash(const ObjectID& object_id, uint8_t* digest) {
+  // Get the plasma object data. We pass in a timeout of 0 to indicate that
+  // the operation should timeout immediately.
+  ObjectBuffer object_buffer;
+  RETURN_NOT_OK(Get(&object_id, 1, 0, &object_buffer));
+  // If the object was not retrieved, return false.
+  if (object_buffer.data_size == -1) {
+    return Status::PlasmaObjectNonexistent("Object not found");
+  }
+  // Compute the hash.
+  uint64_t hash = compute_object_hash(object_buffer);
+  memcpy(digest, &hash, sizeof(hash));
+  // Release the plasma object.
+  return Release(object_id);
+}
+
 Status PlasmaClient::Subscribe(int* fd) {
   int sock[2];
   // Create a non-blocking socket pair. This will only be used to send
@@ -461,6 +461,25 @@ Status PlasmaClient::Subscribe(int* fd) {
   return Status::OK();
 }
 
+Status PlasmaClient::GetNotification(int fd, ObjectID *object_id, int64_t* data_size, int64_t* metadata_size) {
+  uint8_t* notification = read_message_async(fd);
+  if (notification == NULL) {
+    return Status::IOError("Failed to read object notification from Plasma socket");
+  }
+  auto object_info = flatbuffers::GetRoot<ObjectInfo>(notification);
+  ARROW_CHECK(object_info->object_id()->size() == sizeof(ObjectID));
+  memcpy(object_id, object_info->object_id()->data(), sizeof(ObjectID));
+  if (object_info->is_deletion()) {
+    *data_size = -1;
+    *metadata_size = -1;
+  } else {
+    *data_size = object_info->data_size();
+    *metadata_size = object_info->metadata_size();
+  }
+  delete[] notification;
+  return Status::OK();
+}
+
 Status PlasmaClient::Connect(const std::string& store_socket_name,
     const std::string& manager_socket_name, int release_delay) {
   store_conn_ = connect_ipc_sock_retry(store_socket_name, -1, -1);

cpp/src/plasma/client.h

@@ -179,11 +179,19 @@ class ARROW_EXPORT PlasmaClient {
   /// retrieved.
   /// @return The return status.
   Status Evict(int64_t num_bytes, int64_t& num_bytes_evicted);
+  
+  /// Compute the hash of an object in the object store.
+  ///
+  /// @param conn The object containing the connection state.
+  /// @param object_id The ID of the object we want to hash.
+  /// @param digest A pointer at which to return the hash digest of the object.
+  ///        The pointer must have at least kDigestSize bytes allocated.
+  /// @return The return status.
+  Status Hash(const ObjectID& object_id, uint8_t* digest);
 
   /// Subscribe to notifications when objects are sealed in the object store.
   /// Whenever an object is sealed, a message will be written to the client
-  /// socket
-  /// that is returned by this method.
+  /// socket that is returned by this method.
   ///
   /// @param fd Out parameter for the file descriptor the client should use to
   /// read notifications
@@ -191,6 +199,15 @@ class ARROW_EXPORT PlasmaClient {
   /// @return The return status.
   Status Subscribe(int* fd);
 
+  /// Receive next object notification for this client if Subscribe has been called.
+  ///
+  /// @param fd The file descriptor we are reading the notification from.
+  /// @param object_id Out parameter, the object_id of the object that was sealed.
+  /// @param data_size Out parameter, the data size of the object that was sealed.
+  /// @param metadata_size Out parameter, the metadata size of the object that was sealed.
+  /// @return The return status.
+  Status GetNotification(int fd, ObjectID* object_id, int64_t* data_size, int64_t* metadata_size);
+
   /// Disconnect from the local plasma instance, including the local store and
   /// manager.
   ///
@@ -333,14 +350,4 @@ class ARROW_EXPORT PlasmaClient {
   int64_t store_capacity_;
 };
 
-/// Compute the hash of an object in the object store.
-///
-/// @param conn The object containing the connection state.
-/// @param object_id The ID of the object we want to hash.
-/// @param digest A pointer at which to return the hash digest of the object.
-///        The pointer must have at least DIGEST_SIZE bytes allocated.
-/// @return A boolean representing whether the hash operation succeeded.
-bool plasma_compute_object_hash(
-    PlasmaClient* conn, ObjectID object_id, unsigned char* digest);
-
 #endif  // PLASMA_CLIENT_H

cpp/src/plasma/plasma.pyx

@@ -26,8 +26,13 @@ cdef extern from "plasma/common.h" nogil:
     @staticmethod
     CUniqueID from_binary(const c_string& binary)
 
+    c_bool operator==(const CUniqueID& rhs) const
+
     c_string hex()
 
+cdef extern from "plasma/plasma.h":
+    cdef int64_t kDigestSize
+
 cdef extern from "plasma/client.h" nogil:
 
   cdef cppclass CPlasmaClient" PlasmaClient":
@@ -43,8 +48,18 @@ cdef extern from "plasma/client.h" nogil:
 
     CStatus Seal(const CUniqueID& object_id)
 
+    CStatus Evict(int64_t num_bytes, int64_t& num_bytes_evicted)
+
+    CStatus Hash(const CUniqueID& object_id, uint8_t* digest)
+
     CStatus Release(const CUniqueID& object_id)
 
+    CStatus Contains(const CUniqueID& object_id, c_bool* has_object)
+
+    CStatus Subscribe(int* fd)
+
+    CStatus GetNotification(int fd, CUniqueID* object_id, int64_t* data_size, int64_t* metadata_size)
+
     CStatus Disconnect()
 
 cdef extern from "plasma/client.h" nogil:
@@ -63,6 +78,10 @@ cdef class ObjectID:
   def __cinit__(self, object_id):
     self.data = CUniqueID.from_binary(object_id)
 
+  def __richcmp__(ObjectID self, ObjectID object_id, operation):
+    assert operation == 2, "only equality implemented so far"
+    return self.data == object_id.data
+
   def __repr__(self):
     return "ObjectID(" + self.data.hex().decode() + ")"
 
@@ -83,9 +102,11 @@ cdef class PlasmaClient:
 
   cdef:
     shared_ptr[CPlasmaClient] client
+    int notification_fd
 
   def __cinit__(self):
     self.client.reset(new CPlasmaClient())
+    self.notification_fd = -1
 
   cdef _get_object_buffers(self, object_ids, int64_t timeout_ms, c_vector[CObjectBuffer]* result):
     cdef c_vector[CUniqueID] ids
@@ -135,5 +156,30 @@ cdef class PlasmaClient:
   def release(self, ObjectID object_id):
     check_status(self.client.get().Release(object_id.data))
 
+  def contains(self, ObjectID object_id):
+    cdef c_bool is_contained
+    check_status(self.client.get().Contains(object_id.data, &is_contained))
+    return is_contained
+
+  def hash(self, ObjectID object_id):
+    cdef c_vector[uint8_t] digest = c_vector[uint8_t](kDigestSize)
+    check_status(self.client.get().Hash(object_id.data, digest.data()))
+    return bytes(digest[:])
+
+  def evict(self, int64_t num_bytes):
+    cdef int64_t num_bytes_evicted
+    check_status(self.client.get().Evict(num_bytes, num_bytes_evicted))
+    return num_bytes_evicted
+
+  def subscribe(self):
+    check_status(self.client.get().Subscribe(&self.notification_fd))
+
+  def get_next_notification(self):
+    cdef ObjectID object_id = ObjectID(20 * b"\0")
+    cdef int64_t data_size
+    cdef int64_t metadata_size
+    check_status(self.client.get().GetNotification(self.notification_fd, &object_id.data, &data_size, &metadata_size))
+    return object_id, data_size, metadata_size
+
   def disconnect(self):
     check_status(self.client.get().Disconnect())

cpp/src/plasma/test/test.py

@@ -270,6 +270,298 @@ class TestPlasmaClient(unittest.TestCase):
     # Remaining space is 2 * 10 ** 8.
     assert_create_raises_plasma_full(self, 2 * 10 ** 8 + 1)
 
+  def test_contains(self):
+    fake_object_ids = [random_object_id() for _ in range(100)]
+    real_object_ids = [random_object_id() for _ in range(100)]
+    for object_id in real_object_ids:
+      self.assertFalse(self.plasma_client.contains(object_id))
+      self.plasma_client.create(object_id, 100)
+      self.plasma_client.seal(object_id)
+      self.assertTrue(self.plasma_client.contains(object_id))
+    for object_id in fake_object_ids:
+      self.assertFalse(self.plasma_client.contains(object_id))
+    for object_id in real_object_ids:
+      self.assertTrue(self.plasma_client.contains(object_id))
+
+  def test_hash(self):
+    # Check the hash of an object that doesn't exist.
+    object_id1 = random_object_id()
+    try:
+      self.plasma_client.hash(object_id1)
+      # TODO(pcm): Introduce a more specific error type here
+    except pa.lib.ArrowException as e:
+      pass
+    else:
+      self.assertTrue(False)
+
+    length = 1000
+    # Create a random object, and check that the hash function always returns
+    # the same value.
+    metadata = generate_metadata(length)
+    memory_buffer = np.frombuffer(self.plasma_client.create(object_id1, length, metadata), dtype="uint8")
+    for i in range(length):
+      memory_buffer[i] = i % 256
+    self.plasma_client.seal(object_id1)
+    self.assertEqual(self.plasma_client.hash(object_id1),
+                     self.plasma_client.hash(object_id1))
+
+    # Create a second object with the same value as the first, and check that
+    # their hashes are equal.
+    object_id2 = random_object_id()
+    memory_buffer = np.frombuffer(self.plasma_client.create(object_id2, length, metadata), dtype="uint8")
+    for i in range(length):
+      memory_buffer[i] = i % 256
+    self.plasma_client.seal(object_id2)
+    self.assertEqual(self.plasma_client.hash(object_id1),
+                     self.plasma_client.hash(object_id2))
+
+    # Create a third object with a different value from the first two, and
+    # check that its hash is different.
+    object_id3 = random_object_id()
+    metadata = generate_metadata(length)
+    memory_buffer = np.frombuffer(self.plasma_client.create(object_id3, length, metadata), dtype="uint8")
+    for i in range(length):
+      memory_buffer[i] = (i + 1) % 256
+    self.plasma_client.seal(object_id3)
+    self.assertNotEqual(self.plasma_client.hash(object_id1),
+                        self.plasma_client.hash(object_id3))
+
+    # Create a fourth object with the same value as the third, but different
+    # metadata. Check that its hash is different from any of the previous
+    # three.
+    object_id4 = random_object_id()
+    metadata4 = generate_metadata(length)
+    memory_buffer = np.frombuffer(self.plasma_client.create(object_id4, length, metadata4), dtype="uint8")
+    for i in range(length):
+      memory_buffer[i] = (i + 1) % 256
+    self.plasma_client.seal(object_id4)
+    self.assertNotEqual(self.plasma_client.hash(object_id1),
+                        self.plasma_client.hash(object_id4))
+    self.assertNotEqual(self.plasma_client.hash(object_id3),
+                        self.plasma_client.hash(object_id4))
+
+  def test_many_hashes(self):
+    hashes = []
+    length = 2 ** 10
+
+    for i in range(256):
+      object_id = random_object_id()
+      memory_buffer = np.frombuffer(self.plasma_client.create(object_id, length), dtype="uint8")
+      for j in range(length):
+        memory_buffer[j] = i
+      self.plasma_client.seal(object_id)
+      hashes.append(self.plasma_client.hash(object_id))
+
+    # Create objects of varying length. Each pair has two bits different.
+    for i in range(length):
+      object_id = random_object_id()
+      memory_buffer = np.frombuffer(self.plasma_client.create(object_id, length), dtype="uint8")
+      for j in range(length):
+        memory_buffer[j] = 0
+      memory_buffer[i] = 1
+      self.plasma_client.seal(object_id)
+      hashes.append(self.plasma_client.hash(object_id))
+
+    # Create objects of varying length, all with value 0.
+    for i in range(length):
+      object_id = random_object_id()
+      memory_buffer = np.frombuffer(self.plasma_client.create(object_id, i), dtype="uint8")
+      for j in range(i):
+        memory_buffer[j] = 0
+      self.plasma_client.seal(object_id)
+      hashes.append(self.plasma_client.hash(object_id))
+
+    # Check that all hashes were unique.
+    self.assertEqual(len(set(hashes)), 256 + length + length)
+
+  # def test_individual_delete(self):
+  #   length = 100
+  #   # Create an object id string.
+  #   object_id = random_object_id()
+  #   # Create a random metadata string.
+  #   metadata = generate_metadata(100)
+  #   # Create a new buffer and write to it.
+  #   memory_buffer = self.plasma_client.create(object_id, length, metadata)
+  #   for i in range(length):
+  #     memory_buffer[i] = chr(i % 256)
+  #   # Seal the object.
+  #   self.plasma_client.seal(object_id)
+  #   # Check that the object is present.
+  #   self.assertTrue(self.plasma_client.contains(object_id))
+  #   # Delete the object.
+  #   self.plasma_client.delete(object_id)
+  #   # Make sure the object is no longer present.
+  #   self.assertFalse(self.plasma_client.contains(object_id))
+  #
+  # def test_delete(self):
+  #   # Create some objects.
+  #   object_ids = [random_object_id() for _ in range(100)]
+  #   for object_id in object_ids:
+  #     length = 100
+  #     # Create a random metadata string.
+  #     metadata = generate_metadata(100)
+  #     # Create a new buffer and write to it.
+  #     memory_buffer = self.plasma_client.create(object_id, length, metadata)
+  #     for i in range(length):
+  #       memory_buffer[i] = chr(i % 256)
+  #     # Seal the object.
+  #     self.plasma_client.seal(object_id)
+  #     # Check that the object is present.
+  #     self.assertTrue(self.plasma_client.contains(object_id))
+  #
+  #   # Delete the objects and make sure they are no longer present.
+  #   for object_id in object_ids:
+  #     # Delete the object.
+  #     self.plasma_client.delete(object_id)
+  #     # Make sure the object is no longer present.
+  #     self.assertFalse(self.plasma_client.contains(object_id))
+
+  def test_illegal_functionality(self):
+    # Create an object id string.
+    object_id = random_object_id()
+    # Create a new buffer and write to it.
+    length = 1000
+    memory_buffer = self.plasma_client.create(object_id, length)
+    # Make sure we cannot access memory out of bounds.
+    self.assertRaises(Exception, lambda: memory_buffer[length])
+    # Seal the object.
+    self.plasma_client.seal(object_id)
+    # This test is commented out because it currently fails.
+    # # Make sure the object is ready only now.
+    # def illegal_assignment():
+    #   memory_buffer[0] = chr(0)
+    # self.assertRaises(Exception, illegal_assignment)
+    # Get the object.
+    memory_buffer = self.plasma_client.get([object_id])[0]
+
+    # Make sure the object is read only.
+    def illegal_assignment():
+      memory_buffer[0] = chr(0)
+    self.assertRaises(Exception, illegal_assignment)
+
+  def test_evict(self):
+    client = self.plasma_client2
+    object_id1 = random_object_id()
+    b1 = client.create(object_id1, 1000)
+    client.seal(object_id1)
+    del b1
+    self.assertEqual(client.evict(1), 1000)
+
+    object_id2 = random_object_id()
+    object_id3 = random_object_id()
+    b2 = client.create(object_id2, 999)
+    b3 = client.create(object_id3, 998)
+    client.seal(object_id3)
+    del b3
+    self.assertEqual(client.evict(1000), 998)
+
+    object_id4 = random_object_id()
+    b4 = client.create(object_id4, 997)
+    client.seal(object_id4)
+    del b4
+    client.seal(object_id2)
+    del b2
+    self.assertEqual(client.evict(1), 997)
+    self.assertEqual(client.evict(1), 999)
+
+    object_id5 = random_object_id()
+    object_id6 = random_object_id()
+    object_id7 = random_object_id()
+    b5 = client.create(object_id5, 996)
+    b6 = client.create(object_id6, 995)
+    b7 = client.create(object_id7, 994)
+    client.seal(object_id5)
+    client.seal(object_id6)
+    client.seal(object_id7)
+    del b5
+    del b6
+    del b7
+    self.assertEqual(client.evict(2000), 996 + 995 + 994)
+
+  def test_subscribe(self):
+    # Subscribe to notifications from the Plasma Store.
+    self.plasma_client.subscribe()
+    for i in [1, 10, 100, 1000, 10000, 100000]:
+      object_ids = [random_object_id() for _ in range(i)]
+      metadata_sizes = [np.random.randint(1000) for _ in range(i)]
+      data_sizes = [np.random.randint(1000) for _ in range(i)]
+      for j in range(i):
+        self.plasma_client.create(
+            object_ids[j], data_sizes[j],
+            metadata=bytearray(np.random.bytes(metadata_sizes[j])))
+        self.plasma_client.seal(object_ids[j])
+      # Check that we received notifications for all of the objects.
+      for j in range(i):
+        notification_info = self.plasma_client.get_next_notification()
+        recv_objid, recv_dsize, recv_msize = notification_info
+        self.assertEqual(object_ids[j], recv_objid)
+        self.assertEqual(data_sizes[j], recv_dsize)
+        self.assertEqual(metadata_sizes[j], recv_msize)
+
+  def test_subscribe_deletions(self):
+    # Subscribe to notifications from the Plasma Store. We use plasma_client2
+    # to make sure that all used objects will get evicted properly.
+    self.plasma_client2.subscribe()
+    for i in [1, 10, 100, 1000, 10000, 100000]:
+      object_ids = [random_object_id() for _ in range(i)]
+      # Add 1 to the sizes to make sure we have nonzero object sizes.
+      metadata_sizes = [np.random.randint(1000) + 1 for _ in range(i)]
+      data_sizes = [np.random.randint(1000) + 1 for _ in range(i)]
+      for j in range(i):
+        x = self.plasma_client2.create(
+            object_ids[j], data_sizes[j],
+            metadata=bytearray(np.random.bytes(metadata_sizes[j])))
+        self.plasma_client2.seal(object_ids[j])
+      del x
+      # Check that we received notifications for creating all of the objects.
+      for j in range(i):
+        notification_info = self.plasma_client2.get_next_notification()
+        recv_objid, recv_dsize, recv_msize = notification_info
+        self.assertEqual(object_ids[j], recv_objid)
+        self.assertEqual(data_sizes[j], recv_dsize)
+        self.assertEqual(metadata_sizes[j], recv_msize)
+
+      # Check that we receive notifications for deleting all objects, as we
+      # evict them.
+      for j in range(i):
+        self.assertEqual(self.plasma_client2.evict(1),
+                         data_sizes[j] + metadata_sizes[j])
+        notification_info = self.plasma_client2.get_next_notification()
+        recv_objid, recv_dsize, recv_msize = notification_info
+        self.assertEqual(object_ids[j], recv_objid)
+        self.assertEqual(-1, recv_dsize)
+        self.assertEqual(-1, recv_msize)
+
+    # Test multiple deletion notifications. The first 9 object IDs have size 0,
+    # and the last has a nonzero size. When Plasma evicts 1 byte, it will evict
+    # all objects, so we should receive deletion notifications for each.
+    num_object_ids = 10
+    object_ids = [random_object_id() for _ in range(num_object_ids)]
+    metadata_sizes = [0] * (num_object_ids - 1)
+    data_sizes = [0] * (num_object_ids - 1)
+    metadata_sizes.append(np.random.randint(1000))
+    data_sizes.append(np.random.randint(1000))
+    for i in range(num_object_ids):
+      x = self.plasma_client2.create(
+          object_ids[i], data_sizes[i],
+          metadata=bytearray(np.random.bytes(metadata_sizes[i])))
+      self.plasma_client2.seal(object_ids[i])
+    del x
+    for i in range(num_object_ids):
+      notification_info = self.plasma_client2.get_next_notification()
+      recv_objid, recv_dsize, recv_msize = notification_info
+      self.assertEqual(object_ids[i], recv_objid)
+      self.assertEqual(data_sizes[i], recv_dsize)
+      self.assertEqual(metadata_sizes[i], recv_msize)
+    self.assertEqual(self.plasma_client2.evict(1),
+                     data_sizes[-1] + metadata_sizes[-1])
+    for i in range(num_object_ids):
+      notification_info = self.plasma_client2.get_next_notification()
+      recv_objid, recv_dsize, recv_msize = notification_info
+      self.assertEqual(object_ids[i], recv_objid)
+      self.assertEqual(-1, recv_dsize)
+      self.assertEqual(-1, recv_msize)
+
 if __name__ == "__main__":
   if len(sys.argv) > 1:
     # Pop the argument so we don't mess with unittest's own argument parser.