shared ptr and weak ptr

Signed-off-by: Ruiyang Wang <[email protected]>

Author: rynewang

src/ray/common/asio/periodical_runner.cc

@@ -29,7 +29,6 @@ PeriodicalRunner::~PeriodicalRunner() {
 
 void PeriodicalRunner::Clear() {
   absl::MutexLock lock(&mutex_);
-  *stopped_ = true;
   for (const auto &timer : timers_) {
     timer->cancel();
   }
@@ -39,28 +38,23 @@ void PeriodicalRunner::Clear() {
 void PeriodicalRunner::RunFnPeriodically(std::function<void()> fn,
                                          uint64_t period_ms,
                                          const std::string &name) {
-  *stopped_ = false;
   if (period_ms > 0) {
     auto timer = std::make_shared<boost::asio::deadline_timer>(io_service_);
     {
       absl::MutexLock lock(&mutex_);
       timers_.push_back(timer);
     }
+    auto weak_self = weak_from_this();
     io_service_.post(
-        [this,
-         stopped = stopped_,
-         fn = std::move(fn),
-         period_ms,
-         name,
-         timer = std::move(timer)]() {
-          if (*stopped) {
-            return;
-          }
-          if (RayConfig::instance().event_stats()) {
-            DoRunFnPeriodicallyInstrumented(
-                fn, boost::posix_time::milliseconds(period_ms), timer, name);
-          } else {
-            DoRunFnPeriodically(fn, boost::posix_time::milliseconds(period_ms), timer);
+        [weak_self, fn = std::move(fn), period_ms, name, timer = std::move(timer)]() {
+          if (auto self = weak_self.lock(); self) {
+            if (RayConfig::instance().event_stats()) {
+              self->DoRunFnPeriodicallyInstrumented(
+                  std::move(fn), boost::posix_time::milliseconds(period_ms), timer, name);
+            } else {
+              self->DoRunFnPeriodically(
+                  std::move(fn), boost::posix_time::milliseconds(period_ms), timer);
+            }
           }
         },
         "PeriodicalRunner.RunFnPeriodically");
@@ -74,19 +68,19 @@ void PeriodicalRunner::DoRunFnPeriodically(
   fn();
   absl::MutexLock lock(&mutex_);
   timer->expires_from_now(period);
-  timer->async_wait([this, stopped = stopped_, fn, period, timer = std::move(timer)](
+  auto weak_self = weak_from_this();
+  timer->async_wait([weak_self, fn, period, timer = std::move(timer)](
                         const boost::system::error_code &error) {
-    if (*stopped) {
-      return;
-    }
-    if (error == boost::asio::error::operation_aborted) {
-      // `operation_aborted` is set when `timer` is canceled or destroyed.
-      // The Monitor lifetime may be short than the object who use it. (e.g.
-      // gcs_server)
-      return;
+    if (auto self = weak_self.lock(); self) {
+      if (error == boost::asio::error::operation_aborted) {
+        // `operation_aborted` is set when `timer` is canceled or destroyed.
+        // The Monitor lifetime may be short than the object who use it. (e.g.
+        // gcs_server)
+        return;
+      }
+      RAY_CHECK(!error) << error.message();
+      self->DoRunFnPeriodically(fn, period, timer);
     }
-    RAY_CHECK(!error) << error.message();
-    DoRunFnPeriodically(fn, period, timer);
   });
 }
 
@@ -102,31 +96,29 @@ void PeriodicalRunner::DoRunFnPeriodicallyInstrumented(
   // which the handler was elgible to execute on the event loop but was queued by the
   // event loop.
   auto stats_handle = io_service_.stats().RecordStart(name, period.total_nanoseconds());
-  timer->async_wait([this,
+  auto weak_self = weak_from_this();
+  timer->async_wait([weak_self,
                      fn,
-                     stopped = stopped_,
                      period,
                      timer = std::move(timer),
                      stats_handle = std::move(stats_handle),
                      name](const boost::system::error_code &error) {
-    if (*stopped) {
-      return;
+    if (auto self = weak_self.lock(); self) {
+      self->io_service_.stats().RecordExecution(
+          [weak_self, fn, error, period, timer, name]() {
+            if (auto self = weak_self.lock(); self) {
+              if (error == boost::asio::error::operation_aborted) {
+                // `operation_aborted` is set when `timer` is canceled or destroyed.
+                // The Monitor lifetime may be short than the object who use it. (e.g.
+                // gcs_server)
+                return;
+              }
+              RAY_CHECK(!error) << error.message();
+              self->DoRunFnPeriodicallyInstrumented(fn, period, timer, name);
+            }
+          },
+          stats_handle);
     }
-    io_service_.stats().RecordExecution(
-        [this, stopped = stopped, fn, error, period, timer, name]() {
-          if (*stopped) {
-            return;
-          }
-          if (error == boost::asio::error::operation_aborted) {
-            // `operation_aborted` is set when `timer` is canceled or destroyed.
-            // The Monitor lifetime may be short than the object who use it. (e.g.
-            // gcs_server)
-            return;
-          }
-          RAY_CHECK(!error) << error.message();
-          DoRunFnPeriodicallyInstrumented(fn, period, timer, name);
-        },
-        stats_handle);
   });
 }
 

src/ray/common/asio/periodical_runner.h

@@ -28,19 +28,28 @@ namespace ray {
 /// It can run functions with specified period. Each function is triggered by its timer.
 /// To run a function, call `RunFnPeriodically(fn, period_ms)`.
 /// All registered functions will stop running once this object is destructed.
-class PeriodicalRunner {
+//
+// Lifetime: once a PeriodicalRunner is destructed, all its timers are cancelled. The
+// scheduled asio tasks keep a weak_ptr to the PeriodicalRunner, and they won't run after
+// the PeriodicalRunner is destructed.
+class PeriodicalRunner : public std::enable_shared_from_this<PeriodicalRunner> {
  public:
-  explicit PeriodicalRunner(instrumented_io_context &io_service);
+  static std::shared_ptr<PeriodicalRunner> Create(instrumented_io_context &io_service) {
+    // Sadly we can't use std::make_shared because the constructor is private.
+    return std::shared_ptr<PeriodicalRunner>(new PeriodicalRunner(io_service));
+  }
 
   ~PeriodicalRunner();
 
-  void Clear();
-
   void RunFnPeriodically(std::function<void()> fn,
                          uint64_t period_ms,
                          const std::string &name) ABSL_LOCKS_EXCLUDED(mutex_);
 
  private:
+  explicit PeriodicalRunner(instrumented_io_context &io_service);
+
+  void Clear();
+
   void DoRunFnPeriodically(const std::function<void()> &fn,
                            boost::posix_time::milliseconds period,
                            std::shared_ptr<boost::asio::deadline_timer> timer)
@@ -56,9 +65,6 @@ class PeriodicalRunner {
   mutable absl::Mutex mutex_;
   std::vector<std::shared_ptr<boost::asio::deadline_timer>> timers_
       ABSL_GUARDED_BY(mutex_);
-  // `stopped_` is copied to the timer callback, and may outlive `this`.
-  std::shared_ptr<std::atomic<bool>> stopped_ =
-      std::make_shared<std::atomic<bool>>(false);
 };
 
 }  // namespace ray

src/ray/common/file_system_monitor.cc

@@ -33,10 +33,10 @@ FileSystemMonitor::FileSystemMonitor(std::vector<std::string> paths,
         boost::asio::io_service::work io_service_work_(io_context_);
         io_context_.run();
       }),
-      runner_(io_context_) {
-  runner_.RunFnPeriodically([this] { over_capacity_ = CheckIfAnyPathOverCapacity(); },
-                            monitor_interval_ms,
-                            "FileSystemMonitor.CheckIfAnyPathOverCapacity");
+      runner_(PeriodicalRunner::Create(io_context_)) {
+  runner_->RunFnPeriodically([this] { over_capacity_ = CheckIfAnyPathOverCapacity(); },
+                             monitor_interval_ms,
+                             "FileSystemMonitor.CheckIfAnyPathOverCapacity");
 }
 
 FileSystemMonitor::FileSystemMonitor()

src/ray/common/file_system_monitor.h

@@ -70,7 +70,7 @@ class FileSystemMonitor {
   std::atomic<bool> over_capacity_;
   instrumented_io_context io_context_;
   std::thread monitor_thread_;
-  PeriodicalRunner runner_;
+  std::shared_ptr<PeriodicalRunner> runner_;
 };
 
 std::vector<std::string> ParseSpillingPaths(const std::string &spilling_config);

src/ray/common/memory_monitor.cc

@@ -34,7 +34,7 @@ MemoryMonitor::MemoryMonitor(instrumented_io_context &io_service,
     : usage_threshold_(usage_threshold),
       min_memory_free_bytes_(min_memory_free_bytes),
       monitor_callback_(monitor_callback),
-      runner_(io_service) {
+      runner_(PeriodicalRunner::Create(io_service)) {
   RAY_CHECK(monitor_callback_ != nullptr);
   RAY_CHECK_GE(usage_threshold_, 0);
   RAY_CHECK_LE(usage_threshold_, 1);
@@ -48,7 +48,7 @@ MemoryMonitor::MemoryMonitor(instrumented_io_context &io_service,
                   << computed_threshold_bytes_ << " bytes ("
                   << FormatFloat(computed_threshold_fraction_, 2)
                   << " system memory), total system memory bytes: " << total_memory_bytes;
-    runner_.RunFnPeriodically(
+    runner_->RunFnPeriodically(
         [this] {
           auto [used_memory_bytes, total_memory_bytes] = GetMemoryBytes();
           MemorySnapshot system_memory;

src/ray/common/memory_monitor.h

@@ -233,7 +233,7 @@ class MemoryMonitor {
   /// Callback function that executes at each monitoring interval,
   /// on a dedicated thread managed by this class.
   const MemoryUsageRefreshCallback monitor_callback_;
-  PeriodicalRunner runner_;
+  std::shared_ptr<PeriodicalRunner> runner_;
 };
 
 }  // namespace ray

src/ray/common/ray_syncer/ray_syncer.cc

@@ -29,7 +29,7 @@ RaySyncer::RaySyncer(instrumented_io_context &io_context,
     : io_context_(io_context),
       local_node_id_(local_node_id),
       node_state_(std::make_unique<NodeState>()),
-      timer_(io_context) {
+      timer_(PeriodicalRunner::Create(io_context)) {
   stopped_ = std::make_shared<bool>(false);
 }
 
@@ -155,7 +155,7 @@ void RaySyncer::Register(MessageType message_type,
 
         // Set job to pull from reporter periodically
         if (reporter != nullptr && pull_from_reporter_interval_ms > 0) {
-          timer_.RunFnPeriodically(
+          timer_->RunFnPeriodically(
               [this, stopped = stopped_, message_type]() {
                 if (*stopped) {
                   return;

src/ray/common/ray_syncer/ray_syncer.h

@@ -176,7 +176,7 @@ class RaySyncer {
   std::unique_ptr<NodeState> node_state_;
 
   /// Timer is used to do broadcasting.
-  ray::PeriodicalRunner timer_;
+  std::shared_ptr<PeriodicalRunner> timer_;
 
   friend class RaySyncerService;
   /// Test purpose

src/ray/common/test/syncer_service_e2e_test.cc

@@ -28,12 +28,13 @@
 #include "ray/common/ray_syncer/ray_syncer.h"
 using namespace std;
 using namespace ray::syncer;
+using ray::PeriodicalRunner;
 
 class LocalNode : public ReporterInterface {
  public:
   LocalNode(instrumented_io_context &io_context, ray::NodeID node_id)
-      : node_id_(node_id), timer_(io_context) {
-    timer_.RunFnPeriodically(
+      : node_id_(node_id), timer_(PeriodicalRunner::Create(io_context)) {
+    timer_->RunFnPeriodically(
         [this]() {
           auto v = static_cast<double>(std::rand()) / RAND_MAX;
           if (v < 0.3) {
@@ -66,7 +67,7 @@ class LocalNode : public ReporterInterface {
   int64_t version_ = 0;
   int state_ = 0;
   ray::NodeID node_id_;
-  ray::PeriodicalRunner timer_;
+  std::shared_ptr<PeriodicalRunner> timer_;
 };
 
 class RemoteNodes : public ReceiverInterface {

src/ray/core_worker/core_worker.cc

@@ -257,7 +257,7 @@ CoreWorker::CoreWorker(CoreWorkerOptions options, const WorkerID &worker_id)
       worker_context_(options_.worker_type, worker_id, GetProcessJobID(options_)),
       io_work_(io_service_),
       client_call_manager_(new rpc::ClientCallManager(io_service_)),
-      periodical_runner_(io_service_),
+      periodical_runner_(PeriodicalRunner::Create(io_service_)),
       task_queue_length_(0),
       num_executed_tasks_(0),
       resource_ids_(new ResourceMappingType()),
@@ -327,7 +327,7 @@ CoreWorker::CoreWorker(CoreWorkerOptions options, const WorkerID &worker_id)
       options_.raylet_ip_address, options_.node_manager_port, *client_call_manager_);
 
   if (options_.worker_type != WorkerType::DRIVER) {
-    periodical_runner_.RunFnPeriodically(
+    periodical_runner_->RunFnPeriodically(
         [this] { ExitIfParentRayletDies(); },
         RayConfig::instance().raylet_death_check_interval_milliseconds(),
         "CoreWorker.ExitIfParentRayletDies");
@@ -451,11 +451,11 @@ CoreWorker::CoreWorker(CoreWorkerOptions options, const WorkerID &worker_id)
       });
 
   object_info_publisher_ = std::make_unique<pubsub::Publisher>(
-      /*channels=*/std::vector<
-          rpc::ChannelType>{rpc::ChannelType::WORKER_OBJECT_EVICTION,
-                            rpc::ChannelType::WORKER_REF_REMOVED_CHANNEL,
-                            rpc::ChannelType::WORKER_OBJECT_LOCATIONS_CHANNEL},
-      /*periodical_runner=*/&periodical_runner_,
+      /*channels=*/
+      std::vector<rpc::ChannelType>{rpc::ChannelType::WORKER_OBJECT_EVICTION,
+                                    rpc::ChannelType::WORKER_REF_REMOVED_CHANNEL,
+                                    rpc::ChannelType::WORKER_OBJECT_LOCATIONS_CHANNEL},
+      /*periodical_runner=*/*periodical_runner_,
       /*get_time_ms=*/[]() { return absl::GetCurrentTimeNanos() / 1e6; },
       /*subscriber_timeout_ms=*/RayConfig::instance().subscriber_timeout_ms(),
       /*publish_batch_size_=*/RayConfig::instance().publish_batch_size(),
@@ -791,7 +791,7 @@ CoreWorker::CoreWorker(CoreWorkerOptions options, const WorkerID &worker_id)
   const auto event_stats_print_interval_ms =
       RayConfig::instance().event_stats_print_interval_ms();
   if (event_stats_print_interval_ms != -1 && RayConfig::instance().event_stats()) {
-    periodical_runner_.RunFnPeriodically(
+    periodical_runner_->RunFnPeriodically(
         [this] {
           RAY_LOG(INFO) << "Event stats:\n\n"
                         << io_service_.stats().StatsString() << "\n\n"
@@ -811,7 +811,7 @@ CoreWorker::CoreWorker(CoreWorkerOptions options, const WorkerID &worker_id)
       ray::rpc::Event_SourceType::Event_SourceType_CORE_WORKER,
       {{"worker_id", worker_id.Hex()}});
 
-  periodical_runner_.RunFnPeriodically(
+  periodical_runner_->RunFnPeriodically(
       [this] {
         const auto lost_objects = reference_counter_->FlushObjectsToRecover();
         if (!lost_objects.empty()) {
@@ -836,17 +836,17 @@ CoreWorker::CoreWorker(CoreWorkerOptions options, const WorkerID &worker_id)
       100,
       "CoreWorker.RecoverObjects");
 
-  periodical_runner_.RunFnPeriodically(
+  periodical_runner_->RunFnPeriodically(
       [this] { InternalHeartbeat(); },
       RayConfig::instance().core_worker_internal_heartbeat_ms(),
       "CoreWorker.InternalHeartbeat");
 
-  periodical_runner_.RunFnPeriodically(
+  periodical_runner_->RunFnPeriodically(
       [this] { RecordMetrics(); },
       RayConfig::instance().metrics_report_interval_ms() / 2,
       "CoreWorker.RecordMetrics");
 
-  periodical_runner_.RunFnPeriodically(
+  periodical_runner_->RunFnPeriodically(
       [this] { TryDeleteObjectRefStreams(); },
       RayConfig::instance().local_gc_min_interval_s() * 1000,
       "CoreWorker.GCStreamingGeneratorMetadata");
@@ -3014,9 +3014,9 @@ std::unique_ptr<worker::ProfileEvent> CoreWorker::CreateProfileEvent(
 }
 
 void CoreWorker::RunTaskExecutionLoop() {
-  PeriodicalRunner signal_checker(task_execution_service_);
+  auto signal_checker = PeriodicalRunner::Create(task_execution_service_);
   if (options_.check_signals) {
-    signal_checker.RunFnPeriodically(
+    signal_checker->RunFnPeriodically(
         [this] {
           /// The overhead of this is only a single digit microsecond.
           auto status = options_.check_signals();

src/ray/core_worker/core_worker.h

@@ -1695,7 +1695,7 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   std::shared_ptr<rpc::CoreWorkerClientPool> core_worker_client_pool_;
 
   /// The runner to run function periodically.
-  PeriodicalRunner periodical_runner_;
+  std::shared_ptr<PeriodicalRunner> periodical_runner_;
 
   /// RPC server used to receive tasks to execute.
   std::unique_ptr<rpc::GrpcServer> core_worker_server_;

src/ray/core_worker/task_event_buffer.cc

@@ -232,7 +232,7 @@ bool TaskEventBuffer::RecordTaskStatusEventIfNeeded(
 
 TaskEventBufferImpl::TaskEventBufferImpl(std::shared_ptr<gcs::GcsClient> gcs_client)
     : work_guard_(boost::asio::make_work_guard(io_service_)),
-      periodical_runner_(io_service_),
+      periodical_runner_(PeriodicalRunner::Create(io_service_)),
       gcs_client_(std::move(gcs_client)),
       status_events_() {}
 
@@ -283,9 +283,9 @@ Status TaskEventBufferImpl::Start(bool auto_flush) {
   }
 
   RAY_LOG(INFO) << "Reporting task events to GCS every " << report_interval_ms << "ms.";
-  periodical_runner_.RunFnPeriodically([this] { FlushEvents(/* forced */ false); },
-                                       report_interval_ms,
-                                       "CoreWorker.deadline_timer.flush_task_events");
+  periodical_runner_->RunFnPeriodically([this] { FlushEvents(/*forced= */ false); },
+                                        report_interval_ms,
+                                        "CoreWorker.deadline_timer.flush_task_events");
   return Status::OK();
 }
 

src/ray/core_worker/task_event_buffer.h

@@ -430,7 +430,7 @@ class TaskEventBufferImpl : public TaskEventBuffer {
   std::thread io_thread_;
 
   /// The runner to run function periodically.
-  PeriodicalRunner periodical_runner_;
+  std::shared_ptr<PeriodicalRunner> periodical_runner_;
 
   /// Client to the GCS used to push profile events to it.
   std::shared_ptr<gcs::GcsClient> gcs_client_ ABSL_GUARDED_BY(mutex_);