[direct task] For serialized object IDs, check with owner before declaring object unreconstructable (#6286)

* Track borrowed vs owned objects

* Serialize owner address with object ID

* serialize owner task id

* Deserialize object IDs

* Pass direct task ID instead of plasma ID

* it works

* Fix ref count test

* Add unit test

* update warning

* we own ray.put objects

* missing file

* doc

* Fix unit test

* comments

* Fix py2

* lint

* update

python/ray/_raylet.pyx

@@ -41,6 +41,7 @@ from libcpp.vector cimport vector as c_vector
 from cython.operator import dereference, postincrement
 
 from ray.includes.common cimport (
+    CAddress,
     CLanguage,
     CRayObject,
     CRayStatus,
@@ -1052,11 +1053,29 @@ cdef class CoreWorker:
         # Note: faster to not release GIL for short-running op.
         self.core_worker.get().RemoveObjectIDReference(c_object_id)
 
-    def promote_object_to_plasma(self, ObjectID object_id):
+    def serialize_and_promote_object_id(self, ObjectID object_id):
         cdef:
             CObjectID c_object_id = object_id.native()
-        self.core_worker.get().PromoteObjectToPlasma(c_object_id)
-        return object_id.with_plasma_transport_type()
+            CTaskID c_owner_id = CTaskID.Nil()
+            CAddress c_owner_address = CAddress()
+        self.core_worker.get().PromoteToPlasmaAndGetOwnershipInfo(
+                c_object_id, &c_owner_id, &c_owner_address)
+        return (object_id,
+                TaskID(c_owner_id.Binary()),
+                c_owner_address.SerializeAsString())
+
+    def deserialize_and_register_object_id(
+            self, const c_string &object_id_binary, const c_string
+            &owner_id_binary, const c_string &serialized_owner_address):
+        cdef:
+            CObjectID c_object_id = CObjectID.FromBinary(object_id_binary)
+            CTaskID c_owner_id = CTaskID.FromBinary(owner_id_binary)
+            CAddress c_owner_address = CAddress()
+        c_owner_address.ParseFromString(serialized_owner_address)
+        self.core_worker.get().RegisterOwnershipInfoAndResolveFuture(
+                c_object_id,
+                c_owner_id,
+                c_owner_address)
 
     # TODO: handle noreturn better
     cdef store_task_outputs(

python/ray/includes/common.pxd

@@ -132,6 +132,10 @@ cdef extern from "ray/protobuf/common.pb.h" nogil:
         pass
     cdef cppclass CTaskType "ray::TaskType":
         pass
+    cdef cppclass CAddress "ray::rpc::Address":
+        CAddress()
+        const c_string &SerializeAsString()
+        void ParseFromString(const c_string &serialized)
 
 
 # This is a workaround for C++ enum class since Cython has no corresponding

python/ray/includes/libcoreworker.pxd

@@ -17,6 +17,7 @@ from ray.includes.unique_ids cimport (
     CObjectID,
 )
 from ray.includes.common cimport (
+    CAddress,
     CActorCreationOptions,
     CBuffer,
     CRayFunction,
@@ -116,6 +117,12 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
         void AddObjectIDReference(const CObjectID &object_id)
         void RemoveObjectIDReference(const CObjectID &object_id)
         void PromoteObjectToPlasma(const CObjectID &object_id)
+        void PromoteToPlasmaAndGetOwnershipInfo(const CObjectID &object_id,
+                                                CTaskID *owner_id,
+                                                CAddress *owner_address)
+        void RegisterOwnershipInfoAndResolveFuture(
+                const CObjectID &object_id, const CTaskID &owner_id, const
+                CAddress &owner_address)
 
         CRayStatus SetClientOptions(c_string client_name, int64_t limit)
         CRayStatus Put(const CRayObject &object, CObjectID *object_id)

python/ray/includes/task.pxd

@@ -27,7 +27,6 @@ cdef extern from "ray/protobuf/common.pb.h" nogil:
     cdef cppclass RpcTask "ray::rpc::Task":
         RpcTaskSpec *mutable_task_spec()
 
-
 cdef extern from "ray/protobuf/gcs.pb.h" nogil:
     cdef cppclass TaskTableData "ray::rpc::TaskTableData":
         RpcTask *mutable_task()

python/ray/serialization.py

@@ -157,19 +157,52 @@ class SerializationContext(object):
                 serialization_context)
 
             def id_serializer(obj):
-                if isinstance(obj, ray.ObjectID) and obj.is_direct_call_type():
-                    obj = self.worker.core_worker.promote_object_to_plasma(obj)
                 return pickle.dumps(obj)
 
             def id_deserializer(serialized_obj):
                 return pickle.loads(serialized_obj)
 
+            def object_id_serializer(obj):
+                owner_id = ""
+                owner_address = ""
+                if obj.is_direct_call_type():
+                    worker = ray.worker.get_global_worker()
+                    worker.check_connected()
+                    obj, owner_id, owner_address = (
+                        worker.core_worker.serialize_and_promote_object_id(obj)
+                    )
+                obj = obj.__reduce__()
+                owner_id = owner_id.__reduce__() if owner_id else owner_id
+                return pickle.dumps((obj, owner_id, owner_address))
+
+            def object_id_deserializer(serialized_obj):
+                obj_id, owner_id, owner_address = pickle.loads(serialized_obj)
+                # Must deserialize the object in the core worker before we
+                # create the ObjectID to ensure that the reference is added
+                # before we increment its count to 1.
+                if owner_id:
+                    worker = ray.worker.get_global_worker()
+                    worker.check_connected()
+                    # UniqueIDs are serialized as
+                    # (class name, (unique bytes,)).
+                    worker.core_worker.deserialize_and_register_object_id(
+                        obj_id[1][0], owner_id[1][0], owner_address)
+                obj_id = obj_id[0](obj_id[1][0])
+                return obj_id
+
             for id_type in ray._raylet._ID_TYPES:
-                serialization_context.register_type(
-                    id_type,
-                    "{}.{}".format(id_type.__module__, id_type.__name__),
-                    custom_serializer=id_serializer,
-                    custom_deserializer=id_deserializer)
+                if id_type == ray._raylet.ObjectID:
+                    serialization_context.register_type(
+                        id_type,
+                        "{}.{}".format(id_type.__module__, id_type.__name__),
+                        custom_serializer=object_id_serializer,
+                        custom_deserializer=object_id_deserializer)
+                else:
+                    serialization_context.register_type(
+                        id_type,
+                        "{}.{}".format(id_type.__module__, id_type.__name__),
+                        custom_serializer=id_serializer,
+                        custom_deserializer=id_deserializer)
 
             # We register this serializer on each worker instead of calling
             # _register_custom_serializer from the driver so that isinstance
@@ -188,16 +221,46 @@ class SerializationContext(object):
                 custom_deserializer=actor_handle_deserializer)
 
             def id_serializer(obj):
-                if isinstance(obj, ray.ObjectID) and obj.is_direct_call_type():
-                    obj = self.worker.core_worker.promote_object_to_plasma(obj)
                 return obj.__reduce__()
 
             def id_deserializer(serialized_obj):
                 return serialized_obj[0](*serialized_obj[1])
 
+            def object_id_serializer(obj):
+                owner_id = ""
+                owner_address = ""
+                if obj.is_direct_call_type():
+                    worker = ray.worker.get_global_worker()
+                    worker.check_connected()
+                    obj, owner_id, owner_address = (
+                        worker.core_worker.serialize_and_promote_object_id(obj)
+                    )
+                obj = id_serializer(obj)
+                owner_id = id_serializer(owner_id) if owner_id else owner_id
+                return (obj, owner_id, owner_address)
+
+            def object_id_deserializer(serialized_obj):
+                obj_id, owner_id, owner_address = serialized_obj
+                # Must deserialize the object in the core worker before we
+                # create the ObjectID to ensure that the reference is added
+                # before we increment its count to 1.
+                if owner_id:
+                    worker = ray.worker.get_global_worker()
+                    worker.check_connected()
+                    # UniqueIDs are serialized as
+                    # (class name, (unique bytes,)).
+                    worker.core_worker.deserialize_and_register_object_id(
+                        obj_id[1][0], owner_id[1][0], owner_address)
+                obj_id = id_deserializer(obj_id)
+                return obj_id
+
             for id_type in ray._raylet._ID_TYPES:
-                self._register_cloudpickle_serializer(id_type, id_serializer,
-                                                      id_deserializer)
+                if id_type == ray._raylet.ObjectID:
+                    self._register_cloudpickle_serializer(
+                        id_type, object_id_serializer, object_id_deserializer)
+                else:
+                    self._register_cloudpickle_serializer(
+                        id_type, id_serializer, id_deserializer)
 
     def initialize(self):
         """ Register custom serializers """

python/ray/tests/test_failure.py

@@ -951,15 +951,13 @@ def test_direct_call_serialized_id_eviction(ray_start_cluster):
     ray.get(get.remote([obj]))
 
 
-@pytest.mark.skip(
-    "Uncomment once eviction errors for serialized IDs are implemented")
 @pytest.mark.parametrize(
     "ray_start_cluster", [{
         "num_nodes": 2,
-        "num_cpus": 10,
+        "num_cpus": 1,
     }, {
         "num_nodes": 1,
-        "num_cpus": 20,
+        "num_cpus": 2,
     }],
     indirect=True)
 def test_direct_call_serialized_id(ray_start_cluster):
@@ -971,16 +969,33 @@ def test_direct_call_serialized_id(ray_start_cluster):
         return 1
 
     @ray.remote
-    def get(obj_ids):
+    def dependent_task(x):
+        return x
+
+    @ray.remote
+    def get(obj_ids, test_dependent_task):
         print("get", obj_ids)
         obj_id = obj_ids[0]
-        assert ray.get(obj_id) == 1
+        if test_dependent_task:
+            assert ray.get(dependent_task.remote(obj_id)) == 1
+        else:
+            assert ray.get(obj_id) == 1
 
     small_object = small_object.options(is_direct_call=True)
+    dependent_task = dependent_task.options(is_direct_call=True)
     get = get.options(is_direct_call=True)
 
     obj = small_object.remote()
-    ray.get(get.remote([obj]))
+    ray.get(get.remote([obj], False))
+
+    obj = small_object.remote()
+    ray.get(get.remote([obj], True))
+
+    obj = ray.put(1)
+    ray.get(get.remote([obj], False))
+
+    obj = ray.put(1)
+    ray.get(get.remote([obj], True))
 
 
 if __name__ == "__main__":

src/ray/common/ray_object.cc

@@ -2,6 +2,17 @@
 
 namespace ray {
 
+std::shared_ptr<LocalMemoryBuffer> MakeErrorMetadataBuffer(rpc::ErrorType error_type) {
+  std::string meta = std::to_string(static_cast<int>(error_type));
+  auto metadata = const_cast<uint8_t *>(reinterpret_cast<const uint8_t *>(meta.data()));
+  auto meta_buffer =
+      std::make_shared<LocalMemoryBuffer>(metadata, meta.size(), /*copy_data=*/true);
+  return meta_buffer;
+}
+
+RayObject::RayObject(rpc::ErrorType error_type)
+    : RayObject(nullptr, MakeErrorMetadataBuffer(error_type)) {}
+
 bool RayObject::IsException(rpc::ErrorType *error_type) const {
   if (metadata_ == nullptr) {
     return false;

src/ray/common/ray_object.h

@@ -42,6 +42,8 @@ class RayObject {
     RAY_CHECK(data_ || metadata_) << "Data and metadata cannot both be empty.";
   }
 
+  RayObject(rpc::ErrorType error_type);
+
   /// Return the data of the ray object.
   const std::shared_ptr<Buffer> &GetData() const { return data_; };
 

src/ray/core_worker/core_worker.cc

@@ -164,6 +164,7 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
         RAY_CHECK_OK(plasma_store_provider_->Put(obj, obj_id));
       },
       ref_counting_enabled ? reference_counter_ : nullptr, raylet_client_));
+
   task_manager_.reset(new TaskManager(memory_store_));
   resolver_.reset(new LocalDependencyResolver(memory_store_));
 
@@ -208,6 +209,7 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
           },
           memory_store_, task_manager_,
           RayConfig::instance().worker_lease_timeout_milliseconds()));
+  future_resolver_.reset(new FutureResolver(memory_store_, client_factory, io_service_));
 }
 
 CoreWorker::~CoreWorker() {
@@ -281,13 +283,36 @@ void CoreWorker::ReportActiveObjectIDs() {
   heartbeat_timer_.async_wait(boost::bind(&CoreWorker::ReportActiveObjectIDs, this));
 }
 
-void CoreWorker::PromoteObjectToPlasma(const ObjectID &object_id) {
+void CoreWorker::PromoteToPlasmaAndGetOwnershipInfo(const ObjectID &object_id,
+                                                    TaskID *owner_id,
+                                                    rpc::Address *owner_address) {
   RAY_CHECK(object_id.IsDirectCallType());
   auto value = memory_store_->GetOrPromoteToPlasma(object_id);
   if (value != nullptr) {
     RAY_CHECK_OK(
         plasma_store_provider_->Put(*value, object_id.WithPlasmaTransportType()));
   }
+
+  auto has_owner = reference_counter_->GetOwner(object_id, owner_id, owner_address);
+  RAY_CHECK(has_owner)
+      << "Object IDs generated randomly (ObjectID.from_random()) or out-of-band "
+         "(ObjectID.from_binary(...)) cannot be serialized because Ray does not know "
+         "which task will create them. "
+         "If this was not how your object ID was generated, please file an issue "
+         "at https://github.com/ray-project/ray/issues/";
+}
+
+void CoreWorker::RegisterOwnershipInfoAndResolveFuture(
+    const ObjectID &object_id, const TaskID &owner_id,
+    const rpc::Address &owner_address) {
+  // Add the object's owner to the local metadata in case it gets serialized
+  // again.
+  reference_counter_->AddBorrowedObject(object_id, owner_id, owner_address);
+
+  RAY_CHECK(!owner_id.IsNil());
+  // We will ask the owner about the object until the object is
+  // created or we can no longer reach the owner.
+  future_resolver_->ResolveFutureAsync(object_id, owner_id, owner_address);
 }
 
 Status CoreWorker::SetClientOptions(std::string name, int64_t limit_bytes) {
@@ -299,6 +324,8 @@ Status CoreWorker::Put(const RayObject &object, ObjectID *object_id) {
   *object_id = ObjectID::ForPut(worker_context_.GetCurrentTaskID(),
                                 worker_context_.GetNextPutIndex(),
                                 static_cast<uint8_t>(TaskTransportType::RAYLET));
+  reference_counter_->AddOwnedObject(*object_id, GetCallerId(), rpc_address_,
+                                     std::make_shared<std::vector<ObjectID>>());
   return Put(object, *object_id);
 }
 
@@ -557,7 +584,8 @@ void CoreWorker::PinObjectReferences(const TaskSpecification &task_spec,
 
   // Note that we call this even if task_deps.size() == 0, in order to pin the return id.
   for (size_t i = 0; i < num_returns; i++) {
-    reference_counter_->SetDependencies(task_spec.ReturnId(i, transport_type), task_deps);
+    reference_counter_->AddOwnedObject(task_spec.ReturnId(i, transport_type),
+                                       GetCallerId(), rpc_address_, task_deps);
   }
 }
 
@@ -928,6 +956,27 @@ void CoreWorker::HandleDirectActorCallArgWaitComplete(
   });
 }
 
+void CoreWorker::HandleGetObjectStatus(const rpc::GetObjectStatusRequest &request,
+                                       rpc::GetObjectStatusReply *reply,
+                                       rpc::SendReplyCallback send_reply_callback) {
+  ObjectID object_id = ObjectID::FromBinary(request.object_id());
+  TaskID owner_id = TaskID::FromBinary(request.owner_id());
+  if (owner_id != GetCallerId()) {
+    // We may have owned this object in the past, but we are now executing some
+    // other task or actor.
+    reply->set_status(rpc::GetObjectStatusReply::WRONG_OWNER);
+  } else {
+    if (task_manager_->IsTaskPending(object_id.TaskId())) {
+      reply->set_status(rpc::GetObjectStatusReply::PENDING);
+    } else {
+      // TODO: We could probably just send the object value if it is small
+      // enough and we have it local.
+      reply->set_status(rpc::GetObjectStatusReply::CREATED);
+    }
+  }
+  send_reply_callback(Status::OK(), nullptr, nullptr);
+}
+
 void CoreWorker::YieldCurrentFiber(FiberEvent &event) {
   RAY_CHECK(worker_context_.CurrentActorIsAsync());
   boost::this_fiber::yield();

src/ray/core_worker/core_worker.h

@@ -7,6 +7,7 @@
 #include "ray/core_worker/actor_handle.h"
 #include "ray/core_worker/common.h"
 #include "ray/core_worker/context.h"
+#include "ray/core_worker/future_resolver.h"
 #include "ray/core_worker/profiling.h"
 #include "ray/core_worker/reference_count.h"
 #include "ray/core_worker/store_provider/memory_store/memory_store.h"
@@ -25,10 +26,11 @@
 /// 1) Add the rpc to the CoreWorkerService in core_worker.proto, e.g., "ExampleCall"
 /// 2) Add a new handler to the macro below: "RAY_CORE_WORKER_RPC_HANDLER(ExampleCall, 1)"
 /// 3) Add a method to the CoreWorker class below: "CoreWorker::HandleExampleCall"
-#define RAY_CORE_WORKER_RPC_HANDLERS          \
-  RAY_CORE_WORKER_RPC_HANDLER(AssignTask, 5)  \
-  RAY_CORE_WORKER_RPC_HANDLER(PushTask, 9999) \
-  RAY_CORE_WORKER_RPC_HANDLER(DirectActorCallArgWaitComplete, 100)
+#define RAY_CORE_WORKER_RPC_HANDLERS                               \
+  RAY_CORE_WORKER_RPC_HANDLER(AssignTask, 5)                       \
+  RAY_CORE_WORKER_RPC_HANDLER(PushTask, 9999)                      \
+  RAY_CORE_WORKER_RPC_HANDLER(DirectActorCallArgWaitComplete, 100) \
+  RAY_CORE_WORKER_RPC_HANDLER(GetObjectStatus, 100)
 
 namespace ray {
 
@@ -104,7 +106,7 @@ class CoreWorker {
   ///
   /// \param[in] object_id The object ID to increase the reference count for.
   void AddObjectIDReference(const ObjectID &object_id) {
-    reference_counter_->AddReference(object_id);
+    reference_counter_->AddLocalReference(object_id);
   }
 
   /// Decrease the reference count for this object ID.
@@ -112,20 +114,46 @@ class CoreWorker {
   /// \param[in] object_id The object ID to decrease the reference count for.
   void RemoveObjectIDReference(const ObjectID &object_id) {
     std::vector<ObjectID> deleted;
-    reference_counter_->RemoveReference(object_id, &deleted);
+    reference_counter_->RemoveLocalReference(object_id, &deleted);
     if (ref_counting_enabled_) {
       memory_store_->Delete(deleted);
     }
   }
 
-  /// Promote an object to plasma. If it already exists locally, it will be
-  /// put into the plasma store. If it doesn't yet exist, it will be spilled to
-  /// plasma once available.
+  /// Promote an object to plasma and get its owner information. This should be
+  /// called when serializing an object ID, and the returned information should
+  /// be stored with the serialized object ID. For plasma promotion, if the
+  /// object already exists locally, it will be put into the plasma store. If
+  /// it doesn't yet exist, it will be spilled to plasma once available.
+  ///
+  /// This can only be called on object IDs that we created via task
+  /// submission, ray.put, or object IDs that we deserialized. It cannot be
+  /// called on object IDs that were created randomly, e.g.,
+  /// ObjectID::FromRandom.
   ///
   /// Postcondition: Get(object_id.WithPlasmaTransportType()) is valid.
   ///
-  /// \param[in] object_id The object ID to promote to plasma.
-  void PromoteObjectToPlasma(const ObjectID &object_id);
+  /// \param[in] object_id The object ID to serialize.
+  /// \param[out] owner_id The ID of the object's owner. This should be
+  /// appended to the serialized object ID.
+  /// \param[out] owner_address The address of the object's owner. This should
+  /// be appended to the serialized object ID.
+  void PromoteToPlasmaAndGetOwnershipInfo(const ObjectID &object_id, TaskID *owner_id,
+                                          rpc::Address *owner_address);
+
+  /// Add a reference to an ObjectID that was deserialized by the language
+  /// frontend. This will also start the process to resolve the future.
+  /// Specifically, we will periodically contact the owner, until we learn that
+  /// the object has been created or the owner is no longer reachable. This
+  /// will then unblock any Gets or submissions of tasks dependent on the
+  /// object.
+  ///
+  /// \param[in] object_id The object ID to deserialize.
+  /// \param[out] owner_id The ID of the object's owner.
+  /// \param[out] owner_address The address of the object's owner.
+  void RegisterOwnershipInfoAndResolveFuture(const ObjectID &object_id,
+                                             const TaskID &owner_id,
+                                             const rpc::Address &owner_address);
 
   ///
   /// Public methods related to storing and retrieving objects.
@@ -354,6 +382,11 @@ class CoreWorker {
       rpc::DirectActorCallArgWaitCompleteReply *reply,
       rpc::SendReplyCallback send_reply_callback);
 
+  /// Implements gRPC server handler.
+  void HandleGetObjectStatus(const rpc::GetObjectStatusRequest &request,
+                             rpc::GetObjectStatusReply *reply,
+                             rpc::SendReplyCallback send_reply_callback);
+
   ///
   /// Public methods related to async actor call. This should only be used when
   /// the actor is (1) direct actor and (2) using asyncio mode.
@@ -509,6 +542,8 @@ class CoreWorker {
   /// Plasma store interface.
   std::shared_ptr<CoreWorkerPlasmaStoreProvider> plasma_store_provider_;
 
+  std::unique_ptr<FutureResolver> future_resolver_;
+
   ///
   /// Fields related to task submission.
   ///

src/ray/core_worker/future_resolver.cc

@@ -0,0 +1,54 @@
+#include "ray/core_worker/future_resolver.h"
+
+namespace ray {
+
+void FutureResolver::ResolveFutureAsync(const ObjectID &object_id, const TaskID &owner_id,
+                                        const rpc::Address &owner_address) {
+  RAY_CHECK(object_id.IsDirectCallType());
+  absl::MutexLock lock(&mu_);
+  auto it = owner_clients_.find(owner_id);
+  if (it == owner_clients_.end()) {
+    auto client = std::shared_ptr<rpc::CoreWorkerClientInterface>(
+        client_factory_({owner_address.ip_address(), owner_address.port()}));
+    owner_clients_.emplace(owner_id, std::move(client));
+  }
+
+  // This timer will get deallocated once the future has been resolved.
+  auto timer = std::make_shared<boost::asio::deadline_timer>(io_service_);
+  AttemptFutureResolution(object_id, owner_id, std::move(timer));
+}
+
+void FutureResolver::AttemptFutureResolution(
+    const ObjectID &object_id, const TaskID &owner_id,
+    std::shared_ptr<boost::asio::deadline_timer> timer) {
+  auto &owner_client = owner_clients_[owner_id];
+  rpc::GetObjectStatusRequest request;
+  request.set_object_id(object_id.Binary());
+  request.set_owner_id(owner_id.Binary());
+  auto status = owner_client->GetObjectStatus(
+      request, [this, object_id, owner_id, timer](
+                   const Status &status, const rpc::GetObjectStatusReply &reply) {
+        if (!status.ok() || reply.status() != rpc::GetObjectStatusReply::PENDING) {
+          // Either the owner is gone or the owner replied that the object has
+          // been created. In both cases, we can now try to fetch the object via
+          // plasma.
+          RAY_CHECK_OK(in_memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA),
+                                             object_id));
+        } else {
+          // Try again later.
+          timer->expires_from_now(
+              boost::posix_time::milliseconds(wait_future_resolution_milliseconds_));
+          timer->async_wait(
+              [this, object_id, owner_id, timer](const boost::system::error_code &error) {
+                absl::MutexLock lock(&mu_);
+                AttemptFutureResolution(object_id, owner_id, std::move(timer));
+              });
+        }
+      });
+  if (!status.ok()) {
+    RAY_CHECK_OK(
+        in_memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA), object_id));
+  }
+}
+
+}  // namespace ray

src/ray/core_worker/future_resolver.h

@@ -0,0 +1,71 @@
+#ifndef RAY_CORE_WORKER_FUTURE_RESOLVER_H
+#define RAY_CORE_WORKER_FUTURE_RESOLVER_H
+
+#include <memory>
+
+#include "ray/common/id.h"
+#include "ray/core_worker/store_provider/memory_store/memory_store.h"
+#include "ray/protobuf/core_worker.pb.h"
+#include "ray/rpc/worker/core_worker_client.h"
+
+namespace ray {
+
+/// Max time between requests to the owner to check whether the object is still
+/// being computed.
+const int kWaitObjectEvictionMilliseconds = 100;
+
+// Resolve values for futures that were given to us before the value
+// was available. This class is thread-safe.
+class FutureResolver {
+ public:
+  FutureResolver(
+      std::shared_ptr<CoreWorkerMemoryStore> store, rpc::ClientFactoryFn client_factory,
+      boost::asio::io_service &io_service,
+      int wait_future_resolution_milliseconds = kWaitObjectEvictionMilliseconds)
+      : in_memory_store_(store),
+        client_factory_(client_factory),
+        io_service_(io_service),
+        wait_future_resolution_milliseconds_(wait_future_resolution_milliseconds) {}
+
+  /// Resolve the value for a future. This will periodically contact the given
+  /// owner until the owner dies or the owner has finished creating the object.
+  /// In either case, this will put an OBJECT_IN_PLASMA error as the future's
+  /// value.
+  ///
+  /// \param[in] object_id The ID of the future to resolve.
+  /// \param[in] owner_id The ID of the task or actor that owns the future.
+  /// \param[in] owner_address The address of the task or actor that owns the
+  /// future.
+  void ResolveFutureAsync(const ObjectID &object_id, const TaskID &owner_id,
+                          const rpc::Address &owner_address);
+
+ private:
+  // Attempt to contact the owner to ask about the future's current status.
+  void AttemptFutureResolution(const ObjectID &object_id, const TaskID &owner_id,
+                               std::shared_ptr<boost::asio::deadline_timer> timer)
+      EXCLUSIVE_LOCKS_REQUIRED(mu_);
+
+  /// Used to set timers.
+  boost::asio::io_service &io_service_;
+
+  /// Used to store values of resolved futures.
+  std::shared_ptr<CoreWorkerMemoryStore> in_memory_store_;
+
+  /// Factory for producing new core worker clients.
+  const rpc::ClientFactoryFn client_factory_;
+
+  /// The amount of time to wait between requests to a future's owner to get
+  /// the object's current status.
+  const int wait_future_resolution_milliseconds_;
+
+  /// Protects against concurrent access to internal state.
+  absl::Mutex mu_;
+
+  /// Cache of gRPC clients to the objects' owners.
+  absl::flat_hash_map<TaskID, std::shared_ptr<rpc::CoreWorkerClientInterface>>
+      owner_clients_ GUARDED_BY(mu_);
+};
+
+}  // namespace ray
+
+#endif  // RAY_CORE_WORKER_FUTURE_RESOLVER_H

src/ray/core_worker/reference_count.cc

@@ -2,42 +2,48 @@
 
 namespace ray {
 
-void ReferenceCounter::AddReference(const ObjectID &object_id) {
+void ReferenceCounter::AddBorrowedObject(const ObjectID &object_id,
+                                         const TaskID &owner_id,
+                                         const rpc::Address &owner_address) {
   absl::MutexLock lock(&mutex_);
-  AddReferenceInternal(object_id);
+  RAY_CHECK(
+      object_id_refs_.emplace(object_id, Reference(owner_id, owner_address)).second);
 }
 
-void ReferenceCounter::AddReferenceInternal(const ObjectID &object_id) {
-  auto entry = object_id_refs_.find(object_id);
-  if (entry == object_id_refs_.end()) {
-    object_id_refs_[object_id] = std::make_pair(1, nullptr);
-  } else {
-    entry->second.first++;
-  }
-}
-
-void ReferenceCounter::SetDependencies(
-    const ObjectID &object_id, std::shared_ptr<std::vector<ObjectID>> dependencies) {
+void ReferenceCounter::AddOwnedObject(
+    const ObjectID &object_id, const TaskID &owner_id, const rpc::Address &owner_address,
+    std::shared_ptr<std::vector<ObjectID>> dependencies) {
   absl::MutexLock lock(&mutex_);
 
-  auto entry = object_id_refs_.find(object_id);
-  if (entry == object_id_refs_.end()) {
-    // If the entry doesn't exist, we initialize the direct reference count to zero
-    // because this corresponds to a submitted task whose return ObjectID will be created
-    // in the frontend language, incrementing the reference count.
-    object_id_refs_[object_id] = std::make_pair(0, dependencies);
-  } else {
-    RAY_CHECK(!entry->second.second);
-    entry->second.second = dependencies;
-  }
-
   for (const ObjectID &dependency_id : *dependencies) {
-    AddReferenceInternal(dependency_id);
+    AddLocalReferenceInternal(dependency_id);
   }
+
+  RAY_CHECK(object_id_refs_.count(object_id) == 0)
+      << "Cannot create an object that already exists. ObjectID: " << object_id;
+  // If the entry doesn't exist, we initialize the direct reference count to zero
+  // because this corresponds to a submitted task whose return ObjectID will be created
+  // in the frontend language, incrementing the reference count.
+  object_id_refs_.emplace(object_id, Reference(owner_id, owner_address, dependencies));
 }
 
-void ReferenceCounter::RemoveReference(const ObjectID &object_id,
-                                       std::vector<ObjectID> *deleted) {
+void ReferenceCounter::AddLocalReferenceInternal(const ObjectID &object_id) {
+  auto entry = object_id_refs_.find(object_id);
+  if (entry == object_id_refs_.end()) {
+    // TODO: Once ref counting is implemented, we should always know how the
+    // ObjectID was created, so there should always ben an entry.
+    entry = object_id_refs_.emplace(object_id, Reference()).first;
+  }
+  entry->second.local_ref_count++;
+}
+
+void ReferenceCounter::AddLocalReference(const ObjectID &object_id) {
+  absl::MutexLock lock(&mutex_);
+  AddLocalReferenceInternal(object_id);
+}
+
+void ReferenceCounter::RemoveLocalReference(const ObjectID &object_id,
+                                            std::vector<ObjectID> *deleted) {
   absl::MutexLock lock(&mutex_);
   RemoveReferenceRecursive(object_id, deleted);
 }
@@ -50,19 +56,36 @@ void ReferenceCounter::RemoveReferenceRecursive(const ObjectID &object_id,
                      << object_id;
     return;
   }
-  if (--entry->second.first == 0) {
+  if (--entry->second.local_ref_count == 0) {
     // If the reference count reached 0, decrease the reference count for each dependency.
-    if (entry->second.second) {
-      for (const ObjectID &pending_task_object_id : *entry->second.second) {
+    if (entry->second.dependencies) {
+      for (const ObjectID &pending_task_object_id : *entry->second.dependencies) {
         RemoveReferenceRecursive(pending_task_object_id, deleted);
       }
     }
-    object_id_refs_.erase(object_id);
+    object_id_refs_.erase(entry);
     deleted->push_back(object_id);
   }
 }
 
-bool ReferenceCounter::HasReference(const ObjectID &object_id) {
+bool ReferenceCounter::GetOwner(const ObjectID &object_id, TaskID *owner_id,
+                                rpc::Address *owner_address) const {
+  absl::MutexLock lock(&mutex_);
+  auto it = object_id_refs_.find(object_id);
+  if (it == object_id_refs_.end()) {
+    return false;
+  }
+
+  if (it->second.owner.has_value()) {
+    *owner_id = it->second.owner.value().first;
+    *owner_address = it->second.owner.value().second;
+    return true;
+  } else {
+    return false;
+  }
+}
+
+bool ReferenceCounter::HasReference(const ObjectID &object_id) const {
   absl::MutexLock lock(&mutex_);
   return object_id_refs_.find(object_id) != object_id_refs_.end();
 }
@@ -93,12 +116,12 @@ void ReferenceCounter::LogDebugString() const {
 
   for (const auto &entry : object_id_refs_) {
     RAY_LOG(DEBUG) << "\t" << entry.first.Hex();
-    RAY_LOG(DEBUG) << "\t\treference count: " << entry.second.first;
+    RAY_LOG(DEBUG) << "\t\treference count: " << entry.second.local_ref_count;
     RAY_LOG(DEBUG) << "\t\tdependencies: ";
-    if (!entry.second.second) {
+    if (!entry.second.dependencies) {
       RAY_LOG(DEBUG) << "\t\t\tNULL";
     } else {
-      for (const ObjectID &pending_task_object_id : *entry.second.second) {
+      for (const ObjectID &pending_task_object_id : *entry.second.dependencies) {
         RAY_LOG(DEBUG) << "\t\t\t" << pending_task_object_id.Hex();
       }
     }

src/ray/core_worker/reference_count.h

@@ -6,6 +6,7 @@
 #include "absl/synchronization/mutex.h"
 
 #include "ray/common/id.h"
+#include "ray/protobuf/common.pb.h"
 #include "ray/util/logging.h"
 
 namespace ray {
@@ -19,30 +20,57 @@ class ReferenceCounter {
   ~ReferenceCounter() {}
 
   /// Increase the reference count for the ObjectID by one. If there is no
-  /// entry for the ObjectID, one will be created with no dependencies.
-  void AddReference(const ObjectID &object_id) LOCKS_EXCLUDED(mutex_);
+  /// entry for the ObjectID, one will be created. The object ID will not have
+  /// any owner information, since we don't know how it was created.
+  ///
+  /// \param[in] object_id The object to to increment the count for.
+  void AddLocalReference(const ObjectID &object_id) LOCKS_EXCLUDED(mutex_);
 
   /// Decrease the reference count for the ObjectID by one. If the reference count reaches
   /// zero, it will be erased from the map and the reference count for all of its
   /// dependencies will be decreased be one.
   ///
   /// \param[in] object_id The object to to decrement the count for.
-  /// \param[in] deleted List to store objects that hit zero ref count.
-  void RemoveReference(const ObjectID &object_id, std::vector<ObjectID> *deleted)
+  /// \param[out] deleted List to store objects that hit zero ref count.
+  void RemoveLocalReference(const ObjectID &object_id, std::vector<ObjectID> *deleted)
       LOCKS_EXCLUDED(mutex_);
 
-  /// Set the dependencies for the ObjectID. Dependencies for each ObjectID must be
-  /// set at most once. The direct reference count for the ObjectID is set to zero and the
-  /// reference count for each dependency is incremented.
-  void SetDependencies(const ObjectID &object_id,
-                       std::shared_ptr<std::vector<ObjectID>> dependencies)
+  /// Add an object that we own. The object may depend on other objects.
+  /// Dependencies for each ObjectID must be set at most once. The direct
+  /// reference count for the ObjectID is set to zero and the reference count
+  /// for each dependency is incremented.
+  ///
+  /// TODO(swang): We could avoid copying the owner_id and owner_address since
+  /// we are the owner, but it is easier to store a copy for now, since the
+  /// owner ID will change for workers executing normal tasks and it is
+  /// possible to have leftover references after a task has finished.
+  ///
+  /// \param[in] object_id The ID of the object that we own.
+  /// \param[in] owner_id The ID of the object's owner.
+  /// \param[in] owner_address The address of the object's owner.
+  /// \param[in] dependencies The objects that the object depends on.
+  void AddOwnedObject(const ObjectID &object_id, const TaskID &owner_id,
+                      const rpc::Address &owner_address,
+                      std::shared_ptr<std::vector<ObjectID>> dependencies)
       LOCKS_EXCLUDED(mutex_);
 
+  /// Add an object that we are borrowing.
+  ///
+  /// \param[in] object_id The ID of the object that we are borrowing.
+  /// \param[in] owner_id The ID of the owner of the object. This is either the
+  /// task ID (for non-actors) or the actor ID of the owner.
+  /// \param[in] owner_address The owner's address.
+  void AddBorrowedObject(const ObjectID &object_id, const TaskID &owner_id,
+                         const rpc::Address &owner_address) LOCKS_EXCLUDED(mutex_);
+
+  bool GetOwner(const ObjectID &object_id, TaskID *owner_id,
+                rpc::Address *owner_address) const LOCKS_EXCLUDED(mutex_);
+
   /// Returns the total number of ObjectIDs currently in scope.
   size_t NumObjectIDsInScope() const LOCKS_EXCLUDED(mutex_);
 
   /// Returns whether this object has an active reference.
-  bool HasReference(const ObjectID &object_id) LOCKS_EXCLUDED(mutex_);
+  bool HasReference(const ObjectID &object_id) const LOCKS_EXCLUDED(mutex_);
 
   /// Returns a set of all ObjectIDs currently in scope (i.e., nonzero reference count).
   std::unordered_set<ObjectID> GetAllInScopeObjectIDs() const LOCKS_EXCLUDED(mutex_);
@@ -51,9 +79,41 @@ class ReferenceCounter {
   void LogDebugString() const LOCKS_EXCLUDED(mutex_);
 
  private:
+  /// Metadata for an ObjectID reference in the language frontend.
+  struct Reference {
+    /// Constructor for a reference whose origin is unknown.
+    Reference() : owned_by_us(false) {}
+    /// Constructor for a reference that we created.
+    Reference(const TaskID &owner_id, const rpc::Address &owner_address,
+              std::shared_ptr<std::vector<ObjectID>> deps)
+        : dependencies(std::move(deps)),
+          owned_by_us(true),
+          owner({owner_id, owner_address}) {}
+    /// Constructor for a reference that was given to us.
+    Reference(const TaskID &owner_id, const rpc::Address &owner_address)
+        : owned_by_us(false), owner({owner_id, owner_address}) {}
+    /// The local ref count for the ObjectID in the language frontend.
+    size_t local_ref_count = 0;
+    /// The objects that this object depends on. Tracked only by the owner of
+    /// the object. Dependencies are stored as shared_ptrs because the same set
+    /// of dependencies can be shared among multiple entries. For example, when
+    /// a task has multiple return values, the entry for each return ObjectID
+    /// depends on all task dependencies.
+    std::shared_ptr<std::vector<ObjectID>> dependencies;
+    /// Whether we own the object. If we own the object, then we are
+    /// responsible for tracking the state of the task that creates the object
+    /// (see task_manager.h).
+    bool owned_by_us;
+    /// The object's owner, if we know it. This has no value if the object is
+    /// if we do not know the object's owner (because distributed ref counting
+    /// is not yet implemented).
+    const absl::optional<std::pair<TaskID, rpc::Address>> owner;
+  };
+
   /// Helper function with the same semantics as AddReference to allow adding a reference
   /// while already holding mutex_.
-  void AddReferenceInternal(const ObjectID &object_id) EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+  void AddLocalReferenceInternal(const ObjectID &object_id)
+      EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
   /// Recursive helper function for decreasing reference counts. Will recursively call
   /// itself on any dependencies whose reference count reaches zero as a result of
@@ -67,12 +127,8 @@ class ReferenceCounter {
   /// Protects access to the reference counting state.
   mutable absl::Mutex mutex_;
 
-  /// Holds all direct reference counts and dependency information for tracked ObjectIDs.
-  /// Dependencies are stored as shared_ptrs because the same set of dependencies can be
-  /// shared among multiple entries. For example, when a task has multiple return values,
-  /// the entry for each return ObjectID depends on all task dependencies.
-  absl::flat_hash_map<ObjectID, std::pair<size_t, std::shared_ptr<std::vector<ObjectID>>>>
-      object_id_refs_ GUARDED_BY(mutex_);
+  /// Holds all reference counts and dependency information for tracked ObjectIDs.
+  absl::flat_hash_map<ObjectID, Reference> object_id_refs_ GUARDED_BY(mutex_);
 };
 
 }  // namespace ray

src/ray/core_worker/reference_count_test.cc

@@ -20,19 +20,19 @@ class ReferenceCountTest : public ::testing::Test {
 TEST_F(ReferenceCountTest, TestBasic) {
   std::vector<ObjectID> out;
   ObjectID id = ObjectID::FromRandom();
-  rc->AddReference(id);
+  rc->AddLocalReference(id);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 1);
-  rc->AddReference(id);
+  rc->AddLocalReference(id);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 1);
-  rc->AddReference(id);
+  rc->AddLocalReference(id);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 1);
-  rc->RemoveReference(id, &out);
+  rc->RemoveLocalReference(id, &out);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 1);
   ASSERT_EQ(out.size(), 0);
-  rc->RemoveReference(id, &out);
+  rc->RemoveLocalReference(id, &out);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 1);
   ASSERT_EQ(out.size(), 0);
-  rc->RemoveReference(id, &out);
+  rc->RemoveLocalReference(id, &out);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 0);
   ASSERT_EQ(out.size(), 1);
 }
@@ -49,21 +49,21 @@ TEST_F(ReferenceCountTest, TestDependencies) {
   std::shared_ptr<std::vector<ObjectID>> deps = std::make_shared<std::vector<ObjectID>>();
   deps->push_back(id2);
   deps->push_back(id3);
-  rc->SetDependencies(id1, deps);
+  rc->AddOwnedObject(id1, TaskID::Nil(), rpc::Address(), deps);
 
-  rc->AddReference(id1);
-  rc->AddReference(id1);
-  rc->AddReference(id3);
+  rc->AddLocalReference(id1);
+  rc->AddLocalReference(id1);
+  rc->AddLocalReference(id3);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 3);
 
-  rc->RemoveReference(id1, &out);
+  rc->RemoveLocalReference(id1, &out);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 3);
   ASSERT_EQ(out.size(), 0);
-  rc->RemoveReference(id1, &out);
+  rc->RemoveLocalReference(id1, &out);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 1);
   ASSERT_EQ(out.size(), 2);
 
-  rc->RemoveReference(id3, &out);
+  rc->RemoveLocalReference(id3, &out);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 0);
   ASSERT_EQ(out.size(), 3);
 }
@@ -82,22 +82,22 @@ TEST_F(ReferenceCountTest, TestSharedDependencies) {
   std::shared_ptr<std::vector<ObjectID>> deps = std::make_shared<std::vector<ObjectID>>();
   deps->push_back(id3);
   deps->push_back(id4);
-  rc->SetDependencies(id1, deps);
-  rc->SetDependencies(id2, deps);
+  rc->AddOwnedObject(id1, TaskID::Nil(), rpc::Address(), deps);
+  rc->AddOwnedObject(id2, TaskID::Nil(), rpc::Address(), deps);
 
-  rc->AddReference(id1);
-  rc->AddReference(id2);
-  rc->AddReference(id4);
+  rc->AddLocalReference(id1);
+  rc->AddLocalReference(id2);
+  rc->AddLocalReference(id4);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 4);
 
-  rc->RemoveReference(id1, &out);
+  rc->RemoveLocalReference(id1, &out);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 3);
   ASSERT_EQ(out.size(), 1);
-  rc->RemoveReference(id2, &out);
+  rc->RemoveLocalReference(id2, &out);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 1);
   ASSERT_EQ(out.size(), 3);
 
-  rc->RemoveReference(id4, &out);
+  rc->RemoveLocalReference(id4, &out);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 0);
   ASSERT_EQ(out.size(), 4);
 }
@@ -113,34 +113,65 @@ TEST_F(ReferenceCountTest, TestRecursiveDependencies) {
   ObjectID id3 = ObjectID::FromRandom();
   ObjectID id4 = ObjectID::FromRandom();
 
-  std::shared_ptr<std::vector<ObjectID>> deps1 =
-      std::make_shared<std::vector<ObjectID>>();
-  deps1->push_back(id2);
-  rc->SetDependencies(id1, deps1);
-
   std::shared_ptr<std::vector<ObjectID>> deps2 =
       std::make_shared<std::vector<ObjectID>>();
   deps2->push_back(id3);
   deps2->push_back(id4);
-  rc->SetDependencies(id2, deps2);
+  rc->AddOwnedObject(id2, TaskID::Nil(), rpc::Address(), deps2);
 
-  rc->AddReference(id1);
-  rc->AddReference(id2);
-  rc->AddReference(id4);
+  std::shared_ptr<std::vector<ObjectID>> deps1 =
+      std::make_shared<std::vector<ObjectID>>();
+  deps1->push_back(id2);
+  rc->AddOwnedObject(id1, TaskID::Nil(), rpc::Address(), deps1);
+
+  rc->AddLocalReference(id1);
+  rc->AddLocalReference(id2);
+  rc->AddLocalReference(id4);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 4);
 
-  rc->RemoveReference(id2, &out);
+  rc->RemoveLocalReference(id2, &out);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 4);
   ASSERT_EQ(out.size(), 0);
-  rc->RemoveReference(id1, &out);
+  rc->RemoveLocalReference(id1, &out);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 1);
   ASSERT_EQ(out.size(), 3);
 
-  rc->RemoveReference(id4, &out);
+  rc->RemoveLocalReference(id4, &out);
   ASSERT_EQ(rc->NumObjectIDsInScope(), 0);
   ASSERT_EQ(out.size(), 4);
 }
 
+// Tests that we can get the owner address correctly for objects that we own,
+// objects that we borrowed via a serialized object ID, and objects whose
+// origin we do not know.
+TEST_F(ReferenceCountTest, TestOwnerAddress) {
+  auto object_id = ObjectID::FromRandom();
+  TaskID task_id = TaskID::ForFakeTask();
+  rpc::Address address;
+  address.set_ip_address("1234");
+  auto deps = std::make_shared<std::vector<ObjectID>>();
+  rc->AddOwnedObject(object_id, task_id, address, deps);
+
+  TaskID added_id;
+  rpc::Address added_address;
+  ASSERT_TRUE(rc->GetOwner(object_id, &added_id, &added_address));
+  ASSERT_EQ(task_id, added_id);
+  ASSERT_EQ(address.ip_address(), added_address.ip_address());
+
+  auto object_id2 = ObjectID::FromRandom();
+  task_id = TaskID::ForFakeTask();
+  address.set_ip_address("5678");
+  rc->AddOwnedObject(object_id2, task_id, address, deps);
+  ASSERT_TRUE(rc->GetOwner(object_id2, &added_id, &added_address));
+  ASSERT_EQ(task_id, added_id);
+  ASSERT_EQ(address.ip_address(), added_address.ip_address());
+
+  auto object_id3 = ObjectID::FromRandom();
+  ASSERT_FALSE(rc->GetOwner(object_id3, &added_id, &added_address));
+  rc->AddLocalReference(object_id3);
+  ASSERT_FALSE(rc->GetOwner(object_id3, &added_id, &added_address));
+}
+
 // Tests that the ref counts are properly integrated into the local
 // object memory store.
 TEST(MemoryStoreIntegrationTest, TestSimple) {
@@ -157,7 +188,7 @@ TEST(MemoryStoreIntegrationTest, TestSimple) {
   ASSERT_EQ(store.Size(), 0);
 
   // Tests ref counting overrides remove after get option.
-  rc->AddReference(id1);
+  rc->AddLocalReference(id1);
   RAY_CHECK_OK(store.Put(buffer, id1));
   ASSERT_EQ(store.Size(), 1);
   std::vector<std::shared_ptr<RayObject>> results;

src/ray/core_worker/task_manager.cc

@@ -25,12 +25,8 @@ void TaskManager::CompletePendingTask(const TaskID &task_id,
 
     if (return_object.in_plasma()) {
       // Mark it as in plasma with a dummy object.
-      std::string meta =
-          std::to_string(static_cast<int>(rpc::ErrorType::OBJECT_IN_PLASMA));
-      auto metadata =
-          const_cast<uint8_t *>(reinterpret_cast<const uint8_t *>(meta.data()));
-      auto meta_buffer = std::make_shared<LocalMemoryBuffer>(metadata, meta.size());
-      RAY_CHECK_OK(in_memory_store_->Put(RayObject(nullptr, meta_buffer), object_id));
+      RAY_CHECK_OK(
+          in_memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA), object_id));
     } else {
       std::shared_ptr<LocalMemoryBuffer> data_buffer;
       if (return_object.data().size() > 0) {
@@ -70,10 +66,7 @@ void TaskManager::FailPendingTask(const TaskID &task_id, rpc::ErrorType error_ty
     const auto object_id = ObjectID::ForTaskReturn(
         task_id, /*index=*/i + 1,
         /*transport_type=*/static_cast<int>(TaskTransportType::DIRECT));
-    std::string meta = std::to_string(static_cast<int>(error_type));
-    auto metadata = const_cast<uint8_t *>(reinterpret_cast<const uint8_t *>(meta.data()));
-    auto meta_buffer = std::make_shared<LocalMemoryBuffer>(metadata, meta.size());
-    RAY_CHECK_OK(in_memory_store_->Put(RayObject(nullptr, meta_buffer), object_id));
+    RAY_CHECK_OK(in_memory_store_->Put(RayObject(error_type), object_id));
   }
 }
 

src/ray/protobuf/core_worker.proto

@@ -87,6 +87,20 @@ message DirectActorCallArgWaitCompleteRequest {
 message DirectActorCallArgWaitCompleteReply {
 }
 
+message GetObjectStatusRequest {
+  bytes owner_id = 1;
+  bytes object_id = 2;
+}
+
+message GetObjectStatusReply {
+  enum ObjectStatus {
+    PENDING = 0;
+    CREATED = 1;
+    WRONG_OWNER = 2;
+  }
+  ObjectStatus status = 1;
+}
+
 service CoreWorkerService {
   // Push a task to a worker from the raylet.
   rpc AssignTask(AssignTaskRequest) returns (AssignTaskReply);
@@ -95,4 +109,6 @@ service CoreWorkerService {
   // Reply from raylet that wait for direct actor call args has completed.
   rpc DirectActorCallArgWaitComplete(DirectActorCallArgWaitCompleteRequest)
       returns (DirectActorCallArgWaitCompleteReply);
+  // Ask the object's owner about the object's current status.
+  rpc GetObjectStatus(GetObjectStatusRequest) returns (GetObjectStatusReply);
 }

src/ray/rpc/worker/core_worker_client.h

@@ -81,6 +81,13 @@ class CoreWorkerClientInterface {
     return Status::NotImplemented("");
   }
 
+  /// Ask the owner of an object about the object's current status.
+  virtual ray::Status GetObjectStatus(
+      const GetObjectStatusRequest &request,
+      const ClientCallback<GetObjectStatusReply> &callback) {
+    return Status::NotImplemented("");
+  }
+
   virtual ~CoreWorkerClientInterface(){};
 };
 
@@ -148,6 +155,14 @@ class CoreWorkerClient : public std::enable_shared_from_this<CoreWorkerClient>,
     return call->GetStatus();
   }
 
+  virtual ray::Status GetObjectStatus(
+      const GetObjectStatusRequest &request,
+      const ClientCallback<GetObjectStatusReply> &callback) {
+    auto call = client_call_manager_.CreateCall<CoreWorkerService, GetObjectStatusRequest,
+                                                GetObjectStatusReply>(
+        *stub_, &CoreWorkerService::Stub::PrepareAsyncGetObjectStatus, request, callback);
+    return call->GetStatus();
+  }
   /// Send as many pending tasks as possible. This method is thread-safe.
   ///
   /// The client will guarantee no more than kMaxBytesInFlight bytes of RPCs are being