diff --git a/src/server/blocking_controller.cc b/src/server/blocking_controller.cc index 63fb0ed376a3..b82fe3ac4b15 100644 --- a/src/server/blocking_controller.cc +++ b/src/server/blocking_controller.cc @@ -22,17 +22,24 @@ struct WatchItem { Transaction* trans; KeyReadyChecker key_ready_checker; + // True if this waiter should be woken by the missing key. + bool wake_on_missing_key; + Transaction* get() const { return trans; } - WatchItem(Transaction* t, KeyReadyChecker krc) : trans(t), key_ready_checker(std::move(krc)) { + WatchItem(Transaction* t, KeyReadyChecker krc, bool wake_on_missing_key) + : trans(t), key_ready_checker(std::move(krc)), wake_on_missing_key(wake_on_missing_key) { } }; struct BlockingController::WatchQueue { deque items; + // Items with wake_on_missing_key set. + size_t wake_on_missing_key_count = 0; + // Updated by both coordinator and shard threads but at different times. enum State { SUSPENDED, ACTIVE } state = SUSPENDED; @@ -71,6 +78,7 @@ bool BlockingController::DbWatchTable::UnwatchTx(string_view key, Transaction* t bool res = false; if (wq->state == WatchQueue::ACTIVE && wq->items.front().get() == tx) { + wq->wake_on_missing_key_count -= wq->items.front().wake_on_missing_key; wq->items.pop_front(); // We suspend the queue and add keys to re-verification. @@ -87,6 +95,7 @@ bool BlockingController::DbWatchTable::UnwatchTx(string_view key, Transaction* t // This shard has not been awakened and in case this transaction in the queue // we must clean it up. if (auto it = wq->Find(tx); it != wq->items.end()) { + wq->wake_on_missing_key_count -= it->wake_on_missing_key; wq->items.erase(it); } } @@ -203,7 +212,11 @@ void BlockingController::AddWatched(Keys watch_keys, KeyReadyChecker krc, Transa continue; } DVLOG(2) << "Emplace " << trans->DebugId() << " to watch " << key; - res->second->items.emplace_back(trans, krc); + + // Ask the checker itself whether a missing key wakes this waiter + const bool wake_on_missing = krc(key, nullptr) == KeyReadyResult::kReady; + res->second->items.emplace_back(trans, krc, wake_on_missing); + res->second->wake_on_missing_key_count += wake_on_missing; } } @@ -254,16 +267,23 @@ void BlockingController::NotifyWatchQueue(std::string_view key, WatchQueue* wq, auto& queue = wq->items; ShardId sid = owner_->shard_id(); + // Resolve the key once for the whole queue. Waiters check the type themselves, so a missing or + // wrong-type key costs no per-waiter db access. + auto it = context.GetDbSlice(sid).FindReadOnly(context, key); + const CompactObj* obj = IsValid(it) ? &it->second : nullptr; + + // Skipping the sweep if the key is missing and no waiter in queue doesn't need to woke up. + if (obj == nullptr && wq->wake_on_missing_key_count == 0) { + return; + } + // In the most cases we shouldn't have skipped elements at all absl::InlinedVector skipped; while (!queue.empty()) { auto& wi = queue.front(); Transaction* head = wi.get(); - KeyReadyResult result = wi.key_ready_checker(owner_, context, key); - if (result == KeyReadyResult::kKeyNotFound) { - // Key is gone - no tx in this queue can be woken, abort the scan entirely. - break; - } else if (result == KeyReadyResult::kReady) { + KeyReadyResult result = wi.key_ready_checker(key, obj); + if (result == KeyReadyResult::kReady) { DVLOG(2) << "WQ-Pop " << head->DebugId() << " from key " << key << " committed txid " << owner_->committed_txid(); if (head->NotifySuspended(sid, key)) { @@ -273,7 +293,8 @@ void BlockingController::NotifyWatchQueue(std::string_view key, WatchQueue* wq, awakened_transactions_.insert(head); break; } - } else { // kNotReady - key exists but per-tx conditions not met, try next + wq->wake_on_missing_key_count -= wi.wake_on_missing_key; + } else { skipped.push_back(std::move(wi)); } diff --git a/src/server/blocking_controller_test.cc b/src/server/blocking_controller_test.cc index 9111001006c3..ed049ca10612 100644 --- a/src/server/blocking_controller_test.cc +++ b/src/server/blocking_controller_test.cc @@ -98,8 +98,8 @@ TEST_F(BlockingControllerTest, Basic) { // Regression for https://github.com/dragonflydb/dragonfly/pull/7225: // NotifyWatchQueue used to walk every queued waiter (O(N) per notify) when -// the key was absent. The fast path now short-circuits via FindReadOnly. We -// assert the per-waiter checker is never invoked. +// the key was absent. The queue is now skipped outright when no waiter can be +// woken by a missing key. We assert the per-waiter checker is never invoked. TEST_F(BlockingControllerTest, NotifyWatchQueueFastPathOnAbsentKey) { constexpr size_t kWaiters = 64; const std::string_view key = str_vec_[0]; // "x", hashes to shard 0 (verified in SetUp) @@ -119,9 +119,9 @@ TEST_F(BlockingControllerTest, NotifyWatchQueueFastPathOnAbsentKey) { EngineShard* shard = EngineShard::tlocal(); BlockingController bc(shard, &namespaces->GetDefaultNamespace()); - auto checker = [&checker_calls](EngineShard*, const DbContext&, std::string_view) { + auto checker = [&checker_calls](std::string_view, const CompactObj*) { ++checker_calls; - return KeyReadyResult::kKeyNotFound; + return KeyReadyResult::kNotReady; }; for (auto& t : txs) { @@ -129,13 +129,12 @@ TEST_F(BlockingControllerTest, NotifyWatchQueueFastPathOnAbsentKey) { } ASSERT_EQ(1u, bc.NumWatched(0)); // 1 watched key, kWaiters items in its queue + checker_calls = 0; bc.Awaken(0, key); bc.NotifyPending(); }); - // With the enum-based fast path, the first item's checker is called once and returns - // kKeyNotFound, aborting the scan without visiting the remaining kWaiters-1 items. - EXPECT_EQ(1u, checker_calls) << "fast path did not short-circuit"; + EXPECT_EQ(0u, checker_calls) << "fast path did not short-circuit"; } TEST_F(BlockingControllerTest, Timeout) { diff --git a/src/server/container_utils.cc b/src/server/container_utils.cc index 5fd2c722319b..f6f5ccbec677 100644 --- a/src/server/container_utils.cc +++ b/src/server/container_utils.cc @@ -390,15 +390,12 @@ OpResult RunCbOnFirstNonEmptyBlocking(Transaction* trans, int req_obj_ty limit_tp = steady_clock::now() + milliseconds(limit_ms); } - auto* ns = &trans->GetNamespace(); - const auto key_checker = [req_obj_type, ns](EngineShard* owner, const DbContext& context, - std::string_view key) -> KeyReadyResult { - auto res = ns->GetDbSlice(owner->shard_id()).FindReadOnly(context, key, req_obj_type); - if (res.ok()) - return KeyReadyResult::kReady; - if (res.status() == OpStatus::WRONG_TYPE) + const auto key_checker = [req_obj_type](std::string_view /*unused*/, + const CompactObj* obj) -> KeyReadyResult { + if (!obj || obj->ObjType() != CompactObjType(req_obj_type)) { return KeyReadyResult::kNotReady; - return KeyReadyResult::kKeyNotFound; + } + return KeyReadyResult::kReady; }; auto status = diff --git a/src/server/list_family.cc b/src/server/list_family.cc index 600ce96465d8..51a6c5ea0803 100644 --- a/src/server/list_family.cc +++ b/src/server/list_family.cc @@ -982,13 +982,11 @@ void BLMove(CmdArgParser parser, CommandContext* cmd_cntx) { } } -KeyReadyResult ListKeyChecker(EngineShard* owner, const DbContext& context, std::string_view key) { - auto res = context.GetDbSlice(owner->shard_id()).FindReadOnly(context, key, OBJ_LIST); - if (res.ok()) - return KeyReadyResult::kReady; - if (res.status() == OpStatus::WRONG_TYPE) +KeyReadyResult ListKeyChecker(std::string_view /*unused*/, const CompactObj* obj) { + if (!obj || obj->ObjType() != OBJ_LIST) { return KeyReadyResult::kNotReady; - return KeyReadyResult::kKeyNotFound; + } + return KeyReadyResult::kReady; } BPopPusher::BPopPusher(string_view pop_key, string_view push_key, ListDir popdir, ListDir pushdir) diff --git a/src/server/stream_family.cc b/src/server/stream_family.cc index ff3a752b7924..98b3d5426ce9 100644 --- a/src/server/stream_family.cc +++ b/src/server/stream_family.cc @@ -3026,23 +3026,16 @@ void XReadBlock(ReadOpts* opts, Transaction* tx, SinkReplyBuilder* builder, auto tp = (opts->timeout) ? chrono::steady_clock::now() + chrono::milliseconds(opts->timeout) : Transaction::time_point::max(); - const auto key_checker = [opts](EngineShard* owner, const DbContext& context, - std::string_view key) -> KeyReadyResult { - auto& db_slice = context.GetDbSlice(owner->shard_id()); - auto res_it = db_slice.FindReadOnly(context, key, OBJ_STREAM); - if (!res_it.ok()) { + const auto key_checker = [opts](std::string_view key, const CompactObj* obj) -> KeyReadyResult { + if (!obj || obj->ObjType() != OBJ_STREAM) { // A blocked XREADGROUP must revalidate when its key is no longer a stream. The wake path // distinguishes a missing stream (NOGROUP) from a non-stream value (WRONGTYPE). A plain // XREAD keeps waiting for the stream to be recreated. - // TODO: Make kKeyNotFound a per-waiter result in BlockingController. It currently - // short-circuits the whole queue, allowing a blocked XREAD to hide a later XREADGROUP - // waiter. Return kNotReady until the controller handles heterogeneous queues. return opts->read_group ? KeyReadyResult::kReady : KeyReadyResult::kNotReady; } StreamIDsItem& sitem = opts->stream_ids.at(key); - const CompactObj& cobj = (*res_it)->second; - stream* s = GetReadOnlyStream(cobj); + stream* s = GetReadOnlyStream(*obj); // Update group pointer and check it's validity if (opts->read_group) { diff --git a/src/server/stream_family_test.cc b/src/server/stream_family_test.cc index bc403547a4c6..e2248c5ec2e2 100644 --- a/src/server/stream_family_test.cc +++ b/src/server/stream_family_test.cc @@ -891,6 +891,38 @@ TEST_F(StreamFamilyTest, XReadGroupBlockLazyExpireDuringWakeDoesNotCrash) { EXPECT_THAT(resp0, ErrArg("consumer group this client was blocked on no longer exists")); } +// A blocked XREADGROUP gets NOGROUP when its stream is deleted, even with a BLPOP blocked ahead of +// it on the same key. Issue #8067. +TEST_F(StreamFamilyTest, XReadGroupBlockNotHiddenByBlpopWaiter) { + RespExpr blpop_resp; + auto blpop_fb = pp_->at(1)->LaunchFiber(Launch::dispatch, [&] { + blpop_resp = Run({"BLPOP", "k", "0"}); + }); + ASSERT_TRUE(WaitUntilCondition([&] { return IsConnBlocked("IO1"); }, 500ms)); + + Run({"XGROUP", "CREATE", "k", "g", "0", "MKSTREAM"}); + + RespExpr xread_resp; + auto xread_fb = pp_->at(2)->LaunchFiber(Launch::dispatch, [&] { + xread_resp = Run({"XREADGROUP", "GROUP", "g", "c", "BLOCK", "0", "STREAMS", "k", ">"}); + }); + ASSERT_TRUE(WaitUntilCondition([&] { return IsConnBlocked("IO2"); }, 500ms)); + + EXPECT_THAT(Run({"DEL", "k"}), IntArg(1)); + + EXPECT_TRUE(WaitUntilCondition([&] { return !IsConnBlocked("IO2"); }, 500ms)) + << "XREADGROUP stayed blocked after its stream was deleted"; + EXPECT_TRUE(IsConnBlocked("IO1")) << "BLPOP must stay blocked, its key never received a value"; + + // Release BLPOP so both fibers can be joined regardless of the outcome above. + Run({"LPUSH", "k", "v"}); + blpop_fb.Join(); + xread_fb.Join(); + + EXPECT_THAT(xread_resp, ErrArg("consumer group this client was blocked on no longer exists")); + EXPECT_THAT(blpop_resp, RespArray(ElementsAre("k", "v"))); +} + TEST_F(StreamFamilyTest, XReadGroupBlockHonorsCount) { Run({"xgroup", "create", "foo", "group", "0", "MKSTREAM"}); diff --git a/src/server/transaction_test.cc b/src/server/transaction_test.cc index d22b8304d585..18b91043b2db 100644 --- a/src/server/transaction_test.cc +++ b/src/server/transaction_test.cc @@ -240,9 +240,7 @@ TEST_F(TransactionTest, AwakenedPollNotDroppedWhenBlockedTxPresent) { auto tx_bl = MakeTx(&cid_bl, {"x"}); auto tx_z = MakeTx(&cid_z, {"a"}); - auto ready_checker = [](EngineShard*, const DbContext&, std::string_view) { - return KeyReadyResult::kReady; - }; + auto ready_checker = [](std::string_view, const CompactObj*) { return KeyReadyResult::kReady; }; // Phase 1: BL schedules a first hop, then suspends watching key "x". bool bl_blocked = false, bl_paused = false; diff --git a/src/server/tx_base.h b/src/server/tx_base.h index 99c0207661e1..cf87d05e1b32 100644 --- a/src/server/tx_base.h +++ b/src/server/tx_base.h @@ -17,6 +17,8 @@ namespace dfly { using cmn::ArgSlice; +class CompactObj; + struct KeyLockArgs { DbIndex db_index = 0; absl::Span fps; @@ -114,14 +116,12 @@ class LockTag { }; enum class KeyReadyResult { - kKeyNotFound, // key doesn't exist - abort the entire watch queue - kNotReady, // key exists but per-tx conditions not met - skip this tx, try next - kReady, // wake this tx + kNotReady, // key missing or conditions not met keep this tx waiting, check next in queue + kReady, // wake this tx }; // Checks whether the touched key is valid for a blocking transaction watching it. -using KeyReadyChecker = - std::function; +using KeyReadyChecker = std::function; // References arguments in another array. using IndexSlice = std::pair; // [begin, end) diff --git a/src/server/zset_family.cc b/src/server/zset_family.cc index aa327298d56e..5e4ba0bc7074 100644 --- a/src/server/zset_family.cc +++ b/src/server/zset_family.cc @@ -2496,22 +2496,18 @@ void ZMPopGeneric(CmdArgParser parser, CommandContext* cmd_cntx, bool is_blockin if (is_blocking && !key_to_pop.has_value()) { auto trans = cmd_cntx->tx(); auto* cntx = cmd_cntx->server_conn_cntx(); - auto* ns = &trans->GetNamespace(); - auto limit_tp = Transaction::time_point::max(); auto limit_ms = (unsigned)(timeout * 1000); if (limit_ms > 0) { using namespace std::chrono; limit_tp = steady_clock::now() + milliseconds(limit_ms); } - const auto key_checker = [ns](EngineShard* owner, const DbContext& context, - std::string_view key) -> KeyReadyResult { - auto res = ns->GetDbSlice(owner->shard_id()).FindReadOnly(context, key, OBJ_ZSET); - if (res.ok()) - return KeyReadyResult::kReady; - if (res.status() == OpStatus::WRONG_TYPE) + const auto key_checker = [](std::string_view /*unused*/, + const CompactObj* obj) -> KeyReadyResult { + if (!obj || obj->ObjType() != OBJ_ZSET) { return KeyReadyResult::kNotReady; - return KeyReadyResult::kKeyNotFound; + } + return KeyReadyResult::kReady; }; DCHECK(trans->IsScheduled()); // Checking if the transaction is scheduled before calling