Factoring out object_info structure for use in several Ray components (#101)

* change plasma object notifications to carry a struct of information

* factoring out object_info for general use by several Ray components

* fixing a bug in python test

* addressing comments

* handling Robert's comments

* clang format

* Fix valgrind.

src/common/object_info.h

@@ -0,0 +1,19 @@
+#ifndef OBJECT_H
+#define OBJECT_H
+
+#include <stdint.h>
+
+#include "common.h"
+
+/**
+ * Object information data structure.
+ */
+typedef struct {
+  object_id obj_id;
+  int64_t data_size;
+  int64_t metadata_size;
+  int64_t create_time;
+  int64_t construct_duration;
+} object_info;
+
+#endif

src/photon/photon_scheduler.c

@@ -9,6 +9,7 @@
 #include "common.h"
 #include "event_loop.h"
 #include "io.h"
+#include "object_info.h"
 #include "photon.h"
 #include "photon_algorithm.h"
 #include "photon_scheduler.h"
@@ -104,8 +105,9 @@ void process_plasma_notification(event_loop *loop,
                                  int events) {
   local_scheduler_state *s = context;
   /* Read the notification from Plasma. */
-  object_id obj_id;
-  int error = read_bytes(client_sock, (uint8_t *) &obj_id, sizeof(obj_id));
+  object_info object_info;
+  int error =
+      read_bytes(client_sock, (uint8_t *) &object_info, sizeof(object_info));
   if (error < 0) {
     /* The store has closed the socket. */
     LOG_DEBUG(
@@ -115,10 +117,12 @@ void process_plasma_notification(event_loop *loop,
     close(client_sock);
     return;
   }
-  handle_object_available(s, s->algorithm_state, obj_id);
+  handle_object_available(s, s->algorithm_state, object_info.obj_id);
 }
 
-void process_message(event_loop *loop, int client_sock, void *context,
+void process_message(event_loop *loop,
+                     int client_sock,
+                     void *context,
                      int events) {
   local_scheduler_state *s = context;
 
@@ -151,7 +155,9 @@ void process_message(event_loop *loop, int client_sock, void *context,
   }
 }
 
-void new_client_connection(event_loop *loop, int listener_sock, void *context,
+void new_client_connection(event_loop *loop,
+                           int listener_sock,
+                           void *context,
                            int events) {
   local_scheduler_state *s = context;
   int new_socket = accept_client(listener_sock);

src/plasma/lib/python/plasma.py

@@ -211,6 +211,7 @@ class PlasmaClient(object):
     """Subscribe to notifications about sealed objects."""
     fd = libplasma.subscribe(self.conn)
     self.notification_sock = socket.fromfd(fd, socket.AF_UNIX, socket.SOCK_STREAM)
+    self.notification_fd = fd
     # Make the socket non-blocking.
     self.notification_sock.setblocking(0)
 
@@ -221,13 +222,14 @@ class PlasmaClient(object):
     # Loop until we've read PLASMA_ID_SIZE bytes from the socket.
     while True:
       try:
-        message_data = self.notification_sock.recv(PLASMA_ID_SIZE)
+        rv = libplasma.receive_notification(self.notification_fd)
+        obj_id, data_size, metadata_size = rv
       except socket.error:
         time.sleep(0.001)
       else:
-        assert len(message_data) == PLASMA_ID_SIZE
+        assert len(obj_id) == PLASMA_ID_SIZE
         break
-    return message_data
+    return obj_id, data_size, metadata_size
 
 DEFAULT_PLASMA_STORE_MEMORY = 10 ** 9
 

src/plasma/plasma.h

@@ -11,17 +11,11 @@
 #include <unistd.h> /* pid_t */
 
 #include "common.h"
+#include "object_info.h"
 
 #include "utarray.h"
 #include "uthash.h"
 
-typedef struct {
-  int64_t data_size;
-  int64_t metadata_size;
-  int64_t create_time;
-  int64_t construct_duration;
-} plasma_object_info;
-
 /**
  * Object request data structure. Used in the plasma_wait_for_objects()
  * argument.
@@ -199,7 +193,7 @@ typedef struct {
   /** Object id of this object. */
   object_id object_id;
   /** Object info like size, creation time and owner. */
-  plasma_object_info info;
+  object_info info;
   /** Memory mapped file containing the object. */
   int fd;
   /** Size of the underlying map. */

src/plasma/plasma_client.h

@@ -291,17 +291,6 @@ typedef struct {
   uint8_t *metadata;
 } object_buffer;
 
-/**
- * Object information data structure.
- */
-typedef struct {
-  /** The time when the object was created (sealed). */
-  time_t last_access_time;
-  /** The time when the object was last accessed. */
-  time_t creation_date;
-  uint64_t refcount;
-} object_info;
-
 /**
  * Get specified object from the local Plasma Store. This function is
  * non-blocking.

src/plasma/plasma_extension.c

@@ -1,7 +1,9 @@
 #include <Python.h>
 
 #include "common.h"
+#include "io.h"
 #include "plasma_client.h"
+#include "object_info.h"
 
 static int PyObjectToPlasmaConnection(PyObject *object,
                                       plasma_connection **conn) {
@@ -305,10 +307,38 @@ PyObject *PyPlasma_subscribe(PyObject *self, PyObject *args) {
   if (!PyArg_ParseTuple(args, "O&", PyObjectToPlasmaConnection, &conn)) {
     return NULL;
   }
+
   int sock = plasma_subscribe(conn);
   return PyInt_FromLong(sock);
 }
 
+PyObject *PyPlasma_receive_notification(PyObject *self, PyObject *args) {
+  int plasma_sock;
+  object_info object_info;
+
+  if (!PyArg_ParseTuple(args, "i", &plasma_sock)) {
+    return NULL;
+  }
+  /* Receive object notification from the plasma connection socket,
+   * return a tuple of its fields: object_id, data_size, metadata_size. */
+  int nbytes =
+      read_bytes(plasma_sock, (uint8_t *) &object_info, sizeof(object_info));
+
+  if (nbytes < 0) {
+    PyErr_SetString(PyExc_RuntimeError,
+                    "Failed to read object notification from Plasma socket");
+    return NULL;
+  }
+  /* Construct a tuple from object_info and return. */
+  PyObject *t = PyTuple_New(3);
+  PyTuple_SetItem(t, 0, PyString_FromStringAndSize(
+                            (char *) object_info.obj_id.id, UNIQUE_ID_SIZE));
+  PyTuple_SetItem(t, 1, PyInt_FromLong(object_info.data_size));
+  PyTuple_SetItem(t, 2, PyInt_FromLong(object_info.metadata_size));
+
+  return t;
+}
+
 static PyMethodDef plasma_methods[] = {
     {"connect", PyPlasma_connect, METH_VARARGS, "Connect to plasma."},
     {"disconnect", PyPlasma_disconnect, METH_VARARGS,
@@ -332,6 +362,8 @@ static PyMethodDef plasma_methods[] = {
      "Transfer object to another plasma manager."},
     {"subscribe", PyPlasma_subscribe, METH_VARARGS,
      "Subscribe to the plasma notification socket."},
+    {"receive_notification", PyPlasma_receive_notification, METH_VARARGS,
+     "Receive next notification from plasma notification socket."},
     {NULL} /* Sentinel */
 };
 

src/plasma/plasma_manager.c

@@ -1497,13 +1497,15 @@ void process_object_notification(event_loop *loop,
                                  int events) {
   plasma_manager_state *state = context;
   object_id obj_id;
+  object_info object_info;
   retry_info retry = {
       .num_retries = NUM_RETRIES,
       .timeout = MANAGER_TIMEOUT,
       .fail_callback = NULL,
   };
   /* Read the notification from Plasma. */
-  int error = read_bytes(client_sock, (uint8_t *) &obj_id, sizeof(obj_id));
+  int error =
+      read_bytes(client_sock, (uint8_t *) &object_info, sizeof(object_info));
   if (error < 0) {
     /* The store has closed the socket. */
     LOG_DEBUG(
@@ -1513,6 +1515,7 @@ void process_object_notification(event_loop *loop,
     close(client_sock);
     return;
   }
+  obj_id = object_info.obj_id;
   /* Add object to locally available object. */
   /* TODO(pcm): Where is this deallocated? */
   available_object *entry =

src/plasma/plasma_store.c

@@ -164,7 +164,9 @@ bool create_object(client *client_context,
   assert(fd != -1);
 
   entry = malloc(sizeof(object_table_entry));
+  memset(entry, 0, sizeof(object_table_entry));
   memcpy(&entry->object_id, &obj_id, sizeof(entry->object_id));
+  entry->info.obj_id = obj_id;
   entry->info.data_size = data_size;
   entry->info.metadata_size = metadata_size;
   entry->pointer = pointer;
@@ -399,11 +401,21 @@ void send_notifications(event_loop *loop,
    * possible. */
   for (int i = 0; i < utarray_len(queue->object_ids); ++i) {
     object_id *obj_id = (object_id *) utarray_eltptr(queue->object_ids, i);
+    object_table_entry *entry = NULL;
+    /* This object should already exist in plasma store state. */
+    HASH_FIND(handle, plasma_state->plasma_store_info->objects, obj_id,
+              sizeof(object_id), entry);
+    CHECK(entry != NULL);
+
+    object_info object_info = entry->info;
+
     /* Attempt to send a notification about this object ID. */
-    int nbytes = send(client_sock, (char const *) obj_id, sizeof(*obj_id), 0);
+    int nbytes =
+        send(client_sock, (char const *) &object_info, sizeof(object_info), 0);
     if (nbytes >= 0) {
-      CHECK(nbytes == sizeof(*obj_id));
-    } else if (nbytes == -1 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
+      CHECK(nbytes == sizeof(object_info));
+    } else if (nbytes == -1 &&
+               (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)) {
       LOG_DEBUG(
           "The socket's send buffer is full, so we are caching this "
           "notification and will send it later.");

src/plasma/test/test.py

@@ -254,14 +254,18 @@ class TestPlasmaClient(unittest.TestCase):
     sock = self.plasma_client.subscribe()
     for i in [1, 10, 100, 1000, 10000, 100000]:
       object_ids = [random_object_id() for _ in range(i)]
-      for object_id in object_ids:
-        # Create an object and seal it to trigger a notification.
-        self.plasma_client.create(object_id, 1000)
-        self.plasma_client.seal(object_id)
+      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], size=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 object_id in object_ids:
-        message_data = self.plasma_client.get_next_notification()
-        self.assertEqual(object_id, message_data)
+      for j in range(i):
+        recv_objid, recv_dsize, recv_msize = self.plasma_client.get_next_notification()
+        self.assertEqual(object_ids[j], recv_objid)
+        self.assertEqual(data_sizes[j], recv_dsize)
+        self.assertEqual(metadata_sizes[j], recv_msize)
 
 class TestPlasmaManager(unittest.TestCase):