[Core] Remove socket pair exchange in Plasma Store (#9565)

* try use boost::asio for notification processing

src/ray/object_manager/notification/object_store_notification_manager_ipc.cc

@@ -14,17 +14,8 @@
 
 #include "ray/object_manager/notification/object_store_notification_manager_ipc.h"
 
-#include <boost/asio.hpp>
-#include <boost/bind.hpp>
-#include <boost/function.hpp>
-
-#include "ray/common/common_protocol.h"
-#include "ray/common/status.h"
-#include "ray/util/util.h"
-
-#ifdef _WIN32
-#include <win32fd.h>
-#endif
+#include "ray/common/client_connection.h"
+#include "ray/object_manager/plasma/plasma_generated.h"
 
 namespace ray {
 
@@ -32,76 +23,47 @@ ObjectStoreNotificationManagerIPC::ObjectStoreNotificationManagerIPC(
     boost::asio::io_service &io_service, const std::string &store_socket_name,
     bool exit_on_error)
     : ObjectStoreNotificationManager(io_service),
-      store_client_(),
       length_(0),
-      socket_(io_service),
       exit_on_error_(exit_on_error) {
-  RAY_CHECK_OK(store_client_.Connect(store_socket_name.c_str(), "", 0, 300));
-
-  int fd;
-  RAY_CHECK_OK(store_client_.Subscribe(&fd));
-  boost::system::error_code ec;
-#ifdef _WIN32
-  boost::asio::detail::socket_type c_socket = fh_release(fd);
-  WSAPROTOCOL_INFO pi;
-  size_t n = sizeof(pi);
-  char *p = reinterpret_cast<char *>(&pi);
-  const int level = SOL_SOCKET;
-  const int opt = SO_PROTOCOL_INFO;
-  if (boost::asio::detail::socket_ops::getsockopt(c_socket, 0, level, opt, p, &n, ec) !=
-      boost::asio::detail::socket_error_retval) {
-    switch (pi.iAddressFamily) {
-    case AF_INET:
-      socket_.assign(boost::asio::ip::tcp::v4(), c_socket, ec);
-      break;
-    case AF_INET6:
-      socket_.assign(boost::asio::ip::tcp::v6(), c_socket, ec);
-      break;
-    default:
-      ec = boost::system::errc::make_error_code(
-          boost::system::errc::address_family_not_supported);
-      break;
-    }
-  }
-#else
-  socket_.assign(boost::asio::local::stream_protocol(), fd, ec);
-#endif
-  RAY_CHECK(!ec);
+  local_stream_socket socket(io_service);
+  RAY_CHECK_OK(ConnectSocketRetry(socket, store_socket_name));
+  store_client_ = ServerConnection::Create(std::move(socket));
+  // Subscribe messages.
+  flatbuffers::FlatBufferBuilder fbb;
+  auto message = plasma::flatbuf::CreatePlasmaSubscribeRequest(fbb);
+  fbb.Finish(message);
+  RAY_CHECK_OK(store_client_->WriteMessage(
+      static_cast<int64_t>(plasma::flatbuf::MessageType::PlasmaSubscribeRequest),
+      fbb.GetSize(), fbb.GetBufferPointer()));
+  RAY_CHECK_OK(store_client_->SetNonBlocking(true));
   NotificationWait();
 }
 
 ObjectStoreNotificationManagerIPC::~ObjectStoreNotificationManagerIPC() {
-  RAY_CHECK_OK(store_client_.Disconnect());
+  store_client_->Close();
 }
 
-void ObjectStoreNotificationManagerIPC::Shutdown() {
-  RAY_CHECK_OK(store_client_.Disconnect());
-}
+void ObjectStoreNotificationManagerIPC::Shutdown() { store_client_->Close(); }
 
 void ObjectStoreNotificationManagerIPC::NotificationWait() {
-  boost::asio::async_read(
-      socket_, boost::asio::buffer(&length_, sizeof(length_)),
-      boost::bind(&ObjectStoreNotificationManagerIPC::ProcessStoreLength, this,
-                  boost::asio::placeholders::error));
+  store_client_->ReadBufferAsync({boost::asio::buffer(&length_, sizeof(length_))},
+                                 [this](const ray::Status &s) { ProcessStoreLength(s); });
 }
 
-void ObjectStoreNotificationManagerIPC::ProcessStoreLength(
-    const boost::system::error_code &error) {
+void ObjectStoreNotificationManagerIPC::ProcessStoreLength(const ray::Status &s) {
   notification_.resize(length_);
-  if (error) {
+  if (!s.ok()) {
     if (exit_on_error_) {
       // When shutting down a cluster, it's possible that the plasma store is killed
       // earlier than raylet. In this case we don't want raylet to crash, we instead
       // log an error message and exit.
-      RAY_LOG(ERROR) << "Failed to process store length: "
-                     << boost_to_ray_status(error).ToString()
+      RAY_LOG(ERROR) << "Failed to process store length: " << s.ToString()
                      << ", most likely plasma store is down, raylet will exit";
       // Exit raylet process.
       _exit(kRayletStoreErrorExitCode);
     } else {
       // The log level is set to debug so user don't see it on ctrl+c exit.
-      RAY_LOG(DEBUG) << "Failed to process store length: "
-                     << boost_to_ray_status(error).ToString()
+      RAY_LOG(DEBUG) << "Failed to process store length: " << s.ToString()
                      << ", most likely plasma store is down. "
                      << "The error is silenced because exit_on_error_ "
                      << "flag is set.";
@@ -109,24 +71,22 @@ void ObjectStoreNotificationManagerIPC::ProcessStoreLength(
     }
   }
 
-  boost::asio::async_read(
-      socket_, boost::asio::buffer(notification_),
-      boost::bind(&ObjectStoreNotificationManagerIPC::ProcessStoreNotification, this,
-                  boost::asio::placeholders::error));
+  store_client_->ReadBufferAsync(
+      {boost::asio::buffer(notification_)},
+      [this](const ray::Status &s) { ProcessStoreNotification(s); });
 }
 
-void ObjectStoreNotificationManagerIPC::ProcessStoreNotification(
-    const boost::system::error_code &error) {
-  if (error) {
+void ObjectStoreNotificationManagerIPC::ProcessStoreNotification(const ray::Status &s) {
+  if (!s.ok()) {
     if (exit_on_error_) {
       RAY_LOG(FATAL)
           << "Problem communicating with the object store from raylet, check logs or "
-          << "dmesg for previous errors: " << boost_to_ray_status(error).ToString();
+          << "dmesg for previous errors: " << s.ToString();
     } else {
       // The log level is set to debug so user don't see it on ctrl+c exit.
       RAY_LOG(DEBUG)
           << "Problem communicating with the object store from raylet, check logs or "
-          << "dmesg for previous errors: " << boost_to_ray_status(error).ToString()
+          << "dmesg for previous errors: " << s.ToString()
           << " The error is silenced because exit_on_error_ "
           << "flag is set.";
       return;

src/ray/object_manager/notification/object_store_notification_manager_ipc.h

@@ -14,22 +14,16 @@
 
 #pragma once
 
-#include <list>
 #include <memory>
 #include <vector>
 
-#include <boost/asio.hpp>
-#include <boost/asio/error.hpp>
-#include <boost/bind.hpp>
-
-#include "ray/common/client_connection.h"
-#include "ray/common/id.h"
 #include "ray/common/status.h"
 #include "ray/object_manager/notification/object_store_notification_manager.h"
-#include "ray/object_manager/plasma/client.h"
 
 namespace ray {
 
+class ServerConnection;
+
 /// \class ObjectStoreNotificationManagerIPC
 ///
 /// Encapsulates notification handling from the object store.
@@ -53,13 +47,12 @@ class ObjectStoreNotificationManagerIPC : public ObjectStoreNotificationManager
  private:
   /// Async loop for handling object store notifications.
   void NotificationWait();
-  void ProcessStoreLength(const boost::system::error_code &error);
-  void ProcessStoreNotification(const boost::system::error_code &error);
+  void ProcessStoreLength(const ray::Status &s);
+  void ProcessStoreNotification(const ray::Status &s);
 
-  plasma::PlasmaClient store_client_;
+  std::shared_ptr<ServerConnection> store_client_;
   int64_t length_;
   std::vector<uint8_t> notification_;
-  local_stream_socket socket_;
 
   /// Flag to indicate whether or not to exit the process when received socket
   /// error. When it is false, socket error will be ignored. This flag is needed

src/ray/object_manager/plasma/client.cc

@@ -257,15 +257,6 @@ class PlasmaClient::Impl : public std::enable_shared_from_this<PlasmaClient::Imp
 
   Status Refresh(const std::vector<ObjectID>& object_ids);
 
-  Status Subscribe(int* fd);
-
-  Status GetNotification(int fd, ObjectID* object_id, int64_t* data_size,
-                         int64_t* metadata_size);
-
-  Status DecodeNotifications(const uint8_t* buffer, std::vector<ObjectID>* object_ids,
-                             std::vector<int64_t>* data_sizes,
-                             std::vector<int64_t>* metadata_sizes);
-
   Status Disconnect();
 
   std::string DebugString();
@@ -817,93 +808,6 @@ Status PlasmaClient::Impl::Refresh(const std::vector<ObjectID>& object_ids) {
   return ReadRefreshLRUReply(buffer.data(), buffer.size());
 }
 
-Status PlasmaClient::Impl::Subscribe(int* fd) {
-  std::lock_guard<std::recursive_mutex> guard(client_mutex_);
-
-  int sock[2];
-  // Create a non-blocking socket pair. This will only be used to send
-  // notifications from the Plasma store to the client.
-#ifdef _WINSOCKAPI_
-  SOCKET sockets[2] = { INVALID_SOCKET, INVALID_SOCKET };
-  socketpair(AF_INET, SOCK_STREAM, 0, sockets);
-  sock[0] = fh_open(sockets[0], -1);
-  sock[1] = fh_open(sockets[1], -1);
-#else
-  socketpair(AF_UNIX, SOCK_STREAM, 0, sock);
-#endif
-  // Make the socket non-blocking.
-#ifdef _WINSOCKAPI_
-  unsigned long value = 1;
-  RAY_CHECK(ioctlsocket(sock[1], FIONBIO, &value) == 0);
-#else
-  int flags = fcntl(sock[1], F_GETFL, 0);
-  RAY_CHECK(fcntl(sock[1], F_SETFL, flags | O_NONBLOCK) == 0);
-#endif
-  // Tell the Plasma store about the subscription.
-  RAY_RETURN_NOT_OK(SendSubscribeRequest(store_conn_));
-  // Send the file descriptor that the Plasma store should use to push
-  // notifications about sealed objects to this client.
-  RAY_CHECK(send_fd(store_conn_->fd, sock[1]) >= 0);
-  close(sock[1]);
-  // Return the file descriptor that the client should use to read notifications
-  // about sealed objects.
-  *fd = sock[0];
-  return Status::OK();
-}
-
-Status PlasmaClient::Impl::GetNotification(int fd, ObjectID* object_id,
-                                           int64_t* data_size, int64_t* metadata_size) {
-  std::lock_guard<std::recursive_mutex> guard(client_mutex_);
-
-  if (pending_notification_.empty()) {
-    auto message = ReadMessageAsync(fd);
-    if (message == NULL) {
-      return Status::IOError("Failed to read object notification from Plasma socket");
-    }
-
-    std::vector<ObjectID> object_ids;
-    std::vector<int64_t> data_sizes;
-    std::vector<int64_t> metadata_sizes;
-    RAY_RETURN_NOT_OK(
-        DecodeNotifications(message.get(), &object_ids, &data_sizes, &metadata_sizes));
-    for (size_t i = 0; i < object_ids.size(); ++i) {
-      pending_notification_.emplace_back(object_ids[i], data_sizes[i], metadata_sizes[i]);
-    }
-  }
-
-  auto notification = pending_notification_.front();
-  *object_id = std::get<0>(notification);
-  *data_size = std::get<1>(notification);
-  *metadata_size = std::get<2>(notification);
-
-  pending_notification_.pop_front();
-
-  return Status::OK();
-}
-
-Status PlasmaClient::Impl::DecodeNotifications(const uint8_t* buffer,
-                                               std::vector<ObjectID>* object_ids,
-                                               std::vector<int64_t>* data_sizes,
-                                               std::vector<int64_t>* metadata_sizes) {
-  std::lock_guard<std::recursive_mutex> guard(client_mutex_);
-  auto object_info = flatbuffers::GetRoot<ray::object_manager::protocol::PlasmaNotification>(buffer);
-
-  for (size_t i = 0; i < object_info->object_info()->size(); ++i) {
-    auto info = object_info->object_info()->Get(i);
-    ObjectID id = ObjectID::FromBinary(info->object_id()->str());
-    object_ids->push_back(id);
-    if (info->is_deletion()) {
-      data_sizes->push_back(-1);
-      metadata_sizes->push_back(-1);
-    } else {
-      data_sizes->push_back(info->data_size());
-      metadata_sizes->push_back(info->metadata_size());
-    }
-  }
-
-  return Status::OK();
-}
-
 Status PlasmaClient::Impl::Connect(const std::string& store_socket_name,
                                    const std::string& manager_socket_name,
                                    int release_delay, int num_retries) {
@@ -1030,20 +934,6 @@ Status PlasmaClient::Refresh(const std::vector<ObjectID>& object_ids) {
   return impl_->Refresh(object_ids);
 }
 
-Status PlasmaClient::Subscribe(int* fd) { return impl_->Subscribe(fd); }
-
-Status PlasmaClient::GetNotification(int fd, ObjectID* object_id, int64_t* data_size,
-                                     int64_t* metadata_size) {
-  return impl_->GetNotification(fd, object_id, data_size, metadata_size);
-}
-
-Status PlasmaClient::DecodeNotifications(const uint8_t* buffer,
-                                         std::vector<ObjectID>* object_ids,
-                                         std::vector<int64_t>* data_sizes,
-                                         std::vector<int64_t>* metadata_sizes) {
-  return impl_->DecodeNotifications(buffer, object_ids, data_sizes, metadata_sizes);
-}
-
 Status PlasmaClient::Disconnect() { return impl_->Disconnect(); }
 
 std::string PlasmaClient::DebugString() { return impl_->DebugString(); }

src/ray/object_manager/plasma/client.h

@@ -205,30 +205,6 @@ class RAY_EXPORT PlasmaClient {
   /// \return The return status.
   Status Refresh(const std::vector<ObjectID>& object_ids);
 
-  /// Subscribe to notifications when objects are sealed in the object store.
-  /// Whenever an object is sealed, a message will be written to the client
-  /// socket that is returned by this method.
-  ///
-  /// \param fd Out parameter for the file descriptor the client should use to
-  /// read notifications
-  ///         from the object store about sealed objects.
-  /// \return The return status.
-  Status Subscribe(int* fd);
-
-  /// Receive next object notification for this client if Subscribe has been called.
-  ///
-  /// \param fd The file descriptor we are reading the notification from.
-  /// \param object_id Out parameter, the object_id of the object that was sealed.
-  /// \param data_size Out parameter, the data size of the object that was sealed.
-  /// \param metadata_size Out parameter, the metadata size of the object that was sealed.
-  /// \return The return status.
-  Status GetNotification(int fd, ObjectID* object_id, int64_t* data_size,
-                         int64_t* metadata_size);
-
-  Status DecodeNotifications(const uint8_t* buffer, std::vector<ObjectID>* object_ids,
-                             std::vector<int64_t>* data_sizes,
-                             std::vector<int64_t>* metadata_sizes);
-
   /// Disconnect from the local plasma instance, including the local store and
   /// manager.
   ///

src/ray/object_manager/plasma/io.h

@@ -28,6 +28,7 @@
 #include <string>
 #include <vector>
 
+#include "ray/common/ray_config.h"
 #include "ray/common/status.h"
 #include "ray/object_manager/plasma/common.h"
 #include "ray/object_manager/plasma/compat.h"
@@ -43,10 +44,9 @@ enum class MessageType : int64_t;
 
 }  // namespace flatbuf
 
-// TODO(pcm): Replace our own custom message header (message type,
-// message length, plasma protocol version) with one that is serialized
-// using flatbuffers.
-constexpr int64_t kPlasmaProtocolVersion = 0x0000000000000000;
+// TODO(suquark): We temporarily sync this with Ray cookie. This code will
+// be removed soon in later PRs.
+constexpr int64_t kPlasmaProtocolVersion = 0x5241590000000000;
 
 Status WriteBytes(int fd, uint8_t* cursor, size_t length);
 

src/ray/object_manager/plasma/plasma.cc

@@ -65,44 +65,6 @@ int WarnIfSigpipe(int status, int client_sock) {
   return -1;  // This is never reached.
 }
 
-/**
- * This will create a new ObjectInfo buffer. The first sizeof(int64_t) bytes
- * of this buffer are the length of the remaining message and the
- * remaining message is a serialized version of the object info.
- *
- * \param object_info The object info to be serialized
- * \return The object info buffer. It is the caller's responsibility to free
- *         this buffer with "delete" after it has been used.
- */
-std::unique_ptr<uint8_t[]> CreateObjectInfoBuffer(fb::ObjectInfoT* object_info) {
-  flatbuffers::FlatBufferBuilder fbb;
-  auto message = fb::CreateObjectInfo(fbb, object_info);
-  fbb.Finish(message);
-  auto notification =
-      std::unique_ptr<uint8_t[]>(new uint8_t[sizeof(int64_t) + fbb.GetSize()]);
-  *(reinterpret_cast<int64_t*>(notification.get())) = fbb.GetSize();
-  memcpy(notification.get() + sizeof(int64_t), fbb.GetBufferPointer(), fbb.GetSize());
-  return notification;
-}
-
-std::unique_ptr<uint8_t[]> CreatePlasmaNotificationBuffer(
-    const std::vector<fb::ObjectInfoT>& object_info) {
-  flatbuffers::FlatBufferBuilder fbb;
-  std::vector<flatbuffers::Offset<fb::ObjectInfo>> info;
-  for (size_t i = 0; i < object_info.size(); ++i) {
-    info.push_back(fb::CreateObjectInfo(fbb, &object_info[i]));
-  }
-
-  auto info_array = fbb.CreateVector(info);
-  auto message = fb::CreatePlasmaNotification(fbb, info_array);
-  fbb.Finish(message);
-  auto notification =
-      std::unique_ptr<uint8_t[]>(new uint8_t[sizeof(int64_t) + fbb.GetSize()]);
-  *(reinterpret_cast<int64_t*>(notification.get())) = fbb.GetSize();
-  memcpy(notification.get() + sizeof(int64_t), fbb.GetBufferPointer(), fbb.GetSize());
-  return notification;
-}
-
 ObjectTableEntry* GetObjectTableEntry(PlasmaStoreInfo* store_info,
                                       const ObjectID& object_id) {
   auto it = store_info->objects.find(object_id);

src/ray/object_manager/plasma/plasma.h

@@ -70,7 +70,7 @@ constexpr int64_t kBlockSize = 64;
 
 /// Contains all information that is associated with a Plasma store client.
 struct Client {
-  explicit Client(int fd) : fd(fd), notification_fd(-1) {}
+  explicit Client(int fd) : fd(fd) {}
 
   ~Client();
 
@@ -83,11 +83,11 @@ struct Client {
   /// File descriptors that are used by this client.
   std::unordered_set<int> used_fds;
 
-  /// The file descriptor used to push notifications to client. This is only valid
-  /// if client subscribes to plasma store. -1 indicates invalid.
-  int notification_fd;
-
   std::string name = "anonymous_client";
+
+  /// The object notifications for clients. We notify the client about the
+  /// objects in the order that the objects were sealed or deleted.
+  std::deque<std::unique_ptr<uint8_t[]>> object_notifications;
 };
 
 std::ostream &operator<<(std::ostream &os, const std::shared_ptr<Client> &client);
@@ -182,11 +182,6 @@ ObjectTableEntry* GetObjectTableEntry(PlasmaStoreInfo* store_info,
 /// \return The errno set.
 int WarnIfSigpipe(int status, int client_sock);
 
-std::unique_ptr<uint8_t[]> CreateObjectInfoBuffer(ObjectInfoT* object_info);
-
-std::unique_ptr<uint8_t[]> CreatePlasmaNotificationBuffer(
-    const std::vector<ObjectInfoT>& object_info);
-
 /// Globally accessible reference to plasma store configuration.
 extern const PlasmaStoreInfo* plasma_config;
 

src/ray/object_manager/plasma/protocol.cc

@@ -602,14 +602,6 @@ Status ReadGetReply(uint8_t* data, size_t size, ObjectID object_ids[],
   return Status::OK();
 }
 
-// Subscribe messages.
-
-Status SendSubscribeRequest(const std::shared_ptr<StoreConn> &store_conn) {
-  flatbuffers::FlatBufferBuilder fbb;
-  auto message = fb::CreatePlasmaSubscribeRequest(fbb);
-  return PlasmaSend(store_conn, MessageType::PlasmaSubscribeRequest, &fbb, message);
-}
-
 // Data messages.
 
 Status SendDataRequest(const std::shared_ptr<StoreConn> &store_conn, ObjectID object_id, const char* address, int port) {

src/ray/object_manager/plasma/protocol.h

@@ -177,10 +177,6 @@ Status SendEvictReply(const std::shared_ptr<Client> &client, int64_t num_bytes);
 
 Status ReadEvictReply(uint8_t* data, size_t size, int64_t& num_bytes);
 
-/* Plasma Subscribe message functions. */
-
-Status SendSubscribeRequest(const std::shared_ptr<StoreConn> &store_conn);
-
 /* Data messages. */
 
 Status SendDataRequest(const std::shared_ptr<StoreConn> &store_conn, ObjectID object_id, const char* address, int port);

src/ray/object_manager/plasma/store.cc

@@ -742,16 +742,9 @@ void PlasmaStore::DisconnectClient(const std::shared_ptr<Client> &client) {
     RemoveFromClientObjectIds(entry.first, entry.second, client);
   }
 
-  if (client->notification_fd > 0) {
-    // This client has subscribed for notifications.
-    auto notify_fd = client->notification_fd;
-    loop_->RemoveFileEvent(notify_fd);
-    // Close socket.
-    close(notify_fd);
-    // Remove notification queue for this fd from global map.
-    pending_notifications_.erase(notify_fd);
-    // Reset fd.
-    client->notification_fd = -1;
+  if (notification_clients_.find(client) != notification_clients_.end()) {
+    // Remove notification for this client from global map.
+    notification_clients_.erase(client);
   }
 
   // We lose the last borrower of the Client instance here.
@@ -761,16 +754,28 @@ void PlasmaStore::DisconnectClient(const std::shared_ptr<Client> &client) {
 /// Send notifications about sealed objects to the subscribers. This is called
 /// in SealObject. If the socket's send buffer is full, the notification will
 /// be buffered, and this will be called again when the send buffer has room.
-/// Since we call erase on pending_notifications_, all iterators get
-/// invalidated, which is why we return a valid iterator to the next client to
-/// be used in PushNotification.
 ///
-/// \param it Iterator that points to the client to send the notification to.
-/// \return Iterator pointing to the next client.
-PlasmaStore::NotificationMap::iterator PlasmaStore::SendNotifications(
-    PlasmaStore::NotificationMap::iterator it) {
-  int client_fd = it->first;
-  auto& notifications = it->second.object_notifications;
+/// \param client The client to push notifications to.
+/// \param object_info The notifications.
+Status PlasmaStore::SendNotifications(
+    const std::shared_ptr<Client>& client, const std::vector<ObjectInfoT> &object_info) {
+  namespace protocol = ray::object_manager::protocol;
+  flatbuffers::FlatBufferBuilder fbb;
+  std::vector<flatbuffers::Offset<protocol::ObjectInfo>> info;
+  for (size_t i = 0; i < object_info.size(); ++i) {
+    info.push_back(protocol::CreateObjectInfo(fbb, &object_info[i]));
+  }
+  auto info_array = fbb.CreateVector(info);
+  auto message = protocol::CreatePlasmaNotification(fbb, info_array);
+  fbb.Finish(message);
+
+  RAY_LOG(DEBUG) << "Send notifications to fd = " << client->fd;
+  auto& notifications = client->object_notifications;
+  auto new_notifications =
+      std::unique_ptr<uint8_t[]>(new uint8_t[sizeof(int64_t) + fbb.GetSize()]);
+  *(reinterpret_cast<int64_t*>(new_notifications.get())) = fbb.GetSize();
+  memcpy(new_notifications.get() + sizeof(int64_t), fbb.GetBufferPointer(), fbb.GetSize());
+  notifications.emplace_back(std::move(new_notifications));
 
   int num_processed = 0;
   bool closed = false;
@@ -782,7 +787,7 @@ PlasmaStore::NotificationMap::iterator PlasmaStore::SendNotifications(
     int64_t size = *(reinterpret_cast<int64_t*>(notification.get()));
 
     // Attempt to send a notification about this object ID.
-    ssize_t nbytes = send(client_fd, notification.get(), sizeof(int64_t) + size, 0);
+    ssize_t nbytes = send(client->fd, notification.get(), sizeof(int64_t) + size, 0);
     if (nbytes >= 0) {
       RAY_CHECK(nbytes == static_cast<ssize_t>(sizeof(int64_t)) + size);
     } else if (nbytes == -1 &&
@@ -795,12 +800,16 @@ PlasmaStore::NotificationMap::iterator PlasmaStore::SendNotifications(
       // at the end of the method.
       // TODO(pcm): Introduce status codes and check in case the file descriptor
       // is added twice.
-      loop_->AddFileEvent(client_fd, kEventLoopWrite, [this, client_fd](int events) {
-        SendNotifications(pending_notifications_.find(client_fd));
+      loop_->AddFileEvent(client->fd, kEventLoopWrite, [this, client](int events) {
+        Status s = SendNotifications(client, {});
+        if (!s.ok()) {
+          notification_clients_.erase(client);
+        }
       });
       break;
     } else {
-      RAY_LOG(WARNING) << "Failed to send notification to client on fd " << client_fd;
+      RAY_LOG(WARNING) << "Failed to send notification to client on fd " << client->fd
+                       << ", errno = " << errno;
       if (errno == EPIPE) {
         closed = true;
         break;
@@ -813,16 +822,10 @@ PlasmaStore::NotificationMap::iterator PlasmaStore::SendNotifications(
 
   // If we have sent all notifications, remove the fd from the event loop.
   if (notifications.empty()) {
-    loop_->RemoveFileEvent(client_fd);
+    loop_->RemoveFileEvent(client->fd);
   }
-
   // Stop sending notifications if the pipe was broken.
-  if (closed) {
-    close(client_fd);
-    return pending_notifications_.erase(it);
-  } else {
-    return ++it;
-  }
+  return closed ? Status::IOError("Send notifications failed") : Status::OK();
 }
 
 void PlasmaStore::PushNotification(ObjectInfoT* object_info) {
@@ -840,44 +843,31 @@ void PlasmaStore::PushNotifications(const std::vector<ObjectInfoT>& object_info)
     }
   }
 
-  auto it = pending_notifications_.begin();
-  while (it != pending_notifications_.end()) {
-    auto notifications = CreatePlasmaNotificationBuffer(object_info);
-    it->second.object_notifications.emplace_back(std::move(notifications));
-    it = SendNotifications(it);
-  }
-}
-
-void PlasmaStore::PushNotification(ObjectInfoT* object_info, int client_fd) {
-  auto it = pending_notifications_.find(client_fd);
-  if (it != pending_notifications_.end()) {
-    auto notification = CreatePlasmaNotificationBuffer({*object_info});
-    it->second.object_notifications.emplace_back(std::move(notification));
-    SendNotifications(it);
+  auto it = notification_clients_.begin();
+  while (it != notification_clients_.end()) {
+    Status s = SendNotifications(*it, object_info);
+    if (s.ok()) {
+      ++it;
+    } else {
+      it = notification_clients_.erase(it);
+    }
   }
 }
 
 // Subscribe to notifications about sealed objects.
 void PlasmaStore::SubscribeToUpdates(const std::shared_ptr<Client> &client) {
   RAY_LOG(DEBUG) << "subscribing to updates on fd " << client->fd;
-  if (client->notification_fd > 0) {
-    // This client has already subscribed. Return.
-    return;
-  }
-
-  // TODO(rkn): The store could block here if the client doesn't send a file
-  // descriptor.
-  int fd = recv_fd(client->fd);
-  if (fd < 0) {
-    // This may mean that the client died before sending the file descriptor.
-    RAY_LOG(WARNING) << "Failed to receive file descriptor from client on fd "
-                       << client << ".";
-    return;
-  }
-
   // Add this fd to global map, which is needed for this client to receive notifications.
-  pending_notifications_[fd];
-  client->notification_fd = fd;
+  notification_clients_.insert(client);
+
+  // Make the socket non-blocking.
+#ifdef _WINSOCKAPI_
+  unsigned long value = 1;
+  RAY_CHECK(ioctlsocket(client->fd, FIONBIO, &value) == 0);
+#else
+  int flags = fcntl(client->fd, F_GETFL, 0);
+  RAY_CHECK(fcntl(client->fd, F_SETFL, flags | O_NONBLOCK) == 0);
+#endif
 
   // Push notifications to the new subscriber about existing sealed objects.
   for (const auto& entry : store_info_.objects) {
@@ -886,7 +876,10 @@ void PlasmaStore::SubscribeToUpdates(const std::shared_ptr<Client> &client) {
       info.object_id = entry.first.Binary();
       info.data_size = entry.second->data_size;
       info.metadata_size = entry.second->metadata_size;
-      PushNotification(&info, fd);
+      Status s = SendNotifications(client, {info});
+      if (!s.ok()) {
+        notification_clients_.erase(client);
+      }
     }
   }
 }

src/ray/object_manager/plasma/store.h

@@ -47,16 +47,8 @@ using flatbuf::PlasmaError;
 
 struct GetRequest;
 
-struct NotificationQueue {
-  /// The object notifications for clients. We notify the client about the
-  /// objects in the order that the objects were sealed or deleted.
-  std::deque<std::unique_ptr<uint8_t[]>> object_notifications;
-};
-
 class PlasmaStore {
  public:
-  using NotificationMap = std::unordered_map<int, NotificationQueue>;
-
   // TODO: PascalCase PlasmaStore methods.
   PlasmaStore(EventLoop* loop, std::string directory, bool hugepages_enabled,
               const std::string& socket_name,
@@ -167,7 +159,8 @@ class PlasmaStore {
   /// \param client The client that is disconnected.
   void DisconnectClient(const std::shared_ptr<Client> &client);
 
-  NotificationMap::iterator SendNotifications(NotificationMap::iterator it);
+  Status SendNotifications(
+      const std::shared_ptr<Client>& client, const std::vector<ObjectInfoT> &object_info);
 
   Status ProcessMessage(const std::shared_ptr<Client> &client);
 
@@ -193,8 +186,6 @@ class PlasmaStore {
 
   void PushNotifications(const std::vector<ObjectInfoT>& object_notifications);
 
-  void PushNotification(ObjectInfoT* object_notification, int client_fd);
-
   void AddToClientObjectIds(const ObjectID& object_id, ObjectTableEntry* entry,
                             const std::shared_ptr<Client> &client);
 
@@ -239,12 +230,8 @@ class PlasmaStore {
   /// A hash table mapping object IDs to a vector of the get requests that are
   /// waiting for the object to arrive.
   std::unordered_map<ObjectID, std::vector<GetRequest*>> object_get_requests_;
-  /// The pending notifications that have not been sent to subscribers because
-  /// the socket send buffers were full. This is a hash table from client file
-  /// descriptor to an array of object_ids to send to that client.
-  /// TODO(pcm): Consider putting this into the Client data structure and
-  /// reorganize the code slightly.
-  NotificationMap pending_notifications_;
+  /// The registered client for receiving notifications.
+  std::unordered_set<std::shared_ptr<Client>> notification_clients_;
 
   std::unordered_set<ObjectID> deletion_cache_;