dragonflydb / dragonflydb/dragonfly

Design: Redesign `Transaction::shard_data_` — single inline cell + heap array, encapsulated access

Open
#7,791 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
31.5k
Forks
1.3k
Avg merge
1d 10h
Merged PRs (30d)
137

Description

## Context

`Transaction::shard_data_` is an `absl::InlinedVector` (transaction.h:624) where
`PerShardData` is a 64-byte cacheline (transaction.h:411-445). Access goes through
`SidToId(sid) = sid < shard_data_.size() ? sid : 0` (transaction.h:594-596): the array is **only
ever sized 1 or `shard_set->size()`**, so the fold means "index==sid when full, cell 0 when
single". This has two problems:

1. **Footgun / inelegance.** ~24 sites must remember `SidToId`; two raw `shard_data_[sid]`
(GetLockArgs transaction.cc:1183, GetShardArgs transaction.cc:1369) are correct *only* because
they sit in `unique_shard_cnt_ > 1` branches. The invariant lives in the author's head, not the
type. The header even warns `// Index only with SidToId(shard index)!` and `// TODO: explore
dense packing`.
2. **Wasted memory.** Inline capacity 4 bloats *every* `Transaction` by `4 × 64B = 256B`, though
the common case is single-shard.

**Goal:** encapsulate all access behind a typed accessor, and represent storage as a single inline
cell for single-shard transactions plus a heap array for multi-shard — optimizing for the
single-vs-multi split instead of an arbitrary inline capacity.

## Design

### Split PerShardData: small core (both regimes) + full cacheline (multi only)

Today `PerShardData` (64B, `alignas(64)`, transaction.h:411-445) carries fields the single-shard
case never uses: `slice_start/slice_count` (subspan in `args_slices_`) and `fp_start/fp_count`
(subspan in `kv_fp_`) are only read in the `unique_shard_cnt_ > 1` branches of `GetShardArgs`
(transaction.cc:1369) and `GetLockArgs` (transaction.cc:1183) and written by the full-fanout
builder `InitShardData` (transaction.cc:227-245). The `alignas(64)` + `pad` exist only to keep
adjacent array cells from false-sharing across shard fibers — irrelevant for a lone inline cell.

So split into:

```cpp
// Per-shard state needed by BOTH single- and multi-shard transactions. No forced alignment/padding.
struct ShardState {
uint16_t local_mask = 0;
std::atomic_bool is_armed = false; // cross-thread sync point
TxQueue::Iterator pq_pos = TxQueue::kEnd;
uint32_t wake_key_pos = UINT32_MAX;
struct Stats { unsigned total_runs = 0; } stats;
}; // ~16 bytes

// Element of the multi-shard heap array: core (inherited) + arg/fp subspans, one full cache line
// to avoid false sharing between shards' fibers. Deriving (not embedding) keeps core fields flat:
// shards_[sid].local_mask, not .state.local_mask. alignas(64) makes sizeof a multiple of 64, so no
// explicit pad member is needed.
struct alignas(64) PerShardData : ShardState {
uint32_t slice_start = 0, slice_count = 0; // subspan in args_slices_
uint32_t fp_start = 0, fp_count = 0; // subspan in kv_fp_
};
static_assert(sizeof(PerShardData) == 64);
```

### Storage: inline core cell + heap array

```cpp
// Single-shard transactions (the common case) use this small inline cell; multi-shard ones
// allocate shards_ sized to shard_set->size() (index == shard id). Never both at once.
ShardState single_shard_data_;
std::unique_ptr shards_; // non-null == "full, index by shard id"
```

Single-shard inline per-shard cost drops from 256B (`InlinedVector<_,4>`) to ~16B + 8B pointer.

### Central accessors (replace `SidToId` entirely)

```cpp
// Core state, valid in both regimes. shards_[sid] (PerShardData) binds to ShardState& via
// derived->base reference conversion; single_shard_data_ is already ShardState&.
ShardState& ShardData(ShardId sid) { return shards_ ? shards_[sid] : single_shard_data_; }
const ShardState& ShardData(ShardId sid) const { ... }
bool HasFullShardData() const { return bool(shards_); }

// Arg/fp subspans, only meaningful for multi-shard (unique_shard_cnt_ > 1); DCHECKs shards_.
PerShardData& FullShardData(ShardId sid) { DCHECK(shards_); return shards_[sid]; }
```

Sizing helpers to replace the raw `resize(...)` calls:
- `EnsureFullShardData()` — `if (!shards_) shards_ = std::make_unique(shard_set->size());`
(aligned `new[]` honors `alignas(64)`). Used everywhere the code currently does
`shard_data_.resize(shard_set->size())` (transaction.cc:333, 1135) and the full-size branch of
the `NeedsFullShardData() ? ... : 1` resizes (transaction.cc:327, 1128).
- `UseSingleShardData()` — `shards_.reset();` and reset `single_shard_data_`. Used where the code
shrinks to one cell (transaction.cc:356, and the `: 1` branch of 327/1128). This is only ever
reached on `!IsActiveMulti()` paths, preserving the "never shrink an active multi array" rule
(transaction.cc:350-356), so leftover cross-thread callbacks are never invalidated.

Because the full array is always exactly `shard_set->size()` long when present, no separate length
field is needed. The regime for the accessor (`shards_` null-ness) exactly reproduces `SidToId`'s
size-based fold.

**Not a union.** Overlapping the two fields in a union saves only ~8B: the union size is dominated
by the 64B `PerShardData`, so multi mode reclaims nothing and only the 8B pointer is saved in
single mode. It costs manual union lifetime management (placement-new the cell / explicit reset of
the unique_ptr on every regime switch — both types are non-trivial), a separate discriminant bool
(the clean `shards_ ? ... : ...` accessor stops working since you can't read the inactive member),
and reintroduces exactly the footguns we're removing. The ~190B win comes from dropping the
`InlinedVector<_,4>` inline buffer; chasing the last 8B is not worth it. Keep the two fields plain.

## Call-site migration (mechanical)

Core-field access (`local_mask`, `is_armed`, `pq_pos`, `wake_key_pos`, `stats`):
- Every `shard_data_[SidToId(x)]` → `ShardData(x)` (returns `ShardState&`) (the ~24 sites in
section-1a of the inventory: transaction.h:382/386/390/601;
transaction.cc:328/353/560/562/589/902/1191/1197/1215/1219/1433/1448/1507/1542/1566, plus the
SidToId-via-local ones at 588-589, 1234-1235, 1326-1327).
- `shard_data_.front()` (transaction.cc:357, 1129) → `single_shard_data_` after `UseSingleShardData()`.

Arg/fp subspan access (`slice_start/count`, `fp_start/count`) — all multi-only:
- The two raw `shard_data_[sid]` (transaction.cc:1183 GetLockArgs, 1369 GetShardArgs) → `FullShardData(sid)`;
read `.fp_start/.fp_count` / `.slice_start/.slice_count` off it. Both sit in `unique_shard_cnt_ > 1`
branches, so `shards_` is set.
- `InitShardData` (transaction.cc:227-245) builds the full array — index via `FullShardData(i)`;
sets `.local_mask` + the subspan fields flat (core fields inherited, no `.state.` prefix).
- `shard_data_.size()` used as a shard count (transaction.cc:207, 337) → `shard_set->size()`.
- `DCHECK_EQ(shard_data_.size(), shard_set->size())` (transaction.cc:431, 843) → `DCHECK(HasFullShardData())`.
- `DCHECK_LE(shard_data_.size(), 1024u)` (transaction.cc:968) → drop or assert `shard_set->size() <= 1024`.
- Index loops over the full array (transaction.cc:227-228, 432-441, 1116-1117 [dead `#if 0`], and
IterateShards' multi branch transaction.h:603-604) → loop `0 .. shard_set->size()` indexing
`shards_[i]`. `IterateShards` single-shard branch (transaction.h:601) already routes through the
accessor via `unique_shard_cnt_ == 1`, so `IterateShards`/`IterateActiveShards` keep working
unchanged except for the internal loop bound/index source.
- `MultiSwitchCmd` reset loop `for (auto& sd : shard_data_)` (transaction.cc:505) →
`IterateShards` (or loop over the active regime).

`SidToId` is deleted after migration. `shard_data_` no longer exists as a name.

## Cleanup / notes

- With no `InlinedVector`, elements are never moved (array is `new[]` at final size; inline cell is
a stable member). The empty move ctor `PerShardData(PerShardData&&)` (transaction.h:414) becomes
dead — remove it. Keep `alignas(64)` + `static_assert(sizeof(PerShardData)==64)` for the heap
array element (false-sharing avoidance between shard fibers); the old explicit `pad[]` member is
no longer needed since `alignas(64)` rounds the size up to a full cache line.
- **False-sharing trade-off for the inline cell:** dropping alignment means `single_shard_data_`
(incl. its `is_armed`) may share a cache line with neighboring `Transaction` members. For
single-shard txs there is exactly one shard fiber + coordinator touching it (one arm/disarm), not
N fibers on adjacent cells, so contention is minimal — an acceptable trade for the size win.
Revisit only if profiling shows single-shard hop contention.

## Constraints preserved (from inventory)
- Cell address stability for the `is_armed` cross-thread sync point (transaction.cc:588-589,
1190-1205; engine_shard.cc:651-715): inline cell is in the stable `Transaction`; the heap array is
allocated once and never reallocated/shrunk while the multi tx is active.
- Arg/fp subspans are only read when `unique_shard_cnt_ > 1` (verified: GetShardArgs/GetLockArgs
single-shard branches use the whole span), so omitting them from the single cell is safe — even in
the IsActiveMulti single-shard case (full array kept, but those getters take the whole-span branch).
- `shard_set->size() < 1024` (main_service.cc:214); `std::bitset<1024>` in DispatchHop unaffected.
- Sizing regimes unchanged (full only when `NeedsFullShardData()`), so #2458/#7784 invariants hold.

## Verification
- Build debug (DCHECKs on): `cd build-dbg && ninja transaction_test multi_test dragonfly_test string_family_test generic_family_test`
- Run: `./transaction_test && ./multi_test && ./dragonfly_test && ./string_family_test && ./generic_family_test`
(these exercise single-shard, multi-shard, atomic multi, non-atomic squash, and stub paths).
- Rerun the SQUASHED_LOCAL test added earlier: `./multi_test --gtest_filter='MultiTest.SquashShardLocalMultiKey'`.
- ASAN build to stress cross-thread stability of the heap array under concurrent hops:
`./helio/blaze.sh -DWITH_ASAN=ON ...` then run `multi_test`/`dragonfly_test`.
- `pre-commit run --files src/server/transaction.cc src/server/transaction.h`.

## Out of scope
- General dense packing (store only K active cells + a `sid→slot` map). Deferred: high risk
(scheduling/arming/cancel + a second stable cross-thread structure) for marginal gain now that
the inline bloat is removed. This design leaves the door open — a `sid→slot` map could later hide
behind the same `ShardData(sid)` accessor without touching call sites.

Contributor guide

Open the contributing guide

Research direction

Start with src/server/transaction.h and src/server/transaction.cc, tracing the listed shard_data_ and SidToId call sites, then inspect the related synchronization paths in engine_shard.cc. Use the specified debug tests, SQUASHED_LOCAL filter, ASAN run, and pre-commit command to verify that single- and multi-shard transactions retain their behavior and stable cross-thread access.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
backend, databases
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.