[Core] Remove object notification IPC between Plasma and Raylet (initial step) (#8939)

* initial refactoring

redirect notifications to eventloop

implement direct notifications

* protect vector with mutex

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 = [

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"

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];
 }

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 <iostream>
+#include <memory>
+#include <vector>
+
+#include <boost/asio.hpp>
+
+#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<void(const object_manager::protocol::ObjectInfoT &)> 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<void(const ray::ObjectID &)> 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<std::function<void(const object_manager::protocol::ObjectInfoT &)>>
+      add_handlers_;
+  std::vector<std::function<void(const ray::ObjectID &)>> 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

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 <boost/asio.hpp>
 #include <boost/bind.hpp>
 #include <boost/function.hpp>
-#include <future>
-#include <iostream>
 
 #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<void(const object_manager::protocol::ObjectInfoT &)> callback) {
-  add_handlers_.push_back(std::move(callback));
-}
-
-void ObjectStoreNotificationManager::SubscribeObjDeleted(
-    std::function<void(const ObjectID &)> 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

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<void(const object_manager::protocol::ObjectInfoT &)> callback);
-
-  /// Subscribe to notifications of objects deleted from local store.
-  ///
-  /// \param callback A callback expecting an ObjectID.
-  void SubscribeObjDeleted(std::function<void(const ray::ObjectID &)> 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<std::function<void(const object_manager::protocol::ObjectInfoT &)>>
-      add_handlers_;
-  std::vector<std::function<void(const ray::ObjectID &)>> rem_handlers_;
-
   plasma::PlasmaClient store_client_;
   int64_t length_;
-  int64_t num_adds_processed_;
-  int64_t num_removes_processed_;
   std::vector<uint8_t> notification_;
   local_stream_socket socket_;
 

src/ray/object_manager/object_buffer_pool.h

@@ -23,8 +23,6 @@
 #include <boost/asio/error.hpp>
 #include <boost/bind.hpp>
 
-#include "ray/object_manager/plasma/client.h"
-
 #include "ray/common/id.h"
 #include "ray/common/status.h"
 #include "ray/object_manager/plasma/client.h"

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 {
 

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<ObjectStoreNotificationManager>(main_service);
+    plasma::plasma_store_runner->SetNotificationListener(store_notification_);
+  } else {
+    store_notification_ = std::make_shared<ObjectStoreNotificationManagerIPC>(
+        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<void(const object_manager::protocol::ObjectInfoT &)> callback) {
-  store_notification_.SubscribeObjAdded(callback);
+  store_notification_->SubscribeObjAdded(callback);
   return ray::Status::OK();
 }
 
 ray::Status ObjectManager::SubscribeObjDeleted(
     std::function<void(const ObjectID &)> 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();
 }

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<ObjectDirectoryInterface> 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<ObjectStoreNotificationManager> store_notification_;
   ObjectBufferPool buffer_pool_;
 
   /// Weak reference to main service. We ensure this object is destroyed before

src/ray/object_manager/plasma/client.cc

@@ -1035,7 +1035,7 @@ Status PlasmaClient::Impl::DecodeNotifications(const uint8_t* buffer,
                                                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<fb::PlasmaNotification>(buffer);
+  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);

src/ray/object_manager/plasma/common.fbs

@@ -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;
-}

src/ray/object_manager/plasma/plasma.cc

@@ -21,11 +21,11 @@
 #include <sys/types.h>
 #include <unistd.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/protocol.h"
 
-namespace fb = plasma::flatbuf;
+namespace fb = ray::object_manager::protocol;
 
 namespace plasma {
 
@@ -70,9 +70,9 @@ std::unique_ptr<uint8_t[]> CreateObjectInfoBuffer(fb::ObjectInfoT* object_info)
 }
 
 std::unique_ptr<uint8_t[]> CreatePlasmaNotificationBuffer(
-    std::vector<fb::ObjectInfoT>& object_info) {
+    const std::vector<fb::ObjectInfoT>& object_info) {
   flatbuffers::FlatBufferBuilder fbb;
-  std::vector<flatbuffers::Offset<plasma::flatbuf::ObjectInfo>> info;
+  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]));
   }

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;

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<uint8_t[]> CreateObjectInfoBuffer(flatbuf::ObjectInfoT* object_info);
+std::unique_ptr<uint8_t[]> CreateObjectInfoBuffer(ObjectInfoT* object_info);
 
 std::unique_ptr<uint8_t[]> CreatePlasmaNotificationBuffer(
-    std::vector<flatbuf::ObjectInfoT>& object_info);
+    const std::vector<ObjectInfoT>& object_info);
 
 }  // namespace plasma

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(&notification);
@@ -683,7 +683,7 @@ void PlasmaStore::EvictObjects(const std::vector<ObjectID>& 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(&notification);
@@ -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<fb::ObjectInfoT> 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<fb::ObjectInfoT>& object_info) {
+void PlasmaStore::PushNotifications(const std::vector<ObjectInfoT>& 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<fb::ObjectInfoT>& 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<fb::ObjectInfoT> 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);
   }

src/ray/object_manager/plasma/store.h

@@ -24,6 +24,8 @@
 #include <unordered_set>
 #include <vector>
 
+#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<ray::ObjectStoreNotificationManager> &notification_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<ObjectInfoT>& object_notifications);
+  void PushNotifications(const std::vector<ObjectInfoT>& 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<ray::ObjectStoreNotificationManager> notification_listener_;
 };
 
 }  // namespace plasma

src/ray/object_manager/plasma/store_runner.h

@@ -2,6 +2,7 @@
 
 #include <memory>
 
+#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<ray::ObjectStoreNotificationManager> &notification_listener) {
+    store_->SetNotificationListener(notification_listener);
+  }
 
  private:
   std::string socket_name_;
@@ -23,6 +28,7 @@ class PlasmaStoreRunner {
   std::string external_store_endpoint_;
   std::unique_ptr<EventLoop> loop_;
   std::unique_ptr<PlasmaStore> store_;
+  std::shared_ptr<ray::ObjectStoreNotificationManager> listener_;
 };
 
 // We use a global variable for Plasma Store instance here because: