From afa2b40ac99fdfd633ca79c52e1191ece531a8f8 Mon Sep 17 00:00:00 2001
From: Benny Halevy <bhalevy@scylladb.com>
Date: Sun, 1 Jun 2025 10:38:17 +0300
Subject: [PATCH] disk_space_monitor: add space_source_registration

Register the current space_source_fn in an RAII
object that resets monitor._space_source to the
previous function when the RAII object is destroyed.

Use space_source_registration in database_test::
mutation_dump_generated_schema_deterministic_id_version
to prevent use-after-stack-return in the test.

Fixes #24314

Signed-off-by: Benny Halevy <bhalevy@scylladb.com>

Closes scylladb/scylladb#24342

(cherry picked from commit 8b387109fcb28bca1c1427cc56aea6f31c4f9ca7)

Closes scylladb/scylladb#24392
---
 test/boost/database_test.cc |  2 +-
 utils/disk_space_monitor.cc | 10 ++++++++++
 utils/disk_space_monitor.hh | 17 ++++++++++++++++-
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/test/boost/database_test.cc b/test/boost/database_test.cc
index 6477a2f745..a7e71d1da7 100644
--- a/test/boost/database_test.cc
+++ b/test/boost/database_test.cc
@@ -1615,7 +1615,7 @@ SEASTAR_TEST_CASE(test_disk_space_monitor_capacity_override) {
             .free = 12,
             .available = 11,
         };
-        monitor.set_space_source([&] { return make_ready_future<std::filesystem::space_info>(orig_space); });
+        auto reg = monitor.set_space_source([&] { return make_ready_future<std::filesystem::space_info>(orig_space); });
 
         utils::phased_barrier poll_barrier("poll_barrier"); // new operation started whenever monitor calls listeners.
         auto op = poll_barrier.start();
diff --git a/utils/disk_space_monitor.cc b/utils/disk_space_monitor.cc
index 92b8a8c797..79d6eba0d0 100644
--- a/utils/disk_space_monitor.cc
+++ b/utils/disk_space_monitor.cc
@@ -43,6 +43,16 @@ disk_space_monitor::~disk_space_monitor() {
     SCYLLA_ASSERT(_poller_fut.available());
 }
 
+disk_space_monitor::space_source_registration::space_source_registration(disk_space_monitor& m)
+    : _monitor(m)
+    , _prev_space_source(m._space_source)
+{
+}
+
+disk_space_monitor::space_source_registration::~space_source_registration() {
+    _monitor._space_source = _prev_space_source;
+}
+
 future<> disk_space_monitor::start() {
     _space_info = co_await get_filesystem_space();
     _poller_fut = poll();
diff --git a/utils/disk_space_monitor.hh b/utils/disk_space_monitor.hh
index d52f56030c..d5317d64af 100644
--- a/utils/disk_space_monitor.hh
+++ b/utils/disk_space_monitor.hh
@@ -80,9 +80,22 @@ public:
 
     signal_connection_type listen(signal_callback_type callback);
 
+    // Registers a new space source function and returns an object that
+    // restores the previous one when it goes out of scope.
+    class space_source_registration {
+        disk_space_monitor& _monitor;
+        space_source_fn _prev_space_source;
+    public:
+        space_source_registration(disk_space_monitor& m);
+
+        ~space_source_registration();
+    };
+
     // Replaces default way of obtaining file system usage information.
-    void set_space_source(space_source_fn space_source) {
+    space_source_registration set_space_source(space_source_fn space_source) {
+        auto ret = space_source_registration(*this);
         _space_source = std::move(space_source);
+        return ret;
     }
 
     void trigger_poll() noexcept;
@@ -93,6 +106,8 @@ private:
     future<std::filesystem::space_info> get_filesystem_space();
 
     clock_type::duration get_polling_interval() const noexcept;
+
+    friend class space_source_registration;
 };
 
 } // namespace utils