dragonflydb / dragonflydb/dragonfly
Transactions: allow per-shard callback preemption
- Dominant language
- C++
- Stars
- 31.6k
- Forks
- 1.3k
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 137
Description
## Problem
The transaction framework pays two system-wide costs to guard against a rare event (callback preemption during journal writes):
### 1. Inline scheduling disabled during replication
`AllowInlineScheduling()` returns false whenever journal callbacks are registered, and `CanRunInlined()` also checks `DbSlice::HasRegisteredCallbacks()`. This forces ALL single-shard transactions to be posted to the shard queue instead of running inline on the coordinator fiber — adding message dispatch overhead for every command during replication, even though actual preemptions are rare.
### 2. Intent locks acquired for optimistic single-shard transactions
In `ScheduleInShard`, intent locks are **always** acquired — even on the optimistic fast path where the transaction runs its callback inline and concludes immediately without ever entering the tx-queue. The code comment (`transaction.cc`, line 1207-1208) explains why:
```cpp
// We need to acquire the fp locks because the executing callback
// within RunCallback below might preempt.
const bool keys_unlocked = GetDbSlice(shard->shard_id()).Acquire(mode, lock_args);
```
The concern: if an optimistic transaction runs its callback without holding intent locks and that callback preempts, another transaction scheduling on the same shard would check the intent locks, find the keys uncontended, set `OUT_OF_ORDER`, and start executing its own callback — before the preempted one finishes, on the same keys.
The chain of events without locks:
```
1. Tx A: ScheduleInShard (optimistic, no locks acquired)
RunCallback → preempts during journal write
2. Tx B: ScheduleInShard (on a different fiber, same thread)
Check intent locks → keys appear free (Tx A holds no locks)
Set OUT_OF_ORDER
RunCallback → starts executing before Tx A finishes!
3. Tx A: resumes and finishes → operations on same keys interleaved
```
With intent locks held (current behavior), step 2 would see the keys as contended, Tx B would NOT get `OUT_OF_ORDER`, and it would be inserted into the tx-queue to wait. This is correct but expensive: every optimistic single-shard transaction pays the cost of `Acquire()` + `Release()` on the lock table, even though the vast majority never preempt.
## Goal
Track callback execution at the shard level so the system can handle the rare preemption case precisely, rather than pessimistically disabling inlining and acquiring intent locks for all transactions.
## Background: `continuation_trans_` vs `running_tx_`
EngineShard already has a `continuation_trans_` pointer. A natural question is whether it can be extended to cover preemption tracking. The answer is no — the two pointers serve different purposes and are active at different times.
**`continuation_trans_`** is set **between** hops: after a callback returns but before the transaction concludes. The transaction has already been removed from the tx-queue. The pointer persists across PollExecution invocations to give the multi-hop transaction priority.
**`running_tx_`** (proposed) is set **during** a callback: before `RunCallback` enters the command callback, cleared after the callback and journal write complete. The pointer is only non-null while a callback is on the call stack of some fiber.
When an inlined callback preempts, `continuation_trans_` is NOT set — `RunInShard` has not returned yet (we're mid-callback). Reusing `continuation_trans_` would break PollExecution's semantics: it would try to `DisarmInShard` and dispatch the next hop, but the transaction is mid-callback on another fiber, not armed for a new hop.
| | `continuation_trans_` | `running_tx_` |
|---|---|---|
| **When set** | After callback returns, before next hop | During callback execution |
| **Meaning** | "This tx needs another hop" | "A callback is on the call stack" |
| **Fiber** | Same fiber (shard queue) | Any fiber (inline or shard queue) |
| **PollExecution action** | Dispatch next hop | N/A (checked during scheduling) |
| **Duration** | Between hops (can be long) | During single callback (usually short) |
Both can be set simultaneously: a continuation transaction runs its next hop → `running_tx_` is set during that callback → if it preempts, `running_tx_` is visible to scheduling → callback completes → `running_tx_` cleared → `continuation_trans_` remains.
## Solution
### Core idea
Add a per-shard `Transaction* running_tx_` pointer to `EngineShard`, set before entering `RunCallback` and cleared after it returns. This pointer makes callback execution visible to the scheduling algorithm, allowing it to make correct OOO decisions without requiring every optimistic transaction to acquire intent locks upfront.
### Invariants
**Callback exclusivity**: At most one `RunCallback` executes per shard at any point in time. Although Dragonfly uses cooperative scheduling on a single thread per shard, a fiber can yield during journal callbacks. This creates a window where another fiber can enter the scheduling algorithm on the same shard. `running_tx_` makes this window visible.
- `running_tx_ == nullptr`: no callback in flight, the scheduling algorithm can rely on intent locks alone.
- `running_tx_ != nullptr`: a callback is mid-execution (possibly suspended on another fiber). The scheduling algorithm must account for this transaction's keys even though they may not be reflected in the intent lock table.
### Effect on inline scheduling
With `running_tx_`, `AllowInlineScheduling()` no longer needs to check for journal callbacks, and `CanRunInlined()` no longer needs to check `DbSlice::HasRegisteredCallbacks()`. The preemption hazard is handled within the scheduling algorithm itself: when a subsequent transaction schedules on the shard and observes `running_tx_`, it can react appropriately. No system-wide restriction is needed.
The LOADING state check in `AllowInlineScheduling()` remains — it guards a different problem (RdbLoader not using transactions).
### Effect on intent locks and OOO decisions
The key insight is that `running_tx_` is used during **scheduling**, not as a gate in PollExecution. When a transaction enters `ScheduleInShard` and observes `running_tx_ != nullptr`, it knows a callback is in flight that may not have its keys reflected in the lock table. The scheduling algorithm handles this with **lazy locking**:
**Common case** (`running_tx_ == null`): The optimistic path skips intent lock acquisition entirely. The callback runs and completes without touching the lock table. This is safe because no other callback is in flight, so intent locks accurately reflect the state of enqueued transactions (if any).
**Rare case** (`running_tx_ != null`): A callback is suspended on another fiber. The scheduling algorithm acquires intent locks for `running_tx_`'s keys on its behalf (if not already locked), then acquires its own locks, then checks for conflicts as usual. This restores the lock table to an accurate state so the OOO decision is correct:
- If Tx B's keys **don't overlap** with the running Tx A's keys → no conflict → Tx B can proceed OOO. No unnecessary penalty.
- If Tx B's keys **overlap** → conflict detected → Tx B enters the tx-queue and waits.
The cost of lock acquisition is paid only when preemption actually occurs AND another transaction arrives during the preemption window — a rare-on-rare event.
Revisiting the hazard scenario:
```
1. Tx A: ScheduleInShard (optimistic, NO intent locks acquired)
running_tx_ = A
RunCallback → preempts during journal write
2. Tx B: ScheduleInShard (on a different fiber, same thread)
running_tx_ != null (Tx A is in flight)
Acquire intent locks for Tx A's keys (lazy, on its behalf)
Acquire intent locks for Tx B's keys
Check conflicts:
- overlapping keys → Tx B enters tx-queue, waits
- disjoint keys → Tx B proceeds OOO (no penalty)
3. Tx A: resumes, finishes callback
running_tx_ = null
```
Intent locks remain necessary for transactions that enter the tx-queue, where they serve their original VLL purpose.
### Scope
- `EngineShard` — new `running_tx_` member and accessors.
- `Transaction::RunCallback` — set/clear `running_tx_` around the callback and journal write.
- `Transaction::ScheduleInShard` — check `running_tx_` during scheduling; lazy-lock the running transaction's keys when needed; skip upfront lock acquisition on the optimistic path.
- `ServerState::AllowInlineScheduling` — remove journal callback check.
- `Transaction::CanRunInlined` — remove `HasRegisteredCallbacks` check.
## Risks
- **Exception safety**: `running_tx_` must be cleared even if the callback throws. A scope guard ensures this.
- **Lazy locking complexity**: Acquiring locks on behalf of another transaction requires access to its `KeyLockArgs`. The `running_tx_` pointer provides this (via `GetLockArgs`). Since this is single-threaded cooperative scheduling, the pointer is stable as long as we don't yield between the check and the use.
- **LOADING state**: Orthogonal concern. The existing check in `AllowInlineScheduling` stays.
## Design doc
See `docs/running_tx_design.md` for full details including `continuation_trans_` vs `running_tx_` timeline diagrams.
Contributor guide
Research direction
Start with docs/running_tx_design.md, then trace EngineShard, Transaction::RunCallback, Transaction::ScheduleInShard, ServerState::AllowInlineScheduling, and Transaction::CanRunInlined. Compare the design's preemption timeline with transaction.cc and verify that the listed scope is covered, exception safety is preserved, and the LOADING check remains.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- databases
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100