Skip to content

[core] Use shared_ptr for pins_in_flight_ instead of shared_from_this #58744

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
edoakes merged 5 commits into from
Nov 18, 2025
Merged

[core] Use shared_ptr for pins_in_flight_ instead of shared_from_this #58744

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
221ba3e
fix
edoakes Nov 18, 2025
dd56025
fix
edoakes Nov 18, 2025
94d09ac
fix
edoakes Nov 18, 2025
4cf886c
fix
edoakes Nov 18, 2025
5de14fc
fix
edoakes Nov 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions src/ray/raylet_rpc_client/raylet_client.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,8 @@ RayletClient::RayletClient(const rpc::Address &address,
::RayConfig::instance().raylet_rpc_server_reconnect_timeout_max_s(),
/*server_unavailable_timeout_callback=*/
std::move(raylet_unavailable_timeout_callback),
/*server_name=*/std::string("Raylet ") + address.ip_address())) {}
/*server_name=*/std::string("Raylet ") + address.ip_address())),
pins_in_flight_(std::make_shared<std::atomic<int64_t>>(0)) {}

void RayletClient::RequestWorkerLease(
const rpc::LeaseSpec &lease_spec,
Expand Down Expand Up @@ -335,11 +336,10 @@ 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 = [self, callback = std::move(callback)](
pins_in_flight_->fetch_add(1);
auto rpc_callback = [callback, pins_in_flight = pins_in_flight_](
Status status, rpc::PinObjectIDsReply &&reply) {
self->pins_in_flight_--;
pins_in_flight->fetch_sub(1);
callback(status, std::move(reply));
};
INVOKE_RETRYABLE_RPC_CALL(retryable_grpc_client_,
Expand Down
9 changes: 5 additions & 4 deletions src/ray/raylet_rpc_client/raylet_client.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,7 @@ 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,
public std::enable_shared_from_this<RayletClient> {
class RayletClient : public RayletClientInterface {
public:
/// Connect to the raylet.
///
Expand Down Expand Up @@ -160,7 +159,7 @@ class RayletClient : public RayletClientInterface,

const ResourceMappingType &GetResourceIDs() const { return resource_ids_; }

int64_t GetPinsInFlight() const override { return pins_in_flight_.load(); }
int64_t GetPinsInFlight() const override { return pins_in_flight_->load(); }

void GetNodeStats(const rpc::GetNodeStatsRequest &request,
const rpc::ClientCallback<rpc::GetNodeStatsReply> &callback) override;
Expand Down Expand Up @@ -188,7 +187,9 @@ class RayletClient : public RayletClientInterface,
ResourceMappingType resource_ids_;

/// The number of object ID pin RPCs currently in flight.
std::atomic<int64_t> pins_in_flight_ = 0;
/// NOTE: `shared_ptr` because it is captured in a callback that can outlive this
/// instance.
std::shared_ptr<std::atomic<int64_t>> pins_in_flight_;
};

} // namespace rpc
Expand Down