Single Big Object Parallel Transfer. (#1827)

* cache all object info from object added store notification.

* Adds parallel transfer for big objects.

* documentation and clean up.

* compare objects...

* merge buffer_state with chunk vec. Make separate buffer state for get and create.

* use references for Get. Allow partial failure of Create.

* single plasma client.

* changes based on review.

* update documentation and add parameters for object manager in main.cc.

* review feedback.

* use vector consturctor.

* linting

* remove profile visualizations.

* test fixes.

* linting.

* kill specific pids and use less memory.

* linting.

* simplify tests.

* Asynchronous IO for ObjectManager messages and object transfer.

* Revert "Asynchronous IO for ObjectManager messages and object transfer."

This reverts commit 4af43b159babc04daf80d1543e27c2cb46b7b19d.

* update test configuration to reflect changes in #1891

* review feedback.

* linting.

src/ray/CMakeLists.txt

@@ -39,7 +39,7 @@ set(RAY_SRCS
   common/client_connection.cc
   object_manager/object_manager_client_connection.cc
   object_manager/connection_pool.cc
-  object_manager/object_store_client_pool.cc
+  object_manager/object_buffer_pool.cc
   object_manager/object_store_notification_manager.cc
   object_manager/object_directory.cc
   object_manager/transfer_queue.cc

src/ray/object_manager/format/object_manager.fbs

@@ -10,8 +10,12 @@ enum MessageType:int {
 table PushRequestMessage {
   // The object ID being transferred.
   object_id: string;
-  // The size of the object being transferred.
-  object_size: ulong;
+  // The index of the chunk being transferred.
+  chunk_index: ulong;
+  // The total size of the object + metadata.
+  data_size: ulong;
+  // The metadata size.
+  metadata_size: ulong;
 }
 
 table PullRequestMessage {

src/ray/object_manager/object_buffer_pool.cc

@@ -0,0 +1,187 @@
+#include "ray/object_manager/object_buffer_pool.h"
+
+namespace ray {
+
+ObjectBufferPool::ObjectBufferPool(const std::string &store_socket_name,
+                                   uint64_t chunk_size, int release_delay)
+    : chunk_size_(chunk_size) {
+  store_socket_name_ = store_socket_name;
+  ARROW_CHECK_OK(store_client_.Connect(store_socket_name_.c_str(), "", release_delay));
+}
+
+ObjectBufferPool::~ObjectBufferPool() {
+  // Abort everything in progress.
+  auto get_buf_state_copy = get_buffer_state_;
+  for (const auto &pair : get_buf_state_copy) {
+    AbortGet(pair.first);
+  }
+  auto create_buf_state_copy = create_buffer_state_;
+  for (const auto &pair : create_buf_state_copy) {
+    AbortCreate(pair.first);
+  }
+  RAY_CHECK(get_buffer_state_.empty());
+  RAY_CHECK(create_buffer_state_.empty());
+  ARROW_CHECK_OK(store_client_.Disconnect());
+}
+
+uint64_t ObjectBufferPool::GetNumChunks(uint64_t data_size) {
+  return (data_size + chunk_size_ - 1) / chunk_size_;
+}
+
+uint64_t ObjectBufferPool::GetBufferLength(uint64_t chunk_index, uint64_t data_size) {
+  return (chunk_index + 1) * chunk_size_ > data_size ? data_size % chunk_size_
+                                                     : chunk_size_;
+}
+
+std::pair<const ObjectBufferPool::ChunkInfo &, ray::Status> ObjectBufferPool::GetChunk(
+    const ObjectID &object_id, uint64_t data_size, uint64_t metadata_size,
+    uint64_t chunk_index) {
+  std::lock_guard<std::mutex> lock(pool_mutex_);
+  RAY_LOG(DEBUG) << "GetChunk " << object_id << " " << data_size << " " << metadata_size;
+  if (get_buffer_state_.count(object_id) == 0) {
+    plasma::ObjectBuffer object_buffer;
+    plasma::ObjectID plasma_id = ObjectID(object_id).to_plasma_id();
+    ARROW_CHECK_OK(store_client_.Get(&plasma_id, 1, 0, &object_buffer));
+    if (object_buffer.data == nullptr) {
+      RAY_LOG(ERROR) << "Failed to get object";
+      return std::pair<const ObjectBufferPool::ChunkInfo &, ray::Status>(
+          errored_chunk_,
+          ray::Status::IOError("Unable to obtain object chunk, object not local."));
+    }
+    RAY_CHECK(object_buffer.metadata->data() ==
+              object_buffer.data->data() + object_buffer.data->size());
+    RAY_CHECK(data_size == static_cast<uint64_t>(object_buffer.data->size() +
+                                                 object_buffer.metadata_size));
+    auto *data = const_cast<uint8_t *>(object_buffer.data->data());
+    uint64_t num_chunks = GetNumChunks(data_size);
+    get_buffer_state_.emplace(
+        std::piecewise_construct, std::forward_as_tuple(object_id),
+        std::forward_as_tuple(BuildChunks(object_id, data, data_size)));
+    RAY_CHECK(get_buffer_state_[object_id].chunk_info.size() == num_chunks);
+  }
+  get_buffer_state_[object_id].references++;
+  return std::pair<const ObjectBufferPool::ChunkInfo &, ray::Status>(
+      get_buffer_state_[object_id].chunk_info[chunk_index], ray::Status::OK());
+}
+
+void ObjectBufferPool::ReleaseGetChunk(const ObjectID &object_id, uint64_t chunk_index) {
+  std::lock_guard<std::mutex> lock(pool_mutex_);
+  GetBufferState &buffer_state = get_buffer_state_[object_id];
+  buffer_state.references--;
+  RAY_LOG(DEBUG) << "ReleaseBuffer " << object_id << " " << buffer_state.references;
+  if (buffer_state.references == 0) {
+    ARROW_CHECK_OK(store_client_.Release(ObjectID(object_id).to_plasma_id()));
+    get_buffer_state_.erase(object_id);
+  }
+}
+
+void ObjectBufferPool::AbortGet(const ObjectID &object_id) {
+  std::lock_guard<std::mutex> lock(pool_mutex_);
+  ARROW_CHECK_OK(store_client_.Release(ObjectID(object_id).to_plasma_id()));
+  get_buffer_state_.erase(object_id);
+}
+
+std::pair<const ObjectBufferPool::ChunkInfo &, ray::Status> ObjectBufferPool::CreateChunk(
+    const ObjectID &object_id, uint64_t data_size, uint64_t metadata_size,
+    uint64_t chunk_index) {
+  std::lock_guard<std::mutex> lock(pool_mutex_);
+  RAY_LOG(DEBUG) << "CreateChunk " << object_id << " " << data_size << " "
+                 << metadata_size;
+  if (create_buffer_state_.count(object_id) == 0) {
+    const plasma::ObjectID plasma_id = ObjectID(object_id).to_plasma_id();
+    int64_t object_size = data_size - metadata_size;
+    // Try to create shared buffer.
+    std::shared_ptr<Buffer> data;
+    arrow::Status s =
+        store_client_.Create(plasma_id, object_size, NULL, metadata_size, &data);
+    std::vector<boost::asio::mutable_buffer> buffer;
+    if (!s.ok()) {
+      // Create failed. The object may already exist locally. If something else went
+      // wrong, another chunk will succeed in creating the buffer, and this
+      // chunk will eventually make it here via pull requests.
+      return std::pair<const ObjectBufferPool::ChunkInfo &, ray::Status>(
+          errored_chunk_, ray::Status::IOError(s.message()));
+    }
+    // Read object into store.
+    uint8_t *mutable_data = data->mutable_data();
+    uint64_t num_chunks = GetNumChunks(data_size);
+    create_buffer_state_.emplace(
+        std::piecewise_construct, std::forward_as_tuple(object_id),
+        std::forward_as_tuple(BuildChunks(object_id, mutable_data, data_size)));
+    RAY_CHECK(create_buffer_state_[object_id].chunk_info.size() == num_chunks);
+  }
+  if (create_buffer_state_[object_id].chunk_state[chunk_index] !=
+      CreateChunkState::AVAILABLE) {
+    // There can be only one reference to this chunk at any given time.
+    return std::pair<const ObjectBufferPool::ChunkInfo &, ray::Status>(
+        errored_chunk_,
+        ray::Status::IOError("Chunk already referenced by another thread."));
+  }
+  create_buffer_state_[object_id].chunk_state[chunk_index] = CreateChunkState::REFERENCED;
+  return std::pair<const ObjectBufferPool::ChunkInfo &, ray::Status>(
+      create_buffer_state_[object_id].chunk_info[chunk_index], ray::Status::OK());
+}
+
+void ObjectBufferPool::AbortCreateChunk(const ObjectID &object_id,
+                                        const uint64_t chunk_index) {
+  std::lock_guard<std::mutex> lock(pool_mutex_);
+  RAY_CHECK(create_buffer_state_[object_id].chunk_state[chunk_index] ==
+            CreateChunkState::REFERENCED);
+  create_buffer_state_[object_id].chunk_state[chunk_index] = CreateChunkState::AVAILABLE;
+  if (create_buffer_state_[object_id].num_seals_remaining ==
+      create_buffer_state_[object_id].chunk_state.size()) {
+    // If chunk_state is AVAILABLE at every chunk_index and
+    // num_seals_remaining == num_chunks, this is back to the initial state
+    // right before the first CreateChunk.
+    bool abort = true;
+    for (auto chunk_state : create_buffer_state_[object_id].chunk_state) {
+      abort &= chunk_state == CreateChunkState::AVAILABLE;
+    }
+    if (abort) {
+      AbortCreate(object_id);
+    }
+  }
+}
+
+void ObjectBufferPool::SealChunk(const ObjectID &object_id, const uint64_t chunk_index) {
+  std::lock_guard<std::mutex> lock(pool_mutex_);
+  RAY_CHECK(create_buffer_state_[object_id].chunk_state[chunk_index] ==
+            CreateChunkState::REFERENCED);
+  create_buffer_state_[object_id].chunk_state[chunk_index] = CreateChunkState::SEALED;
+  create_buffer_state_[object_id].num_seals_remaining--;
+  RAY_LOG(DEBUG) << "SealChunk" << object_id << " "
+                 << create_buffer_state_[object_id].num_seals_remaining;
+  if (create_buffer_state_[object_id].num_seals_remaining == 0) {
+    const plasma::ObjectID plasma_id = ObjectID(object_id).to_plasma_id();
+    ARROW_CHECK_OK(store_client_.Seal(plasma_id));
+    ARROW_CHECK_OK(store_client_.Release(plasma_id));
+    create_buffer_state_.erase(object_id);
+  }
+}
+
+void ObjectBufferPool::AbortCreate(const ObjectID &object_id) {
+  const plasma::ObjectID plasma_id = ObjectID(object_id).to_plasma_id();
+  ARROW_CHECK_OK(store_client_.Release(plasma_id));
+  ARROW_CHECK_OK(store_client_.Abort(plasma_id));
+  create_buffer_state_.erase(object_id);
+}
+
+std::vector<ObjectBufferPool::ChunkInfo> ObjectBufferPool::BuildChunks(
+    const ObjectID &object_id, uint8_t *data, uint64_t data_size) {
+  uint64_t space_remaining = data_size;
+  std::vector<ChunkInfo> chunks;
+  int64_t position = 0;
+  while (space_remaining) {
+    position = data_size - space_remaining;
+    if (space_remaining < chunk_size_) {
+      chunks.emplace_back(chunks.size(), data + position, space_remaining);
+      space_remaining = 0;
+    } else {
+      chunks.emplace_back(chunks.size(), data + position, chunk_size_);
+      space_remaining -= chunk_size_;
+    }
+  }
+  return chunks;
+}
+
+}  // namespace ray

src/ray/object_manager/object_buffer_pool.h

@@ -0,0 +1,196 @@
+#ifndef RAY_OBJECT_MANAGER_OBJECT_BUFFER_POOL_H
+#define RAY_OBJECT_MANAGER_OBJECT_BUFFER_POOL_H
+
+#include <list>
+#include <memory>
+#include <mutex>
+#include <vector>
+
+#include <boost/asio.hpp>
+#include <boost/asio/error.hpp>
+#include <boost/bind.hpp>
+
+#include "plasma/client.h"
+#include "plasma/events.h"
+#include "plasma/plasma.h"
+
+#include "ray/id.h"
+#include "ray/status.h"
+
+namespace ray {
+
+/// \class ObjectBufferPool Exposes chunks of object buffers for use by the ObjectManager.
+class ObjectBufferPool {
+ public:
+  /// Information needed to read or write an object chunk.
+  /// This is the structure returned whenever an object chunk is
+  /// accessed via Get and Create.
+  struct ChunkInfo {
+    ChunkInfo(uint64_t chunk_index, uint8_t *data, uint64_t buffer_length)
+        : chunk_index(chunk_index), data(data), buffer_length(buffer_length){};
+    /// A pointer to the start position of this object chunk.
+    uint64_t chunk_index;
+    /// A pointer to the start position of this object chunk.
+    uint8_t *data;
+    /// The size of this object chunk.
+    uint64_t buffer_length;
+  };
+
+  /// Constructor.
+  ///
+  /// \param store_socket_name The socket name of the store to which plasma clients
+  /// connect.
+  /// \param chunk_size The chunk size into which objects are to be split.
+  /// \param release_delay The number of release calls before objects are released
+  /// from the store client (FIFO).
+  ObjectBufferPool(const std::string &store_socket_name, const uint64_t chunk_size,
+                   const int release_delay);
+
+  ~ObjectBufferPool();
+
+  /// This object cannot be copied due to pool_mutex.
+  RAY_DISALLOW_COPY_AND_ASSIGN(ObjectBufferPool);
+
+  /// Computes the number of chunks needed to transfer an object and its metadata.
+  ///
+  /// \param data_size The size of the object + metadata.
+  /// \return The number of chunks into which the object will be split.
+  uint64_t GetNumChunks(uint64_t data_size);
+
+  /// Computes the buffer length of a chunk of an object.
+  ///
+  /// \param chunk_index The chunk index for which to obtain the buffer length.
+  /// \param data_size The size of the object + metadata.
+  /// \return The buffer length of the chunk at chunk_index.
+  uint64_t GetBufferLength(uint64_t chunk_index, uint64_t data_size);
+
+  /// Returns a chunk of an object at the given chunk_index. The object chunk serves
+  /// as the data that is to be written to a connection as part of sending an object to
+  /// a remote node.
+  ///
+  /// \param object_id The ObjectID.
+  /// \param data_size The sum of the object size and metadata size.
+  /// \param metadata_size The size of the metadata.
+  /// \param chunk_index The index of the chunk.
+  /// \return A pair consisting of a ChunkInfo and status of invoking this method.
+  /// An IOError status is returned if the Get call on the plasma store fails.
+  std::pair<const ObjectBufferPool::ChunkInfo &, ray::Status> GetChunk(
+      const ObjectID &object_id, uint64_t data_size, uint64_t metadata_size,
+      uint64_t chunk_index);
+
+  /// When a chunk is done being used as part of a get, this method releases the chunk.
+  /// If all chunks of an object are released, the object buffer will be released.
+  ///
+  /// \param object_id The object_id of the buffer to release.
+  /// \param chunk_index The index of the chunk.
+  void ReleaseGetChunk(const ObjectID &object_id, uint64_t chunk_index);
+
+  /// Returns a chunk of an empty object at the given chunk_index. The object chunk
+  /// serves as the buffer that is to be written to by a connection receiving an object
+  /// from a remote node. Only one thread is permitted to create the object chunk at
+  /// chunk_index. Multiple threads attempting to create the same object chunk will
+  /// result in one succeeding. The ObjectManager is responsible for handling
+  /// create failures. This method will fail if it's invoked on a chunk_index on which
+  /// SealChunk has already been invoked.
+  ///
+  /// \param object_id The ObjectID.
+  /// \param data_size The sum of the object size and metadata size.
+  /// \param metadata_size The size of the metadata.
+  /// \param chunk_index The index of the chunk.
+  /// \return A pair consisting of ChunkInfo and status of invoking this method.
+  /// An IOError status is returned if object creation on the store client fails,
+  /// or if create is invoked consecutively on the same chunk
+  /// (with no intermediate AbortCreateChunk).
+  std::pair<const ObjectBufferPool::ChunkInfo &, ray::Status> CreateChunk(
+      const ObjectID &object_id, uint64_t data_size, uint64_t metadata_size,
+      uint64_t chunk_index);
+
+  /// Abort the create operation associated with a chunk at chunk_index.
+  /// This method will fail if it's invoked on a chunk_index on which
+  /// CreateChunk was not first invoked, or a chunk_index on which
+  /// SealChunk has already been invoked.
+  ///
+  /// \param object_id The ObjectID.
+  /// \param chunk_index The index of the chunk.
+  void AbortCreateChunk(const ObjectID &object_id, uint64_t chunk_index);
+
+  /// Seal the object associated with a create operation. This is invoked whenever
+  /// a chunk is successfully written to.
+  /// This method will fail if it's invoked on a chunk_index on which
+  /// CreateChunk was not first invoked, or a chunk_index on which
+  /// SealChunk or AbortCreateChunk has already been invoked.
+  ///
+  /// \param object_id The ObjectID.
+  /// \param chunk_index The index of the chunk.
+  void SealChunk(const ObjectID &object_id, uint64_t chunk_index);
+
+ private:
+  /// Abort the create operation associated with an object. This destroys the buffer
+  /// state, including create operations in progress for all chunks of the object.
+  void AbortCreate(const ObjectID &object_id);
+
+  /// Abort the get operation associated with an object.
+  void AbortGet(const ObjectID &object_id);
+
+  /// Splits an object into ceil(data_size/chunk_size) chunks, which will
+  /// either be read or written to in parallel.
+  std::vector<ChunkInfo> BuildChunks(const ObjectID &object_id, uint8_t *data,
+                                     uint64_t data_size);
+
+  /// Holds the state of a get buffer.
+  struct GetBufferState {
+    GetBufferState() {}
+    GetBufferState(std::vector<ChunkInfo> chunk_info) : chunk_info(chunk_info) {}
+    /// A vector maintaining information about the chunks which comprise
+    /// an object.
+    std::vector<ChunkInfo> chunk_info;
+    /// The number of references that currently rely on this buffer.
+    /// Once this reaches 0, the buffer is released and this object is erased
+    /// from get_buffer_state_.
+    uint64_t references = 0;
+  };
+
+  /// The state of a chunk associated with a create operation.
+  enum class CreateChunkState : uint { AVAILABLE = 0, REFERENCED, SEALED };
+
+  /// Holds the state of a create buffer.
+  struct CreateBufferState {
+    CreateBufferState() {}
+    CreateBufferState(std::vector<ChunkInfo> chunk_info)
+        : chunk_info(chunk_info),
+          chunk_state(chunk_info.size(), CreateChunkState::AVAILABLE),
+          num_seals_remaining(chunk_info.size()) {}
+    /// A vector maintaining information about the chunks which comprise
+    /// an object.
+    std::vector<ChunkInfo> chunk_info;
+    /// The state of each chunk, which is used to enforce strict state
+    /// transitions of each chunk.
+    std::vector<CreateChunkState> chunk_state;
+    /// The number of chunks left to seal before the buffer is sealed.
+    uint64_t num_seals_remaining;
+  };
+
+  /// Returned when GetChunk or CreateChunk fails.
+  const ChunkInfo errored_chunk_ = {0, nullptr, 0};
+
+  /// Mutex on public methods for thread-safe operations on
+  /// get_buffer_state_, create_buffer_state_, and store_client_.
+  std::mutex pool_mutex_;
+  /// Determines the maximum chunk size to be transferred by a single thread.
+  const uint64_t chunk_size_;
+  /// The state of a buffer that's currently being used.
+  std::unordered_map<ray::ObjectID, GetBufferState, ray::UniqueIDHasher>
+      get_buffer_state_;
+  /// The state of a buffer that's currently being used.
+  std::unordered_map<ray::ObjectID, CreateBufferState, ray::UniqueIDHasher>
+      create_buffer_state_;
+
+  /// Plasma client pool.
+  plasma::PlasmaClient store_client_;
+  /// Socket name of plasma store.
+  std::string store_socket_name_;
+};
+
+}  // namespace ray
+
+#endif  // RAY_OBJECT_MANAGER_OBJECT_BUFFER_POOL_H

src/ray/object_manager/object_directory.h

@@ -13,6 +13,7 @@
 
 namespace ray {
 
+/// Connection information for remote object managers.
 struct RemoteConnectionInfo {
   RemoteConnectionInfo() = default;
   RemoteConnectionInfo(const ClientID &id, const std::string &ip_address,
@@ -23,7 +24,6 @@ struct RemoteConnectionInfo {
   uint16_t port;
 };
 
-/// Connection information for remote object managers.
 class ObjectDirectoryInterface {
  public:
   ObjectDirectoryInterface() = default;

src/ray/object_manager/object_manager.cc

@@ -14,7 +14,10 @@ ObjectManager::ObjectManager(asio::io_service &main_service,
     : client_id_(gcs_client->client_table().GetLocalClientId()),
       object_directory_(new ObjectDirectory(gcs_client)),
       store_notification_(main_service, config.store_socket_name),
-      store_pool_(config.store_socket_name),
+      // release_delay of 2 * config.max_sends is to ensure the pool does not release
+      // an object prematurely whenever we reach the maximum number of sends.
+      buffer_pool_(config.store_socket_name, config.object_chunk_size,
+                   /*release_delay=*/2 * config.max_sends),
       object_manager_service_(std::move(object_manager_service)),
       work_(*object_manager_service_),
       connection_pool_(),
@@ -39,7 +42,10 @@ ObjectManager::ObjectManager(asio::io_service &main_service,
                              std::unique_ptr<ObjectDirectoryInterface> od)
     : object_directory_(std::move(od)),
       store_notification_(main_service, config.store_socket_name),
-      store_pool_(config.store_socket_name),
+      // release_delay of 2 * config.max_sends is to ensure the pool does not release
+      // an object prematurely whenever we reach the maximum number of sends.
+      buffer_pool_(config.store_socket_name, config.object_chunk_size,
+                   /*release_delay=*/2 * config.max_sends),
       object_manager_service_(std::move(object_manager_service)),
       work_(*object_manager_service_),
       connection_pool_(),
@@ -59,12 +65,7 @@ ObjectManager::ObjectManager(asio::io_service &main_service,
   StartIOService();
 }
 
-ObjectManager::~ObjectManager() {
-  object_manager_service_->stop();
-  for (int i = 0; i < num_threads_; ++i) {
-    io_threads_[i].join();
-  }
-}
+ObjectManager::~ObjectManager() { StopIOService(); }
 
 void ObjectManager::StartIOService() {
   for (int i = 0; i < num_threads_; ++i) {
@@ -74,6 +75,13 @@ void ObjectManager::StartIOService() {
 
 void ObjectManager::IOServiceLoop() { object_manager_service_->run(); }
 
+void ObjectManager::StopIOService() {
+  object_manager_service_->stop();
+  for (int i = 0; i < num_threads_; ++i) {
+    io_threads_[i].join();
+  }
+}
+
 void ObjectManager::NotifyDirectoryObjectAdd(const ObjectInfoT &object_info) {
   ObjectID object_id = ObjectID::from_binary(object_info.object_id);
   local_objects_[object_id] = object_info;
@@ -197,13 +205,28 @@ ray::Status ObjectManager::PullSendRequest(const ObjectID &object_id,
 }
 
 ray::Status ObjectManager::Push(const ObjectID &object_id, const ClientID &client_id) {
-  // TODO(hme): Cache this data in ObjectDirectory.
-  // Okay for now since the GCS client caches this data.
+  if (local_objects_.count(object_id) == 0) {
+    // TODO(hme): Do not retry indefinitely...
+    main_service_->post(
+        [this, object_id, client_id]() { RAY_CHECK_OK(Push(object_id, client_id)); });
+    return ray::Status::OK();
+  }
+
   main_service_->dispatch([this, object_id, client_id]() {
+    // TODO(hme): Cache this data in ObjectDirectory.
+    // Okay for now since the GCS client caches this data.
     Status status = object_directory_->GetInformation(
         client_id,
         [this, object_id, client_id](const RemoteConnectionInfo &info) {
-          transfer_queue_.QueueSend(client_id, object_id, info);
+          ObjectInfoT object_info = local_objects_[object_id];
+          uint64_t data_size =
+              static_cast<uint64_t>(object_info.data_size + object_info.metadata_size);
+          uint64_t metadata_size = static_cast<uint64_t>(object_info.metadata_size);
+          uint64_t num_chunks = buffer_pool_.GetNumChunks(data_size);
+          for (uint64_t chunk_index = 0; chunk_index < num_chunks; ++chunk_index) {
+            transfer_queue_.QueueSend(client_id, object_id, data_size, metadata_size,
+                                      chunk_index, info);
+          }
           RAY_CHECK_OK(DequeueTransfers());
         },
         [](const Status &status) {
@@ -226,8 +249,9 @@ ray::Status ObjectManager::DequeueTransfers() {
         object_manager_service_->dispatch([this, req]() {
           RAY_LOG(DEBUG) << "DequeueSend " << client_id_ << " " << req.object_id << " "
                          << num_transfers_send_ << "/" << config_.max_sends;
-          RAY_CHECK_OK(
-              ExecuteSendObject(req.object_id, req.client_id, req.connection_info));
+          RAY_CHECK_OK(ExecuteSendObject(req.client_id, req.object_id, req.data_size,
+                                         req.metadata_size, req.chunk_index,
+                                         req.connection_info));
         });
       } else {
         std::atomic_fetch_sub(&num_transfers_send_, 1);
@@ -248,7 +272,8 @@ ray::Status ObjectManager::DequeueTransfers() {
         object_manager_service_->dispatch([this, req]() {
           RAY_LOG(DEBUG) << "DequeueReceive " << client_id_ << " " << req.object_id << " "
                          << num_transfers_receive_ << "/" << config_.max_receives;
-          RAY_CHECK_OK(ExecuteReceiveObject(req.client_id, req.object_id, req.object_size,
+          RAY_CHECK_OK(ExecuteReceiveObject(req.client_id, req.object_id, req.data_size,
+                                            req.metadata_size, req.chunk_index,
                                             req.conn));
         });
       } else {
@@ -273,100 +298,75 @@ ray::Status ObjectManager::TransferCompleted(TransferQueue::TransferType type) {
 };
 
 ray::Status ObjectManager::ExecuteSendObject(
-    const ObjectID &object_id, const ClientID &client_id,
+    const ClientID &client_id, const ObjectID &object_id, uint64_t data_size,
+    uint64_t metadata_size, uint64_t chunk_index,
     const RemoteConnectionInfo &connection_info) {
+  RAY_LOG(DEBUG) << "ExecuteSendObject " << client_id << " " << object_id << " "
+                 << chunk_index;
   ray::Status status;
   std::shared_ptr<SenderConnection> conn;
   status = connection_pool_.GetSender(ConnectionPool::ConnectionType::TRANSFER, client_id,
                                       &conn);
-  if (!status.ok()) {
-    // TODO(hme): Keep track of retries,
-    // and only retry on object not local
-    // for now.
-    RAY_CHECK_OK(Push(object_id, conn->GetClientID()));
-    return Status::OK();
-  }
   if (conn == nullptr) {
     conn =
         CreateSenderConnection(ConnectionPool::ConnectionType::TRANSFER, connection_info);
     connection_pool_.RegisterSender(ConnectionPool::ConnectionType::TRANSFER, client_id,
                                     conn);
   }
-  status = SendObjectHeaders(object_id, conn);
-  if (!status.ok()) {
-    RAY_CHECK_OK(Push(object_id, conn->GetClientID()));
-    return Status::OK();
-  }
+  status = SendObjectHeaders(object_id, data_size, metadata_size, chunk_index, conn);
   return Status::OK();
 }
 
-ray::Status ObjectManager::SendObjectHeaders(const ObjectID &object_id_const,
+ray::Status ObjectManager::SendObjectHeaders(const ObjectID &object_id,
+                                             uint64_t data_size, uint64_t metadata_size,
+                                             uint64_t chunk_index,
                                              std::shared_ptr<SenderConnection> conn) {
-  ObjectID object_id = ObjectID(object_id_const);
-  // Allocate and append the request to the transfer queue.
-  plasma::ObjectBuffer object_buffer;
-  plasma::ObjectID plasma_id = object_id.to_plasma_id();
-  std::shared_ptr<plasma::PlasmaClient> store_client = store_pool_.GetObjectStore();
-  ARROW_CHECK_OK(store_client->Get(&plasma_id, 1, 0, &object_buffer));
-  if (object_buffer.data == nullptr) {
-    RAY_LOG(ERROR) << "Failed to get object";
-    // If the object wasn't locally available, exit immediately. If the object
-    // later appears locally, the requesting plasma manager should request the
-    // transfer again.
-    RAY_CHECK_OK(
-        connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::TRANSFER, conn));
-    store_pool_.ReleaseObjectStore(store_client);
+  std::pair<const ObjectBufferPool::ChunkInfo &, ray::Status> chunk_status =
+      buffer_pool_.GetChunk(object_id, data_size, metadata_size, chunk_index);
+  ObjectBufferPool::ChunkInfo chunk_info = chunk_status.first;
+
+  if (!chunk_status.second.ok()) {
+    // This is the first thread to invoke GetChunk => Get failed on the
+    // plasma client.
+    // No reference is acquired for this chunk, so no need to release the chunk.
+    // TODO(hme): Retry send here? If so, store RemoteConnectionInfo in SenderConnection.
     RAY_CHECK_OK(TransferCompleted(TransferQueue::TransferType::SEND));
-    return ray::Status::IOError(
-        "Unable to transfer object to requesting plasma manager, object not local.");
+    return chunk_status.second;
   }
-  RAY_CHECK(object_buffer.metadata->data() ==
-            object_buffer.data->data() + object_buffer.data->size());
-
-  TransferQueue::SendContext context;
-  context.client_id = conn->GetClientID();
-  context.object_id = object_id;
-  context.object_size = static_cast<uint64_t>(object_buffer.data->size());
-  context.data = const_cast<uint8_t *>(object_buffer.data->data());
-  UniqueID context_id = transfer_queue_.AddContext(context);
-
   // Create buffer.
   flatbuffers::FlatBufferBuilder fbb;
   // TODO(hme): use to_flatbuf
   auto message = object_manager_protocol::CreatePushRequestMessage(
-      fbb, fbb.CreateString(object_id.binary()), context.object_size);
+      fbb, fbb.CreateString(object_id.binary()), chunk_index, data_size, metadata_size);
   fbb.Finish(message);
   ray::Status status =
       conn->WriteMessage(object_manager_protocol::MessageType_PushRequest, fbb.GetSize(),
                          fbb.GetBufferPointer());
   RAY_CHECK_OK(status);
-  // TODO(hme): Make this async.
-  return SendObjectData(conn, context_id, store_client);
+  return SendObjectData(object_id, chunk_info, conn);
 }
 
-ray::Status ObjectManager::SendObjectData(
-    std::shared_ptr<SenderConnection> conn, const UniqueID &context_id,
-    std::shared_ptr<plasma::PlasmaClient> store_client) {
-  TransferQueue::SendContext context = transfer_queue_.GetContext(context_id);
+ray::Status ObjectManager::SendObjectData(const ObjectID &object_id,
+                                          const ObjectBufferPool::ChunkInfo &chunk_info,
+                                          std::shared_ptr<SenderConnection> conn) {
+  // TransferQueue::SendContext context = transfer_queue_.GetContext(context_id);
   boost::system::error_code ec;
   std::vector<asio::const_buffer> buffer;
-  buffer.push_back(asio::buffer(context.data, context.object_size));
+  buffer.push_back(asio::buffer(chunk_info.data, chunk_info.buffer_length));
   conn->WriteBuffer(buffer, ec);
 
   ray::Status status = ray::Status::OK();
   if (ec.value() != 0) {
-    // push failed.
-    // TODO(hme): Trash sender.
+    // Push failed. Deal with partial objects on the receiving end.
+    // TODO(hme): Try to invoke disconnect on sender connection, then remove it.
     status = ray::Status::IOError(ec.message());
   }
 
   // Do this regardless of whether it failed or succeeded.
-  ARROW_CHECK_OK(store_client->Release(context.object_id.to_plasma_id()));
-  store_pool_.ReleaseObjectStore(store_client);
+  buffer_pool_.ReleaseGetChunk(object_id, chunk_info.chunk_index);
   RAY_CHECK_OK(
       connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::TRANSFER, conn));
-  RAY_CHECK_OK(transfer_queue_.RemoveContext(context_id));
-  RAY_LOG(DEBUG) << "SendCompleted " << client_id_ << " " << context.object_id << " "
+  RAY_LOG(DEBUG) << "SendCompleted " << client_id_ << " " << object_id << " "
                  << num_transfers_send_ << "/" << config_.max_sends;
   RAY_CHECK_OK(TransferCompleted(TransferQueue::TransferType::SEND));
   return status;
@@ -469,45 +469,53 @@ void ObjectManager::ReceivePushRequest(std::shared_ptr<TcpClientConnection> conn
   auto object_header =
       flatbuffers::GetRoot<object_manager_protocol::PushRequestMessage>(message);
   ObjectID object_id = ObjectID::from_binary(object_header->object_id()->str());
-  int64_t object_size = (int64_t)object_header->object_size();
-  transfer_queue_.QueueReceive(conn->GetClientID(), object_id, object_size, conn);
+  uint64_t chunk_index = object_header->chunk_index();
+  uint64_t data_size = object_header->data_size();
+  uint64_t metadata_size = object_header->metadata_size();
+  transfer_queue_.QueueReceive(conn->GetClientID(), object_id, data_size, metadata_size,
+                               chunk_index, conn);
+  RAY_LOG(DEBUG) << "ReceivePushRequest " << conn->GetClientID() << " " << object_id
+                 << " " << chunk_index;
   RAY_CHECK_OK(DequeueTransfers());
 }
 
 ray::Status ObjectManager::ExecuteReceiveObject(
-    const ClientID &client_id, const ObjectID &object_id, uint64_t object_size,
+    const ClientID &client_id, const ObjectID &object_id, uint64_t data_size,
+    uint64_t metadata_size, uint64_t chunk_index,
     std::shared_ptr<TcpClientConnection> conn) {
-  boost::system::error_code ec;
-  int64_t metadata_size = 0;
-  const plasma::ObjectID plasma_id = ObjectID(object_id).to_plasma_id();
-  // Try to create shared buffer.
-  std::shared_ptr<Buffer> data;
-  std::shared_ptr<plasma::PlasmaClient> store_client = store_pool_.GetObjectStore();
-  arrow::Status s =
-      store_client->Create(plasma_id, object_size, NULL, metadata_size, &data);
-  std::vector<boost::asio::mutable_buffer> buffer;
-  if (s.ok()) {
-    // Read object into store.
-    uint8_t *mutable_data = data->mutable_data();
-    buffer.push_back(asio::buffer(mutable_data, object_size));
+  RAY_LOG(DEBUG) << "ExecuteReceiveObject " << client_id << " " << object_id << " "
+                 << chunk_index;
+
+  std::pair<const ObjectBufferPool::ChunkInfo &, ray::Status> chunk_status =
+      buffer_pool_.CreateChunk(object_id, data_size, metadata_size, chunk_index);
+  ObjectBufferPool::ChunkInfo chunk_info = chunk_status.first;
+  if (chunk_status.second.ok()) {
+    // Avoid handling this chunk if it's already being handled by another process.
+    std::vector<boost::asio::mutable_buffer> buffer;
+    buffer.push_back(asio::buffer(chunk_info.data, chunk_info.buffer_length));
+    boost::system::error_code ec;
     conn->ReadBuffer(buffer, ec);
-    if (!ec.value()) {
-      ARROW_CHECK_OK(store_client->Seal(plasma_id));
-      ARROW_CHECK_OK(store_client->Release(plasma_id));
+    if (ec.value() == 0) {
+      buffer_pool_.SealChunk(object_id, chunk_index);
     } else {
-      ARROW_CHECK_OK(store_client->Release(plasma_id));
-      ARROW_CHECK_OK(store_client->Abort(plasma_id));
-      RAY_LOG(ERROR) << "Receive Failed";
+      buffer_pool_.AbortCreateChunk(object_id, chunk_index);
+      // TODO(hme): This chunk failed, so create a pull request for this chunk.
     }
   } else {
-    RAY_LOG(ERROR) << "Buffer Create Failed: " << s.message();
+    RAY_LOG(ERROR) << "Buffer Create Failed: " << chunk_status.second.message();
     // Read object into empty buffer.
-    std::vector<uint8_t> mutable_data;
-    mutable_data.resize(object_size + metadata_size);
-    buffer.push_back(asio::buffer(mutable_data, object_size + metadata_size));
+    uint64_t buffer_length = buffer_pool_.GetBufferLength(chunk_index, data_size);
+    std::vector<uint8_t> mutable_vec;
+    mutable_vec.resize(buffer_length);
+    std::vector<boost::asio::mutable_buffer> buffer;
+    buffer.push_back(asio::buffer(mutable_vec, buffer_length));
+    boost::system::error_code ec;
     conn->ReadBuffer(buffer, ec);
+    if (ec.value() != 0) {
+      RAY_LOG(ERROR) << ec.message();
+    }
+    // TODO(hme): If the object isn't local, create a pull request for this chunk.
   }
-  store_pool_.ReleaseObjectStore(store_client);
   conn->ProcessMessages();
   RAY_LOG(DEBUG) << "ReceiveCompleted " << client_id_ << " " << object_id << " "
                  << num_transfers_receive_ << "/" << config_.max_receives;

src/ray/object_manager/object_manager.h

@@ -22,9 +22,9 @@
 
 #include "ray/object_manager/connection_pool.h"
 #include "ray/object_manager/format/object_manager_generated.h"
+#include "ray/object_manager/object_buffer_pool.h"
 #include "ray/object_manager/object_directory.h"
 #include "ray/object_manager/object_manager_client_connection.h"
-#include "ray/object_manager/object_store_client_pool.h"
 #include "ray/object_manager/object_store_notification_manager.h"
 #include "ray/object_manager/transfer_queue.h"
 
@@ -38,6 +38,8 @@ struct ObjectManagerConfig {
   int max_sends = 2;
   /// Maximum number of receives allowed.
   int max_receives = 2;
+  /// Object chunk size, in bytes
+  uint64_t object_chunk_size = std::pow(10, 8);
   // TODO(hme): Implement num retries (to avoid infinite retries).
   std::string store_socket_name;
 };
@@ -153,7 +155,7 @@ class ObjectManager {
   ObjectManagerConfig config_;
   std::unique_ptr<ObjectDirectoryInterface> object_directory_;
   ObjectStoreNotificationManager store_notification_;
-  ObjectStoreClientPool store_pool_;
+  ObjectBufferPool buffer_pool_;
 
   /// An io service for creating connections to other object managers.
   /// This runs on a thread pool.
@@ -252,19 +254,22 @@ class ObjectManager {
 
   /// Begin executing a send.
   /// Executes on object_manager_service_ thread pool.
-  ray::Status ExecuteSendObject(const ObjectID &object_id, const ClientID &client_id,
+  ray::Status ExecuteSendObject(const ClientID &client_id, const ObjectID &object_id,
+                                uint64_t data_size, uint64_t metadata_size,
+                                uint64_t chunk_index,
                                 const RemoteConnectionInfo &connection_info);
   /// This method synchronously sends the object id and object size
   /// to the remote object manager.
   /// Executes on object_manager_service_ thread pool.
-  ray::Status SendObjectHeaders(const ObjectID &object_id,
-                                std::shared_ptr<SenderConnection> client);
+  ray::Status SendObjectHeaders(const ObjectID &object_id, uint64_t data_size,
+                                uint64_t metadata_size, uint64_t chunk_index,
+                                std::shared_ptr<SenderConnection> conn);
 
   /// This method initiates the actual object transfer.
   /// Executes on object_manager_service_ thread pool.
-  ray::Status SendObjectData(std::shared_ptr<SenderConnection> conn,
-                             const UniqueID &context_id,
-                             std::shared_ptr<plasma::PlasmaClient> store_client);
+  ray::Status SendObjectData(const ObjectID &object_id,
+                             const ObjectBufferPool::ChunkInfo &chunk_info,
+                             std::shared_ptr<SenderConnection> conn);
 
   /// Invoked when a remote object manager pushes an object to this object manager.
   /// This will queue the receive.
@@ -272,7 +277,8 @@ class ObjectManager {
                           const uint8_t *message);
   /// Execute a receive that was in the queue.
   ray::Status ExecuteReceiveObject(const ClientID &client_id, const ObjectID &object_id,
-                                   uint64_t object_size,
+                                   uint64_t data_size, uint64_t metadata_size,
+                                   uint64_t chunk_index,
                                    std::shared_ptr<TcpClientConnection> conn);
 
   /// Handles receiving a pull request message.

src/ray/object_manager/object_store_client_pool.cc

@@ -1,36 +0,0 @@
-#include "object_store_client_pool.h"
-
-namespace ray {
-
-ObjectStoreClientPool::ObjectStoreClientPool(const std::string &store_socket_name)
-    : store_socket_name_(store_socket_name) {}
-
-ObjectStoreClientPool::~ObjectStoreClientPool() {
-  for (const auto &client : clients) {
-    ARROW_CHECK_OK(client->Disconnect());
-  }
-}
-
-std::shared_ptr<plasma::PlasmaClient> ObjectStoreClientPool::GetObjectStore() {
-  std::lock_guard<std::mutex> lock(pool_mutex);
-  if (available_clients.empty()) {
-    Add();
-  }
-  std::shared_ptr<plasma::PlasmaClient> client = available_clients.back();
-  available_clients.pop_back();
-  return client;
-}
-
-void ObjectStoreClientPool::ReleaseObjectStore(
-    std::shared_ptr<plasma::PlasmaClient> client) {
-  std::lock_guard<std::mutex> lock(pool_mutex);
-  available_clients.push_back(client);
-}
-
-void ObjectStoreClientPool::Add() {
-  clients.emplace_back(new plasma::PlasmaClient());
-  ARROW_CHECK_OK(clients.back()->Connect(store_socket_name_.c_str(), "",
-                                         PLASMA_DEFAULT_RELEASE_DELAY));
-  available_clients.push_back(clients.back());
-}
-}  // namespace ray

src/ray/object_manager/object_store_client_pool.h

@@ -1,64 +0,0 @@
-#ifndef RAY_OBJECT_MANAGER_OBJECT_STORE_CLIENT_POOL_H
-#define RAY_OBJECT_MANAGER_OBJECT_STORE_CLIENT_POOL_H
-
-#include <list>
-#include <memory>
-#include <mutex>
-#include <vector>
-
-#include <boost/asio.hpp>
-#include <boost/asio/error.hpp>
-#include <boost/bind.hpp>
-
-#include "plasma/client.h"
-#include "plasma/events.h"
-#include "plasma/plasma.h"
-
-#include "object_directory.h"
-#include "ray/id.h"
-#include "ray/status.h"
-
-namespace ray {
-
-/// \class ObjectStoreClientPool
-///
-/// Provides connections to the object store. Enables concurrent communication with
-/// the object store.
-class ObjectStoreClientPool {
- public:
-  /// Constructor.
-  ///
-  /// \param store_socket_name The object store socket name.
-  ObjectStoreClientPool(const std::string &store_socket_name);
-
-  ~ObjectStoreClientPool();
-
-  /// This object cannot be copied due to pool_mutex.
-  RAY_DISALLOW_COPY_AND_ASSIGN(ObjectStoreClientPool);
-
-  /// Provides a connection to the object store from the object store pool.
-  /// This removes the object store client from the pool of available clients.
-  ///
-  /// \return A connection to the object store.
-  std::shared_ptr<plasma::PlasmaClient> GetObjectStore();
-
-  /// Releases a client object and puts it back into the object store pool
-  /// for reuse.
-  /// Once a client is released, it is assumed that it is not being used.
-  /// \param client The client to return.
-  /// \param client
-  void ReleaseObjectStore(std::shared_ptr<plasma::PlasmaClient> client);
-
- private:
-  /// Adds a client to the client pool and mark it as available.
-  void Add();
-
-  std::mutex pool_mutex;
-  std::vector<std::shared_ptr<plasma::PlasmaClient>> available_clients;
-  std::vector<std::shared_ptr<plasma::PlasmaClient>> clients;
-  std::string store_socket_name_;
-};
-
-}  // namespace ray
-
-#endif  // RAY_OBJECT_MANAGER_OBJECT_STORE_CLIENT_POOL_H

src/ray/object_manager/test/object_manager_stress_test.cc

@@ -96,14 +96,25 @@ class TestObjectManagerBase : public ::testing::Test {
   std::string StartStore(const std::string &id) {
     std::string store_id = "/tmp/store";
     store_id = store_id + id;
+    std::string store_pid = store_id + ".pid";
     std::string plasma_command = store_executable + " -m 1000000000 -s " + store_id +
-                                 " 1> /dev/null 2> /dev/null &";
+                                 " 1> /dev/null 2> /dev/null &" + " echo $! > " +
+                                 store_pid;
+
     RAY_LOG(DEBUG) << plasma_command;
     int ec = system(plasma_command.c_str());
     RAY_CHECK(ec == 0);
+    sleep(1);
     return store_id;
   }
 
+  void StopStore(std::string store_id) {
+    std::string store_pid = store_id + ".pid";
+    std::string kill_1 = "kill -9 `cat " + store_pid + "`";
+    int s = system(kill_1.c_str());
+    ASSERT_TRUE(!s);
+  }
+
   void SetUp() {
     flushall_redis();
 
@@ -111,30 +122,32 @@ class TestObjectManagerBase : public ::testing::Test {
     object_manager_service_2.reset(new boost::asio::io_service());
 
     // start store
-    std::string store_sock_1 = StartStore("1");
-    std::string store_sock_2 = StartStore("2");
+    store_id_1 = StartStore(UniqueID::from_random().hex());
+    store_id_2 = StartStore(UniqueID::from_random().hex());
 
     // start first server
     gcs_client_1 = std::shared_ptr<gcs::AsyncGcsClient>(new gcs::AsyncGcsClient());
     ObjectManagerConfig om_config_1;
-    om_config_1.store_socket_name = store_sock_1;
+    om_config_1.store_socket_name = store_id_1;
     om_config_1.max_sends = 2;
     om_config_1.max_receives = 2;
+    om_config_1.object_chunk_size = static_cast<uint64_t>(std::pow(10, 3));
     server1.reset(new MockServer(main_service, std::move(object_manager_service_1),
                                  om_config_1, gcs_client_1));
 
     // start second server
     gcs_client_2 = std::shared_ptr<gcs::AsyncGcsClient>(new gcs::AsyncGcsClient());
     ObjectManagerConfig om_config_2;
-    om_config_2.store_socket_name = store_sock_2;
+    om_config_2.store_socket_name = store_id_2;
     om_config_2.max_sends = 2;
     om_config_2.max_receives = 2;
+    om_config_2.object_chunk_size = static_cast<uint64_t>(std::pow(10, 3));
     server2.reset(new MockServer(main_service, std::move(object_manager_service_2),
                                  om_config_2, gcs_client_2));
 
     // connect to stores.
-    ARROW_CHECK_OK(client1.Connect(store_sock_1, "", PLASMA_DEFAULT_RELEASE_DELAY));
-    ARROW_CHECK_OK(client2.Connect(store_sock_2, "", PLASMA_DEFAULT_RELEASE_DELAY));
+    ARROW_CHECK_OK(client1.Connect(store_id_1, "", PLASMA_DEFAULT_RELEASE_DELAY));
+    ARROW_CHECK_OK(client2.Connect(store_id_2, "", PLASMA_DEFAULT_RELEASE_DELAY));
   }
 
   void TearDown() {
@@ -145,8 +158,8 @@ class TestObjectManagerBase : public ::testing::Test {
     this->server1.reset();
     this->server2.reset();
 
-    int s = system("killall plasma_store &");
-    ASSERT_TRUE(!s);
+    StopStore(store_id_1);
+    StopStore(store_id_2);
   }
 
   ObjectID WriteDataToClient(plasma::PlasmaClient &client, int64_t data_size) {
@@ -179,6 +192,9 @@ class TestObjectManagerBase : public ::testing::Test {
   plasma::PlasmaClient client2;
   std::vector<ObjectID> v1;
   std::vector<ObjectID> v2;
+
+  std::string store_id_1;
+  std::string store_id_2;
 };
 
 class StressTestObjectManager : public TestObjectManagerBase {
@@ -257,7 +273,7 @@ class StressTestObjectManager : public TestObjectManagerBase {
     async_loop_index += 1;
     if ((uint)async_loop_index < async_loop_patterns.size()) {
       TransferPattern pattern = async_loop_patterns[async_loop_index];
-      TransferTestExecute(100, 100, pattern);
+      TransferTestExecute(100, 3 * std::pow(10, 3) - 1, pattern);
     } else {
       main_service.stop();
     }
@@ -279,11 +295,14 @@ class StressTestObjectManager : public TestObjectManagerBase {
 
   void CompareObjects(ObjectID &object_id_1, ObjectID &object_id_2) {
     plasma::ObjectBuffer object_buffer_1 = GetObject(client1, object_id_1);
-    plasma::ObjectBuffer object_buffer_2 = GetObject(client1, object_id_1);
+    plasma::ObjectBuffer object_buffer_2 = GetObject(client2, object_id_2);
     uint8_t *data_1 = const_cast<uint8_t *>(object_buffer_1.data->data());
     uint8_t *data_2 = const_cast<uint8_t *>(object_buffer_2.data->data());
-    ASSERT_EQ(object_buffer_1.data->size(), object_buffer_2.data->size());
-    for (int i = -1; ++i < object_buffer_1.data->size();) {
+    ASSERT_EQ(object_buffer_1.data->size(), object_buffer_2.data_size);
+    ASSERT_EQ(object_buffer_1.metadata_size, object_buffer_2.metadata_size);
+    int64_t total_size = object_buffer_1.data->size() + object_buffer_1.metadata_size;
+    RAY_LOG(DEBUG) << "total_size " << total_size;
+    for (int i = -1; ++i < total_size;) {
       ASSERT_TRUE(data_1[i] == data_2[i]);
     }
   }

src/ray/object_manager/test/object_manager_test.cc

@@ -87,14 +87,24 @@ class TestObjectManager : public ::testing::Test {
   std::string StartStore(const std::string &id) {
     std::string store_id = "/tmp/store";
     store_id = store_id + id;
+    std::string store_pid = store_id + ".pid";
     std::string plasma_command = store_executable + " -m 1000000000 -s " + store_id +
-                                 " 1> /dev/null 2> /dev/null &";
+                                 " 1> /dev/null 2> /dev/null &" + " echo $! > " +
+                                 store_pid;
+
     RAY_LOG(DEBUG) << plasma_command;
     int ec = system(plasma_command.c_str());
     RAY_CHECK(ec == 0);
+    sleep(1);
     return store_id;
   }
 
+  void StopStore(std::string store_id) {
+    std::string store_pid = store_id + ".pid";
+    std::string kill_1 = "kill -9 `cat " + store_pid + "`";
+    ASSERT_TRUE(!system(kill_1.c_str()));
+  }
+
   void SetUp() {
     flushall_redis();
 
@@ -102,26 +112,26 @@ class TestObjectManager : public ::testing::Test {
     object_manager_service_2.reset(new boost::asio::io_service());
 
     // start store
-    std::string store_sock_1 = StartStore("1");
-    std::string store_sock_2 = StartStore("2");
+    store_id_1 = StartStore(UniqueID::from_random().hex());
+    store_id_2 = StartStore(UniqueID::from_random().hex());
 
     // start first server
     gcs_client_1 = std::shared_ptr<gcs::AsyncGcsClient>(new gcs::AsyncGcsClient());
     ObjectManagerConfig om_config_1;
-    om_config_1.store_socket_name = store_sock_1;
+    om_config_1.store_socket_name = store_id_1;
     server1.reset(new MockServer(main_service, std::move(object_manager_service_1),
                                  om_config_1, gcs_client_1));
 
     // start second server
     gcs_client_2 = std::shared_ptr<gcs::AsyncGcsClient>(new gcs::AsyncGcsClient());
     ObjectManagerConfig om_config_2;
-    om_config_2.store_socket_name = store_sock_2;
+    om_config_2.store_socket_name = store_id_2;
     server2.reset(new MockServer(main_service, std::move(object_manager_service_2),
                                  om_config_2, gcs_client_2));
 
     // connect to stores.
-    ARROW_CHECK_OK(client1.Connect(store_sock_1, "", PLASMA_DEFAULT_RELEASE_DELAY));
-    ARROW_CHECK_OK(client2.Connect(store_sock_2, "", PLASMA_DEFAULT_RELEASE_DELAY));
+    ARROW_CHECK_OK(client1.Connect(store_id_1, "", PLASMA_DEFAULT_RELEASE_DELAY));
+    ARROW_CHECK_OK(client2.Connect(store_id_2, "", PLASMA_DEFAULT_RELEASE_DELAY));
   }
 
   void TearDown() {
@@ -132,8 +142,8 @@ class TestObjectManager : public ::testing::Test {
     this->server1.reset();
     this->server2.reset();
 
-    int s = system("killall plasma_store &");
-    ASSERT_TRUE(!s);
+    StopStore(store_id_1);
+    StopStore(store_id_2);
   }
 
   ObjectID WriteDataToClient(plasma::PlasmaClient &client, int64_t data_size) {
@@ -166,6 +176,9 @@ class TestObjectManager : public ::testing::Test {
   plasma::PlasmaClient client2;
   std::vector<ObjectID> v1;
   std::vector<ObjectID> v2;
+
+  std::string store_id_1;
+  std::string store_id_2;
 };
 
 class TestObjectManagerCommands : public TestObjectManager {

src/ray/object_manager/transfer_queue.cc

@@ -3,9 +3,10 @@
 namespace ray {
 
 void TransferQueue::QueueSend(const ClientID &client_id, const ObjectID &object_id,
-                              const RemoteConnectionInfo &info) {
-  WriteLock guard(send_mutex);
-  SendRequest req = {client_id, object_id, info};
+                              uint64_t data_size, uint64_t metadata_size,
+                              uint64_t chunk_index, const RemoteConnectionInfo &info) {
+  std::unique_lock<std::mutex> guard(send_mutex_);
+  SendRequest req = {client_id, object_id, data_size, metadata_size, chunk_index, info};
   // TODO(hme): Use a set to speed this up.
   if (std::find(send_queue_.begin(), send_queue_.end(), req) != send_queue_.end()) {
     // already queued.
@@ -15,10 +16,12 @@ void TransferQueue::QueueSend(const ClientID &client_id, const ObjectID &object_
 }
 
 void TransferQueue::QueueReceive(const ClientID &client_id, const ObjectID &object_id,
-                                 uint64_t object_size,
+                                 uint64_t data_size, uint64_t metadata_size,
+                                 uint64_t chunk_index,
                                  std::shared_ptr<TcpClientConnection> conn) {
-  WriteLock guard(receive_mutex);
-  ReceiveRequest req = {client_id, object_id, object_size, conn};
+  std::unique_lock<std::mutex> guard(receive_mutex_);
+  ReceiveRequest req = {client_id,     object_id,   data_size,
+                        metadata_size, chunk_index, conn};
   if (std::find(receive_queue_.begin(), receive_queue_.end(), req) !=
       receive_queue_.end()) {
     // already queued.
@@ -28,7 +31,7 @@ void TransferQueue::QueueReceive(const ClientID &client_id, const ObjectID &obje
 }
 
 bool TransferQueue::DequeueSendIfPresent(TransferQueue::SendRequest *send_ptr) {
-  WriteLock guard(send_mutex);
+  std::unique_lock<std::mutex> guard(send_mutex_);
   if (send_queue_.empty()) {
     return false;
   }
@@ -38,7 +41,7 @@ bool TransferQueue::DequeueSendIfPresent(TransferQueue::SendRequest *send_ptr) {
 }
 
 bool TransferQueue::DequeueReceiveIfPresent(TransferQueue::ReceiveRequest *receive_ptr) {
-  WriteLock guard(receive_mutex);
+  std::unique_lock<std::mutex> guard(receive_mutex_);
   if (receive_queue_.empty()) {
     return false;
   }
@@ -47,21 +50,4 @@ bool TransferQueue::DequeueReceiveIfPresent(TransferQueue::ReceiveRequest *recei
   return true;
 }
 
-UniqueID TransferQueue::AddContext(SendContext &context) {
-  WriteLock guard(context_mutex);
-  UniqueID id = UniqueID::from_random();
-  send_context_set_.emplace(id, context);
-  return id;
-}
-
-TransferQueue::SendContext &TransferQueue::GetContext(const UniqueID &id) {
-  ReadLock guard(context_mutex);
-  return send_context_set_[id];
-}
-
-ray::Status TransferQueue::RemoveContext(const UniqueID &id) {
-  WriteLock guard(context_mutex);
-  send_context_set_.erase(id);
-  return Status::OK();
-}
 }  // namespace ray

src/ray/object_manager/transfer_queue.h

@@ -26,21 +26,17 @@ class TransferQueue {
  public:
   enum TransferType { SEND = 1, RECEIVE };
 
-  /// Context maintained during an object send.
-  struct SendContext {
-    ClientID client_id;
-    ObjectID object_id;
-    uint64_t object_size;
-    uint8_t *data;
-  };
-
   /// The structure used in the send queue.
   struct SendRequest {
     ClientID client_id;
     ObjectID object_id;
+    uint64_t data_size;
+    uint64_t metadata_size;
+    uint64_t chunk_index;
     RemoteConnectionInfo connection_info;
     bool operator==(const SendRequest &rhs) const {
-      return client_id == rhs.client_id && object_id == rhs.object_id;
+      return client_id == rhs.client_id && object_id == rhs.object_id &&
+             chunk_index == rhs.chunk_index;
     }
   };
 
@@ -48,10 +44,14 @@ class TransferQueue {
   struct ReceiveRequest {
     ClientID client_id;
     ObjectID object_id;
-    uint64_t object_size;
+    uint64_t data_size;
+    uint64_t metadata_size;
+    uint64_t chunk_index;
     std::shared_ptr<TcpClientConnection> conn;
     bool operator==(const ReceiveRequest &rhs) const {
-      return client_id == rhs.client_id && object_id == rhs.object_id;
+      return client_id == rhs.client_id && object_id == rhs.object_id &&
+             chunk_index == rhs.chunk_index;
+      ;
     }
   };
 
@@ -61,11 +61,13 @@ class TransferQueue {
   ///
   /// \param client_id The ClientID to which the object needs to be sent.
   /// \param object_id The ObjectID of the object to be sent.
-  void QueueSend(const ClientID &client_id, const ObjectID &object_id,
+  void QueueSend(const ClientID &client_id, const ObjectID &object_id, uint64_t data_size,
+                 uint64_t metadata_size, uint64_t chunk_index,
                  const RemoteConnectionInfo &info);
 
   /// If send_queue_ is not empty, removes a SendRequest from send_queue_ and assigns
   /// it to send_ptr. The queue is FIFO.
+  ///
   /// \param send_ptr A pointer to an empty SendRequest.
   /// \return A bool indicating whether the queue was empty at the time this method
   /// was invoked.
@@ -76,51 +78,27 @@ class TransferQueue {
   /// \param client_id The ClientID from which the object is being received.
   /// \param object_id The ObjectID of the object to be received.
   void QueueReceive(const ClientID &client_id, const ObjectID &object_id,
-                    uint64_t object_size, std::shared_ptr<TcpClientConnection> conn);
+                    uint64_t data_size, uint64_t metadata_size, uint64_t chunk_index,
+                    std::shared_ptr<TcpClientConnection> conn);
 
   /// If receive_queue_ is not empty, removes a ReceiveRequest from receive_queue_ and
-  /// assigns
-  /// it to receive_ptr. The queue is FIFO.
+  /// assigns it to receive_ptr. The queue is FIFO.
+  ///
   /// \param receive_ptr A pointer to an empty ReceiveRequest.
   /// \return A bool indicating whether the queue was empty at the time this method
   /// was invoked.
   bool DequeueReceiveIfPresent(TransferQueue::ReceiveRequest *receive_ptr);
 
-  /// Maintain ownership over SendContext for sends in transit.
-  ///
-  /// \param context The context to maintain.
-  /// \return A unique identifier identifying the context that was added.
-  UniqueID AddContext(SendContext &context);
-
-  /// Gets the SendContext associated with the given id.
-  ///
-  /// \param id The unique identifier of the context.
-  /// \return The context.
-  SendContext &GetContext(const UniqueID &id);
-
-  /// Removes the context associated with the given id.
-  ///
-  /// \param id The unique identifier of the context.
-  /// \return The status of invoking this method.
-  ray::Status RemoveContext(const UniqueID &id);
-
   /// This object cannot be copied for thread-safety.
   RAY_DISALLOW_COPY_AND_ASSIGN(TransferQueue);
 
  private:
-  // TODO(hme): make this a shared mutex.
-  typedef std::mutex Lock;
-  typedef std::unique_lock<Lock> WriteLock;
-  // TODO(hme): make this a shared lock.
-  typedef std::unique_lock<Lock> ReadLock;
-  Lock send_mutex;
-  Lock receive_mutex;
-  Lock context_mutex;
-
+  std::mutex send_mutex_;
+  std::mutex receive_mutex_;
   std::deque<SendRequest> send_queue_;
   std::deque<ReceiveRequest> receive_queue_;
-  std::unordered_map<ray::UniqueID, SendContext, ray::UniqueIDHasher> send_context_set_;
 };
+
 }  // namespace ray
 
 #endif  // RAY_OBJECT_MANAGER_TRANSFER_QUEUE_H

src/ray/raylet/main.cc

@@ -47,6 +47,14 @@ int main(int argc, char *argv[]) {
   // Configuration for the object manager.
   ray::ObjectManagerConfig object_manager_config;
   object_manager_config.store_socket_name = store_socket_name;
+  // Time out in milliseconds to wait before retrying a failed pull.
+  object_manager_config.pull_timeout_ms = 100;
+  // Maximum number of sends allowed.
+  object_manager_config.max_sends = 2;
+  // Maximum number of receives allowed.
+  object_manager_config.max_receives = 2;
+  // Object chunk size, in bytes.
+  object_manager_config.object_chunk_size = static_cast<uint64_t>(std::pow(10, 8));
 
   //  initialize mock gcs & object directory
   auto gcs_client = std::make_shared<ray::gcs::AsyncGcsClient>();

src/ray/raylet/raylet.cc

@@ -37,16 +37,13 @@ Raylet::Raylet(boost::asio::io_service &main_service,
   DoAcceptNodeManager();
 
   RAY_CHECK_OK(RegisterGcs(node_ip_address, socket_name_,
-                           object_manager_config.store_socket_name,
-                           redis_address, redis_port, main_service,
-                           node_manager_config));
+                           object_manager_config.store_socket_name, redis_address,
+                           redis_port, main_service, node_manager_config));
 
   RAY_CHECK_OK(RegisterPeriodicTimer(main_service));
 }
 
-Raylet::~Raylet() {
-  RAY_CHECK_OK(gcs_client_->client_table().Disconnect());
-}
+Raylet::~Raylet() { RAY_CHECK_OK(gcs_client_->client_table().Disconnect()); }
 
 ray::Status Raylet::RegisterPeriodicTimer(boost::asio::io_service &io_service) {
   boost::posix_time::milliseconds timer_period_ms(100);

src/ray/test/run_object_manager_tests.sh

@@ -30,12 +30,11 @@ echo "$REDIS_DIR/redis-server --loglevel warning --loadmodule $REDIS_MODULE --po
 echo "$REDIS_DIR/redis-cli -p 6379 shutdown"
 
 # Allow cleanup commands to fail.
-killall plasma_store || true
-$REDIS_DIR/redis-cli -p 6379 shutdown || true
+# killall plasma_store || true
+# $REDIS_DIR/redis-cli -p 6379 shutdown || true
 sleep 1s
 $REDIS_DIR/redis-server --loglevel warning --loadmodule $REDIS_MODULE --port 6379 &
 sleep 1s
-
 # Run tests.
 $CORE_DIR/src/ray/object_manager/object_manager_stress_test $STORE_EXEC
 sleep 1s