diff --git a/BUILD.bazel b/BUILD.bazel index 43740df7e..04ba84223 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -278,9 +278,10 @@ cc_library( "src/ray/object_manager/plasma/protocol.cc", ], hdrs = [ + "src/ray/object_manager/format/object_manager_generated.h", + "src/ray/object_manager/notification/object_store_notification_manager.h", "src/ray/object_manager/plasma/client.h", "src/ray/object_manager/plasma/common.h", - "src/ray/object_manager/plasma/common_generated.h", "src/ray/object_manager/plasma/compat.h", "src/ray/object_manager/plasma/external_store.h", "src/ray/object_manager/plasma/fling.h", @@ -298,7 +299,7 @@ cc_library( linkopts = PLASMA_LINKOPTS, strip_include_prefix = "src", deps = [ - ":common_fbs", + ":object_manager_fbs", ":plasma_fbs", ":platform_shims", ":ray_common", @@ -423,18 +424,11 @@ FLATC_ARGS = [ "--scoped-enums", ] -flatbuffer_cc_library( - name = "common_fbs", - srcs = ["src/ray/object_manager/plasma/common.fbs"], - flatc_args = FLATC_ARGS, - out_prefix = "src/ray/object_manager/plasma/", -) - flatbuffer_cc_library( name = "plasma_fbs", srcs = ["src/ray/object_manager/plasma/plasma.fbs"], flatc_args = FLATC_ARGS, - includes = ["src/ray/object_manager/plasma/common.fbs"], + includes = ["src/ray/object_manager/format/object_manager.fbs"], out_prefix = "src/ray/object_manager/plasma/", ) @@ -1197,9 +1191,11 @@ cc_library( name = "object_manager", srcs = glob([ "src/ray/object_manager/*.cc", + "src/ray/object_manager/notification/*.cc", ]), hdrs = glob([ "src/ray/object_manager/*.h", + "src/ray/object_manager/notification/*.h", ]), copts = COPTS, includes = [ diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 816b41439..e5bf4db1e 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -32,7 +32,6 @@ #include "ray/core_worker/transport/raylet_transport.h" #include "ray/gcs/redis_gcs_client.h" #include "ray/gcs/subscription_executor.h" -#include "ray/object_manager/object_store_notification_manager.h" #include "ray/raylet/raylet_client.h" #include "ray/rpc/node_manager/node_manager_client.h" #include "ray/rpc/worker/core_worker_client.h" diff --git a/src/ray/object_manager/format/object_manager.fbs b/src/ray/object_manager/format/object_manager.fbs index 50ed3d341..887c87bbe 100644 --- a/src/ray/object_manager/format/object_manager.fbs +++ b/src/ray/object_manager/format/object_manager.fbs @@ -16,9 +16,6 @@ namespace ray.object_manager.protocol; // Object information data structure. -// NOTE(pcm): This structure is replicated in -// https://github.com/apache/arrow/blob/master/cpp/src/plasma/common.fbs, -// so if you modify it, you should also modify that one. table ObjectInfo { // Object ID of this object. object_id: string; @@ -36,9 +33,6 @@ table ObjectInfo { is_deletion: bool; } -// NOTE(pcm): This structure is replicated in -// https://github.com/apache/arrow/blob/master/cpp/src/plasma/plasma.fbs -// so if you modify it, you should also modify that one. table PlasmaNotification { object_info: [ObjectInfo]; } diff --git a/src/ray/object_manager/notification/object_store_notification_manager.h b/src/ray/object_manager/notification/object_store_notification_manager.h new file mode 100644 index 000000000..623da175b --- /dev/null +++ b/src/ray/object_manager/notification/object_store_notification_manager.h @@ -0,0 +1,103 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef RAY_OBJECT_STORE_NOTIFICATION_MANAGER_H +#define RAY_OBJECT_STORE_NOTIFICATION_MANAGER_H + +#include +#include +#include + +#include + +#include "absl/synchronization/mutex.h" + +#include "ray/common/id.h" +#include "ray/common/status.h" +#include "ray/object_manager/format/object_manager_generated.h" + +namespace ray { + +/// \class ObjectStoreNotificationManager +/// +/// Encapsulates notification handling from the object store. +class ObjectStoreNotificationManager { + public: + ObjectStoreNotificationManager(boost::asio::io_service &io_service) + : main_service_(&io_service), num_adds_processed_(0), num_removes_processed_(0) {} + virtual ~ObjectStoreNotificationManager() {} + /// Subscribe to notifications of objects added to local store. + /// Upon subscribing, the callback will be invoked for all objects that + /// already exist in the local store + /// + /// \param callback A callback expecting an ObjectID. + void SubscribeObjAdded( + std::function callback) { + absl::MutexLock lock(&store_add_mutex_); + add_handlers_.push_back(std::move(callback)); + } + + /// Subscribe to notifications of objects deleted from local store. + /// + /// \param callback A callback expecting an ObjectID. + void SubscribeObjDeleted(std::function callback) { + absl::MutexLock lock(&store_remove_mutex_); + rem_handlers_.push_back(std::move(callback)); + } + + /// Support for rebroadcasting object add/rem events. + void ProcessStoreAdd(const object_manager::protocol::ObjectInfoT &object_info) { + // TODO(suquark): Use strand in boost asio to enforce sequential execution. + absl::MutexLock lock(&store_add_mutex_); + for (auto &handler : add_handlers_) { + main_service_->post([handler, object_info]() { handler(object_info); }); + } + num_adds_processed_++; + } + + void ProcessStoreRemove(const ObjectID &object_id) { + absl::MutexLock lock(&store_remove_mutex_); + for (auto &handler : rem_handlers_) { + main_service_->post([handler, object_id]() { handler(object_id); }); + } + num_removes_processed_++; + } + + /// Returns debug string for class. + /// + /// \return string. + std::string DebugString() const { + std::stringstream result; + result << "ObjectStoreNotificationManager:"; + result << "\n- num adds processed: " << num_adds_processed_; + result << "\n- num removes processed: " << num_removes_processed_; + return result.str(); + } + + private: + /// Weak reference to main service. We ensure this object is destroyed before + /// main_service_ is stopped. + boost::asio::io_service *main_service_; + std::vector> + add_handlers_; + std::vector> rem_handlers_; + absl::Mutex store_add_mutex_; + absl::Mutex store_remove_mutex_; + int64_t num_adds_processed_; + int64_t num_removes_processed_; +}; + +} // namespace ray + +#endif // RAY_OBJECT_STORE_NOTIFICATION_MANAGER_H diff --git a/src/ray/object_manager/object_store_notification_manager.cc b/src/ray/object_manager/notification/object_store_notification_manager_ipc.cc similarity index 70% rename from src/ray/object_manager/object_store_notification_manager.cc rename to src/ray/object_manager/notification/object_store_notification_manager_ipc.cc index 9d612b4c1..5419ee289 100644 --- a/src/ray/object_manager/object_store_notification_manager.cc +++ b/src/ray/object_manager/notification/object_store_notification_manager_ipc.cc @@ -12,13 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "ray/object_manager/object_store_notification_manager.h" +#include "ray/object_manager/notification/object_store_notification_manager_ipc.h" #include #include #include -#include -#include #include "ray/common/common_protocol.h" #include "ray/common/status.h" @@ -30,13 +28,12 @@ namespace ray { -ObjectStoreNotificationManager::ObjectStoreNotificationManager( +ObjectStoreNotificationManagerIPC::ObjectStoreNotificationManagerIPC( boost::asio::io_service &io_service, const std::string &store_socket_name, bool exit_on_error) - : store_client_(), + : ObjectStoreNotificationManager(io_service), + store_client_(), length_(0), - num_adds_processed_(0), - num_removes_processed_(0), socket_(io_service), exit_on_error_(exit_on_error) { RAY_ARROW_CHECK_OK(store_client_.Connect(store_socket_name.c_str(), "", 0, 300)); @@ -73,21 +70,22 @@ ObjectStoreNotificationManager::ObjectStoreNotificationManager( NotificationWait(); } -ObjectStoreNotificationManager::~ObjectStoreNotificationManager() { +ObjectStoreNotificationManagerIPC::~ObjectStoreNotificationManagerIPC() { RAY_ARROW_CHECK_OK(store_client_.Disconnect()); } -void ObjectStoreNotificationManager::Shutdown() { +void ObjectStoreNotificationManagerIPC::Shutdown() { RAY_ARROW_CHECK_OK(store_client_.Disconnect()); } -void ObjectStoreNotificationManager::NotificationWait() { - boost::asio::async_read(socket_, boost::asio::buffer(&length_, sizeof(length_)), - boost::bind(&ObjectStoreNotificationManager::ProcessStoreLength, - this, boost::asio::placeholders::error)); +void ObjectStoreNotificationManagerIPC::NotificationWait() { + boost::asio::async_read( + socket_, boost::asio::buffer(&length_, sizeof(length_)), + boost::bind(&ObjectStoreNotificationManagerIPC::ProcessStoreLength, this, + boost::asio::placeholders::error)); } -void ObjectStoreNotificationManager::ProcessStoreLength( +void ObjectStoreNotificationManagerIPC::ProcessStoreLength( const boost::system::error_code &error) { notification_.resize(length_); if (error) { @@ -113,11 +111,11 @@ void ObjectStoreNotificationManager::ProcessStoreLength( boost::asio::async_read( socket_, boost::asio::buffer(notification_), - boost::bind(&ObjectStoreNotificationManager::ProcessStoreNotification, this, + boost::bind(&ObjectStoreNotificationManagerIPC::ProcessStoreNotification, this, boost::asio::placeholders::error)); } -void ObjectStoreNotificationManager::ProcessStoreNotification( +void ObjectStoreNotificationManagerIPC::ProcessStoreNotification( const boost::system::error_code &error) { if (error) { if (exit_on_error_) { @@ -152,37 +150,4 @@ void ObjectStoreNotificationManager::ProcessStoreNotification( NotificationWait(); } -void ObjectStoreNotificationManager::ProcessStoreAdd( - const object_manager::protocol::ObjectInfoT &object_info) { - for (auto &handler : add_handlers_) { - handler(object_info); - } - num_adds_processed_++; -} - -void ObjectStoreNotificationManager::ProcessStoreRemove(const ObjectID &object_id) { - for (auto &handler : rem_handlers_) { - handler(object_id); - } - num_removes_processed_++; -} - -void ObjectStoreNotificationManager::SubscribeObjAdded( - std::function callback) { - add_handlers_.push_back(std::move(callback)); -} - -void ObjectStoreNotificationManager::SubscribeObjDeleted( - std::function callback) { - rem_handlers_.push_back(std::move(callback)); -} - -std::string ObjectStoreNotificationManager::DebugString() const { - std::stringstream result; - result << "ObjectStoreNotificationManager:"; - result << "\n- num adds processed: " << num_adds_processed_; - result << "\n- num removes processed: " << num_removes_processed_; - return result.str(); -} - } // namespace ray diff --git a/src/ray/object_manager/object_store_notification_manager.h b/src/ray/object_manager/notification/object_store_notification_manager_ipc.h similarity index 58% rename from src/ray/object_manager/object_store_notification_manager.h rename to src/ray/object_manager/notification/object_store_notification_manager_ipc.h index 204cb8e96..885fa0c03 100644 --- a/src/ray/object_manager/object_store_notification_manager.h +++ b/src/ray/object_manager/notification/object_store_notification_manager_ipc.h @@ -25,15 +25,15 @@ #include "ray/common/client_connection.h" #include "ray/common/id.h" #include "ray/common/status.h" -#include "ray/object_manager/object_directory.h" +#include "ray/object_manager/notification/object_store_notification_manager.h" #include "ray/object_manager/plasma/client.h" namespace ray { -/// \class ObjectStoreClientPool +/// \class ObjectStoreNotificationManagerIPC /// /// Encapsulates notification handling from the object store. -class ObjectStoreNotificationManager { +class ObjectStoreNotificationManagerIPC : public ObjectStoreNotificationManager { public: /// Constructor. /// @@ -41,51 +41,23 @@ class ObjectStoreNotificationManager { /// \param store_socket_name The store socket to connect to. /// \param exit_on_error The manager will exit with error when it fails /// to process messages from socket. - ObjectStoreNotificationManager(boost::asio::io_service &io_service, - const std::string &store_socket_name, - bool exit_on_error = true); + ObjectStoreNotificationManagerIPC(boost::asio::io_service &io_service, + const std::string &store_socket_name, + bool exit_on_error = true); - ~ObjectStoreNotificationManager(); - - /// Subscribe to notifications of objects added to local store. - /// Upon subscribing, the callback will be invoked for all objects that - /// already exist in the local store - /// - /// \param callback A callback expecting an ObjectID. - void SubscribeObjAdded( - std::function callback); - - /// Subscribe to notifications of objects deleted from local store. - /// - /// \param callback A callback expecting an ObjectID. - void SubscribeObjDeleted(std::function callback); + ~ObjectStoreNotificationManagerIPC() override; /// Explicitly shutdown the manager. void Shutdown(); - /// Returns debug string for class. - /// - /// \return string. - std::string DebugString() const; - 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); - /// Support for rebroadcasting object add/rem events. - void ProcessStoreAdd(const object_manager::protocol::ObjectInfoT &object_info); - void ProcessStoreRemove(const ObjectID &object_id); - - std::vector> - add_handlers_; - std::vector> rem_handlers_; - plasma::PlasmaClient store_client_; int64_t length_; - int64_t num_adds_processed_; - int64_t num_removes_processed_; std::vector notification_; local_stream_socket socket_; diff --git a/src/ray/object_manager/object_buffer_pool.h b/src/ray/object_manager/object_buffer_pool.h index 851645469..e5c036795 100644 --- a/src/ray/object_manager/object_buffer_pool.h +++ b/src/ray/object_manager/object_buffer_pool.h @@ -23,8 +23,6 @@ #include #include -#include "ray/object_manager/plasma/client.h" - #include "ray/common/id.h" #include "ray/common/status.h" #include "ray/object_manager/plasma/client.h" diff --git a/src/ray/object_manager/object_directory.h b/src/ray/object_manager/object_directory.h index ad391bfd0..426992f57 100644 --- a/src/ray/object_manager/object_directory.h +++ b/src/ray/object_manager/object_directory.h @@ -24,7 +24,6 @@ #include "ray/common/status.h" #include "ray/gcs/redis_gcs_client.h" #include "ray/object_manager/format/object_manager_generated.h" -#include "ray/object_manager/plasma/client.h" namespace ray { diff --git a/src/ray/object_manager/object_manager.cc b/src/ray/object_manager/object_manager.cc index 53f7fc378..3c20987c6 100644 --- a/src/ray/object_manager/object_manager.cc +++ b/src/ray/object_manager/object_manager.cc @@ -55,7 +55,6 @@ ObjectManager::ObjectManager(asio::io_service &main_service, const ClientID &sel config_(config), object_directory_(std::move(object_directory)), object_store_internal_(config), - store_notification_(main_service, config_.store_socket_name), buffer_pool_(config_.store_socket_name, config_.object_chunk_size), rpc_work_(rpc_service_), gen_(std::chrono::high_resolution_clock::now().time_since_epoch().count()), @@ -65,11 +64,20 @@ ObjectManager::ObjectManager(asio::io_service &main_service, const ClientID &sel client_call_manager_(main_service, config_.rpc_service_threads_number) { RAY_CHECK(config_.rpc_service_threads_number > 0); main_service_ = &main_service; - store_notification_.SubscribeObjAdded( + + if (plasma::plasma_store_runner) { + store_notification_ = std::make_shared(main_service); + plasma::plasma_store_runner->SetNotificationListener(store_notification_); + } else { + store_notification_ = std::make_shared( + main_service, config_.store_socket_name); + } + + store_notification_->SubscribeObjAdded( [this](const object_manager::protocol::ObjectInfoT &object_info) { HandleObjectAdded(object_info); }); - store_notification_.SubscribeObjDeleted( + store_notification_->SubscribeObjDeleted( [this](const ObjectID &oid) { NotifyDirectoryObjectDeleted(oid); }); // Start object manager rpc server and send & receive request threads @@ -144,13 +152,13 @@ void ObjectManager::NotifyDirectoryObjectDeleted(const ObjectID &object_id) { ray::Status ObjectManager::SubscribeObjAdded( std::function callback) { - store_notification_.SubscribeObjAdded(callback); + store_notification_->SubscribeObjAdded(callback); return ray::Status::OK(); } ray::Status ObjectManager::SubscribeObjDeleted( std::function callback) { - store_notification_.SubscribeObjDeleted(callback); + store_notification_->SubscribeObjDeleted(callback); return ray::Status::OK(); } @@ -875,7 +883,7 @@ std::string ObjectManager::DebugString() const { result << "\n- num pull requests: " << pull_requests_.size(); result << "\n- num buffered profile events: " << profile_events_.size(); result << "\n" << object_directory_->DebugString(); - result << "\n" << store_notification_.DebugString(); + result << "\n" << store_notification_->DebugString(); result << "\n" << buffer_pool_.DebugString(); return result.str(); } diff --git a/src/ray/object_manager/object_manager.h b/src/ray/object_manager/object_manager.h index a193639cd..4928ea7af 100644 --- a/src/ray/object_manager/object_manager.h +++ b/src/ray/object_manager/object_manager.h @@ -32,10 +32,9 @@ #include "ray/common/ray_config.h" #include "ray/common/status.h" #include "ray/object_manager/format/object_manager_generated.h" +#include "ray/object_manager/notification/object_store_notification_manager_ipc.h" #include "ray/object_manager/object_buffer_pool.h" #include "ray/object_manager/object_directory.h" -#include "ray/object_manager/object_store_notification_manager.h" -#include "ray/object_manager/plasma/client.h" #include "ray/object_manager/plasma/store_runner.h" #include "ray/rpc/object_manager/object_manager_client.h" #include "ray/rpc/object_manager/object_manager_server.h" @@ -391,7 +390,9 @@ class ObjectManager : public ObjectManagerInterface, std::shared_ptr object_directory_; // Object store runner. ObjectStoreRunner object_store_internal_; - ObjectStoreNotificationManager store_notification_; + // Process notifications from Plasma. We make it a shared pointer because + // we will decide its type at runtime, and we would pass it to Plasma Store. + std::shared_ptr store_notification_; ObjectBufferPool buffer_pool_; /// Weak reference to main service. We ensure this object is destroyed before diff --git a/src/ray/object_manager/plasma/client.cc b/src/ray/object_manager/plasma/client.cc index f816792f8..edbed32ae 100644 --- a/src/ray/object_manager/plasma/client.cc +++ b/src/ray/object_manager/plasma/client.cc @@ -1035,7 +1035,7 @@ Status PlasmaClient::Impl::DecodeNotifications(const uint8_t* buffer, std::vector* data_sizes, std::vector* metadata_sizes) { std::lock_guard guard(client_mutex_); - auto object_info = flatbuffers::GetRoot(buffer); + 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); diff --git a/src/ray/object_manager/plasma/common.fbs b/src/ray/object_manager/plasma/common.fbs deleted file mode 100644 index 299282774..000000000 --- a/src/ray/object_manager/plasma/common.fbs +++ /dev/null @@ -1,36 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -namespace plasma.flatbuf; - -// Object information data structure. -table ObjectInfo { - // Object ID of this object. - object_id: string; - // Number of bytes the content of this object occupies in memory. - data_size: long; - // Number of bytes the metadata of this object occupies in memory. - metadata_size: long; - // Number of clients using the objects. - ref_count: int; - // Unix epoch of when this object was created. - create_time: long; - // How long creation of this object took. - construct_duration: long; - // Specifies if this object was deleted or added. - is_deletion: bool; -} diff --git a/src/ray/object_manager/plasma/plasma.cc b/src/ray/object_manager/plasma/plasma.cc index b65588b6a..55afa22e0 100644 --- a/src/ray/object_manager/plasma/plasma.cc +++ b/src/ray/object_manager/plasma/plasma.cc @@ -21,11 +21,11 @@ #include #include +#include "ray/object_manager/format/object_manager_generated.h" #include "ray/object_manager/plasma/common.h" -#include "ray/object_manager/plasma/common_generated.h" #include "ray/object_manager/plasma/protocol.h" -namespace fb = plasma::flatbuf; +namespace fb = ray::object_manager::protocol; namespace plasma { @@ -70,9 +70,9 @@ std::unique_ptr CreateObjectInfoBuffer(fb::ObjectInfoT* object_info) } std::unique_ptr CreatePlasmaNotificationBuffer( - std::vector& object_info) { + const std::vector& object_info) { flatbuffers::FlatBufferBuilder fbb; - std::vector> info; + std::vector> info; for (size_t i = 0; i < object_info.size(); ++i) { info.push_back(fb::CreateObjectInfo(fbb, &object_info[i])); } diff --git a/src/ray/object_manager/plasma/plasma.fbs b/src/ray/object_manager/plasma/plasma.fbs index 10475364b..41de16556 100644 --- a/src/ray/object_manager/plasma/plasma.fbs +++ b/src/ray/object_manager/plasma/plasma.fbs @@ -15,8 +15,6 @@ // specific language governing permissions and limitations // under the License. -include "common.fbs"; - // Plasma protocol specification namespace plasma.flatbuf; @@ -311,10 +309,6 @@ table PlasmaEvictReply { table PlasmaSubscribeRequest { } -table PlasmaNotification { - object_info: [ObjectInfo]; -} - table PlasmaDataRequest { // ID of the object that is requested. object_id: string; diff --git a/src/ray/object_manager/plasma/plasma.h b/src/ray/object_manager/plasma/plasma.h index 4971b574c..c9e83a91b 100644 --- a/src/ray/object_manager/plasma/plasma.h +++ b/src/ray/object_manager/plasma/plasma.h @@ -35,6 +35,7 @@ #include "ray/object_manager/plasma/compat.h" #include "arrow/status.h" +#include "ray/object_manager/format/object_manager_generated.h" #include "ray/object_manager/plasma/common.h" #include "ray/util/logging.h" @@ -44,9 +45,7 @@ using arrow::cuda::CudaIpcMemHandle; namespace plasma { -namespace flatbuf { -struct ObjectInfoT; -} // namespace flatbuf +using ray::object_manager::protocol::ObjectInfoT; #define HANDLE_SIGPIPE(s, fd_) \ do { \ @@ -166,9 +165,9 @@ ObjectTableEntry* GetObjectTableEntry(PlasmaStoreInfo* store_info, /// \return The errno set. int WarnIfSigpipe(int status, int client_sock); -std::unique_ptr CreateObjectInfoBuffer(flatbuf::ObjectInfoT* object_info); +std::unique_ptr CreateObjectInfoBuffer(ObjectInfoT* object_info); std::unique_ptr CreatePlasmaNotificationBuffer( - std::vector& object_info); + const std::vector& object_info); } // namespace plasma diff --git a/src/ray/object_manager/plasma/store.cc b/src/ray/object_manager/plasma/store.cc index 76bda0f27..e5e169b4a 100644 --- a/src/ray/object_manager/plasma/store.cc +++ b/src/ray/object_manager/plasma/store.cc @@ -43,8 +43,8 @@ #include "arrow/status.h" +#include "ray/object_manager/format/object_manager_generated.h" #include "ray/object_manager/plasma/common.h" -#include "ray/object_manager/plasma/common_generated.h" #include "ray/object_manager/plasma/fling.h" #include "ray/object_manager/plasma/io.h" #include "ray/object_manager/plasma/malloc.h" @@ -644,7 +644,7 @@ PlasmaError PlasmaStore::DeleteObject(ObjectID& object_id) { eviction_policy_.RemoveObject(object_id); EraseFromObjectTable(object_id); // Inform all subscribers that the object has been deleted. - fb::ObjectInfoT notification; + ObjectInfoT notification; notification.object_id = object_id.Binary(); notification.is_deletion = true; PushNotification(¬ification); @@ -683,7 +683,7 @@ void PlasmaStore::EvictObjects(const std::vector& object_ids) { // and send a deletion notification. EraseFromObjectTable(object_id); // Inform all subscribers that the object has been deleted. - fb::ObjectInfoT notification; + ObjectInfoT notification; notification.object_id = object_id.Binary(); notification.is_deletion = true; PushNotification(¬ification); @@ -835,18 +835,21 @@ PlasmaStore::NotificationMap::iterator PlasmaStore::SendNotifications( } } -void PlasmaStore::PushNotification(fb::ObjectInfoT* object_info) { - auto it = pending_notifications_.begin(); - while (it != pending_notifications_.end()) { - std::vector info; - info.push_back(*object_info); - auto notification = CreatePlasmaNotificationBuffer(info); - it->second.object_notifications.emplace_back(std::move(notification)); - it = SendNotifications(it); - } +void PlasmaStore::PushNotification(ObjectInfoT* object_info) { + PushNotifications({*object_info}); } -void PlasmaStore::PushNotifications(std::vector& object_info) { +void PlasmaStore::PushNotifications(const std::vector& object_info) { + if (notification_listener_) { + for (const auto& info : object_info) { + if (!info.is_deletion) { + notification_listener_->ProcessStoreAdd(info); + } else { + notification_listener_->ProcessStoreRemove(ObjectID::FromBinary(info.object_id)); + } + } + } + auto it = pending_notifications_.begin(); while (it != pending_notifications_.end()) { auto notifications = CreatePlasmaNotificationBuffer(object_info); @@ -855,12 +858,10 @@ void PlasmaStore::PushNotifications(std::vector& object_info) { } } -void PlasmaStore::PushNotification(fb::ObjectInfoT* object_info, int client_fd) { +void PlasmaStore::PushNotification(ObjectInfoT* object_info, int client_fd) { auto it = pending_notifications_.find(client_fd); if (it != pending_notifications_.end()) { - std::vector info; - info.push_back(*object_info); - auto notification = CreatePlasmaNotificationBuffer(info); + auto notification = CreatePlasmaNotificationBuffer({*object_info}); it->second.object_notifications.emplace_back(std::move(notification)); SendNotifications(it); } diff --git a/src/ray/object_manager/plasma/store.h b/src/ray/object_manager/plasma/store.h index cbc5f33cd..bc06a78a6 100644 --- a/src/ray/object_manager/plasma/store.h +++ b/src/ray/object_manager/plasma/store.h @@ -24,6 +24,8 @@ #include #include +#include "ray/object_manager/format/object_manager_generated.h" +#include "ray/object_manager/notification/object_store_notification_manager.h" #include "ray/object_manager/plasma/common.h" #include "ray/object_manager/plasma/events.h" #include "ray/object_manager/plasma/external_store.h" @@ -38,11 +40,10 @@ class Status; namespace plasma { namespace flatbuf { -struct ObjectInfoT; enum class PlasmaError; } // namespace flatbuf -using flatbuf::ObjectInfoT; +using ray::object_manager::protocol::ObjectInfoT; using flatbuf::PlasmaError; struct GetRequest; @@ -171,10 +172,27 @@ class PlasmaStore { arrow::Status ProcessMessage(Client* client); + void SetNotificationListener( + const std::shared_ptr ¬ification_listener) { + notification_listener_ = notification_listener; + if (notification_listener_) { + // Push notifications to the new subscriber about existing sealed objects. + for (const auto& entry : store_info_.objects) { + if (entry.second->state == ObjectState::PLASMA_SEALED) { + ObjectInfoT info; + info.object_id = entry.first.Binary(); + info.data_size = entry.second->data_size; + info.metadata_size = entry.second->metadata_size; + notification_listener_->ProcessStoreAdd(info); + } + } + } + } + private: void PushNotification(ObjectInfoT* object_notification); - void PushNotifications(std::vector& object_notifications); + void PushNotifications(const std::vector& object_notifications); void PushNotification(ObjectInfoT* object_notification, int client_fd); @@ -239,6 +257,7 @@ class PlasmaStore { #ifdef PLASMA_CUDA arrow::cuda::CudaDeviceManager* manager_; #endif + std::shared_ptr notification_listener_; }; } // namespace plasma diff --git a/src/ray/object_manager/plasma/store_runner.h b/src/ray/object_manager/plasma/store_runner.h index 942c1e4d3..a9b1b07a3 100644 --- a/src/ray/object_manager/plasma/store_runner.h +++ b/src/ray/object_manager/plasma/store_runner.h @@ -2,6 +2,7 @@ #include +#include "ray/object_manager/notification/object_store_notification_manager.h" #include "ray/object_manager/plasma/store.h" namespace plasma { @@ -14,6 +15,10 @@ class PlasmaStoreRunner { void Start(); void Stop(); void Shutdown(); + void SetNotificationListener( + const std::shared_ptr ¬ification_listener) { + store_->SetNotificationListener(notification_listener); + } private: std::string socket_name_; @@ -23,6 +28,7 @@ class PlasmaStoreRunner { std::string external_store_endpoint_; std::unique_ptr loop_; std::unique_ptr store_; + std::shared_ptr listener_; }; // We use a global variable for Plasma Store instance here because: