Convert local scheduler messages to flatbuffers (#340)

* use flatbuffer messages for local scheduler

* make sure constructor gets called for C++ object ObjectInfoT

* fix typo

* fix Robert's comments

* Small change to actor test.

* fix valgrind error

* linting

* free notification

* fix

* valgrind

* fix valgrind

* fix other bugs

* valgrind fix

* fixes

* more fixes

* Small changes to comments.

src/plasma/format/plasma.fbs

@@ -38,8 +38,15 @@ enum MessageType:int {
   // Unsubscribe.
   PlasmaUnsubscribeRequest,
   // Sending and receiving data.
+  // PlasmaDataRequest initiates sending the data, there will be one
+  // such message per data transfer.
   PlasmaDataRequest,
-  PlasmaDataReply
+  // PlasmaDataReply contains the actual data and is sent back to the
+  // object store that requested the data. For each transfer, multiple
+  // reply messages get sent. Each one contains a fixed number of bytes.
+  PlasmaDataReply,
+  // Object notifications.
+  PlasmaNotification
 }
 
 enum PlasmaError:int {

src/plasma/plasma.cc

@@ -20,3 +20,22 @@ void warn_if_sigpipe(int status, int client_sock) {
   }
   LOG_FATAL("Failed to write message to client on fd %d.", client_sock);
 }
+
+/**
+ * This will create a new ObjectInfo buffer. The first sizeof(int64_t) bytes
+ * of this buffer are the length of the remaining message and the
+ * remaining message is a serialized version of the object info.
+ *
+ * @param object_info The object info to be serialized
+ * @return The object info buffer. It is the caller's responsibility to free
+ *         this buffer with "free" after it has been used.
+ */
+uint8_t *create_object_info_buffer(ObjectInfoT *object_info) {
+  flatbuffers::FlatBufferBuilder fbb;
+  auto message = CreateObjectInfo(fbb, object_info);
+  fbb.Finish(message);
+  uint8_t *notification = (uint8_t *) malloc(sizeof(int64_t) + fbb.GetSize());
+  *((int64_t *) notification) = fbb.GetSize();
+  memcpy(notification + sizeof(int64_t), fbb.GetBufferPointer(), fbb.GetSize());
+  return notification;
+}

src/plasma/plasma.h

@@ -11,7 +11,7 @@
 #include <unistd.h> /* pid_t */
 
 #include "common.h"
-#include "object_info.h"
+#include "format/common_generated.h"
 
 #include "utarray.h"
 #include "uthash.h"
@@ -89,7 +89,7 @@ typedef struct {
   /** Object id of this object. */
   ObjectID object_id;
   /** Object info like size, creation time and owner. */
-  ObjectInfo info;
+  ObjectInfoT info;
   /** Memory mapped file containing the object. */
   int fd;
   /** Size of the underlying map. */
@@ -134,4 +134,6 @@ typedef struct {
  */
 void warn_if_sigpipe(int status, int client_sock);
 
+uint8_t *create_object_info_buffer(ObjectInfoT *object_info);
+
 #endif /* PLASMA_H */

src/plasma/plasma_extension.cc

@@ -6,7 +6,6 @@
 #include "io.h"
 #include "plasma_protocol.h"
 #include "plasma_client.h"
-#include "object_info.h"
 
 PyObject *PlasmaOutOfMemoryError;
 PyObject *PlasmaObjectExistsError;
@@ -348,7 +347,6 @@ PyObject *PyPlasma_subscribe(PyObject *self, PyObject *args) {
 
 PyObject *PyPlasma_receive_notification(PyObject *self, PyObject *args) {
   int plasma_sock;
-  ObjectInfo object_info;
 
   if (!PyArg_ParseTuple(args, "i", &plasma_sock)) {
     return NULL;
@@ -357,26 +355,27 @@ PyObject *PyPlasma_receive_notification(PyObject *self, PyObject *args) {
    * object was added, return a tuple of its fields: ObjectID, data_size,
    * metadata_size. If the object was deleted, data_size and metadata_size will
    * be set to -1. */
-  int nbytes =
-      read_bytes(plasma_sock, (uint8_t *) &object_info, sizeof(object_info));
-
-  if (nbytes < 0) {
+  uint8_t *notification = read_message_async(NULL, plasma_sock);
+  if (notification == NULL) {
     PyErr_SetString(PyExc_RuntimeError,
                     "Failed to read object notification from Plasma socket");
     return NULL;
   }
+  auto object_info = flatbuffers::GetRoot<ObjectInfo>(notification);
   /* Construct a tuple from object_info and return. */
   PyObject *t = PyTuple_New(3);
-  PyTuple_SetItem(t, 0, PyBytes_FromStringAndSize(
-                            (char *) object_info.obj_id.id, UNIQUE_ID_SIZE));
-  if (object_info.is_deletion) {
+  PyTuple_SetItem(t, 0,
+                  PyBytes_FromStringAndSize(object_info->object_id()->data(),
+                                            object_info->object_id()->size()));
+  if (object_info->is_deletion()) {
     PyTuple_SetItem(t, 1, PyLong_FromLong(-1));
     PyTuple_SetItem(t, 2, PyLong_FromLong(-1));
   } else {
-    PyTuple_SetItem(t, 1, PyLong_FromLong(object_info.data_size));
-    PyTuple_SetItem(t, 2, PyLong_FromLong(object_info.metadata_size));
+    PyTuple_SetItem(t, 1, PyLong_FromLong(object_info->data_size()));
+    PyTuple_SetItem(t, 2, PyLong_FromLong(object_info->metadata_size()));
   }
 
+  free(notification);
   return t;
 }
 

src/plasma/plasma_manager.cc

@@ -25,6 +25,7 @@
 #include "utlist.h"
 #include "utarray.h"
 #include "utstring.h"
+#include "common_protocol.h"
 #include "common.h"
 #include "io.h"
 #include "net.h"
@@ -1252,10 +1253,10 @@ void process_status_request(ClientConnection *client_conn, ObjectID object_id) {
 }
 
 void process_delete_object_notification(PlasmaManagerState *state,
-                                        ObjectInfo object_info) {
-  ObjectID obj_id = object_info.obj_id;
+                                        ObjectID object_id) {
   AvailableObject *entry;
-  HASH_FIND(hh, state->local_available_objects, &obj_id, sizeof(obj_id), entry);
+  HASH_FIND(hh, state->local_available_objects, &object_id, sizeof(object_id),
+            entry);
   if (entry != NULL) {
     HASH_DELETE(hh, state->local_available_objects, entry);
     free(entry);
@@ -1263,7 +1264,7 @@ void process_delete_object_notification(PlasmaManagerState *state,
 
   /* Remove this object from the (redis) object table. */
   if (state->db) {
-    object_table_remove(state->db, obj_id, NULL, NULL, NULL, NULL);
+    object_table_remove(state->db, object_id, NULL, NULL, NULL, NULL);
   }
 
   /* NOTE: There could be pending wait requests for this object that will now
@@ -1311,33 +1312,35 @@ void log_object_hash_mismatch_error_object_callback(ObjectID object_id,
 }
 
 void process_add_object_notification(PlasmaManagerState *state,
-                                     ObjectInfo object_info) {
-  ObjectID obj_id = object_info.obj_id;
+                                     ObjectID object_id,
+                                     int64_t data_size,
+                                     int64_t metadata_size,
+                                     unsigned char *digest) {
   AvailableObject *entry = (AvailableObject *) malloc(sizeof(AvailableObject));
-  entry->object_id = obj_id;
+  entry->object_id = object_id;
   HASH_ADD(hh, state->local_available_objects, object_id, sizeof(ObjectID),
            entry);
 
   /* Add this object to the (redis) object table. */
   if (state->db) {
-    object_table_add(
-        state->db, obj_id, object_info.data_size + object_info.metadata_size,
-        object_info.digest, NULL,
-        log_object_hash_mismatch_error_object_callback, (void *) state);
+    object_table_add(state->db, object_id, data_size + metadata_size, digest,
+                     NULL, log_object_hash_mismatch_error_object_callback,
+                     (void *) state);
   }
 
   /* If we were trying to fetch this object, finish up the fetch request. */
   FetchRequest *fetch_req;
-  HASH_FIND(hh, state->fetch_requests, &obj_id, sizeof(obj_id), fetch_req);
+  HASH_FIND(hh, state->fetch_requests, &object_id, sizeof(object_id),
+            fetch_req);
   if (fetch_req != NULL) {
     remove_fetch_request(state, fetch_req);
     /* TODO(rkn): We also really should unsubscribe from the object table. */
   }
 
   /* Update the in-progress local and remote wait requests. */
-  update_object_wait_requests(state, obj_id, PLASMA_QUERY_LOCAL,
+  update_object_wait_requests(state, object_id, PLASMA_QUERY_LOCAL,
                               ObjectStatus_Local);
-  update_object_wait_requests(state, obj_id, PLASMA_QUERY_ANYWHERE,
+  update_object_wait_requests(state, object_id, PLASMA_QUERY_ANYWHERE,
                               ObjectStatus_Local);
 }
 
@@ -1346,25 +1349,22 @@ void process_object_notification(event_loop *loop,
                                  void *context,
                                  int events) {
   PlasmaManagerState *state = (PlasmaManagerState *) context;
-  ObjectInfo object_info;
-  /* Read the notification from Plasma. */
-  int error =
-      read_bytes(client_sock, (uint8_t *) &object_info, sizeof(object_info));
-  if (error < 0) {
-    /* The store has closed the socket. */
-    LOG_DEBUG(
-        "The plasma store has closed the object notification socket, or some "
-        "other error has occurred.");
-    event_loop_remove_file(loop, client_sock);
-    close(client_sock);
+  uint8_t *notification = read_message_async(loop, client_sock);
+  if (notification == NULL) {
     return;
   }
+  auto object_info = flatbuffers::GetRoot<ObjectInfo>(notification);
   /* Add object to locally available object. */
-  if (object_info.is_deletion) {
-    process_delete_object_notification(state, object_info);
+  ObjectID object_id = from_flatbuf(object_info->object_id());
+  if (object_info->is_deletion()) {
+    process_delete_object_notification(state, object_id);
   } else {
-    process_add_object_notification(state, object_info);
+    process_add_object_notification(
+        state, object_id, object_info->data_size(),
+        object_info->metadata_size(),
+        (unsigned char *) object_info->digest()->data());
   }
+  free(notification);
 }
 
 void process_message(event_loop *loop,

src/plasma/plasma_store.cc

@@ -26,6 +26,7 @@
 #include <poll.h>
 
 #include "common.h"
+#include "format/common_generated.h"
 #include "event_loop.h"
 #include "eviction_policy.h"
 #include "io.h"
@@ -56,7 +57,7 @@ UT_icd client_icd = {sizeof(Client *), NULL, NULL, NULL};
 
 /* This is used to define the queue of object notifications for plasma
  * subscribers. */
-UT_icd object_info_icd = {sizeof(ObjectInfo), NULL, NULL, NULL};
+UT_icd object_info_icd = {sizeof(uint8_t *), NULL, NULL, NULL};
 
 typedef struct {
   /** Client file descriptor. This is used as a key for the hash table. */
@@ -124,6 +125,8 @@ struct PlasmaStoreState {
   protocol_builder *builder;
 };
 
+PlasmaStoreState *g_state;
+
 UT_icd byte_icd = {sizeof(uint8_t), NULL, NULL, NULL};
 
 PlasmaStoreState *PlasmaStoreState_init(event_loop *loop,
@@ -145,8 +148,30 @@ PlasmaStoreState *PlasmaStoreState_init(event_loop *loop,
   return state;
 }
 
+void PlasmaStoreState_free(PlasmaStoreState *state) {
+  /* Here we only clean up objects that need to be cleaned
+   * up to make the valgrind warnings go away. Objects that
+   * are still reachable are not cleaned up. */
+  object_table_entry *entry, *tmp;
+  HASH_ITER(handle, state->plasma_store_info->objects, entry, tmp) {
+    HASH_DELETE(handle, state->plasma_store_info->objects, entry);
+    utarray_free(entry->clients);
+    delete entry;
+  }
+  NotificationQueue *queue, *temp_queue;
+  HASH_ITER(hh, state->pending_notifications, queue, temp_queue) {
+    for (int i = 0; i < utarray_len(queue->object_notifications); ++i) {
+      uint8_t **notification =
+          (uint8_t **) utarray_eltptr(queue->object_notifications, i);
+      uint8_t *data = *notification;
+      free(data);
+    }
+    utarray_free(queue->object_notifications);
+  }
+}
+
 void push_notification(PlasmaStoreState *state,
-                       ObjectInfo *object_notification);
+                       ObjectInfoT *object_notification);
 
 /* If this client is not already using the object, add the client to the
  * object's list of clients, otherwise do nothing. */
@@ -213,10 +238,9 @@ int create_object(Client *client_context,
   get_malloc_mapinfo(pointer, &fd, &map_size, &offset);
   assert(fd != -1);
 
-  entry = (object_table_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 = new object_table_entry();
+  entry->object_id = obj_id;
+  entry->info.object_id = std::string((char *) &obj_id.id[0], sizeof(obj_id));
   entry->info.data_size = data_size;
   entry->info.metadata_size = metadata_size;
   entry->pointer = pointer;
@@ -546,7 +570,7 @@ void seal_object(Client *client_context,
   /* Set the state of object to SEALED. */
   entry->state = PLASMA_SEALED;
   /* Set the object digest. */
-  memcpy(entry->info.digest, digest, DIGEST_SIZE);
+  entry->info.digest = std::string((char *) &digest[0], DIGEST_SIZE);
   /* Inform all subscribers that a new object has been sealed. */
   push_notification(plasma_state, &entry->info);
 
@@ -573,13 +597,11 @@ void delete_object(PlasmaStoreState *plasma_state, ObjectID object_id) {
   HASH_DELETE(handle, plasma_state->plasma_store_info->objects, entry);
   dlfree(pointer);
   utarray_free(entry->clients);
-  free(entry);
+  delete entry;
   /* Inform all subscribers that the object has been deleted. */
-  ObjectInfo notification;
-  /* We memset the struct here because we have to initialize the full struct.
-   * However, we do not use most of the fields. */
-  memset(&notification, 0, sizeof(notification));
-  notification.obj_id = object_id;
+  ObjectInfoT notification;
+  notification.object_id =
+      std::string((char *) &object_id.id[0], sizeof(object_id));
   notification.is_deletion = true;
   push_notification(plasma_state, &notification);
 }
@@ -598,12 +620,15 @@ void remove_objects(PlasmaStoreState *plasma_state,
 }
 
 void push_notification(PlasmaStoreState *plasma_state,
-                       ObjectInfo *notification) {
+                       ObjectInfoT *object_info) {
   NotificationQueue *queue, *temp_queue;
   HASH_ITER(hh, plasma_state->pending_notifications, queue, temp_queue) {
-    utarray_push_back(queue->object_notifications, notification);
+    uint8_t *notification = create_object_info_buffer(object_info);
+    utarray_push_back(queue->object_notifications, &notification);
     send_notifications(plasma_state->loop, queue->subscriber_fd, plasma_state,
                        0);
+    /* The notification gets freed in send_notifications when the notification
+     * is sent over the socket. */
   }
 }
 
@@ -622,14 +647,16 @@ void send_notifications(event_loop *loop,
   /* Loop over the array of pending notifications and send as many of them as
    * possible. */
   for (int i = 0; i < utarray_len(queue->object_notifications); ++i) {
-    ObjectInfo *notification =
-        (ObjectInfo *) utarray_eltptr(queue->object_notifications, i);
+    uint8_t **notification =
+        (uint8_t **) utarray_eltptr(queue->object_notifications, i);
+    uint8_t *data = *notification;
+    /* Decode the length, which is the first bytes of the message. */
+    int64_t size = *((int64_t *) data);
 
     /* Attempt to send a notification about this object ID. */
-    int nbytes = send(client_sock, (char const *) notification,
-                      sizeof(*notification), 0);
+    int nbytes = send(client_sock, data, sizeof(int64_t) + size, 0);
     if (nbytes >= 0) {
-      CHECK(nbytes == sizeof(*notification));
+      CHECK(nbytes == sizeof(int64_t) + size);
     } else if (nbytes == -1 &&
                (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)) {
       LOG_DEBUG(
@@ -650,6 +677,9 @@ void send_notifications(event_loop *loop,
       }
     }
     num_processed += 1;
+    /* The corresponding malloc happened in create_object_info_buffer
+     * within push_notification. */
+    free(data);
   }
   /* Remove the sent notifications from the array. */
   utarray_erase(queue->object_notifications, 0, num_processed);
@@ -694,7 +724,7 @@ void subscribe_to_updates(Client *client_context, int conn) {
   object_table_entry *entry, *temp_entry;
   HASH_ITER(handle, plasma_state->plasma_store_info->objects, entry,
             temp_entry) {
-    utarray_push_back(queue->object_notifications, &entry->info);
+    push_notification(plasma_state, &entry->info);
   }
   send_notifications(plasma_state->loop, queue->subscriber_fd, plasma_state, 0);
 }
@@ -832,6 +862,7 @@ void new_client_connection(event_loop *loop,
 /* Report "success" to valgrind. */
 void signal_handler(int signal) {
   if (signal == SIGTERM) {
+    PlasmaStoreState_free(g_state);
     exit(0);
   }
 }
@@ -847,6 +878,7 @@ void start_server(char *socket_name, int64_t system_memory) {
   CHECK(socket >= 0);
   event_loop_add_file(loop, socket, EVENT_LOOP_READ, new_client_connection,
                       state);
+  g_state = state;
   event_loop_run(loop);
 }
 

src/plasma/test/manager_tests.cc

@@ -23,7 +23,7 @@ SUITE(plasma_manager_tests);
 const char *plasma_store_socket_name = "/tmp/plasma_store_socket_1";
 const char *plasma_manager_socket_name_format = "/tmp/plasma_manager_socket_%d";
 const char *manager_addr = "127.0.0.1";
-ObjectID oid;
+ObjectID object_id;
 
 void wait_for_pollin(int fd) {
   struct pollfd poll_list[1];
@@ -126,7 +126,7 @@ TEST request_transfer_test(void) {
   utstring_new(addr);
   utstring_printf(addr, "127.0.0.1:%d", remote_mock->port);
   manager_vector[0] = utstring_body(addr);
-  call_request_transfer(oid, 1, manager_vector, local_mock->state);
+  call_request_transfer(object_id, 1, manager_vector, local_mock->state);
   free(manager_vector);
   event_loop_add_timer(local_mock->loop, MANAGER_TIMEOUT, test_done_handler,
                        local_mock->state);
@@ -134,11 +134,11 @@ TEST request_transfer_test(void) {
   int read_fd = get_client_sock(remote_mock->read_conn);
   uint8_t *request_data =
       plasma_receive(read_fd, MessageType_PlasmaDataRequest);
-  ObjectID oid2;
+  ObjectID object_id2;
   char *address;
   int port;
-  plasma_read_DataRequest(request_data, &oid2, &address, &port);
-  ASSERT(ObjectID_equal(oid, oid2));
+  plasma_read_DataRequest(request_data, &object_id2, &address, &port);
+  ASSERT(ObjectID_equal(object_id, object_id2));
   free(address);
   /* Clean up. */
   utstring_free(addr);
@@ -173,7 +173,7 @@ TEST request_transfer_retry_test(void) {
   utstring_new(addr1);
   utstring_printf(addr1, "127.0.0.1:%d", remote_mock2->port);
   manager_vector[1] = utstring_body(addr1);
-  call_request_transfer(oid, 2, manager_vector, local_mock->state);
+  call_request_transfer(object_id, 2, manager_vector, local_mock->state);
   free(manager_vector);
   event_loop_add_timer(local_mock->loop, MANAGER_TIMEOUT * 2, test_done_handler,
                        local_mock->state);
@@ -187,12 +187,12 @@ TEST request_transfer_retry_test(void) {
   int read_fd = get_client_sock(remote_mock2->read_conn);
   uint8_t *request_data =
       plasma_receive(read_fd, MessageType_PlasmaDataRequest);
-  ObjectID oid2;
+  ObjectID object_id2;
   char *address;
   int port;
-  plasma_read_DataRequest(request_data, &oid2, &address, &port);
+  plasma_read_DataRequest(request_data, &object_id2, &address, &port);
   free(address);
-  ASSERT(ObjectID_equal(oid, oid2));
+  ASSERT(ObjectID_equal(object_id, object_id2));
   /* Clean up. */
   utstring_free(addr0);
   utstring_free(addr1);
@@ -219,13 +219,13 @@ TEST read_write_object_chunk_test(void) {
   const int metadata_size = 0;
   PlasmaRequestBuffer remote_buf;
   remote_buf.type = MessageType_PlasmaDataReply;
-  remote_buf.object_id = oid;
+  remote_buf.object_id = object_id;
   remote_buf.data = (uint8_t *) data;
   remote_buf.data_size = data_size;
   remote_buf.metadata = (uint8_t *) data + data_size;
   remote_buf.metadata_size = metadata_size;
   PlasmaRequestBuffer local_buf;
-  local_buf.object_id = oid;
+  local_buf.object_id = object_id;
   local_buf.data_size = data_size;
   local_buf.metadata_size = metadata_size;
   local_buf.data = (uint8_t *) malloc(data_size);
@@ -257,32 +257,39 @@ TEST object_notifications_test(void) {
   int flags = fcntl(fd[1], F_GETFL, 0);
   CHECK(fcntl(fd[1], F_SETFL, flags | O_NONBLOCK) == 0);
 
-  ObjectID oid = globally_unique_id();
-  ObjectInfo info = {.obj_id = oid,
-                     .data_size = 10,
-                     .metadata_size = 1,
-                     .create_time = 0,
-                     .construct_duration = 0,
-                     .digest = {0},
-                     .is_deletion = false};
+  ObjectID object_id = globally_unique_id();
+  ObjectInfoT info;
+  info.object_id = std::string((char *) &object_id.id[0], sizeof(object_id));
+  info.data_size = 10;
+  info.metadata_size = 1;
+  info.create_time = 0;
+  info.construct_duration = 0;
+  info.digest = std::string("0");
+  info.is_deletion = false;
 
   /* Check that the object is not local at first. */
-  bool is_local = is_object_local(local_mock->state, oid);
+  bool is_local = is_object_local(local_mock->state, object_id);
   ASSERT(!is_local);
 
   /* Check that the object is local after receiving an object notification. */
-  send(fd[1], (char const *) &info, sizeof(info), 0);
+  uint8_t *notification = create_object_info_buffer(&info);
+  int64_t size = *((int64_t *) notification);
+  send(fd[1], notification, sizeof(int64_t) + size, 0);
   process_object_notification(local_mock->loop, fd[0], local_mock->state, 0);
-  is_local = is_object_local(local_mock->state, oid);
+  is_local = is_object_local(local_mock->state, object_id);
   ASSERT(is_local);
+  free(notification);
 
   /* Check that the object is not local after receiving a notification about
    * the object deletion. */
   info.is_deletion = true;
-  send(fd[1], (char const *) &info, sizeof(info), 0);
+  notification = create_object_info_buffer(&info);
+  size = *((int64_t *) notification);
+  send(fd[1], notification, sizeof(int64_t) + size, 0);
   process_object_notification(local_mock->loop, fd[0], local_mock->state, 0);
-  is_local = is_object_local(local_mock->state, oid);
+  is_local = is_object_local(local_mock->state, object_id);
   ASSERT(!is_local);
+  free(notification);
 
   /* Clean up. */
   close(fd[0]);
@@ -292,7 +299,7 @@ TEST object_notifications_test(void) {
 }
 
 SUITE(plasma_manager_tests) {
-  memset(&oid, 1, sizeof(oid));
+  memset(&object_id, 1, sizeof(object_id));
   RUN_TEST(request_transfer_test);
   RUN_TEST(request_transfer_retry_test);
   RUN_TEST(read_write_object_chunk_test);