Fix the another pre-emption bug with inline scheduling

Author: royjacobson

src/server/journal/journal_slice.cc

@@ -11,6 +11,7 @@
 #include <filesystem>
 
 #include "base/logging.h"
+#include "src/server/server_state.h"
 
 namespace dfly {
 namespace journal {
@@ -145,6 +146,7 @@ uint32_t JournalSlice::RegisterOnChange(ChangeCallback cb) {
   lock_guard lk(cb_mu_);
   uint32_t id = next_cb_id_++;
   change_cb_arr_.emplace_back(id, std::move(cb));
+  ServerState::tlocal()->SetJournalCallbacksActive(true);
   return id;
 }
 
@@ -154,6 +156,9 @@ void JournalSlice::UnregisterOnChange(uint32_t id) {
                     [id](const auto& e) { return e.first == id; });
   CHECK(it != change_cb_arr_.end());
   change_cb_arr_.erase(it);
+
+  if (change_cb_arr_.empty())
+    ServerState::tlocal()->SetJournalCallbacksActive(false);
 }
 
 }  // namespace journal

src/server/server_state.h

@@ -130,6 +130,22 @@ class ServerState {  // public struct - to allow initialization.
     gstate_ = s;
   }
 
+  bool AllowInlineScheduling() const {
+    // We can't allow inline scheduling during a full sync, because then journaling transactions
+    // will be scheduled before RdbLoader::LoadItemsBuffer is finished. We can't use the regular
+    // locking mechanism because RdbLoader is not using transactions.
+
+    // Journal callbacks can preempt; This means we have to disallow inline scheduling
+    // because then we might interleave the callbacks loop from an inlined-scheduled command
+    // and a normally-scheduled command.
+    // The problematic loop is in JournalSlice::AddLogRecord, going over all the callbacks.
+    return !journal_callbacks_active_ && gstate_ != GlobalState::LOADING;
+  }
+
+  void SetJournalCallbacksActive(bool is_allowed) {
+    journal_callbacks_active_ = is_allowed;
+  }
+
   // Borrow interpreter from internal manager. Return int with ReturnInterpreter.
   Interpreter* BorrowInterpreter();
 
@@ -206,6 +222,11 @@ class ServerState {  // public struct - to allow initialization.
   mi_heap_t* data_heap_;
   journal::Journal* journal_ = nullptr;
 
+  // Inline scheduling is an optimization for running in low thread count situations.
+  // Since it bypasses the regular transactions mechanism, it has to be disabled in
+  // some situations.
+  bool journal_callbacks_active_ = false;
+
   InterpreterManager interpreter_mgr_;
   absl::flat_hash_map<ScriptMgr::ScriptKey, ScriptMgr::ScriptParams> cached_script_params_;
 

src/server/transaction.cc

@@ -689,11 +689,7 @@ OpStatus Transaction::ScheduleSingleHop(RunnableType cb) {
       }
     };
 
-    // We can't allow inline scheduling during a full sync, because then journaling transactions
-    // will be scheduled before RdbLoader::LoadItemsBuffer is finished. We can't use the regular
-    // locking mechanism because RdbLoader is not using transactions.
-    if (coordinator_index_ == unique_shard_id_ &&
-        ServerState::tlocal()->gstate() != GlobalState::LOADING) {
+    if (coordinator_index_ == unique_shard_id_ && ServerState::tlocal()->AllowInlineScheduling()) {
       DVLOG(2) << "Inline scheduling a transaction";
       schedule_cb();
     } else {