From 7edd1e6694690dade0ea3013f2ac1e5eb70c2552 Mon Sep 17 00:00:00 2001 From: "Siyuan (Ryans) Zhuang" Date: Sat, 18 Jul 2020 15:47:52 -0700 Subject: [PATCH] [Core] Remove socket pair exchange in Plasma Store (#9565) * try use boost::asio for notification processing --- .../object_store_notification_manager_ipc.cc | 96 +++++--------- .../object_store_notification_manager_ipc.h | 17 +-- src/ray/object_manager/plasma/client.cc | 110 ---------------- src/ray/object_manager/plasma/client.h | 24 ---- src/ray/object_manager/plasma/io.h | 8 +- src/ray/object_manager/plasma/plasma.cc | 38 ------ src/ray/object_manager/plasma/plasma.h | 15 +-- src/ray/object_manager/plasma/protocol.cc | 8 -- src/ray/object_manager/plasma/protocol.h | 4 - src/ray/object_manager/plasma/store.cc | 119 +++++++++--------- src/ray/object_manager/plasma/store.h | 21 +--- 11 files changed, 102 insertions(+), 358 deletions(-) diff --git a/src/ray/object_manager/notification/object_store_notification_manager_ipc.cc b/src/ray/object_manager/notification/object_store_notification_manager_ipc.cc index dbd96ce9f..b43600cb5 100644 --- a/src/ray/object_manager/notification/object_store_notification_manager_ipc.cc +++ b/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 -#include -#include - -#include "ray/common/common_protocol.h" -#include "ray/common/status.h" -#include "ray/util/util.h" - -#ifdef _WIN32 -#include -#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(&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(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; diff --git a/src/ray/object_manager/notification/object_store_notification_manager_ipc.h b/src/ray/object_manager/notification/object_store_notification_manager_ipc.h index 885fa0c03..d7475cfd0 100644 --- a/src/ray/object_manager/notification/object_store_notification_manager_ipc.h +++ b/src/ray/object_manager/notification/object_store_notification_manager_ipc.h @@ -14,22 +14,16 @@ #pragma once -#include #include #include -#include -#include -#include - -#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 store_client_; int64_t length_; std::vector 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 diff --git a/src/ray/object_manager/plasma/client.cc b/src/ray/object_manager/plasma/client.cc index 440f71072..f4af036c9 100644 --- a/src/ray/object_manager/plasma/client.cc +++ b/src/ray/object_manager/plasma/client.cc @@ -257,15 +257,6 @@ class PlasmaClient::Impl : public std::enable_shared_from_this& 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* object_ids, - std::vector* data_sizes, - std::vector* metadata_sizes); - Status Disconnect(); std::string DebugString(); @@ -817,93 +808,6 @@ Status PlasmaClient::Impl::Refresh(const std::vector& object_ids) { return ReadRefreshLRUReply(buffer.data(), buffer.size()); } -Status PlasmaClient::Impl::Subscribe(int* fd) { - std::lock_guard 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 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 object_ids; - std::vector data_sizes; - std::vector 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* object_ids, - std::vector* data_sizes, - std::vector* metadata_sizes) { - std::lock_guard guard(client_mutex_); - auto object_info = flatbuffers::GetRoot(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& 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* object_ids, - std::vector* data_sizes, - std::vector* 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(); } diff --git a/src/ray/object_manager/plasma/client.h b/src/ray/object_manager/plasma/client.h index 34bc6d4ab..f67b52d2f 100644 --- a/src/ray/object_manager/plasma/client.h +++ b/src/ray/object_manager/plasma/client.h @@ -205,30 +205,6 @@ class RAY_EXPORT PlasmaClient { /// \return The return status. Status Refresh(const std::vector& 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* object_ids, - std::vector* data_sizes, - std::vector* metadata_sizes); - /// Disconnect from the local plasma instance, including the local store and /// manager. /// diff --git a/src/ray/object_manager/plasma/io.h b/src/ray/object_manager/plasma/io.h index 45b54cc76..d9d38edb6 100644 --- a/src/ray/object_manager/plasma/io.h +++ b/src/ray/object_manager/plasma/io.h @@ -28,6 +28,7 @@ #include #include +#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); diff --git a/src/ray/object_manager/plasma/plasma.cc b/src/ray/object_manager/plasma/plasma.cc index 5a5ad49d3..fd866dcee 100644 --- a/src/ray/object_manager/plasma/plasma.cc +++ b/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 CreateObjectInfoBuffer(fb::ObjectInfoT* object_info) { - flatbuffers::FlatBufferBuilder fbb; - auto message = fb::CreateObjectInfo(fbb, object_info); - fbb.Finish(message); - auto notification = - std::unique_ptr(new uint8_t[sizeof(int64_t) + fbb.GetSize()]); - *(reinterpret_cast(notification.get())) = fbb.GetSize(); - memcpy(notification.get() + sizeof(int64_t), fbb.GetBufferPointer(), fbb.GetSize()); - return notification; -} - -std::unique_ptr CreatePlasmaNotificationBuffer( - const std::vector& object_info) { - flatbuffers::FlatBufferBuilder fbb; - std::vector> 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(new uint8_t[sizeof(int64_t) + fbb.GetSize()]); - *(reinterpret_cast(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); diff --git a/src/ray/object_manager/plasma/plasma.h b/src/ray/object_manager/plasma/plasma.h index 428259786..4f2ff68a0 100644 --- a/src/ray/object_manager/plasma/plasma.h +++ b/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 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> object_notifications; }; std::ostream &operator<<(std::ostream &os, const std::shared_ptr &client); @@ -182,11 +182,6 @@ ObjectTableEntry* GetObjectTableEntry(PlasmaStoreInfo* store_info, /// \return The errno set. int WarnIfSigpipe(int status, int client_sock); -std::unique_ptr CreateObjectInfoBuffer(ObjectInfoT* object_info); - -std::unique_ptr CreatePlasmaNotificationBuffer( - const std::vector& object_info); - /// Globally accessible reference to plasma store configuration. extern const PlasmaStoreInfo* plasma_config; diff --git a/src/ray/object_manager/plasma/protocol.cc b/src/ray/object_manager/plasma/protocol.cc index fea821001..cbfe53a80 100644 --- a/src/ray/object_manager/plasma/protocol.cc +++ b/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 &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 &store_conn, ObjectID object_id, const char* address, int port) { diff --git a/src/ray/object_manager/plasma/protocol.h b/src/ray/object_manager/plasma/protocol.h index 070f41b11..f16d38d38 100644 --- a/src/ray/object_manager/plasma/protocol.h +++ b/src/ray/object_manager/plasma/protocol.h @@ -177,10 +177,6 @@ Status SendEvictReply(const std::shared_ptr &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 &store_conn); - /* Data messages. */ Status SendDataRequest(const std::shared_ptr &store_conn, ObjectID object_id, const char* address, int port); diff --git a/src/ray/object_manager/plasma/store.cc b/src/ray/object_manager/plasma/store.cc index 6aa983294..0d65ffdb6 100644 --- a/src/ray/object_manager/plasma/store.cc +++ b/src/ray/object_manager/plasma/store.cc @@ -742,16 +742,9 @@ void PlasmaStore::DisconnectClient(const std::shared_ptr &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) { /// 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, const std::vector &object_info) { + namespace protocol = ray::object_manager::protocol; + flatbuffers::FlatBufferBuilder fbb; + std::vector> 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(new uint8_t[sizeof(int64_t) + fbb.GetSize()]); + *(reinterpret_cast(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(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(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& 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) { 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) { 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); + } } } } diff --git a/src/ray/object_manager/plasma/store.h b/src/ray/object_manager/plasma/store.h index cf037ecb8..ee6dc9bb6 100644 --- a/src/ray/object_manager/plasma/store.h +++ b/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> object_notifications; -}; - class PlasmaStore { public: - using NotificationMap = std::unordered_map; - // 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); - NotificationMap::iterator SendNotifications(NotificationMap::iterator it); + Status SendNotifications( + const std::shared_ptr& client, const std::vector &object_info); Status ProcessMessage(const std::shared_ptr &client); @@ -193,8 +186,6 @@ class PlasmaStore { void PushNotifications(const std::vector& object_notifications); - void PushNotification(ObjectInfoT* object_notification, int client_fd); - void AddToClientObjectIds(const ObjectID& object_id, ObjectTableEntry* entry, const std::shared_ptr &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> 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> notification_clients_; std::unordered_set deletion_cache_;