make plasma robust wrt to number of open file descriptors (ray-projec…

…t#25)

lib/python/plasma.py

@@ -40,15 +40,9 @@ def __init__(self, socket_name, addr=None, port=None):
     plasma_client_library = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../../build/plasma_client.so")
     self.client = ctypes.cdll.LoadLibrary(plasma_client_library)
 
-    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_uint8), ctypes.c_int64, ctypes.POINTER(ctypes.c_void_p)]
+    self.client.plasma_store_connect.restype = 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), 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]
     self.client.plasma_seal.restype = None
 
     self.buffer_from_memory = ctypes.pythonapi.PyBuffer_FromMemory
@@ -59,7 +53,7 @@ def __init__(self, socket_name, addr=None, port=None):
     self.buffer_from_read_write_memory.argtypes = [ctypes.c_void_p, ctypes.c_int64]
     self.buffer_from_read_write_memory.restype = ctypes.py_object
 
-    self.sock = self.client.plasma_store_connect(socket_name)
+    self.store_conn = ctypes.c_void_p(self.client.plasma_store_connect(socket_name))
 
     if addr is not None and port is not None:
       self.manager_conn = self.client.plasma_manager_connect(addr, port)
@@ -82,7 +76,7 @@ def create(self, object_id, size, metadata=None):
     # 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))
+    self.client.plasma_create(self.store_conn, make_plasma_id(object_id), size, ctypes.cast(metadata, ctypes.POINTER(ctypes.c_ubyte * len(metadata))), len(metadata), ctypes.byref(data))
     return self.buffer_from_read_write_memory(data, size)
 
   def get(self, object_id):
@@ -98,7 +92,7 @@ def get(self, object_id):
     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))
+    buf = self.client.plasma_get(self.store_conn, 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):
@@ -114,7 +108,7 @@ def get_metadata(self, object_id):
     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))
+    buf = self.client.plasma_get(self.store_conn, 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):
@@ -126,7 +120,7 @@ def seal(self, object_id):
     Args:
       object_id (str): A string used to identify an object.
     """
-    self.client.plasma_seal(self.sock, make_plasma_id(object_id))
+    self.client.plasma_seal(self.store_conn, make_plasma_id(object_id))
 
   def transfer(self, addr, port, object_id):
     """Transfer local object with id object_id to another plasma instance

src/plasma.h

@@ -7,6 +7,8 @@
 #include <stddef.h>
 #include <string.h>
 
+#include "uthash.h"
+
 #ifdef NDEBUG
 #define LOG_DEBUG(M, ...)
 #else
@@ -79,6 +81,8 @@ typedef struct {
   int64_t data_size;
   /* The size of the metadata. */
   int64_t metadata_size;
+  /* Numerical value of the fd of the memory mapped file in the store. */
+  int store_fd_val;
 } plasma_reply;
 
 typedef struct {
@@ -90,25 +94,23 @@ typedef struct {
   int writable;
 } plasma_buffer;
 
-/* Connect to the local plasma store UNIX domain socket */
-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 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);
+typedef struct {
+  /* Key that uniquely identifies the  memory mapped file. In practice, we
+   * take the numerical value of the file descriptor in the object store. */
+  int key;
+  /* The result of mmap for this file descriptor. */
+  uint8_t *pointer;
+  /* Handle for the uthash table. */
+  UT_hash_handle hh;
+} client_mmap_table_entry;
+
+/* A client connection with a plasma store */
+typedef struct {
+  /* File descriptor of the Unix domain socket that connects to the store. */
+  int conn;
+  /* Table of dlmalloc buffer files that have been memory mapped so far. */
+  client_mmap_table_entry *mmap_table;
+} plasma_store_conn;
 
 void plasma_send(int conn, plasma_request *req);
 

src/plasma_client.c

@@ -13,17 +13,46 @@
 #include <netdb.h>
 
 #include "plasma.h"
+#include "plasma_client.h"
 #include "fling.h"
 
 void plasma_send(int fd, plasma_request *req) {
   int req_count = sizeof(plasma_request);
   if (write(fd, req, req_count) != req_count) {
-    LOG_ERR("write error");
+    LOG_ERR("write error, fd = %d", fd);
     exit(-1);
   }
 }
 
-void plasma_create(int conn,
+/* If the file descriptor fd has been mmapped in this client process before,
+ * return the pointer that was returned by mmap, otherwise mmap it and store the
+ * pointer in a hash table. */
+uint8_t *lookup_or_mmap(plasma_store_conn *conn,
+                        int fd,
+                        int store_fd_val,
+                        int64_t map_size) {
+  client_mmap_table_entry *entry;
+  HASH_FIND_INT(conn->mmap_table, &store_fd_val, entry);
+  if (entry) {
+    close(fd);
+    return entry->pointer;
+  } else {
+    uint8_t *result =
+        mmap(NULL, map_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+    if (result == MAP_FAILED) {
+      LOG_ERR("mmap failed");
+      exit(-1);
+    }
+    close(fd);
+    entry = malloc(sizeof(client_mmap_table_entry));
+    entry->key = store_fd_val;
+    entry->pointer = result;
+    HASH_ADD_INT(conn->mmap_table, key, entry);
+    return result;
+  }
+}
+
+void plasma_create(plasma_store_conn *conn,
                    plasma_id object_id,
                    int64_t data_size,
                    uint8_t *metadata,
@@ -37,54 +66,37 @@ void plasma_create(int conn,
                         .object_id = object_id,
                         .data_size = data_size,
                         .metadata_size = metadata_size};
-  plasma_send(conn, &req);
+  plasma_send(conn->conn, &req);
   plasma_reply reply;
-  int fd = recv_fd(conn, (char *) &reply, sizeof(plasma_reply));
+  int fd = recv_fd(conn->conn, (char *) &reply, sizeof(plasma_reply));
   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);
-
-  // TOOD(rshin): Don't call mmap if this fd has already been mapepd.
-  *data = ((uint8_t *) mmap(NULL, reply.map_size, PROT_READ | PROT_WRITE,
-                            MAP_SHARED, fd, 0)) +
+  *data = lookup_or_mmap(conn, fd, reply.store_fd_val, reply.map_size) +
           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);
 }
 
 /* This method is used to get both the data and the metadata. */
-void plasma_get(int conn,
+void plasma_get(plasma_store_conn *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_send(conn->conn, &req);
   plasma_reply reply;
-  /* The following loop is run at most twice. */
-  int fd = recv_fd(conn, (char *) &reply, sizeof(plasma_reply));
-
-  // TOOD(rshin): Don't call mmap if this fd has already been mapepd.
-  *data =
-      ((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);
+  int fd = recv_fd(conn->conn, (char *) &reply, sizeof(plasma_reply));
+  *data = lookup_or_mmap(conn, fd, reply.store_fd_val, reply.map_size) +
+          reply.data_offset;
   *size = reply.data_size;
   /* If requested, return the metadata as well. */
   if (metadata != NULL) {
@@ -93,12 +105,12 @@ void plasma_get(int conn,
   }
 }
 
-void plasma_seal(int fd, plasma_id object_id) {
+void plasma_seal(plasma_store_conn *conn, plasma_id object_id) {
   plasma_request req = {.type = PLASMA_SEAL, .object_id = object_id};
-  plasma_send(fd, &req);
+  plasma_send(conn->conn, &req);
 }
 
-int plasma_store_connect(const char *socket_name) {
+plasma_store_conn *plasma_store_connect(const char *socket_name) {
   assert(socket_name);
   struct sockaddr_un addr;
   int fd;
@@ -125,7 +137,11 @@ int plasma_store_connect(const char *socket_name) {
     LOG_ERR("could not connect to store %s", socket_name);
     exit(-1);
   }
-  return fd;
+  /* Initialize the store connection struct */
+  plasma_store_conn *result = malloc(sizeof(plasma_store_conn));
+  result->conn = fd;
+  result->mmap_table = NULL;
+  return result;
 }
 
 #define h_addr h_addr_list[0]

src/plasma_client.h

@@ -0,0 +1,27 @@
+#ifndef PLASMA_CLIENT_H
+#define PLASMA_CLIENT_H
+
+/* Connect to the local plasma store UNIX domain socket with path socket_name
+ * and return the resulting connection. */
+plasma_store_conn *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(plasma_store_conn *conn,
+                   plasma_id object_id,
+                   int64_t size,
+                   uint8_t *metadata,
+                   int64_t metadata_size,
+                   uint8_t **data);
+
+void plasma_get(plasma_store_conn *conn,
+                plasma_id object_id,
+                int64_t *size,
+                uint8_t **data,
+                int64_t *metadata_size,
+                uint8_t **metadata);
+
+void plasma_seal(plasma_store_conn *conn, plasma_id object_id);
+
+#endif