From fabab2f46f0a201ee9dfcf28c1424ec6dae4bf18 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Chojnowski?= <michal.chojnowski@scylladb.com>
Date: Wed, 7 Feb 2024 11:55:51 +0100
Subject: [PATCH] utils: preempt: add preemption_source

While `preemption_check` can be passed to functions to control
their preemption points, there is no way to inspect the
state of the system after the preemption results in a yield.

`preemption_source` is a superset of `preemption_check`,
which also allows for customizing the yield, not just the preemption
check. An implementation passed by a test can hook the yield to
put the tested function to sleep, run some code, and then wake the
function up.

We use the preprocessor to minimize the impact on release builds.
Only dev-mode preemption_source is hookable. When it's used in other
modes, it should compile to direct reactor calls, as if it wasn't used.
---
 configure.py     |  2 +-
 utils/preempt.hh | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/configure.py b/configure.py
index bbe7c96ce6..38665fe27a 100755
--- a/configure.py
+++ b/configure.py
@@ -420,7 +420,7 @@ modes = {
         'description': 'a mode with optimizations and no debug checks, used for production builds',
     },
     'dev': {
-        'cxxflags': '-DDEVEL -DSEASTAR_ENABLE_ALLOC_FAILURE_INJECTION -DSCYLLA_ENABLE_ERROR_INJECTION',
+        'cxxflags': '-DDEVEL -DSEASTAR_ENABLE_ALLOC_FAILURE_INJECTION -DSCYLLA_ENABLE_ERROR_INJECTION -DSCYLLA_ENABLE_PREEMPTION_SOURCE',
         'cxx_ld_flags': '',
         'stack-usage-threshold': 1024*21,
         'optimization-level': '2',
diff --git a/utils/preempt.hh b/utils/preempt.hh
index 066a5734d1..513cd01295 100644
--- a/utils/preempt.hh
+++ b/utils/preempt.hh
@@ -11,6 +11,7 @@
 #include <seastar/util/bool_class.hh>
 #include <seastar/util/noncopyable_function.hh>
 #include <seastar/core/preempt.hh>
+#include <seastar/core/thread.hh>
 
 #include "seastarx.hh"
 
@@ -41,3 +42,36 @@ preemption_check always_preempt() {
         return true;
     };
 }
+
+struct basic_preemption_source {
+    bool should_preempt() {
+        return seastar::need_preempt();
+    }
+    void thread_yield() {
+        seastar::thread::yield();
+    }
+};
+
+struct custom_preemption_source {
+    struct impl {
+        virtual bool should_preempt() = 0;
+        virtual void thread_yield() = 0;
+        virtual ~impl() = default;
+    };
+    std::unique_ptr<impl> _impl;
+    bool should_preempt() {
+        return _impl ? _impl->should_preempt() : seastar::need_preempt();
+    }
+    void thread_yield() {
+        _impl ? _impl->thread_yield() : seastar::thread::yield();
+    }
+    custom_preemption_source() = default;
+    custom_preemption_source(std::unique_ptr<impl> i) : _impl(std::move(i)) {}
+};
+
+#ifdef SCYLLA_ENABLE_PREEMPTION_SOURCE
+using preemption_source = custom_preemption_source;
+#else
+using preemption_source = basic_preemption_source;
+#endif
+