From aba475fe1d24d5c2581bc5763f40924bbf688500 Mon Sep 17 00:00:00 2001
From: Pavel Emelyanov <xemul@scylladb.com>
Date: Tue, 9 Nov 2021 18:48:04 +0300
Subject: [PATCH] storage_service: Remove bool from do_drain

The do_drain() today tells shutdown drain from API drain. The reason
is that compaction manager subscribes on the main's abort signal and
drains itself early. Thus, on regular drain it needs this extra kick
that would crash if called from shutdown drain.

This differentiation should sit in the compaction manager itself.

Signed-off-by: Pavel Emelyanov <xemul@scylladb.com>
---
 compaction/compaction_manager.cc |  4 ++++
 service/storage_service.cc       | 18 ++++++++----------
 service/storage_service.hh       |  2 +-
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/compaction/compaction_manager.cc b/compaction/compaction_manager.cc
index f456433530..86ffe3e51e 100644
--- a/compaction/compaction_manager.cc
+++ b/compaction/compaction_manager.cc
@@ -508,6 +508,10 @@ future<> compaction_manager::stop_ongoing_compactions(sstring reason) {
 }
 
 future<> compaction_manager::drain() {
+    if (!*_early_abort_subscription) {
+        return make_ready_future<>();
+    }
+
     _state = state::disabled;
     return stop_ongoing_compactions("drain");
 }
diff --git a/service/storage_service.cc b/service/storage_service.cc
index ba9d9ad3be..852e8d2e95 100644
--- a/service/storage_service.cc
+++ b/service/storage_service.cc
@@ -1304,7 +1304,7 @@ future<> storage_service::stop_transport() {
 future<> storage_service::drain_on_shutdown() {
     assert(this_shard_id() == 0);
     return (_operation_mode == mode::DRAINING || _operation_mode == mode::DRAINED) ?
-        _drain_finished.get_future() : do_drain(true);
+        _drain_finished.get_future() : do_drain();
 }
 
 future<> storage_service::init_messaging_service_part() {
@@ -2645,25 +2645,23 @@ future<> storage_service::drain() {
         }
 
         ss.set_mode(mode::DRAINING, "starting drain process", true);
-        return ss.do_drain(false).then([&ss] {
+        return ss.do_drain().then([&ss] {
             ss._drain_finished.set_value();
             ss.set_mode(mode::DRAINED, true);
         });
     });
 }
 
-future<> storage_service::do_drain(bool on_shutdown) {
-    return seastar::async([this, on_shutdown] {
+future<> storage_service::do_drain() {
+    return seastar::async([this] {
         stop_transport().get();
 
         tracing::tracing::tracing_instance().invoke_on_all(&tracing::tracing::shutdown).get();
 
-        if (!on_shutdown) {
-            // Interrupt on going compaction and shutdown to prevent further compaction
-            _db.invoke_on_all([] (auto& db) {
-                return db.get_compaction_manager().drain();
-            }).get();
-        }
+        // Interrupt on going compaction and shutdown to prevent further compaction
+        _db.invoke_on_all([] (auto& db) {
+            return db.get_compaction_manager().drain();
+        }).get();
 
         set_mode(mode::DRAINING, "flushing column families", false);
         flush_column_families();
diff --git a/service/storage_service.hh b/service/storage_service.hh
index 77d1a8273c..79940e2d31 100644
--- a/service/storage_service.hh
+++ b/service/storage_service.hh
@@ -802,7 +802,7 @@ public:
 private:
     promise<> _drain_finished;
     std::optional<shared_promise<>> _transport_stopped;
-    future<> do_drain(bool on_shutdown);
+    future<> do_drain();
     /**
      * Seed data to the endpoints that will be responsible for it at the future
      *