Skip to content

fix: avoid bthread mutex waits in shard completion paths #492

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
starrysky9959 wants to merge 2 commits into from
Closed

fix: avoid bthread mutex waits in shard completion paths #492

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6a151d2
fix: avoid bthread mutex waits in shard completion paths
starrysky9959 Jun 5, 2026
a8ec8e9
Potential fix for pull request finding
thweetkomputer Jun 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 35 additions & 47 deletions tx_service/include/cc/cc_req_misc.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -871,11 +871,11 @@ struct WaitableCc : public RunOnTxProcessorCc
void Reset(std::function<bool(CcShard &ccs)> task = {},
uint16_t core_cnt = 1)
{
std::lock_guard<bthread::Mutex> lk(mux_);
RunOnTxProcessorCc::Reset(std::move(task));

unfinished_cnt_ = core_cnt;
error_code_ = CcErrorCode::NO_ERROR;
unfinished_cnt_.store(core_cnt, std::memory_order_relaxed);
error_code_.store(CcErrorCode::NO_ERROR, std::memory_order_relaxed);
waiting_.store(false, std::memory_order_relaxed);
}

void SetCoroCallbacks(const std::function<void()> *yield_fn,
Expand All @@ -887,10 +887,13 @@ struct WaitableCc : public RunOnTxProcessorCc

void Wait()
{
std::unique_lock<bthread::Mutex> lk(mux_);
while (unfinished_cnt_)
uint64_t interval_us = 100;
constexpr uint64_t kMaxIntervalUs = 100000;
while (unfinished_cnt_.load(std::memory_order_acquire) > 0)
{
cv_.wait(lk);
bthread_usleep(interval_us);
if ((interval_us << 1) < kMaxIntervalUs)
interval_us <<= 1;
}
}

Expand All @@ -902,53 +905,47 @@ struct WaitableCc : public RunOnTxProcessorCc
Wait();
return;
}
std::unique_lock<bthread::Mutex> lk(mux_);
while (unfinished_cnt_)
while (unfinished_cnt_.load(std::memory_order_acquire) > 0)
{
waiting_.store(true, std::memory_order_release);
lk.unlock();
if (unfinished_cnt_.load(std::memory_order_acquire) == 0)
{
waiting_.store(false, std::memory_order_release);

@thweetkomputer thweetkomputer Jun 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The waiter enters the loop after observing unfinished_cnt_ == 1 (W1) and sets waiting_ = true (W2). At this point the completer on the TxProcessor thread decrements the counter to zero (C1) and then exchanges waiting_ from true to false (C2) — since it observed true, it commits to calling resume_fn. Meanwhile, the waiter's re-check (W3) observes unfinished_cnt_ == 0 and takes the early-exit branch: it stores waiting_ = false (W4, which is now a no-op because C2 already cleared the flag) and breaks out of the loop without ever calling yield_fn (W5). Finally the completer invokes resume_fn() (C3), which unconditionally enqueues the coroutine's CoroCtx into resume_queue_.
▎
▎ The result is an orphan resume: a resume has been published with no matching yield to consume it. resume_fn is not a level-triggered notification like a condition variable — every enqueued entry will eventually be popped by the flush worker, which calls ctx->coro_.resume(). Since the coroutine never suspended (and may have already run to completion by then), the worker ends up resuming a running or finished boost::context continuation, which is undefined behavior.
▎
▎ Note the bug is timing-dependent: if W4 wins the race against C2 instead, the completer's exchange reads false, no resume is published, and everything works — the failure window is only the gap between W2 and W4.

break;
Comment on lines 910 to +914
}
(*yield_fn)();
lk.lock();
waiting_.store(false, std::memory_order_release);
}
waiting_.store(false, std::memory_order_release);
}
Comment on lines +908 to 919

@coderabbitai coderabbitai Bot Jun 12, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

waiting_ remains true outside the actual suspended window, so completion can issue a stale resume.

After (*yield_fn)() returns, waiting_ is still true until loop exit. If final completion happens in that runnable window, AbortCcRequest/Execute can call resume_fn_ even though the waiter is not suspended.

Suggested fix 
 while (unfinished_cnt_.load(std::memory_order_acquire) > 0)
 {
     waiting_.store(true, std::memory_order_release);
     if (unfinished_cnt_.load(std::memory_order_acquire) == 0)
     {
         waiting_.store(false, std::memory_order_release);
         break;
     }
     (*yield_fn)();
+    waiting_.store(false, std::memory_order_release);
 }
 waiting_.store(false, std::memory_order_release);

Also applies to: 939-949, 957-965

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tx_service/include/cc/cc_req_misc.h` around lines 908 - 919, The loop leaves
waiting_ true after (*yield_fn)() returns, allowing a concurrent completion to
call resume_fn_ while the waiter is no longer suspended; change the loop so
waiting_ is cleared immediately after the yield returns (i.e., store(false,
std::memory_order_release) right after (*yield_fn)() inside the loop), keeping
the existing early-exit check (if unfinished_cnt_.load(...) == 0) that also
clears waiting_, so waiting_ is only true for the actual suspended window;
references: waiting_, unfinished_cnt_, (*yield_fn)(), AbortCcRequest/Execute and
resume_fn_.


bool IsFinished() const
{
std::lock_guard<bthread::Mutex> lk(mux_);
return unfinished_cnt_ == 0;
return unfinished_cnt_.load(std::memory_order_acquire) == 0;
}

bool IsError() const
{
std::lock_guard<bthread::Mutex> lk(mux_);
return error_code_ != CcErrorCode::NO_ERROR;
return error_code_.load(std::memory_order_acquire) !=
CcErrorCode::NO_ERROR;
}

CcErrorCode ErrorCode() const
{
std::lock_guard<bthread::Mutex> lk(mux_);
return error_code_;
return error_code_.load(std::memory_order_acquire);
}

void AbortCcRequest(CcErrorCode error_code) override
{
std::unique_lock<bthread::Mutex> lk(mux_);
unfinished_cnt_--;
error_code_ = error_code;
if (unfinished_cnt_ == 0)
error_code_.store(error_code, std::memory_order_release);
if (unfinished_cnt_.fetch_sub(1, std::memory_order_acq_rel) == 1)
{
if (resume_fn_ != nullptr &&
waiting_.load(std::memory_order_acquire))
if (resume_fn_ != nullptr)
{
waiting_.store(false, std::memory_order_release);
auto *fn = resume_fn_;
lk.unlock();
(*fn)();
}
else if (resume_fn_ == nullptr)
{
cv_.notify_one();
if (waiting_.exchange(false, std::memory_order_acq_rel))
{
(*fn)();
}
}
}
}
Expand All @@ -957,21 +954,15 @@ struct WaitableCc : public RunOnTxProcessorCc
{
if (RunOnTxProcessorCc::Execute(ccs))
{
std::unique_lock<bthread::Mutex> lk(mux_);
error_code_ = CcErrorCode::NO_ERROR;
if (--unfinished_cnt_ == 0)
if (unfinished_cnt_.fetch_sub(1, std::memory_order_acq_rel) == 1)
{
if (resume_fn_ != nullptr &&
waiting_.load(std::memory_order_acquire))
if (resume_fn_ != nullptr)
{
waiting_.store(false, std::memory_order_release);
auto *fn = resume_fn_;
lk.unlock();
(*fn)();
}
else if (resume_fn_ == nullptr)
{
cv_.notify_one();
if (waiting_.exchange(false, std::memory_order_acq_rel))
{
(*fn)();
}
}
}
}
Expand All @@ -989,11 +980,8 @@ struct WaitableCc : public RunOnTxProcessorCc
}

private:
mutable bthread::Mutex mux_;
bthread::ConditionVariable cv_;

uint32_t unfinished_cnt_{0};
CcErrorCode error_code_;
std::atomic<uint32_t> unfinished_cnt_{0};
std::atomic<CcErrorCode> error_code_;

// Coroutine yield/resume support
const std::function<void()> *yield_fn_{nullptr};
Expand Down
92 changes: 46 additions & 46 deletions tx_service/include/cc/cc_request.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3159,11 +3159,7 @@ struct ActiveTxMaxTsCc : public CcRequestBase
{
public:
ActiveTxMaxTsCc(size_t shard_cnt, NodeGroupId ng_id)
: active_tx_max_ts_(0),
mux_(),
cv_(),
unfinish_cnt_(shard_cnt),
cc_ng_id_(ng_id)
: active_tx_max_ts_(0), unfinish_cnt_(shard_cnt), cc_ng_id_(ng_id)
{
}

Expand All @@ -3181,11 +3177,7 @@ struct ActiveTxMaxTsCc : public CcRequestBase
old_val, shard_active_tx_max_ts, std::memory_order_acq_rel))
;

std::unique_lock lk(mux_);
if (--unfinish_cnt_ == 0)
{
cv_.notify_one();
}
unfinish_cnt_.fetch_sub(1, std::memory_order_acq_rel);

// return false since ActiveTxMaxTsCc is not reused and does not need
// to call CcRequestBase::Free
Expand All @@ -3194,10 +3186,13 @@ struct ActiveTxMaxTsCc : public CcRequestBase

void Wait()
{
std::unique_lock lk(mux_);
while (unfinish_cnt_ > 0)
uint64_t interval_us = 100;
constexpr uint64_t kMaxIntervalUs = 100000;
while (unfinish_cnt_.load(std::memory_order_acquire) > 0)
{
cv_.wait(lk);
bthread_usleep(interval_us);
if ((interval_us << 1) < kMaxIntervalUs)
interval_us <<= 1;
}
}

Expand All @@ -3208,9 +3203,7 @@ struct ActiveTxMaxTsCc : public CcRequestBase

private:
std::atomic<uint64_t> active_tx_max_ts_;
bthread::Mutex mux_;
bthread::ConditionVariable cv_;
size_t unfinish_cnt_;
std::atomic_size_t unfinish_cnt_;
NodeGroupId cc_ng_id_;
};

Expand Down Expand Up @@ -8410,8 +8403,9 @@ struct DbSizeCc : public CcRequestBase
Clear();
table_names_ = table_names;

total_ref_cnt_ = local_ref_cnt + remote_ref_cnt;
remote_ref_cnt_ = remote_ref_cnt;
total_ref_cnt_.store(local_ref_cnt + remote_ref_cnt,
std::memory_order_relaxed);
remote_ref_cnt_.store(remote_ref_cnt, std::memory_order_relaxed);
total_obj_sizes_.resize(table_names_->size(), 0);
}

Expand All @@ -8433,13 +8427,7 @@ struct DbSizeCc : public CcRequestBase
}
}

std::unique_lock lk(mux_);
if (--total_ref_cnt_ == 0)
{
cv_.notify_one();
}

return false;
return OnLocalRefFinished();

@thweetkomputer thweetkomputer Jun 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this should always return false? or DbSizeCc is recycled and reused by another caller.

OnLocalRefFinished();
return false;

}

std::vector<int64_t> GetTotalObjSizes()
Expand Down Expand Up @@ -8473,13 +8461,7 @@ struct DbSizeCc : public CcRequestBase
idx, total_obj_sizes[idx], std::memory_order_relaxed);
}

std::unique_lock lk(mux_);
--remote_ref_cnt_;
--total_ref_cnt_;
if (total_ref_cnt_ == 0)
{
cv_.notify_one();
}
OnRemoteRefFinished();

@coderabbitai coderabbitai Bot Jun 12, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Guard atomic ref counters against underflow on late DBSize callbacks.

OnRemoteRefFinished() / OnLocalRefFinished() unconditionally fetch_sub(1) on unsigned atomics. Given the known timeout race in the DBSize response flow (tx_service/src/remote/cc_stream_receiver.cpp TODO near DBSize timeout handling), late callbacks can arrive after counters were reset/cleared and wrap to SIZE_MAX, corrupting completion state.

🔧 Suggested fix 
 protected:
+    static bool TryDec(std::atomic_size_t &cnt, bool *is_last = nullptr)
+    {
+        size_t cur = cnt.load(std::memory_order_acquire);
+        while (cur > 0)
+        {
+            if (cnt.compare_exchange_weak(cur,
+                                          cur - 1,
+                                          std::memory_order_acq_rel,
+                                          std::memory_order_acquire))
+            {
+                if (is_last != nullptr)
+                {
+                    *is_last = (cur == 1);
+                }
+                return true;
+            }
+        }
+        if (is_last != nullptr)
+        {
+            *is_last = false;
+        }
+        return false;
+    }
+
     bool OnLocalRefFinished()
     {
-        return total_ref_cnt_.fetch_sub(1, std::memory_order_acq_rel) == 1;
+        bool is_last = false;
+        return TryDec(total_ref_cnt_, &is_last) && is_last;
     }

     bool OnRemoteRefFinished()
     {
-        remote_ref_cnt_.fetch_sub(1, std::memory_order_acq_rel);
-        return total_ref_cnt_.fetch_sub(1, std::memory_order_acq_rel) == 1;
+        if (!TryDec(remote_ref_cnt_))
+        {
+            return false;  // stale/duplicate callback
+        }
+        bool is_last = false;
+        return TryDec(total_ref_cnt_, &is_last) && is_last;
     }

Also applies to: 8481-8482, 8537-8546

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tx_service/include/cc/cc_request.h` at line 8464, Guard the unsigned atomic
refcount decrements in OnRemoteRefFinished and OnLocalRefFinished to avoid
underflow when late DBSize callbacks arrive: before calling fetch_sub(1) on the
atomic counters used in cc_request.h, read/check the current value and only
decrement if it is > 0 (use an atomic compare-and-swap loop or fetch-update API
to do this atomically), and ensure the functions return or no-op when the
counter is already zero so they cannot wrap to SIZE_MAX and corrupt completion
state; apply the same guarded-decrement pattern to the other similar sites
referenced (the other calls around the OnRemoteRefFinished/OnLocalRefFinished
usages).

}

int32_t GetTerm()
Expand All @@ -8496,25 +8478,35 @@ struct DbSizeCc : public CcRequestBase
total_obj_sizes_.clear();
total_obj_sizes_.shrink_to_fit();

total_ref_cnt_ = 0;
remote_ref_cnt_ = 0;
total_ref_cnt_.store(0, std::memory_order_relaxed);
remote_ref_cnt_.store(0, std::memory_order_relaxed);
table_names_ = nullptr;
vct_ng_id_.clear();
}

void Wait()
{
const uint64_t MAX_WAIT_TS = 2000000;
std::unique_lock lk(mux_);
uint64_t remaining_wait_us = 2000000;
uint64_t interval_us = 100;
constexpr uint64_t kMaxIntervalUs = 100000;

while (total_ref_cnt_ > 0)
while (total_ref_cnt_.load(std::memory_order_acquire) > 0)
{
int wait_res = cv_.wait_for(lk, MAX_WAIT_TS);
if (wait_res == ETIMEDOUT && total_ref_cnt_ <= remote_ref_cnt_)
bthread_usleep(interval_us);
if (total_ref_cnt_.load(std::memory_order_acquire) <=
remote_ref_cnt_.load(std::memory_order_acquire))
{
LOG(WARNING) << "Waitting timeout for dbsize";
break;
remaining_wait_us = remaining_wait_us > interval_us
? remaining_wait_us - interval_us
: 0;
if (remaining_wait_us == 0)
{
LOG(WARNING) << "Waiting timeout for dbsize";
break;
}
}
if ((interval_us << 1) < kMaxIntervalUs)
interval_us <<= 1;
}
}

Expand All @@ -8538,15 +8530,23 @@ struct DbSizeCc : public CcRequestBase
return total_obj_sizes_.size();
}

bthread::Mutex mux_;
bthread::ConditionVariable cv_;

private:
std::vector<int64_t /*atomic*/> total_obj_sizes_;

protected:
size_t total_ref_cnt_{0};
size_t remote_ref_cnt_{0};
bool OnLocalRefFinished()
{
return total_ref_cnt_.fetch_sub(1, std::memory_order_acq_rel) == 1;
}

bool OnRemoteRefFinished()
{
remote_ref_cnt_.fetch_sub(1, std::memory_order_acq_rel);
return total_ref_cnt_.fetch_sub(1, std::memory_order_acq_rel) == 1;
}

std::atomic_size_t total_ref_cnt_{0};
std::atomic_size_t remote_ref_cnt_{0};
int32_t term_{0};
std::vector<uint32_t> vct_ng_id_;
std::vector<TableName> *table_names_{nullptr};
Expand Down
12 changes: 5 additions & 7 deletions tx_service/src/remote/remote_cc_request.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2317,7 +2317,7 @@ txservice::remote::RemoteDbSizeCc::RemoteDbSizeCc()

post_lambda_ = [this]()
{
assert(total_ref_cnt_ == 0);
assert(total_ref_cnt_.load(std::memory_order_acquire) == 0);
output_msg_.set_handler_addr(input_msg_->handler_addr());
output_msg_.set_txm_addr(input_msg_->txm_addr());

Expand Down Expand Up @@ -2362,8 +2362,8 @@ void txservice::remote::RemoteDbSizeCc::Reset(

DbSizeCc::Reset(&redis_table_names_, core_cnt, 0);
assert(table_names_ == &redis_table_names_);
assert(total_ref_cnt_ == core_cnt);
assert(remote_ref_cnt_ == 0);
assert(total_ref_cnt_.load(std::memory_order_relaxed) == core_cnt);
assert(remote_ref_cnt_.load(std::memory_order_relaxed) == 0);

AddLocalNodeGroupId(cmds_req.node_group_id());

Expand All @@ -2388,10 +2388,8 @@ bool txservice::remote::RemoteDbSizeCc::Execute(CcShard &ccs)
}
}

std::unique_lock lk(mux_);
assert(remote_ref_cnt_ == 0);
--total_ref_cnt_;
if (total_ref_cnt_ == 0)
assert(remote_ref_cnt_.load(std::memory_order_relaxed) == 0);
if (OnLocalRefFinished())
{
table_names_ = nullptr;
redis_table_names_.clear();
Expand Down