From c9fb62437133ddcd5d17c548aa9f0fa5eab3f59d Mon Sep 17 00:00:00 2001
From: art-gor <artgor4github@gmail.com>
Date: Tue, 20 Apr 2021 13:54:35 +0300
Subject: [PATCH] workaround in old scheduler due to possible lifetime issues
 (#135)

---
 .../protocol/common/asio/asio_scheduler.hpp   |  5 +++--
 src/protocol/common/asio/asio_scheduler.cpp   | 19 ++++++++++++-------
 .../common/asio/asio_scheduler_test.cpp       |  2 +-
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/include/libp2p/protocol/common/asio/asio_scheduler.hpp b/include/libp2p/protocol/common/asio/asio_scheduler.hpp
index d31d4a5e..81ad616e 100644
--- a/include/libp2p/protocol/common/asio/asio_scheduler.hpp
+++ b/include/libp2p/protocol/common/asio/asio_scheduler.hpp
@@ -14,7 +14,8 @@ namespace libp2p::protocol {
 
   class AsioScheduler : public Scheduler {
    public:
-    AsioScheduler(boost::asio::io_context &io, SchedulerConfig config);
+    AsioScheduler(const std::shared_ptr<boost::asio::io_context> &io,
+                  SchedulerConfig config);
 
     ~AsioScheduler() override;
 
@@ -27,7 +28,7 @@ namespace libp2p::protocol {
 
     void onImmediate();
 
-    boost::asio::io_context &io_;
+    std::weak_ptr<boost::asio::io_context> io_;
     Ticks interval_;
     boost::asio::deadline_timer timer_;
     boost::posix_time::ptime started_;
diff --git a/src/protocol/common/asio/asio_scheduler.cpp b/src/protocol/common/asio/asio_scheduler.cpp
index c155bbc6..9b9ce67a 100644
--- a/src/protocol/common/asio/asio_scheduler.cpp
+++ b/src/protocol/common/asio/asio_scheduler.cpp
@@ -9,11 +9,11 @@
 
 namespace libp2p::protocol {
 
-  AsioScheduler::AsioScheduler(boost::asio::io_context &io,
+  AsioScheduler::AsioScheduler(const std::shared_ptr<boost::asio::io_context> &io,
                                SchedulerConfig config)
       : io_(io),
         interval_(config.period_msec),
-        timer_(io, boost::posix_time::milliseconds(interval_)),
+        timer_(*io, boost::posix_time::milliseconds(interval_)),
         started_(boost::posix_time::microsec_clock::local_time()),
         canceled_(std::make_shared<bool>(false)),
         timer_cb_([this, canceled{canceled_}](
@@ -28,7 +28,9 @@ namespace libp2p::protocol {
 
   AsioScheduler::~AsioScheduler() {
     *canceled_ = true;
-    timer_.cancel();
+    if (!io_.expired()) {
+      timer_.cancel();
+    }
   };
 
   Scheduler::Ticks AsioScheduler::now() const {
@@ -39,15 +41,18 @@ namespace libp2p::protocol {
   void AsioScheduler::scheduleImmediate() {
     if (!immediate_cb_scheduled_) {
       immediate_cb_scheduled_ = true;
-      io_.post(immediate_cb_);
+      auto io = io_.lock();
+      if (io) io->post(immediate_cb_);
     }
   }
 
   void AsioScheduler::onTimer() {
     pulse(false);
-    timer_.expires_at(timer_.expires_at()
-                      + boost::posix_time::milliseconds(interval_));
-    timer_.async_wait(timer_cb_);
+    if (!io_.expired()) {
+      timer_.expires_at(timer_.expires_at()
+                        + boost::posix_time::milliseconds(interval_));
+      timer_.async_wait(timer_cb_);
+    }
   }
 
   void AsioScheduler::onImmediate() {
diff --git a/test/libp2p/protocol/common/asio/asio_scheduler_test.cpp b/test/libp2p/protocol/common/asio/asio_scheduler_test.cpp
index 79667b81..ace90727 100644
--- a/test/libp2p/protocol/common/asio/asio_scheduler_test.cpp
+++ b/test/libp2p/protocol/common/asio/asio_scheduler_test.cpp
@@ -21,7 +21,7 @@ TEST(AsioScheduler, Construct) {
   auto context = injector.create<std::shared_ptr<boost::asio::io_context>>();
   std::shared_ptr<Scheduler> scheduler =
       std::make_shared<libp2p::protocol::AsioScheduler>(
-          *context, libp2p::protocol::SchedulerConfig{});
+          context, libp2p::protocol::SchedulerConfig{});
   scheduler.reset();
   context->run();
 }