diff --git a/src/ray/raylet_rpc_client/raylet_client.cc b/src/ray/raylet_rpc_client/raylet_client.cc index eee711936a02..baac6411b6f5 100644 --- a/src/ray/raylet_rpc_client/raylet_client.cc +++ b/src/ray/raylet_rpc_client/raylet_client.cc @@ -335,10 +335,11 @@ void RayletClient::PinObjectIDs( if (!generator_id.IsNil()) { request.set_generator_id(generator_id.Binary()); } + auto self = shared_from_this(); pins_in_flight_++; - auto rpc_callback = [this, callback = std::move(callback)]( + auto rpc_callback = [self, callback = std::move(callback)]( Status status, rpc::PinObjectIDsReply &&reply) { - pins_in_flight_--; + self->pins_in_flight_--; callback(status, std::move(reply)); }; INVOKE_RETRYABLE_RPC_CALL(retryable_grpc_client_, diff --git a/src/ray/raylet_rpc_client/raylet_client.h b/src/ray/raylet_rpc_client/raylet_client.h index 78a9f52e4726..39f8b97b4a22 100644 --- a/src/ray/raylet_rpc_client/raylet_client.h +++ b/src/ray/raylet_rpc_client/raylet_client.h @@ -39,7 +39,8 @@ namespace rpc { /// Raylet client is responsible for communication with raylet. It implements /// [RayletClientInterface] and works on worker registration, lease management, etc. -class RayletClient : public RayletClientInterface { +class RayletClient : public RayletClientInterface, + public std::enable_shared_from_this { public: /// Connect to the raylet. ///