remove C struct duplication and python plasma manager

lib/python/plasma.py

@@ -9,77 +9,38 @@ ID = ctypes.c_ubyte * 20
 class PlasmaID(ctypes.Structure):
   _fields_ = [("plasma_id", ID)]
 
-# these must be in sync with plasma_request_type in plasma.h (can we have a test for that?)
-PLASMA_CREATE = 0
-PLASMA_GET = 1
-PLASMA_SEAL = 2
-PLASMA_TRANSFER = 3
-PLASMA_DATA = 4
-PLASMA_REGISTER = 5
-
-class PlasmaRequest(ctypes.Structure):
-  _fields_ = [("type", ctypes.c_int),
-              ("object_id", PlasmaID),
-              ("size", ctypes.c_int64),
-              ("addr", Addr),
-              ("port", ctypes.c_int)]
-
-class PlasmaBuffer(ctypes.Structure):
-  _fields_ = [("plasma_id", PlasmaID),
-              ("data", ctypes.c_void_p),
-              ("size", ctypes.c_int64),
-              ("writable", ctypes.c_int)]
-
 def make_plasma_id(string):
   if len(string) != 20:
     raise Exception("PlasmaIDs must be 20 characters long")
   object_id = map(ord, string)
   return PlasmaID(plasma_id=ID(*object_id))
 
-class PlasmaManager(object):
-  """The PlasmaManager is used to manage a PlasmaStore.
-
-  There should be one PlasmaManager per PlasmaStore. The PlasmaManager is
-  responsible for interfacing with other PlasmaManagers in order to transfer
-  objects between PlasmaStores. This class sends commands to the C
-  implementation of the PlasmaManager using sockets.
-
-  Attributes:
-    sock: The socket used to communicate with the C implementation of the
-      PlasmaManager.
-  """
-
-  def __init__(self, addr, port):
-    """Initialize the PlasmaManager."""
-    self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    self.sock.connect((addr, port))
-
-  def transfer(self, addr, port, object_id):
-    """Transfer local object with id object_id to manager with id manager_id."""
-    req = PlasmaRequest(type=PLASMA_TRANSFER, object_id=make_plasma_id(object_id),
-                        addr=Addr(*map(int, addr.split("."))), port=port)
-    print "sending port", port
-    self.sock.send(buffer(req)[:])
-
 class PlasmaClient(object):
-  """The PlasmaClient is used to interface with a PlasmaStore.
+  """The PlasmaClient is used to interface with a plasma store and a plasma manager.
 
   The PlasmaClient can ask the PlasmaStore to allocate a new buffer, seal a
   buffer, and get a buffer. Buffers are referred to by object IDs, which are
   strings.
   """
 
-  def __init__(self, socket_name):
+  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.
+      port (int): Port number of the plasma manager attached to the plasma store.
+    """
     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]
-    self.client.plasma_create.restype = PlasmaBuffer
+    self.client.plasma_create.argtypes = [ctypes.c_int, PlasmaID, ctypes.c_int64, ctypes.POINTER(ctypes.c_void_p)]
+    self.client.plasma_create.restype = None
 
-    self.client.plasma_get.argtypes = [ctypes.c_int, PlasmaID]
-    self.client.plasma_get.restype = PlasmaBuffer
+    self.client.plasma_get.argtypes = [ctypes.c_int, PlasmaID, 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
@@ -94,6 +55,11 @@ class PlasmaClient(object):
 
     self.sock = 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)
+    else:
+      self.manager_conn = -1 # not connected
+
   def create(self, object_id, size):
     """Create a new buffer in the PlasmaStore for a particular object ID.
 
@@ -103,8 +69,9 @@ class PlasmaClient(object):
       object_id (str): A string used to identify an object.
       size (int): The size in bytes of the created buffer.
     """
-    buf = self.client.plasma_create(self.sock, make_plasma_id(object_id), size)
-    return self.buffer_from_read_write_memory(buf.data, buf.size)
+    data = ctypes.c_void_p()
+    self.client.plasma_create(self.sock, make_plasma_id(object_id), size, 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.
@@ -115,8 +82,10 @@ class PlasmaClient(object):
     Args:
       object_id (str): A string used to identify an object.
     """
-    buf = self.client.plasma_get(self.sock, make_plasma_id(object_id))
-    return self.buffer_from_memory(buf.data, buf.size)
+    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))
+    return self.buffer_from_memory(data, size)
 
   def seal(self, object_id):
     """Seal the buffer in the PlasmaStore for a particular object ID.
@@ -128,3 +97,16 @@ class PlasmaClient(object):
       object_id (str): A string used to identify an object.
     """
     self.client.plasma_seal(self.sock, make_plasma_id(object_id))
+
+  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.
+      object_id (str): A string used to identify an 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))
+