Add metadata handling (#23)

* Automatic whitespace fixes.

* Add metadata handling.

* Make create take a buffer instead of a string for the metadata.

* Small fixes.

lib/python/plasma.py

@@ -25,7 +25,7 @@ class PlasmaClient(object):
 
   def __init__(self, socket_name, addr=None, port=None):
     """Initialize the PlasmaClient.
-    
+
     Args:
       socket_name (str): Name of the socket the plasma store is listening at.
       addr (str): IPv4 address of plasma manager attached to the plasma store.
@@ -36,10 +36,10 @@ class PlasmaClient(object):
 
     self.client.plasma_store_connect.restype = ctypes.c_int
 
-    self.client.plasma_create.argtypes = [ctypes.c_int, PlasmaID, ctypes.c_int64, ctypes.POINTER(ctypes.c_void_p)]
+    self.client.plasma_create.argtypes = [ctypes.c_int, PlasmaID, ctypes.c_int64, ctypes.POINTER(ctypes.c_uint8), ctypes.c_int64, ctypes.POINTER(ctypes.c_void_p)]
     self.client.plasma_create.restype = None
 
-    self.client.plasma_get.argtypes = [ctypes.c_int, PlasmaID, ctypes.POINTER(ctypes.c_int64), ctypes.POINTER(ctypes.c_void_p)]
+    self.client.plasma_get.argtypes = [ctypes.c_int, PlasmaID, ctypes.POINTER(ctypes.c_int64), ctypes.POINTER(ctypes.c_void_p), ctypes.POINTER(ctypes.c_int64), ctypes.POINTER(ctypes.c_void_p)]
     self.client.plasma_get.restype = None
 
     self.client.plasma_seal.argtypes = [ctypes.c_int, PlasmaID]
@@ -60,7 +60,7 @@ class PlasmaClient(object):
     else:
       self.manager_conn = -1 # not connected
 
-  def create(self, object_id, size):
+  def create(self, object_id, size, metadata=None):
     """Create a new buffer in the PlasmaStore for a particular object ID.
 
     The returned buffer is mutable until seal is called.
@@ -68,25 +68,49 @@ class PlasmaClient(object):
     Args:
       object_id (str): A string used to identify an object.
       size (int): The size in bytes of the created buffer.
+      metadata (buffer): An optional buffer encoding whatever metadata the user
+        wishes to encode.
     """
+    # This is used to hold the address of the buffer.
     data = ctypes.c_void_p()
-    self.client.plasma_create(self.sock, make_plasma_id(object_id), size, ctypes.byref(data))
+    # Turn the metadata into the right type.
+    metadata = buffer("") if metadata is None else metadata
+    metadata = (ctypes.c_ubyte * len(metadata)).from_buffer_copy(metadata)
+    self.client.plasma_create(self.sock, make_plasma_id(object_id), size, metadata, len(metadata), ctypes.byref(data))
     return self.buffer_from_read_write_memory(data, size)
 
   def get(self, object_id):
     """Create a buffer from the PlasmaStore based on object ID.
 
-    This method can only be called after the buffer has been sealed. The
-    retrieved buffer is immutable.
+    If the object has not been sealed yet, this call will block. The retrieved
+    buffer is immutable.
 
     Args:
       object_id (str): A string used to identify an object.
     """
     size = ctypes.c_int64()
     data = ctypes.c_void_p()
-    buf = self.client.plasma_get(self.sock, make_plasma_id(object_id), ctypes.byref(size), ctypes.byref(data))
+    metadata_size = ctypes.c_int64()
+    metadata = ctypes.c_void_p()
+    buf = self.client.plasma_get(self.sock, make_plasma_id(object_id), ctypes.byref(size), ctypes.byref(data), ctypes.byref(metadata_size), ctypes.byref(metadata))
     return self.buffer_from_memory(data, size)
 
+  def get_metadata(self, object_id):
+    """Create a buffer from the PlasmaStore based on object ID.
+
+    If the object has not been sealed yet, this call will block until the object
+    has been sealed. The retrieved buffer is immutable.
+
+    Args:
+      object_id (str): A string used to identify an object.
+    """
+    size = ctypes.c_int64()
+    data = ctypes.c_void_p()
+    metadata_size = ctypes.c_int64()
+    metadata = ctypes.c_void_p()
+    buf = self.client.plasma_get(self.sock, make_plasma_id(object_id), ctypes.byref(size), ctypes.byref(data), ctypes.byref(metadata_size), ctypes.byref(metadata))
+    return self.buffer_from_memory(metadata, metadata_size)
+
   def seal(self, object_id):
     """Seal the buffer in the PlasmaStore for a particular object ID.
 
@@ -100,7 +124,7 @@ class PlasmaClient(object):
 
   def transfer(self, addr, port, object_id):
     """Transfer local object with id object_id to another plasma instance
-    
+
     Args:
       addr (str): IPv4 address of the plasma instance the object is sent to.
       port (int): Port number of the plasma instance the object is sent to.
@@ -109,4 +133,3 @@ class PlasmaClient(object):
     if self.manager_conn == -1:
       raise Exception("Not connected to the plasma manager socket")
     self.client.plasma_transfer(self.manager_conn, addr, port, make_plasma_id(object_id))
-    

src/example.c

@@ -17,7 +17,7 @@
 int main(int argc, char *argv[]) {
   int conn = -1;
   int64_t size;
-  void *data;
+  uint8_t *data;
   int c;
   plasma_id id = {{255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
                    255, 255, 255, 255, 255, 255, 255, 255, 255, 255}};
@@ -28,14 +28,14 @@ int main(int argc, char *argv[]) {
       break;
     case 'c':
       assert(conn != -1);
-      plasma_create(conn, id, 100, &data);
+      plasma_create(conn, id, 100, NULL, 0, &data);
       break;
     case 'f':
       assert(conn != -1);
       plasma_seal(conn, id);
       break;
     case 'g':
-      plasma_get(conn, id, &size, &data);
+      plasma_get(conn, id, &size, &data, NULL, NULL);
       break;
     default:
       abort();

src/plasma.h

@@ -35,7 +35,8 @@
   } while (0)
 
 typedef struct {
-  int64_t size;
+  int64_t data_size;
+  int64_t metadata_size;
   int64_t create_time;
   int64_t construct_duration;
 } plasma_object_info;
@@ -59,21 +60,33 @@ enum plasma_request_type {
 typedef struct {
   int type;
   plasma_id object_id;
-  int64_t size;
+  /* The size of the data. */
+  int64_t data_size;
+  /* The size of the metadata. */
+  int64_t metadata_size;
   uint8_t addr[4];
   int port;
 } plasma_request;
 
 typedef struct {
-  ptrdiff_t offset;
+  /* The offset in the memory mapped file of the data. */
+  ptrdiff_t data_offset;
+  /* The offset in the memory mapped file of the metadata. */
+  ptrdiff_t metadata_offset;
+  /* The size of the memory mapped file. */
   int64_t map_size;
-  int64_t object_size;
+  /* The size of the data. */
+  int64_t data_size;
+  /* The size of the metadata. */
+  int64_t metadata_size;
 } plasma_reply;
 
 typedef struct {
   plasma_id object_id;
-  void *data;
-  int64_t size;
+  uint8_t *data;
+  int64_t data_size;
+  uint8_t *metadata;
+  int64_t metadata_size;
   int writable;
 } plasma_buffer;
 
@@ -83,8 +96,18 @@ int plasma_store_connect(const char *socket_name);
 /* Connect to a possibly remote plasma manager */
 int plasma_manager_connect(const char *addr, int port);
 
-void plasma_create(int store, plasma_id object_id, int64_t size, void **data);
-void plasma_get(int store, plasma_id object_id, int64_t *size, void **data);
+void plasma_create(int conn,
+                   plasma_id object_id,
+                   int64_t size,
+                   uint8_t *metadata,
+                   int64_t metadata_size,
+                   uint8_t **data);
+void plasma_get(int conn,
+                plasma_id object_id,
+                int64_t *size,
+                uint8_t **data,
+                int64_t *metadata_size,
+                uint8_t **metadata);
 void plasma_seal(int store, plasma_id object_id);
 
 void plasma_send(int conn, plasma_request *req);

src/plasma_client.c

@@ -23,38 +23,70 @@ void plasma_send(int fd, plasma_request *req) {
   }
 }
 
-void plasma_create(int conn, plasma_id object_id, int64_t size, void **data) {
-  LOG_INFO("called plasma_create on conn %d with size %" PRId64, conn, size);
-  plasma_request req = {
-      .type = PLASMA_CREATE, .object_id = object_id, .size = size};
+void plasma_create(int conn,
+                   plasma_id object_id,
+                   int64_t data_size,
+                   uint8_t *metadata,
+                   int64_t metadata_size,
+                   uint8_t **data) {
+  LOG_INFO(
+      "called plasma_create on conn %d with size %d and metadata size "
+      "%d" PRId64,
+      conn, size, metadata_size);
+  plasma_request req = {.type = PLASMA_CREATE,
+                        .object_id = object_id,
+                        .data_size = data_size,
+                        .metadata_size = metadata_size};
   plasma_send(conn, &req);
   plasma_reply reply;
   int fd = recv_fd(conn, (char *) &reply, sizeof(plasma_reply));
-  assert(reply.object_size == size);
-  *data =
-      mmap(NULL, reply.map_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0) +
-      reply.offset;
+  assert(reply.data_size == data_size);
+  assert(reply.metadata_size == metadata_size);
+  /* The metadata should come right after the data. */
+  assert(reply.metadata_offset == reply.data_offset + data_size);
+  *data = ((uint8_t *) mmap(NULL, reply.map_size, PROT_READ | PROT_WRITE,
+                            MAP_SHARED, fd, 0)) +
+          reply.data_offset;
   if (*data == MAP_FAILED) {
     LOG_ERR("mmap failed");
     exit(-1);
   }
+  /* If plasma_create is being called from a transfer, then we will not copy the
+   * metadata here. The metadata will be written along with the data streamed
+   * from the transfer. */
+  if (metadata != NULL) {
+    /* Copy the metadata to the buffer. */
+    memcpy(*data + reply.data_size, metadata, metadata_size);
+  }
   close(fd);
 }
 
-void plasma_get(int conn, plasma_id object_id, int64_t *size, void **data) {
+/* This method is used to get both the data and the metadata. */
+void plasma_get(int conn,
+                plasma_id object_id,
+                int64_t *size,
+                uint8_t **data,
+                int64_t *metadata_size,
+                uint8_t **metadata) {
   plasma_request req = {.type = PLASMA_GET, .object_id = object_id};
   plasma_send(conn, &req);
   plasma_reply reply;
   /* The following loop is run at most twice. */
   int fd = recv_fd(conn, (char *) &reply, sizeof(plasma_reply));
   *data =
-      mmap(NULL, reply.map_size, PROT_READ, MAP_SHARED, fd, 0) + reply.offset;
+      ((uint8_t *) mmap(NULL, reply.map_size, PROT_READ, MAP_SHARED, fd, 0)) +
+      reply.data_offset;
   if (*data == MAP_FAILED) {
     LOG_ERR("mmap failed");
     exit(-1);
   }
   close(fd);
-  *size = reply.object_size;
+  *size = reply.data_size;
+  /* If requested, return the metadata as well. */
+  if (metadata != NULL) {
+    *metadata = *data + reply.data_size;
+    *metadata_size = reply.metadata_size;
+  }
 }
 
 void plasma_seal(int fd, plasma_id object_id) {

src/plasma_manager.c

@@ -46,9 +46,19 @@ void init_plasma_manager(plasma_manager_state* s,
  * the data header to the other object manager. */
 void initiate_transfer(plasma_manager_state* s, plasma_request* req) {
   int store_conn = plasma_store_connect(s->store_socket_name);
-  plasma_buffer buf = {.object_id = req->object_id, .writable = 0};
-  plasma_get(store_conn, req->object_id, &buf.size, &buf.data);
-
+  uint8_t* data;
+  int64_t data_size;
+  uint8_t* metadata;
+  int64_t metadata_size;
+  plasma_get(store_conn, req->object_id, &data_size, &data, &metadata_size,
+             &metadata);
+  assert(metadata == data + data_size);
+  plasma_buffer buf = {.object_id = req->object_id,
+                       .data = data, /* We treat this as a pointer to the
+                                        concatenated data and metadata. */
+                       .data_size = data_size,
+                       .metadata_size = metadata_size,
+                       .writable = 0};
   char ip_addr[32];
   snprintf(ip_addr, 32, "%d.%d.%d.%d", req->addr[0], req->addr[1], req->addr[2],
            req->addr[3]);
@@ -59,9 +69,10 @@ void initiate_transfer(plasma_manager_state* s, plasma_request* req) {
                           .buf = buf,
                           .cursor = 0};
   event_loop_attach(s->loop, CONNECTION_DATA, &conn, fd, POLLOUT);
-
-  plasma_request manager_req = {
-      .type = PLASMA_DATA, .object_id = req->object_id, .size = buf.size};
+  plasma_request manager_req = {.type = PLASMA_DATA,
+                                .object_id = req->object_id,
+                                .data_size = buf.data_size,
+                                .metadata_size = buf.metadata_size};
   plasma_send(fd, &manager_req);
 }
 
@@ -72,9 +83,12 @@ void start_reading_data(int64_t index,
                         plasma_manager_state* s,
                         plasma_request* req) {
   int store_conn = plasma_store_connect(s->store_socket_name);
-  plasma_buffer buf = {
-      .object_id = req->object_id, .size = req->size, .writable = 1};
-  plasma_create(store_conn, req->object_id, req->size, &buf.data);
+  plasma_buffer buf = {.object_id = req->object_id,
+                       .data_size = req->data_size,
+                       .metadata_size = req->metadata_size,
+                       .writable = 1};
+  plasma_create(store_conn, req->object_id, req->data_size, NULL,
+                req->metadata_size, &buf.data);
   data_connection conn = {.type = DATA_CONNECTION_READ,
                           .store_conn = store_conn,
                           .buf = buf,
@@ -140,7 +154,7 @@ void read_from_socket(plasma_manager_state* state,
     break;
   case DATA_CONNECTION_WRITE:
     LOG_DEBUG("polled DATA_CONNECTION_WRITE");
-    s = conn->buf.size - conn->cursor;
+    s = conn->buf.data_size + conn->buf.metadata_size - conn->cursor;
     if (s > BUFSIZE)
       s = BUFSIZE;
     r = write(waiting->fd, conn->buf.data + conn->cursor, s);

src/plasma_store.c

@@ -85,7 +85,7 @@ void create_object(int conn, plasma_request* req) {
   HASH_FIND(handle, open_objects, &req->object_id, sizeof(plasma_id), entry);
   PLASMA_CHECK(entry == NULL, "Cannot create object twice.");
 
-  void* pointer = dlmalloc(req->size);
+  uint8_t* pointer = dlmalloc(req->data_size + req->metadata_size);
   int fd;
   int64_t map_size;
   ptrdiff_t offset;
@@ -94,7 +94,8 @@ void create_object(int conn, plasma_request* req) {
 
   entry = malloc(sizeof(object_table_entry));
   memcpy(&entry->object_id, &req->object_id, 20);
-  entry->info.size = req->size;
+  entry->info.data_size = req->data_size;
+  entry->info.metadata_size = req->metadata_size;
   /* TODO(pcm): set the other fields */
   entry->fd = fd;
   entry->map_size = map_size;
@@ -102,9 +103,11 @@ void create_object(int conn, plasma_request* req) {
   HASH_ADD(handle, open_objects, object_id, sizeof(plasma_id), entry);
   plasma_reply reply;
   memset(&reply, 0, sizeof(reply));
-  reply.offset = offset;
+  reply.data_offset = offset;
+  reply.metadata_offset = offset + req->data_size;
   reply.map_size = map_size;
-  reply.object_size = req->size;
+  reply.data_size = req->data_size;
+  reply.metadata_size = req->metadata_size;
   send_fd(conn, fd, (char*) &reply, sizeof(reply));
 }
 
@@ -115,9 +118,10 @@ void get_object(int conn, plasma_request* req) {
   if (entry) {
     plasma_reply reply;
     memset(&reply, 0, sizeof(plasma_reply));
-    reply.offset = entry->offset;
+    reply.data_offset = entry->offset;
     reply.map_size = entry->map_size;
-    reply.object_size = entry->info.size;
+    reply.data_size = entry->info.data_size;
+    reply.metadata_size = entry->info.metadata_size;
     send_fd(conn, entry->fd, (char*) &reply, sizeof(plasma_reply));
   } else {
     object_notify_entry* notify_entry;
@@ -156,9 +160,9 @@ void seal_object(int conn, plasma_request* req) {
   if (!notify_entry) {
     return;
   }
-  plasma_reply reply = {.offset = entry->offset,
+  plasma_reply reply = {.data_offset = entry->offset,
                         .map_size = entry->map_size,
-                        .object_size = entry->info.size};
+                        .data_size = entry->info.data_size};
   for (int i = 0; i < notify_entry->num_waiting; ++i) {
     send_fd(notify_entry->conn[i], entry->fd, (char*) &reply,
             sizeof(plasma_reply));

test/test.py

@@ -14,6 +14,31 @@ import plasma
 def random_object_id():
   return "".join([chr(random.randint(0, 255)) for _ in range(20)])
 
+def generate_metadata(length):
+  metadata = length * ["\x00"]
+  if length > 0:
+    metadata[0] = chr(random.randint(0, 255))
+    metadata[-1] = chr(random.randint(0, 255))
+    for _ in range(100):
+      metadata[random.randint(0, length - 1)] = chr(random.randint(0, 255))
+  return buffer("".join(metadata))
+
+def write_to_data_buffer(buff, length):
+  if length > 0:
+    buff[0] = chr(random.randint(0, 255))
+    buff[-1] = chr(random.randint(0, 255))
+    for _ in range(100):
+      buff[random.randint(0, length - 1)] = chr(random.randint(0, 255))
+
+def create_object(client, data_size, metadata_size, seal=True):
+  object_id = random_object_id()
+  metadata = generate_metadata(metadata_size)
+  memory_buffer = client.create(object_id, data_size, metadata)
+  write_to_data_buffer(memory_buffer, data_size)
+  if seal:
+    client.seal(object_id)
+  return object_id, memory_buffer, metadata
+
 class TestPlasmaClient(unittest.TestCase):
 
   def setUp(self):
@@ -32,7 +57,7 @@ class TestPlasmaClient(unittest.TestCase):
     # Create an object id string.
     object_id = random_object_id()
     # Create a new buffer and write to it.
-    length = 1000
+    length = 50
     memory_buffer = self.plasma_client.create(object_id, length)
     for i in range(length):
       memory_buffer[i] = chr(i % 256)
@@ -43,6 +68,28 @@ class TestPlasmaClient(unittest.TestCase):
     for i in range(length):
       self.assertEqual(memory_buffer[i], chr(i % 256))
 
+  def test_create_with_metadata(self):
+    for length in range(1000):
+      # Create an object id string.
+      object_id = random_object_id()
+      # Create a random metadata string.
+      metadata = generate_metadata(length)
+      # 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)
+      # Get the object.
+      memory_buffer = self.plasma_client.get(object_id)
+      for i in range(length):
+        self.assertEqual(memory_buffer[i], chr(i % 256))
+      # Get the metadata.
+      metadata_buffer = self.plasma_client.get_metadata(object_id)
+      self.assertEqual(len(metadata), len(metadata_buffer))
+      for i in range(len(metadata)):
+        self.assertEqual(metadata[i], metadata_buffer[i])
+
   def test_illegal_functionality(self):
     # Create an object id string.
     object_id = random_object_id()
@@ -95,34 +142,30 @@ class TestPlasmaManager(unittest.TestCase):
 
   def test_transfer(self):
     for _ in range(100):
-      # Create an object id string.
-      object_id1 = random_object_id()
-      # Create a new buffer and set the first and last entries.
-      memory_buffer = self.client1.create(object_id1, 20000)
-      memory_buffer[0] = chr(1)
-      memory_buffer[-1] = chr(2)
-      # Seal the buffer.
-      self.client1.seal(object_id1)
+      # Create an object.
+      object_id1, memory_buffer1, metadata1 = create_object(self.client1, 2000, 2000)
       # Transfer the buffer to the the other PlasmaStore.
       self.client1.transfer("127.0.0.1", self.port2, object_id1)
       # Compare the two buffers.
+      self.assertEqual(memory_buffer1[:], self.client2.get(object_id1)[:])
       self.assertEqual(self.client1.get(object_id1)[:], self.client2.get(object_id1)[:])
+      self.assertEqual(metadata1[:], self.client2.get_metadata(object_id1)[:])
+      self.assertEqual(self.client1.get_metadata(object_id1)[:], self.client2.get_metadata(object_id1)[:])
       # Transfer the buffer again.
       self.client1.transfer("127.0.0.1", self.port2, object_id1)
+      self.assertEqual(metadata1[:], self.client2.get_metadata(object_id1)[:])
       # Compare the two buffers.
       self.assertEqual(self.client1.get(object_id1)[:], self.client2.get(object_id1)[:])
-      # Create a new object id string.
-      object_id2 = random_object_id()
-      # Create a new buffer and set the first and last entries.
-      memory_buffer = self.client2.create(object_id2, 20000)
-      memory_buffer[0] = chr(3)
-      memory_buffer[-1] = chr(4)
-      # Seal the buffer.
-      self.client2.seal(object_id2)
+
+      # Create an object.
+      object_id2, memory_buffer2, metadata2 = create_object(self.client2, 20000, 20000)
       # Transfer the buffer to the the other PlasmaStore.
       self.client2.transfer("127.0.0.1", self.port1, object_id2)
       # Compare the two buffers.
+      self.assertEqual(memory_buffer2[:], self.client2.get(object_id2)[:])
       self.assertEqual(self.client1.get(object_id2)[:], self.client2.get(object_id2)[:])
+      self.assertEqual(metadata2[:], self.client2.get_metadata(object_id2)[:])
+      self.assertEqual(self.client1.get_metadata(object_id2)[:], self.client2.get_metadata(object_id2)[:])
 
   def test_illegal_functionality(self):
     # Create an object id string.