Clean up blocking (#5836)

Author: dranikpg

src/server/blocking_controller.cc

@@ -31,16 +31,10 @@ struct WatchItem {
 
 struct BlockingController::WatchQueue {
   deque<WatchItem> items;
-  TxId notify_txid = UINT64_MAX;
 
   // Updated  by both coordinator and shard threads but at different times.
   enum State { SUSPENDED, ACTIVE } state = SUSPENDED;
 
-  void Suspend() {
-    state = SUSPENDED;
-    notify_txid = UINT64_MAX;
-  }
-
   auto Find(Transaction* tx) const {
     return find_if(items.begin(), items.end(),
                    [tx](const WatchItem& wi) { return wi.get() == tx; });
@@ -49,7 +43,8 @@ struct BlockingController::WatchQueue {
 
 // Watch state per db.
 struct BlockingController::DbWatchTable {
-  WatchQueueMap queue_map;
+  // Watch queues per key
+  absl::flat_hash_map<std::string, std::unique_ptr<WatchQueue>> queue_map;
 
   // awakened keys point to blocked keys that can potentially be unblocked.
   absl::flat_hash_set<std::string> awakened_keys;
@@ -163,32 +158,16 @@ void BlockingController::NotifyPending() {
       continue;
 
     context.db_index = index;
-    DbWatchTable& wt = *dbit->second;
-    for (const auto& key : wt.awakened_keys) {
-      string_view sv_key = key;
-      DVLOG(1) << "Processing awakened key " << sv_key;
-      auto w_it = wt.queue_map.find(sv_key);
-      if (w_it == wt.queue_map.end()) {
-        // This should not happen because we remove keys from awakened_keys every type we remove
-        // the entry from queue_map. TODO: to make it a CHECK after Dec 2024
-        LOG(ERROR) << "Internal error: Key " << sv_key
-                   << " was not found in the watch queue, wt.awakened_keys len is "
-                   << wt.awakened_keys.size() << " wt.queue_map len is " << wt.queue_map.size();
-        for (const auto& item : wt.awakened_keys) {
-          LOG(ERROR) << "Awakened key: " << item;
-        }
-
-        continue;
-      }
-
+    DbWatchTable& wt = *dbit->second;  // pointer stability due to node_hash_map
+    for (string_view key : wt.awakened_keys) {
+      DVLOG(1) << "Processing awakened key " << key;
+      auto w_it = wt.queue_map.find(key);
       CHECK(w_it != wt.queue_map.end());
-      DVLOG(1) << "Notify WQ: [" << owner_->shard_id() << "] " << key;
+
       WatchQueue* wq = w_it->second.get();
-      NotifyWatchQueue(sv_key, wq, context);
-      if (wq->items.empty()) {
-        // we erase awakened_keys right after this loop finishes running.
+      NotifyWatchQueue(key, wq, context);
+      if (wq->items.empty())
         wt.queue_map.erase(w_it);
-      }
     }
     wt.awakened_keys.clear();
 
@@ -202,16 +181,15 @@ void BlockingController::NotifyPending() {
 void BlockingController::AddWatched(Keys watch_keys, KeyReadyChecker krc, Transaction* trans) {
   auto [dbit, added] = watched_dbs_.emplace(trans->GetDbIndex(), nullptr);
   if (added) {
-    dbit->second.reset(new DbWatchTable);
+    dbit->second = make_unique<DbWatchTable>();
   }
 
   DbWatchTable& wt = *dbit->second;
 
   for (auto key : watch_keys) {
     auto [res, inserted] = wt.queue_map.emplace(key, nullptr);
-    if (inserted) {
-      res->second.reset(new WatchQueue);
-    }
+    if (inserted)
+      res->second = make_unique<WatchQueue>();
 
     if (!res->second->items.empty()) {
       Transaction* last = res->second->items.back().get();
@@ -227,7 +205,7 @@ void BlockingController::AddWatched(Keys watch_keys, KeyReadyChecker krc, Transa
 }
 
 // Called from commands like lpush.
-void BlockingController::AwakeWatched(DbIndex db_index, string_view db_key) {
+void BlockingController::Awaken(DbIndex db_index, string_view db_key) {
   auto it = watched_dbs_.find(db_index);
   if (it == watched_dbs_.end())
     return;
@@ -236,8 +214,7 @@ void BlockingController::AwakeWatched(DbIndex db_index, string_view db_key) {
   DCHECK(!wt.queue_map.empty());
 
   if (wt.AddAwakeEvent(db_key)) {
-    VLOG(1) << "AwakeWatched: db(" << db_index << ") " << db_key;
-
+    VLOG(1) << "Touch: db(" << db_index << ") " << db_key;
     awakened_indices_.insert(db_index);
   }
 }
@@ -263,7 +240,6 @@ void BlockingController::NotifyWatchQueue(std::string_view key, WatchQueue* wq,
         wq->state = WatchQueue::ACTIVE;
         // We deliberately keep the notified transaction in the queue to know which queue
         // must handled when this transaction finished.
-        wq->notify_txid = owner_->committed_txid();
         awakened_transactions_.insert(head);
         break;
       }

src/server/blocking_controller.h

@@ -17,6 +17,9 @@ namespace dfly {
 class Transaction;
 class Namespace;
 
+// Used for tracking keys of blocking transactions and properly notifying them.
+// First, keys are marked as watched and associated with an owner transaction. A mutating
+// transaction marks them as touched, and once it concludes, the watching transactions are notified.
 class BlockingController {
  public:
   explicit BlockingController(EngineShard* owner, Namespace* ns);
@@ -32,20 +35,17 @@ class BlockingController {
     return awakened_transactions_;
   }
 
-  // Removes transaction from watching these keys.
-  void RemovedWatched(Keys keys, Transaction* tx);
+  // Associate given keys with transaction, checked via the krc checker
+  void AddWatched(Keys watch_keys, KeyReadyChecker krc, Transaction* me);
 
-  // go over potential wakened keys, verify them and activate watch queues.
-  void NotifyPending();
+  // Remove transaction from watching these keys
+  void RemovedWatched(Keys keys, Transaction* tx);
 
-  // Blocking API
-  // TODO: consider moving all watched functions to
-  // EngineShard with separate per db map.
-  //! AddWatched adds a transaction to the blocking queue.
-  void AddWatched(Keys watch_keys, KeyReadyChecker krc, Transaction* me);
+  // Mark given key as awakened. Called by commands mutating this key.
+  void Awaken(DbIndex db_index, std::string_view key);
 
-  // Called from operations that create keys like lpush, rename etc.
-  void AwakeWatched(DbIndex db_index, std::string_view db_key);
+  // Notify transactions of awakened keys
+  void NotifyPending();
 
   // Used in tests and debugging functions.
   size_t NumWatched(DbIndex db_indx) const;
@@ -55,24 +55,17 @@ class BlockingController {
   struct WatchQueue;
   struct DbWatchTable;
 
-  using WatchQueueMap = absl::flat_hash_map<std::string, std::unique_ptr<WatchQueue>>;
-
   void NotifyWatchQueue(std::string_view key, WatchQueue* wqm, const DbContext& context);
 
-  // void NotifyConvergence(Transaction* tx);
-
   EngineShard* owner_;
   Namespace* ns_;
 
-  absl::flat_hash_map<DbIndex, std::unique_ptr<DbWatchTable>> watched_dbs_;
-
-  // serves as a temporary queue that aggregates all the possible awakened dbs.
-  // flushed by RunStep().
-  absl::flat_hash_set<DbIndex> awakened_indices_;
+  // TODO: check if unique_ptr indirection is required
+  absl::flat_hash_map<DbIndex, std::unique_ptr<DbWatchTable>> watched_dbs_;  // watched keys
+  absl::flat_hash_set<DbIndex> awakened_indices_;  // watched_dbs_ with awakened keys
 
-  // tracks currently notified and awaked transactions.
-  // There can be multiple transactions like this because a transaction
-  // could awaken arbitrary number of keys.
+  // Transactions that got awakened with NotifySuspended
+  // TODO: Used only for one DCHECK
   absl::flat_hash_set<Transaction*> awakened_transactions_;
 };
 }  // namespace dfly

src/server/generic_family.cc

@@ -507,7 +507,7 @@ OpStatus Renamer::DeserializeDest(Transaction* t, EngineShard* shard) {
       << "Unexpected override for key " << dest_key_ << " " << dest_found_;
   auto bc = op_args.db_cntx.ns->GetBlockingController(op_args.shard->shard_id());
   if (bc) {
-    bc->AwakeWatched(t->GetDbIndex(), dest_key_);
+    bc->Awaken(t->GetDbIndex(), dest_key_);
   }
 
   if (shard->journal()) {
@@ -886,7 +886,7 @@ OpStatus OpMove(const OpArgs& op_args, string_view key, DbIndex target_db) {
 
   auto bc = op_args.db_cntx.ns->GetBlockingController(op_args.shard->shard_id());
   if (add_res.it->second.ObjType() == OBJ_LIST && bc) {
-    bc->AwakeWatched(target_db, key);
+    bc->Awaken(target_db, key);
   }
 
   return OpStatus::OK;
@@ -956,7 +956,7 @@ OpResult<void> OpRen(const OpArgs& op_args, string_view from_key, string_view to
 
   auto bc = op_args.db_cntx.ns->GetBlockingController(es->shard_id());
   if (!is_prior_list && to_res.it->second.ObjType() == OBJ_LIST && bc) {
-    bc->AwakeWatched(op_args.db_cntx.db_index, to_key);
+    bc->Awaken(op_args.db_cntx.db_index, to_key);
   }
   return OpStatus::OK;
 }

src/server/list_family.cc

@@ -177,7 +177,7 @@ OpResult<string> OpMoveSingleShard(const OpArgs& op_args, string_view src, strin
     dest_res.it->second.InitRobj(OBJ_LIST, kEncodingQL2, destql_v2);
     auto blocking_controller = op_args.db_cntx.ns->GetBlockingController(op_args.shard->shard_id());
     if (blocking_controller) {
-      blocking_controller->AwakeWatched(op_args.db_cntx.db_index, dest);
+      blocking_controller->Awaken(op_args.db_cntx.db_index, dest);
     }
   } else {
     destql_v2 = GetQLV2(dest_res.it->second);
@@ -257,7 +257,7 @@ OpResult<uint32_t> OpPush(const OpArgs& op_args, std::string_view key, ListDir d
   if (res.is_new) {
     auto blocking_controller = op_args.db_cntx.ns->GetBlockingController(es->shard_id());
     if (blocking_controller) {
-      blocking_controller->AwakeWatched(op_args.db_cntx.db_index, key);
+      blocking_controller->Awaken(op_args.db_cntx.db_index, key);
     }
   }
 

src/server/stream_family.cc

@@ -765,7 +765,7 @@ OpResult<streamID> OpAdd(const OpArgs& op_args, string_view key, const AddOpts&
 
   auto blocking_controller = op_args.db_cntx.ns->GetBlockingController(op_args.shard->shard_id());
   if (blocking_controller) {
-    blocking_controller->AwakeWatched(op_args.db_cntx.db_index, key);
+    blocking_controller->Awaken(op_args.db_cntx.db_index, key);
   }
 
   return result_id;
@@ -1439,7 +1439,7 @@ OpStatus OpDestroyGroup(const OpArgs& op_args, string_view key, string_view gnam
   // Awake readers blocked on this group
   auto blocking_controller = op_args.db_cntx.ns->GetBlockingController(op_args.shard->shard_id());
   if (blocking_controller) {
-    blocking_controller->AwakeWatched(op_args.db_cntx.db_index, key);
+    blocking_controller->Awaken(op_args.db_cntx.db_index, key);
   }
 
   return OpStatus::OK;

src/server/zset_family.cc

@@ -182,7 +182,7 @@ OpResult<DbSlice::ItAndUpdater> PrepareZEntry(const ZSetFamily::ZParams& zparams
 
   auto* blocking_controller = op_args.db_cntx.ns->GetBlockingController(op_args.shard->shard_id());
   if (add_res.is_new && blocking_controller) {
-    blocking_controller->AwakeWatched(op_args.db_cntx.db_index, key);
+    blocking_controller->Awaken(op_args.db_cntx.db_index, key);
   }
 
   return DbSlice::ItAndUpdater{add_res.it, add_res.exp_it, std::move(add_res.post_updater)};