Cache object info from store notification. (#1815)

Cache all object info from object added store notification & submit to GCS via object directory.

src/ray/object_manager/object_directory.cc

@@ -7,12 +7,14 @@ ObjectDirectory::ObjectDirectory(std::shared_ptr<gcs::AsyncGcsClient> gcs_client
 };
 
 ray::Status ObjectDirectory::ReportObjectAdded(const ObjectID &object_id,
-                                               const ClientID &client_id) {
+                                               const ClientID &client_id,
+                                               const ObjectInfoT &object_info) {
   // TODO(hme): Determine whether we need to do lookup to append.
   JobID job_id = JobID::from_random();
   auto data = std::make_shared<ObjectTableDataT>();
   data->manager = client_id.binary();
   data->is_eviction = false;
+  data->object_size = object_info.data_size;
   ray::Status status = gcs_client_->object_table().Append(
       job_id, object_id, data, [](gcs::AsyncGcsClient *client, const UniqueID &id,
                                   const std::shared_ptr<ObjectTableDataT> data) {

src/ray/object_manager/object_directory.h

@@ -69,9 +69,11 @@ class ObjectDirectoryInterface {
   ///
   /// \param object_id The object id that was put into the store.
   /// \param client_id The client id corresponding to this node.
+  /// \param object_info Additional information about the object.
   /// \return Status of whether this method succeeded.
   virtual ray::Status ReportObjectAdded(const ObjectID &object_id,
-                                        const ClientID &client_id) = 0;
+                                        const ClientID &client_id,
+                                        const ObjectInfoT &object_info) = 0;
 
   /// Report objects removed from this client's store to the object directory.
   ///
@@ -101,8 +103,8 @@ class ObjectDirectory : public ObjectDirectoryInterface {
                            const OnLocationsFailure &fail_callback) override;
   ray::Status Cancel(const ObjectID &object_id) override;
   ray::Status Terminate() override;
-  ray::Status ReportObjectAdded(const ObjectID &object_id,
-                                const ClientID &client_id) override;
+  ray::Status ReportObjectAdded(const ObjectID &object_id, const ClientID &client_id,
+                                const ObjectInfoT &object_info) override;
   ray::Status ReportObjectRemoved(const ObjectID &object_id,
                                   const ClientID &client_id) override;
   /// Ray only (not part of the OD interface).

src/ray/object_manager/object_manager.cc

@@ -24,7 +24,7 @@ ObjectManager::ObjectManager(asio::io_service &main_service,
   main_service_ = &main_service;
   config_ = config;
   store_notification_.SubscribeObjAdded(
-      [this](const ObjectID &oid) { NotifyDirectoryObjectAdd(oid); });
+      [this](const ObjectInfoT &object_info) { NotifyDirectoryObjectAdd(object_info); });
   store_notification_.SubscribeObjDeleted(
       [this](const ObjectID &oid) { NotifyDirectoryObjectDeleted(oid); });
   StartIOService();
@@ -47,7 +47,7 @@ ObjectManager::ObjectManager(asio::io_service &main_service,
   main_service_ = &main_service;
   config_ = config;
   store_notification_.SubscribeObjAdded(
-      [this](const ObjectID &oid) { NotifyDirectoryObjectAdd(oid); });
+      [this](const ObjectInfoT &object_info) { NotifyDirectoryObjectAdd(object_info); });
   store_notification_.SubscribeObjDeleted(
       [this](const ObjectID &oid) { NotifyDirectoryObjectDeleted(oid); });
   StartIOService();
@@ -68,9 +68,11 @@ void ObjectManager::StopIOService() {
   }
 }
 
-void ObjectManager::NotifyDirectoryObjectAdd(const ObjectID &object_id) {
-  local_objects_.insert(object_id);
-  ray::Status status = object_directory_->ReportObjectAdded(object_id, client_id_);
+void ObjectManager::NotifyDirectoryObjectAdd(const ObjectInfoT &object_info) {
+  ObjectID object_id = ObjectID::from_binary(object_info.object_id);
+  local_objects_[object_id] = object_info;
+  ray::Status status =
+      object_directory_->ReportObjectAdded(object_id, client_id_, object_info);
 }
 
 void ObjectManager::NotifyDirectoryObjectDeleted(const ObjectID &object_id) {
@@ -88,7 +90,7 @@ ray::Status ObjectManager::Terminate() {
 }
 
 ray::Status ObjectManager::SubscribeObjAdded(
-    std::function<void(const ObjectID &)> callback) {
+    std::function<void(const ObjectInfoT &)> callback) {
   store_notification_.SubscribeObjAdded(callback);
   return ray::Status::OK();
 }

src/ray/object_manager/object_manager.h

@@ -12,14 +12,14 @@
 #include <boost/asio/error.hpp>
 #include <boost/bind.hpp>
 
-#include "ray/common/client_connection.h"
-#include "ray/id.h"
-#include "ray/status.h"
-
 #include "plasma/client.h"
 #include "plasma/events.h"
 #include "plasma/plasma.h"
 
+#include "ray/common/client_connection.h"
+#include "ray/id.h"
+#include "ray/status.h"
+
 #include "ray/object_manager/connection_pool.h"
 #include "ray/object_manager/format/object_manager_generated.h"
 #include "ray/object_manager/object_directory.h"
@@ -77,7 +77,7 @@ class ObjectManager {
   /// already exist in the local store.
   /// \param callback The callback to invoke when objects are added to the local store.
   /// \return Status of whether adding the subscription succeeded.
-  ray::Status SubscribeObjAdded(std::function<void(const ray::ObjectID &)> callback);
+  ray::Status SubscribeObjAdded(std::function<void(const ObjectInfoT &)> callback);
 
   /// Subscribe to notifications of objects deleted from local store.
   ///
@@ -192,7 +192,7 @@ class ObjectManager {
   std::atomic<int> num_transfers_receive_;
 
   /// Cache of locally available objects.
-  std::unordered_set<ObjectID, UniqueIDHasher> local_objects_;
+  std::unordered_map<ObjectID, ObjectInfoT, UniqueIDHasher> local_objects_;
 
   /// Handle starting, running, and stopping asio io_service.
   void StartIOService();
@@ -200,7 +200,7 @@ class ObjectManager {
   void StopIOService();
 
   /// Register object add with directory.
-  void NotifyDirectoryObjectAdd(const ObjectID &object_id);
+  void NotifyDirectoryObjectAdd(const ObjectInfoT &object_info);
 
   /// Register object remove with directory.
   void NotifyDirectoryObjectDeleted(const ObjectID &object_id);

src/ray/object_manager/object_store_notification_manager.cc

@@ -55,19 +55,16 @@ void ObjectStoreNotificationManager::ProcessStoreNotification(
   if (object_info->is_deletion()) {
     ProcessStoreRemove(object_id);
   } else {
-    ProcessStoreAdd(object_id);
-    // TODO(hme): Determine what data is actually needed by consumer of this notification.
-    //    ProcessStoreAdd(
-    //        object_id, object_info->data_size(),
-    //        object_info->metadata_size(),
-    //        (unsigned char *) object_info->digest()->data());
+    ObjectInfoT result;
+    object_info->UnPackTo(&result);
+    ProcessStoreAdd(result);
   }
   NotificationWait();
 }
 
-void ObjectStoreNotificationManager::ProcessStoreAdd(const ObjectID &object_id) {
+void ObjectStoreNotificationManager::ProcessStoreAdd(const ObjectInfoT &object_info) {
   for (auto handler : add_handlers_) {
-    handler(object_id);
+    handler(object_info);
   }
 }
 
@@ -78,7 +75,7 @@ void ObjectStoreNotificationManager::ProcessStoreRemove(const ObjectID &object_i
 }
 
 void ObjectStoreNotificationManager::SubscribeObjAdded(
-    std::function<void(const ObjectID &)> callback) {
+    std::function<void(const ObjectInfoT &)> callback) {
   add_handlers_.push_back(callback);
 }
 

src/ray/object_manager/object_store_notification_manager.h

@@ -37,7 +37,7 @@ class ObjectStoreNotificationManager {
   /// already exist in the local store
   ///
   /// \param callback A callback expecting an ObjectID.
-  void SubscribeObjAdded(std::function<void(const ray::ObjectID &)> callback);
+  void SubscribeObjAdded(std::function<void(const ObjectInfoT &)> callback);
 
   /// Subscribe to notifications of objects deleted from local store.
   ///
@@ -54,10 +54,10 @@ class ObjectStoreNotificationManager {
   void ProcessStoreNotification(const boost::system::error_code &error);
 
   /// Support for rebroadcasting object add/rem events.
-  void ProcessStoreAdd(const ObjectID &object_id);
+  void ProcessStoreAdd(const ObjectInfoT &object_info);
   void ProcessStoreRemove(const ObjectID &object_id);
 
-  std::vector<std::function<void(const ray::ObjectID &)>> add_handlers_;
+  std::vector<std::function<void(const ObjectInfoT &)>> add_handlers_;
   std::vector<std::function<void(const ray::ObjectID &)>> rem_handlers_;
 
   plasma::PlasmaClient store_client_;

src/ray/object_manager/test/object_manager_stress_test.cc

@@ -242,17 +242,17 @@ class StressTestObjectManager : public TestObjectManagerBase {
 
   void AddTransferTestHandlers() {
     ray::Status status = ray::Status::OK();
-    status =
-        server1->object_manager_.SubscribeObjAdded([this](const ObjectID &object_id) {
-          object_added_handler_1(object_id);
+    status = server1->object_manager_.SubscribeObjAdded(
+        [this](const ObjectInfoT &object_info) {
+          object_added_handler_1(ObjectID::from_binary(object_info.object_id));
           if (v1.size() == num_expected_objects && v1.size() == v2.size()) {
             TransferTestComplete();
           }
         });
     RAY_CHECK_OK(status);
-    status =
-        server2->object_manager_.SubscribeObjAdded([this](const ObjectID &object_id) {
-          object_added_handler_2(object_id);
+    status = server2->object_manager_.SubscribeObjAdded(
+        [this](const ObjectInfoT &object_info) {
+          object_added_handler_2(ObjectID::from_binary(object_info.object_id));
           if (v2.size() == num_expected_objects && v1.size() == v2.size()) {
             TransferTestComplete();
           }

src/ray/object_manager/test/object_manager_test.cc

@@ -204,11 +204,12 @@ class TestObjectManagerCommands : public TestObjectManager {
 
   void TestNotifications() {
     ray::Status status = ray::Status::OK();
-    status =
-        server1->object_manager_.SubscribeObjAdded([this](const ObjectID &object_id) {
-          object_added_handler_1(object_id);
+    status = server1->object_manager_.SubscribeObjAdded(
+        [this](const ObjectInfoT &object_info) {
+          object_added_handler_1(ObjectID::from_binary(object_info.object_id));
           if (v1.size() == num_expected_objects) {
-            NotificationTestComplete(created_object_id, object_id);
+            NotificationTestComplete(created_object_id,
+                                     ObjectID::from_binary(object_info.object_id));
           }
         });
     RAY_CHECK_OK(status);

src/ray/raylet/object_manager_integration_test.cc

@@ -157,17 +157,17 @@ class TestObjectManagerIntegration : public TestObjectManagerBase {
 
   void AddTransferTestHandlers() {
     ray::Status status = ray::Status::OK();
-    status =
-        server1->object_manager_.SubscribeObjAdded([this](const ObjectID &object_id) {
-          v1.push_back(object_id);
+    status = server1->object_manager_.SubscribeObjAdded(
+        [this](const ObjectInfoT &object_info) {
+          v1.push_back(ObjectID::from_binary(object_info.object_id));
           if (v1.size() == num_expected_objects && v1.size() == v2.size()) {
             TestPushComplete();
           }
         });
     RAY_CHECK_OK(status);
-    status =
-        server2->object_manager_.SubscribeObjAdded([this](const ObjectID &object_id) {
-          v2.push_back(object_id);
+    status = server2->object_manager_.SubscribeObjAdded(
+        [this](const ObjectInfoT &object_info) {
+          v2.push_back(ObjectID::from_binary(object_info.object_id));
           if (v2.size() == num_expected_objects && v1.size() == v2.size()) {
             TestPushComplete();
           }

src/ray/raylet/task_dependency_manager.cc

@@ -12,8 +12,10 @@ TaskDependencyManager::TaskDependencyManager(
       // reconstruction_policy_(reconstruction_policy),
       task_ready_callback_(handler) {
   // TODO(swang): Check return status.
-  ray::Status status = object_manager_.SubscribeObjAdded(
-      [this](const ObjectID &object_id) { handleObjectReady(object_id); });
+  ray::Status status =
+      object_manager_.SubscribeObjAdded([this](const ObjectInfoT &object_info) {
+        handleObjectReady(ObjectID::from_binary(object_info.object_id));
+      });
   // TODO(swang): Subscribe to object removed notifications.
 }