Use shared_ptr for gcs client in profiler (#6150)

src/ray/core_worker/core_worker.cc

@@ -78,7 +78,6 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
       client_call_manager_(new rpc::ClientCallManager(io_service_)),
       heartbeat_timer_(io_service_),
       core_worker_server_(WorkerTypeString(worker_type), 0 /* let grpc choose a port */),
-      gcs_client_(gcs_options),
       memory_store_(std::make_shared<CoreWorkerMemoryStore>()),
       memory_store_provider_(memory_store_),
       task_execution_service_work_(task_execution_service_),
@@ -95,7 +94,8 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
   }
 
   // Initialize gcs client.
-  RAY_CHECK_OK(gcs_client_.Connect(io_service_));
+  gcs_client_ = std::make_shared<gcs::RedisGcsClient>(gcs_options);
+  RAY_CHECK_OK(gcs_client_->Connect(io_service_));
 
   // Initialize profiler.
   profiler_ = std::make_shared<worker::Profiler>(worker_context_, node_ip_address,
@@ -171,7 +171,7 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
 
     std::shared_ptr<gcs::TaskTableData> data = std::make_shared<gcs::TaskTableData>();
     data->mutable_task()->mutable_task_spec()->CopyFrom(builder.Build().GetMessage());
-    RAY_CHECK_OK(gcs_client_.raylet_task_table().Add(job_id, task_id, data, nullptr));
+    RAY_CHECK_OK(gcs_client_->raylet_task_table().Add(job_id, task_id, data, nullptr));
     SetCurrentTaskId(task_id);
   }
 
@@ -209,7 +209,7 @@ void CoreWorker::Shutdown() {
 
 void CoreWorker::Disconnect() {
   io_service_.stop();
-  gcs_client_.Disconnect();
+  gcs_client_->Disconnect();
   if (raylet_client_) {
     RAY_IGNORE_EXPR(raylet_client_->Disconnect());
   }
@@ -232,7 +232,7 @@ void CoreWorker::SetCurrentTaskId(const TaskID &task_id) {
   // Clear all actor handles at the end of each non-actor task.
   if (actor_id_.IsNil() && task_id.IsNil()) {
     for (const auto &handle : actor_handles_) {
-      RAY_CHECK_OK(gcs_client_.Actors().AsyncUnsubscribe(handle.first, nullptr));
+      RAY_CHECK_OK(gcs_client_->Actors().AsyncUnsubscribe(handle.first, nullptr));
     }
     actor_handles_.clear();
   }
@@ -642,7 +642,7 @@ bool CoreWorker::AddActorHandle(std::unique_ptr<ActorHandle> actor_handle) {
           it->second->Reset();
         }
       } else if (actor_data.state() == gcs::ActorTableData::DEAD) {
-        RAY_CHECK_OK(gcs_client_.Actors().AsyncUnsubscribe(actor_id, nullptr));
+        RAY_CHECK_OK(gcs_client_->Actors().AsyncUnsubscribe(actor_id, nullptr));
         // We cannot erase the actor handle here because clients can still
         // submit tasks to dead actors.
       }
@@ -655,7 +655,7 @@ bool CoreWorker::AddActorHandle(std::unique_ptr<ActorHandle> actor_handle) {
                     << ", port: " << actor_data.port();
     };
 
-    RAY_CHECK_OK(gcs_client_.Actors().AsyncSubscribe(
+    RAY_CHECK_OK(gcs_client_->Actors().AsyncSubscribe(
         actor_id, actor_notification_callback, nullptr));
   }
   return inserted;

src/ray/core_worker/core_worker.h

@@ -1,9 +1,7 @@
 #ifndef RAY_CORE_WORKER_CORE_WORKER_H
 #define RAY_CORE_WORKER_CORE_WORKER_H
 
-#include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
-#include "absl/synchronization/mutex.h"
 
 #include "ray/common/buffer.h"
 #include "ray/core_worker/actor_handle.h"
@@ -321,11 +319,11 @@ class CoreWorker {
                                const std::vector<std::shared_ptr<Buffer>> &metadatas,
                                std::vector<std::shared_ptr<RayObject>> *return_objects);
 
-  /**
-   * The following methods are handlers for the core worker's gRPC server, which follow
-   * a macro-generated call convention. These are executed on the io_service_ and
-   * post work to the appropriate event loop.
-   */
+  ///
+  /// The following methods are handlers for the core worker's gRPC server, which follow
+  /// a macro-generated call convention. These are executed on the io_service_ and
+  /// post work to the appropriate event loop.
+  ///
 
   /// Implements gRPC server handler.
   void HandleAssignTask(const rpc::AssignTaskRequest &request,
@@ -462,7 +460,7 @@ class CoreWorker {
   rpc::GrpcServer core_worker_server_;
 
   // Client to the GCS shared by core worker interfaces.
-  gcs::RedisGcsClient gcs_client_;
+  std::shared_ptr<gcs::RedisGcsClient> gcs_client_;
 
   // Client to the raylet shared by core worker interfaces.
   std::unique_ptr<RayletClient> raylet_client_;

src/ray/core_worker/profiling.cc

@@ -6,7 +6,7 @@ namespace ray {
 
 namespace worker {
 
-ProfileEvent::ProfileEvent(const std::shared_ptr<Profiler> profiler,
+ProfileEvent::ProfileEvent(const std::shared_ptr<Profiler> &profiler,
                            const std::string &event_type)
     : profiler_(profiler) {
   rpc_event_.set_event_type(event_type);
@@ -14,11 +14,11 @@ ProfileEvent::ProfileEvent(const std::shared_ptr<Profiler> profiler,
 }
 
 Profiler::Profiler(WorkerContext &worker_context, const std::string &node_ip_address,
-                   boost::asio::io_service &io_service, gcs::RedisGcsClient &gcs_client)
+                   boost::asio::io_service &io_service,
+                   const std::shared_ptr<gcs::RedisGcsClient> &gcs_client)
     : io_service_(io_service),
       timer_(io_service_, boost::asio::chrono::seconds(1)),
       gcs_client_(gcs_client) {
-  absl::MutexLock l(&mu_);
   rpc_profile_data_.set_component_type(WorkerTypeString(worker_context.GetWorkerType()));
   rpc_profile_data_.set_component_id(worker_context.GetWorkerID().Binary());
   rpc_profile_data_.set_node_ip_address(node_ip_address);
@@ -26,16 +26,16 @@ Profiler::Profiler(WorkerContext &worker_context, const std::string &node_ip_add
 }
 
 void Profiler::AddEvent(const rpc::ProfileTableData::ProfileEvent &event) {
-  absl::MutexLock l(&mu_);
+  absl::MutexLock lock(&mutex_);
   rpc_profile_data_.add_profile_events()->CopyFrom(event);
 }
 
 void Profiler::FlushEvents() {
-  absl::MutexLock l(&mu_);
+  absl::MutexLock lock(&mutex_);
   if (rpc_profile_data_.profile_events_size() != 0) {
     // TODO(edoakes): this should be migrated to use the new GCS client interface
     // instead of the raw table interface once it's ready.
-    if (!gcs_client_.profile_table().AddProfileEventBatch(rpc_profile_data_).ok()) {
+    if (!gcs_client_->profile_table().AddProfileEventBatch(rpc_profile_data_).ok()) {
       RAY_LOG(WARNING) << "Failed to push profile events to GCS.";
     } else {
       RAY_LOG(DEBUG) << "Pushed " << rpc_profile_data_.profile_events_size()

src/ray/core_worker/profiling.h

@@ -1,6 +1,7 @@
 #ifndef RAY_CORE_WORKER_PROFILING_H
 #define RAY_CORE_WORKER_PROFILING_H
 
+#include "absl/base/thread_annotations.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/time/clock.h"
 
@@ -14,14 +15,18 @@ namespace worker {
 class Profiler {
  public:
   Profiler(WorkerContext &worker_context, const std::string &node_ip_address,
-           boost::asio::io_service &io_service, gcs::RedisGcsClient &gcs_client);
+           boost::asio::io_service &io_service,
+           const std::shared_ptr<gcs::RedisGcsClient> &gcs_client);
 
   // Add an event to the queue to be flushed periodically.
-  void AddEvent(const rpc::ProfileTableData::ProfileEvent &event);
+  void AddEvent(const rpc::ProfileTableData::ProfileEvent &event) LOCKS_EXCLUDED(mutex_);
 
  private:
   // Flush all of the events that have been added since last flush to the GCS.
-  void FlushEvents();
+  void FlushEvents() LOCKS_EXCLUDED(mutex_);
+
+  // Mutex guarding rpc_profile_data_.
+  absl::Mutex mutex_;
 
   // ASIO IO service event loop. Must be started by the caller.
   boost::asio::io_service &io_service_;
@@ -31,28 +36,32 @@ class Profiler {
 
   // RPC message containing profiling data. Holds the queue of profile events
   // until they are flushed.
-  rpc::ProfileTableData rpc_profile_data_ GUARDED_BY(mu_);
+  rpc::ProfileTableData rpc_profile_data_ GUARDED_BY(mutex_);
 
-  gcs::RedisGcsClient &gcs_client_;
-
-  absl::Mutex mu_;
+  // Client to the GCS used to push profile events to it.
+  std::shared_ptr<gcs::RedisGcsClient> gcs_client_;
 };
 
 class ProfileEvent {
  public:
-  ProfileEvent(const std::shared_ptr<Profiler> profiler, const std::string &event_type);
+  ProfileEvent(const std::shared_ptr<Profiler> &profiler, const std::string &event_type);
 
+  // Set the end time for the event and add it to the profiler.
   ~ProfileEvent() {
     rpc_event_.set_end_time(absl::GetCurrentTimeNanos() / 1e9);
     profiler_->AddEvent(rpc_event_);
   }
 
+  // Set extra metadata for the event, which could change during the event.
   void SetExtraData(const std::string &extra_data) {
     rpc_event_.set_extra_data(extra_data);
   }
 
  private:
-  const std::shared_ptr<Profiler> profiler_;
+  // shared_ptr to the profiler that this event will be added to when it is destructed.
+  std::shared_ptr<Profiler> profiler_;
+
+  // Underlying proto data structure that holds the event data.
   rpc::ProfileTableData::ProfileEvent rpc_event_;
 };