diff --git a/src/ray/ray_syncer/ray_syncer_bidi_reactor.h b/src/ray/ray_syncer/ray_syncer_bidi_reactor.h index 8fa83c67136d..8f0f0fb183e3 100644 --- a/src/ray/ray_syncer/ray_syncer_bidi_reactor.h +++ b/src/ray/ray_syncer/ray_syncer_bidi_reactor.h @@ -120,8 +120,11 @@ class RaySyncerBidiReactor { /// See MessagesReceivedCallback in common.h. MessagesReceivedCallback on_messages_received_; - /// Self-owned shared_ptr for lifetime guarding in async callbacks (timer/dispatch). Set - /// by caller/factory, released in OnDone. + /// Self-owned shared_ptr for lifetime guarding in async callbacks. Set by the + /// caller/factory before any gRPC callback can fire and released in OnDone. Every + /// handler dispatched to io_context_ captures a copy so that it can never outlive + /// the reactor; the batch timer holds a weak_ptr instead so that a pending timer + /// does not keep a finished reactor alive. std::shared_ptr self_ref_; private: diff --git a/src/ray/ray_syncer/ray_syncer_bidi_reactor_base.h b/src/ray/ray_syncer/ray_syncer_bidi_reactor_base.h index 32356d8abfa3..ed39ed345ad7 100644 --- a/src/ray/ray_syncer/ray_syncer_bidi_reactor_base.h +++ b/src/ray/ray_syncer/ray_syncer_bidi_reactor_base.h @@ -233,9 +233,14 @@ class RaySyncerBidiReactorBase : public RaySyncerBidiReactor, public T { using T::StartRead; using T::StartWrite; + // NOTE: Every handler dispatched to io_context_ below captures a strong + // reference to the reactor (self_ref_). OnDone() releases self_ref_ from a + // handler on the same io_context_, and gRPC may invoke OnReadDone/OnWriteDone + // before that handler runs; without the capture, a handler queued earlier + // could run after the reactor has been destroyed (use-after-free). void OnWriteDone(bool ok) override { io_context_.dispatch( - [this, disconnected = IsDisconnected(), ok]() { + [this, self = self_ref_, disconnected = IsDisconnected(), ok]() { if (*disconnected) { return; } @@ -253,7 +258,7 @@ class RaySyncerBidiReactorBase : public RaySyncerBidiReactor, public T { void OnReadDone(bool ok) override { if (!ok) { io_context_.dispatch( - [this]() { + [this, self = self_ref_]() { RAY_LOG_EVERY_MS(INFO, 1000) << "Failed to read a message from node: " << NodeID::FromBinary(GetRemoteNodeID()); Disconnect(); @@ -309,7 +314,7 @@ class RaySyncerBidiReactorBase : public RaySyncerBidiReactor, public T { // task is queued and this read must post one. If not, a task is already queued and // will pick up all messages in the buffer later. if (buffer_was_empty) { - io_context_.dispatch([this]() { ProcessBufferedSyncMessages(); }, + io_context_.dispatch([this, self = self_ref_]() { ProcessBufferedSyncMessages(); }, "RaySyncer.ProcessBufferedSyncMessages"); } } diff --git a/src/ray/ray_syncer/ray_syncer_client.cc b/src/ray/ray_syncer/ray_syncer_client.cc index 7fcfa4e67851..8559ffc049de 100644 --- a/src/ray/ray_syncer/ray_syncer_client.cc +++ b/src/ray/ray_syncer/ray_syncer_client.cc @@ -55,8 +55,11 @@ RayClientBidiReactor::RayClientBidiReactor( } void RayClientBidiReactor::OnDone(const grpc::Status &status) { + // `self` keeps the reactor alive until this handler has finished, including + // through cleanup_cb_; resetting self_ref_ then drops the last owning + // reference once the handler returns. io_context_.dispatch( - [this, status]() { + [this, self = self_ref_, status]() { cleanup_cb_(this, !status.ok()); self_ref_.reset(); }, @@ -65,7 +68,7 @@ void RayClientBidiReactor::OnDone(const grpc::Status &status) { void RayClientBidiReactor::DoDisconnect() { io_context_.dispatch( - [this]() { + [this, self = self_ref_]() { StartWritesDone(); // Free the hold to allow OnDone being called. RemoveHold(); diff --git a/src/ray/ray_syncer/ray_syncer_server.cc b/src/ray/ray_syncer/ray_syncer_server.cc index b01724eac240..2f95b8dcac53 100644 --- a/src/ray/ray_syncer/ray_syncer_server.cc +++ b/src/ray/ray_syncer/ray_syncer_server.cc @@ -109,16 +109,20 @@ RayServerBidiReactor::RayServerBidiReactor( } void RayServerBidiReactor::DoDisconnect() { - io_context_.dispatch([this]() { Finish(grpc::Status::OK); }, ""); + io_context_.dispatch([this, self = self_ref_]() { Finish(grpc::Status::OK); }, ""); } void RayServerBidiReactor::OnCancel() { - io_context_.dispatch([this]() { Disconnect(); }, ""); + io_context_.dispatch([this, self = self_ref_]() { Disconnect(); }, ""); } void RayServerBidiReactor::OnDone() { + // See RayClientBidiReactor::OnDone for why `self` is captured here. io_context_.dispatch( - [this, cleanup_cb = cleanup_cb_, remote_node_id = GetRemoteNodeID()]() { + [this, + self = self_ref_, + cleanup_cb = cleanup_cb_, + remote_node_id = GetRemoteNodeID()]() { cleanup_cb(this, false); self_ref_.reset(); },