tsoracle_driver_openraft/state_machine.rs
1//
2// ░▀█▀░█▀▀░█▀█░█▀▄░█▀█░█▀▀░█░░░█▀▀
3// ░░█░░▀▀█░█░█░█▀▄░█▀█░█░░░█░░░█▀▀
4// ░░▀░░▀▀▀░▀▀▀░▀░▀░▀░▀░▀▀▀░▀▀▀░▀▀▀
5//
6// tsoracle — Distributed Timestamp Oracle
7// https://www.tsoracle.rs
8//
9// Copyright (c) 2026 Prisma Risk
10//
11// Licensed under the Apache License, Version 2.0 (the "License");
12// you may not use this file except in compliance with the License.
13// You may obtain a copy of the License at
14//
15// https://www.apache.org/licenses/LICENSE-2.0
16//
17// Unless required by applicable law or agreed to in writing, software
18// distributed under the License is distributed on an "AS IS" BASIS,
19// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20// See the License for the specific language governing permissions and
21// limitations under the License.
22//
23
24// #[PerformanceCriticalPath]
25//! `RaftStateMachine` for the high-water counter with pluggable snapshot
26//! persistence.
27//!
28//! State is one `u64` plus the openraft-required apply-progress metadata
29//! (`last_applied`, `last_membership`). Snapshots are encoded with the toolkit's
30//! version-prefixed codec ([`tsoracle_openraft_toolkit::encode`] /
31//! [`tsoracle_openraft_toolkit::decode`] at the active write version) — a leading
32//! version byte followed by the postcard body — written through to a
33//! [`SnapshotStore`] so they survive a
34//! process restart. The default store is in-memory
35//! ([`HighWaterStateMachine::new`]); production deployments construct via
36//! [`HighWaterStateMachine::with_store`] and supply a durable backend such as
37//! [`crate::snapshot_store::RocksdbSnapshotStore`].
38//!
39//! All state lives inside an `Arc<Mutex<Core>>` so the state machine is
40//! cheaply cloneable — required because openraft's
41//! `RaftStateMachine::SnapshotBuilder = Self` design hands out clones to drive
42//! `build_snapshot` concurrently with `apply`. The `Arc<dyn SnapshotStore>` is
43//! shared the same way, so every clone sees the same persisted snapshot.
44
45use std::io;
46use std::io::Cursor;
47use std::sync::Arc;
48
49use futures::StreamExt;
50use openraft::EntryPayload;
51use openraft::RaftSnapshotBuilder;
52use openraft::StoredMembership;
53use openraft::storage::EntryResponder;
54use openraft::storage::RaftStateMachine;
55use openraft::storage::Snapshot;
56use openraft::type_config::alias::{LogIdOf, SnapshotMetaOf, SnapshotOf, StoredMembershipOf};
57use parking_lot::Mutex;
58use serde::{Deserialize, Serialize};
59
60use tsoracle_codec::{
61 VersionedCodec, decode_framed, decode_postcard_exact, encode_framed, encode_postcard,
62};
63use tsoracle_openraft_toolkit::{
64 ActiveWriteVersion, BASELINE_WRITE_VERSION, BATCH_WRITE_VERSION, DENSE_WRITE_VERSION,
65 MAX_READABLE_VERSION, MIN_READABLE_VERSION, codec_io_error,
66};
67
68use crate::log_entry::{HighWaterCommand, SetFormatVersionPayload};
69use crate::snapshot_store::{InMemorySnapshotStore, SnapshotStore};
70use crate::type_config::{ApplyOutcome, HighWaterApplied, TypeConfig};
71
72type LogId = LogIdOf<TypeConfig>;
73type SnapMeta = SnapshotMetaOf<TypeConfig>;
74type SnapOf = SnapshotOf<TypeConfig, Cursor<Vec<u8>>>;
75type SnapData = Cursor<Vec<u8>>;
76type StoredMem = StoredMembershipOf<TypeConfig>;
77
78/// Snapshot payload. The persisted/streamed bytes are version-framed as
79/// `[active write version | postcard(Self)]`, so decode them through the
80/// toolkit `decode_framed` over the readable range
81/// `[MIN_READABLE_VERSION, MAX_READABLE_VERSION]` rather than raw
82/// `postcard::from_bytes`.
83///
84/// Exposed at the crate root so callers building tooling around the snapshot
85/// format (e.g. inspectors, migration tools) can decode it without re-deriving
86/// the layout.
87#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
88pub struct HighWaterStateMachineSnapshot {
89 pub current_value: u64,
90 pub last_applied: Option<LogId>,
91 pub last_membership: StoredMem,
92 /// Per-key dense counters (v5+). Empty in a v4 snapshot (pre-dense); the
93 /// restore path keeps the genesis cap seeded from the constructor.
94 pub dense: std::collections::BTreeMap<String, u64>,
95 /// Genesis cardinality cap (v5+). Zero in a v4 snapshot ("absent
96 /// sentinel"); the restore path keeps the cap seeded from the constructor.
97 pub dense_cap: u64,
98}
99
100/// The v4 on-disk snapshot layout, frozen. Used only to decode
101/// `BASELINE_WRITE_VERSION` bytes written before the dense fields existed;
102/// converted into the current payload (empty dense map, sentinel cap 0 =
103/// "absent") on decode. Do not edit its field set — it must remain
104/// byte-identical to the pre-dense layout.
105#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
106struct HighWaterStateMachineSnapshotV4 {
107 current_value: u64,
108 last_applied: Option<LogId>,
109 last_membership: StoredMem,
110}
111
112/// On-disk envelope written to the [`SnapshotStore`]: pairs the openraft
113/// snapshot meta with the version-framed payload, where `data` holds
114/// `[active write version | postcard(HighWaterStateMachineSnapshot)]`. Kept
115/// private — embedders that need to inspect persisted snapshots decode the
116/// inner `data` blob through the toolkit `decode_framed` over the readable
117/// range.
118#[derive(Debug, Clone, Serialize, Deserialize)]
119struct PersistedSnapshot {
120 meta: SnapMeta,
121 data: Vec<u8>,
122}
123
124impl VersionedCodec for HighWaterStateMachineSnapshot {
125 fn decode_version(version: u8, body: &[u8]) -> Result<Self, tsoracle_codec::CodecError> {
126 match version {
127 // v4 (BASELINE_WRITE_VERSION): decode the frozen old layout, lift
128 // into the current payload. Pre-dense snapshots carry no dense
129 // bytes; dense is empty and cap is the sentinel 0 ("absent") — the
130 // restore path keeps the genesis cap seeded from the constructor.
131 v if v == BASELINE_WRITE_VERSION => {
132 let old: HighWaterStateMachineSnapshotV4 = decode_postcard_exact(body)?;
133 Ok(HighWaterStateMachineSnapshot {
134 current_value: old.current_value,
135 last_applied: old.last_applied,
136 last_membership: old.last_membership,
137 dense: std::collections::BTreeMap::new(),
138 dense_cap: 0,
139 })
140 }
141 // v5 (DENSE_WRITE_VERSION): the current layout decodes directly.
142 // Ungated — production reads dense snapshots from v5-activated
143 // clusters without any feature flag.
144 v if v == DENSE_WRITE_VERSION => decode_postcard_exact(body),
145 // v6 (BATCH_WRITE_VERSION): snapshot layout is byte-identical to v5;
146 // batch adds a log command, not snapshot state, so the same decoder
147 // applies.
148 v if v == BATCH_WRITE_VERSION => decode_postcard_exact(body),
149 other => Err(tsoracle_codec::CodecError::VersionUnsupported {
150 min: MIN_READABLE_VERSION,
151 max: MAX_READABLE_VERSION,
152 actual: other,
153 }),
154 }
155 }
156
157 fn encode_version(&self, version: u8) -> Result<Vec<u8>, tsoracle_codec::CodecError> {
158 match version {
159 // v4: project to the frozen old layout (dense is always empty before
160 // activation, so dropping it is lossless). This is what
161 // build_snapshot emits while the active write version is still 4.
162 v if v == BASELINE_WRITE_VERSION => {
163 // Fail loud in ALL build profiles (not just debug): a non-empty
164 // dense map at this point means write-version 5 data is being
165 // encoded into a v4 snapshot, which would silently drop it.
166 // The invariant (dense only populated after activation flips the
167 // write version to 5) is maintained by the rollout gate, so
168 // reaching here with a non-empty map indicates a protocol
169 // violation that must be surfaced rather than silently lost.
170 if !self.dense.is_empty() {
171 return Err(tsoracle_codec::CodecError::NotRepresentable { version });
172 }
173 encode_postcard(&HighWaterStateMachineSnapshotV4 {
174 current_value: self.current_value,
175 last_applied: self.last_applied,
176 last_membership: self.last_membership.clone(),
177 })
178 }
179 // v5: the current (dense) layout encodes directly. Ungated.
180 v if v == DENSE_WRITE_VERSION => encode_postcard(self),
181 // v6 (BATCH_WRITE_VERSION): snapshot layout is byte-identical to v5;
182 // batch adds a log command, not snapshot state, so the same encoder
183 // applies.
184 v if v == BATCH_WRITE_VERSION => encode_postcard(self),
185 other => Err(tsoracle_codec::CodecError::VersionUnsupported {
186 min: MIN_READABLE_VERSION,
187 max: MAX_READABLE_VERSION,
188 actual: other,
189 }),
190 }
191 }
192}
193
194impl VersionedCodec for PersistedSnapshot {
195 fn decode_version(version: u8, body: &[u8]) -> Result<Self, tsoracle_codec::CodecError> {
196 match version {
197 // v4: envelope layout unchanged — opaque `data` bytes hold the
198 // version-framed payload.
199 v if v == BASELINE_WRITE_VERSION => decode_postcard_exact(body),
200 // v5: envelope layout still unchanged (only the payload data blob
201 // inside gains dense bytes). Ungated — production reads v5
202 // envelopes once activated.
203 v if v == DENSE_WRITE_VERSION => decode_postcard_exact(body),
204 // v6 (BATCH_WRITE_VERSION): envelope layout is byte-identical to v5;
205 // batch adds a log command, not snapshot state, so the same decoder
206 // applies.
207 v if v == BATCH_WRITE_VERSION => decode_postcard_exact(body),
208 other => Err(tsoracle_codec::CodecError::VersionUnsupported {
209 min: MIN_READABLE_VERSION,
210 max: MAX_READABLE_VERSION,
211 actual: other,
212 }),
213 }
214 }
215
216 fn encode_version(&self, version: u8) -> Result<Vec<u8>, tsoracle_codec::CodecError> {
217 match version {
218 v if v == BASELINE_WRITE_VERSION => encode_postcard(self),
219 // v5: envelope layout unchanged; real ungated arm for production.
220 v if v == DENSE_WRITE_VERSION => encode_postcard(self),
221 // v6 (BATCH_WRITE_VERSION): envelope layout is byte-identical to v5;
222 // batch adds a log command, not snapshot state, so the same encoder
223 // applies.
224 v if v == BATCH_WRITE_VERSION => encode_postcard(self),
225 other => Err(tsoracle_codec::CodecError::VersionUnsupported {
226 min: MIN_READABLE_VERSION,
227 max: MAX_READABLE_VERSION,
228 actual: other,
229 }),
230 }
231 }
232}
233
234/// Mutable core guarded by a single mutex. `parking_lot::Mutex` because the
235/// apply path is hot and the state is fully re-derivable from the raft log
236/// (poison semantics would just mask the originating panic).
237struct Core {
238 current_value: u64,
239 /// Per-key dense counters. Deterministic iteration (BTreeMap) so snapshots
240 /// are byte-identical across replicas.
241 dense: std::collections::BTreeMap<String, u64>,
242 /// Immutable, cluster-wide genesis cardinality cap. Seeded from the SM
243 /// constructor (identical config on every node) or restored from a
244 /// snapshot; never mutated at runtime. Stored in the replicated snapshot so
245 /// replay/restore reproduces the same accept/reject decisions (spec §6.2).
246 dense_cap: u64,
247 last_applied: Option<LogId>,
248 last_membership: StoredMem,
249 /// Snapshot index counter, used to make snapshot ids unique even when two
250 /// snapshots are produced at the same `last_applied` log id.
251 snapshot_idx: u64,
252 /// The most-recently built or installed snapshot, retained in memory so
253 /// `get_current_snapshot` does not need to rebuild on every call.
254 current_snapshot: Option<StoredSnapshot>,
255}
256
257#[derive(Clone)]
258struct StoredSnapshot {
259 meta: SnapMeta,
260 data: Vec<u8>,
261}
262
263/// Whether a snapshot at `incoming` may replace one currently published at
264/// `published` without regressing the applied log id.
265///
266/// Equal ids may replace (a fresh rebuild at the same `last_applied` carries
267/// the same state, only a new `snapshot_id`). `None` — no prior snapshot, or a
268/// pre-genesis snapshot — is the minimum, so a `None` incoming never displaces
269/// a `Some` published. This is the monotone guard that closes the
270/// `build_snapshot`/`install_snapshot` publish TOCTOU.
271fn supersedes_published(incoming: Option<LogId>, published: Option<LogId>) -> bool {
272 incoming >= published
273}
274
275/// `RaftStateMachine` for the high-water counter, with pluggable snapshot
276/// persistence.
277///
278/// Clone-cheap: the `Arc<Mutex<Core>>`, the `Arc<dyn SnapshotStore>`, and the
279/// `Arc<Mutex<()>>` persist lock are all shared by clone. Required by
280/// openraft's `get_snapshot_builder` contract which uses
281/// `SnapshotBuilder = Self`.
282pub struct HighWaterStateMachine {
283 core: Arc<Mutex<Core>>,
284 store: Arc<dyn SnapshotStore>,
285 /// Serializes the persist-then-publish sequence of `build_snapshot` and
286 /// `install_snapshot` across all clones. openraft drives `build_snapshot`
287 /// on a clone concurrently with `install_snapshot` on the main handle;
288 /// without this lock a slow build could write the store and publish
289 /// `current_snapshot` *after* a newer install, rolling both the durable
290 /// and in-memory snapshot back to a stale `last_log_id`. Distinct from
291 /// `core` precisely so the hot `apply` path never serializes against
292 /// snapshot I/O — `core` is never held across `store.save`.
293 persist: Arc<Mutex<()>>,
294 /// Shared active write version: the version this SM stamps onto snapshots
295 /// it builds and installs. Shares the one cell with the log store (and,
296 /// in a later phase, the wire sender) so all writers emit the same
297 /// version. Mutated only by a successful activation apply (later phase);
298 /// defaults to BASELINE.
299 active_write_version: ActiveWriteVersion,
300}
301
302impl Clone for HighWaterStateMachine {
303 fn clone(&self) -> Self {
304 Self {
305 core: Arc::clone(&self.core),
306 store: Arc::clone(&self.store),
307 persist: Arc::clone(&self.persist),
308 active_write_version: self.active_write_version.clone(),
309 }
310 }
311}
312
313impl Default for HighWaterStateMachine {
314 fn default() -> Self {
315 Self::new()
316 }
317}
318
319impl HighWaterStateMachine {
320 /// Build a state machine backed by an in-memory snapshot store.
321 ///
322 /// Equivalent to
323 /// `HighWaterStateMachine::with_store(Arc::new(InMemorySnapshotStore::new()))`.
324 /// Snapshots survive within a process but not across restarts; for that
325 /// use [`HighWaterStateMachine::with_store`] with a durable backend such
326 /// as [`crate::snapshot_store::RocksdbSnapshotStore`].
327 #[expect(
328 clippy::expect_used,
329 reason = "`with_store` only fails when `store.load()` returns Err; \
330 `InMemorySnapshotStore::load` is `Ok(None)` for a fresh \
331 store, so this branch is unreachable."
332 )]
333 pub fn new() -> Self {
334 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
335 Self::with_store(store).expect("InMemorySnapshotStore::load is infallible")
336 }
337
338 /// Build a state machine with a custom genesis cardinality cap, backed by
339 /// an in-memory snapshot store. The cap is immutable after construction and
340 /// must be identical on every cluster member. Intended for tests that need
341 /// a small cap to exercise `DenseCardinalityExceeded`; production callers
342 /// use [`new`](Self::new) which defaults to
343 /// [`DEFAULT_DENSE_CARDINALITY_CAP`](crate::DEFAULT_DENSE_CARDINALITY_CAP).
344 #[cfg(test)]
345 #[expect(
346 clippy::expect_used,
347 reason = "`new_with_dense_cap` only fails when `store.load()` returns Err; \
348 `InMemorySnapshotStore::load` is `Ok(None)` for a fresh \
349 store, so this branch is unreachable."
350 )]
351 fn new_with_dense_cap(dense_cap: u64) -> Self {
352 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
353 let sm = Self::with_store(store).expect("InMemorySnapshotStore::load is infallible");
354 sm.core.lock().dense_cap = dense_cap;
355 sm
356 }
357
358 /// Build a state machine backed by `store` and rehydrated from whatever
359 /// snapshot the store currently holds.
360 ///
361 /// If `store.load()` returns a snapshot, the state machine's
362 /// `current_value`, `last_applied`, `last_membership`, dense state, and
363 /// `get_current_snapshot()` all reflect that snapshot on return. This is
364 /// the contract that lets openraft re-enable the default snapshot policy:
365 /// after a restart the log store's `last_purged_log_id` may sit above
366 /// index 0, and the state machine must already cover those purged entries
367 /// or openraft will panic during recovery.
368 ///
369 /// The active write version is recovered to at least the persisted
370 /// snapshot's leading version byte, so a store holding a v5/dense snapshot
371 /// does not come back at the baseline (which would trip the fail-loud
372 /// `NotRepresentable` guard on the next `build_snapshot`). For full
373 /// max-of-evidence recovery — folding in the highest version byte among
374 /// durable *log* records as well — seed the cell via
375 /// [`recover_active_write_version`](tsoracle_openraft_toolkit::recover_active_write_version)
376 /// and pass it to [`with_store_and_active_version`](Self::with_store_and_active_version);
377 /// `with_store` only sees the snapshot store.
378 pub fn with_store(store: Arc<dyn SnapshotStore>) -> io::Result<Self> {
379 Self::with_store_and_active_version(store, ActiveWriteVersion::default())
380 }
381
382 /// Build a state machine backed by `store`, sharing `active_write_version`
383 /// with the log store (and, in a later phase, the wire sender). Bootstrap
384 /// constructs and seeds the cell once, then threads the same clone here
385 /// and into the log store; non-bootstrap callers use
386 /// [`with_store`](Self::with_store), which supplies a fresh BASELINE cell.
387 pub fn with_store_and_active_version(
388 store: Arc<dyn SnapshotStore>,
389 active_write_version: ActiveWriteVersion,
390 ) -> io::Result<Self> {
391 let mut core = Core {
392 current_value: 0,
393 dense: std::collections::BTreeMap::new(),
394 dense_cap: crate::DEFAULT_DENSE_CARDINALITY_CAP,
395 last_applied: None,
396 last_membership: StoredMembership::default(),
397 snapshot_idx: 0,
398 current_snapshot: None,
399 };
400 if let Some(bytes) = store.load()? {
401 let persisted: PersistedSnapshot =
402 decode_framed(MIN_READABLE_VERSION, MAX_READABLE_VERSION, &bytes)
403 .map_err(|e| codec_io_error("persisted snapshot envelope decode", e))?;
404 let payload: HighWaterStateMachineSnapshot =
405 decode_framed(MIN_READABLE_VERSION, MAX_READABLE_VERSION, &persisted.data)
406 .map_err(|e| codec_io_error("persisted snapshot payload decode", e))?;
407 core.current_value = payload.current_value;
408 core.last_applied = payload.last_applied;
409 core.last_membership = payload.last_membership;
410 core.dense = payload.dense;
411 // A v4 snapshot carries dense_cap == 0 (absent sentinel); keep the
412 // constructor-seeded genesis cap. A v5 snapshot carries the real cap.
413 if payload.dense_cap != 0 {
414 core.dense_cap = payload.dense_cap;
415 }
416 core.current_snapshot = Some(StoredSnapshot {
417 meta: persisted.meta,
418 data: persisted.data,
419 });
420 // Recover the active write version from the persisted envelope's
421 // leading version byte, mirroring `install_snapshot`'s
422 // max-of-evidence rule (and `recover_active_write_version`'s).
423 // Without this, the bare `with_store` path — which hands in a fresh
424 // BASELINE cell — would rehydrate a populated dense map (from a v5
425 // snapshot) while leaving the cell at v4, then trip the fail-loud
426 // `NotRepresentable` guard on its next `build_snapshot`. The byte is
427 // a safe floor: a v5 snapshot proves a `SetFormatVersion` to at
428 // least that version already committed cluster-wide. Only ever
429 // advances, so a caller (e.g. standalone bootstrap) that already
430 // seeded the cell from richer log+snapshot evidence is unaffected.
431 // `decode_framed` accepted the byte above, so it is within
432 // `[MIN, MAX_READABLE_VERSION]` and never exceeds what this binary
433 // can write.
434 if let Some(&snapshot_version) = bytes.first() {
435 if snapshot_version > active_write_version.get() {
436 active_write_version.set(snapshot_version);
437 }
438 }
439 }
440 let state_machine = Self {
441 core: Arc::new(Mutex::new(core)),
442 store,
443 persist: Arc::new(Mutex::new(())),
444 active_write_version,
445 };
446 // Format-migration observability: publish the compile-time readable
447 // bounds and the recovered active write version so a freshly-started
448 // node reports its version state before any apply runs.
449 crate::observability::record_readable_bounds(MIN_READABLE_VERSION, MAX_READABLE_VERSION);
450 crate::observability::record_active_write_version(state_machine.active_write_version());
451 Ok(state_machine)
452 }
453
454 /// The version this SM currently stamps onto snapshots.
455 pub fn active_write_version(&self) -> u8 {
456 self.active_write_version.get()
457 }
458
459 /// Read the current high-water value without going through raft.
460 ///
461 /// Returns the value most-recently written by `apply` or
462 /// `install_snapshot`. This is a state-machine-local read; callers that
463 /// need linearizability must coordinate a read barrier through `Raft`
464 /// before calling.
465 pub async fn current_value(&self) -> u64 {
466 self.core.lock().current_value
467 }
468
469 /// Read a dense counter by key, returning 0 if absent. State-machine-local
470 /// read; callers needing linearizability must issue a read barrier first.
471 pub fn dense_value(&self, key: &str) -> u64 {
472 self.core.lock().dense.get(key).copied().unwrap_or(0)
473 }
474
475 fn snapshot_id_for(last_applied: Option<&LogId>, idx: u64) -> String {
476 let log_index = last_applied.map(|l| l.index).unwrap_or(0);
477 format!("{log_index}-{idx}")
478 }
479
480 /// Durably persist `envelope`, then publish `(meta, data)` as the current
481 /// in-memory snapshot, running `on_adopt` against the core under the same
482 /// publish lock — but only if `meta.last_log_id` does not regress the
483 /// snapshot already published.
484 ///
485 /// The whole sequence is serialized across clones by `self.persist` so a
486 /// slow `build_snapshot` cannot interleave its store write or its publish
487 /// behind a newer `install_snapshot` (or vice-versa) and roll the durable
488 /// and in-memory snapshot back to a stale `last_log_id`. The published
489 /// `last_log_id` is re-read *inside* the lock, immediately before the
490 /// store write, so the decision can't be invalidated by a concurrent
491 /// commit. `self.core` is taken only in brief bursts and never held across
492 /// `store.save`, so `apply` never serializes against snapshot I/O.
493 ///
494 /// Returns `Ok(true)` if the snapshot was adopted, `Ok(false)` if it was
495 /// dropped as stale because a newer snapshot is already durable and
496 /// published.
497 fn commit_snapshot(
498 &self,
499 meta: SnapMeta,
500 data: Vec<u8>,
501 envelope: &[u8],
502 on_adopt: impl FnOnce(&mut Core),
503 ) -> io::Result<bool> {
504 let _persist = self.persist.lock();
505
506 let published = self
507 .core
508 .lock()
509 .current_snapshot
510 .as_ref()
511 .and_then(|s| s.meta.last_log_id);
512 if !supersedes_published(meta.last_log_id, published) {
513 // `debug!`, not `warn!`: this file is on the per-entry apply hot
514 // path (`#[PerformanceCriticalPath]`), where info-or-higher logging
515 // is banned. A discarded stale publish is also a benign, expected
516 // race resolution — the monotone gate doing its job — not an
517 // operational fault, so debug is the right level on the merits too.
518 tracing::debug!(
519 incoming.last_log_id = ?meta.last_log_id,
520 published.last_log_id = ?published,
521 snapshot_id = %meta.snapshot_id,
522 "discarding stale snapshot publish: a newer snapshot is \
523 already durable and published",
524 );
525 return Ok(false);
526 }
527
528 self.store.save(envelope)?;
529
530 let mut core = self.core.lock();
531 on_adopt(&mut core);
532 core.current_snapshot = Some(StoredSnapshot { meta, data });
533 Ok(true)
534 }
535}
536
537impl RaftSnapshotBuilder<TypeConfig> for HighWaterStateMachine {
538 type SnapshotData = Cursor<Vec<u8>>;
539
540 async fn build_snapshot(&mut self) -> Result<SnapOf, io::Error> {
541 // Build the payload + meta under the lock, then release it before
542 // calling the store: `SnapshotStore::save` may do disk I/O (rocksdb
543 // sync write), and holding a `parking_lot::Mutex` across that would
544 // serialize `apply` against snapshot persistence unnecessarily.
545 let (snapshot_payload, meta) = {
546 let mut core = self.core.lock();
547 core.snapshot_idx += 1;
548 let payload = HighWaterStateMachineSnapshot {
549 current_value: core.current_value,
550 last_applied: core.last_applied,
551 last_membership: core.last_membership.clone(),
552 dense: core.dense.clone(),
553 dense_cap: core.dense_cap,
554 };
555 let snapshot_id = Self::snapshot_id_for(core.last_applied.as_ref(), core.snapshot_idx);
556 let meta = SnapMeta {
557 last_log_id: core.last_applied,
558 last_membership: core.last_membership.clone(),
559 snapshot_id,
560 };
561 let bytes = encode_framed(self.active_write_version.get(), &payload)
562 .map_err(|e| codec_io_error("snapshot payload serialize", e))?;
563 (bytes, meta)
564 };
565
566 let persisted = PersistedSnapshot {
567 meta: meta.clone(),
568 data: snapshot_payload.clone(),
569 };
570 let envelope = encode_framed(self.active_write_version.get(), &persisted)
571 .map_err(|e| codec_io_error("snapshot envelope serialize", e))?;
572 // Persist + publish through the monotone, serialized commit path. The
573 // store write happens BEFORE `current_snapshot` is updated, so a
574 // failed write leaves the prior (already-persisted) snapshot in place
575 // and a follower streaming via `get_current_snapshot` never observes a
576 // snapshot we could not durably write. A build whose `last_applied`
577 // was captured before a newer install bumped the published snapshot is
578 // dropped here rather than rolling the durable + in-memory snapshot
579 // back to its stale `last_log_id`. We still return the snapshot we
580 // built — openraft asked for one, and the bytes faithfully capture the
581 // state we observed.
582 self.commit_snapshot(meta.clone(), snapshot_payload.clone(), &envelope, |_| {})?;
583
584 Ok(Snapshot {
585 meta,
586 snapshot: Cursor::new(snapshot_payload),
587 })
588 }
589}
590
591impl RaftStateMachine<TypeConfig> for HighWaterStateMachine {
592 type SnapshotData = Cursor<Vec<u8>>;
593 type SnapshotBuilder = Self;
594
595 async fn applied_state(&mut self) -> Result<(Option<LogId>, StoredMem), io::Error> {
596 let core = self.core.lock();
597 Ok((core.last_applied, core.last_membership.clone()))
598 }
599
600 async fn apply<Strm>(&mut self, mut entries: Strm) -> Result<(), io::Error>
601 where
602 Strm: futures::Stream<Item = Result<EntryResponder<TypeConfig>, io::Error>>
603 + Unpin
604 + openraft::OptionalSend,
605 {
606 while let Some(item) = entries.next().await {
607 let (entry, responder_opt) = item?;
608 let log_id = entry.log_id;
609
610 let applied = match &entry.payload {
611 EntryPayload::Blank => {
612 let mut core = self.core.lock();
613 core.last_applied = Some(log_id);
614 HighWaterApplied {
615 value: core.current_value,
616 outcome: ApplyOutcome::Advanced,
617 }
618 }
619 EntryPayload::Normal(HighWaterCommand::Advance(advance)) => {
620 let mut core = self.core.lock();
621 core.current_value = advance.merge(core.current_value);
622 core.last_applied = Some(log_id);
623 HighWaterApplied {
624 value: core.current_value,
625 outcome: ApplyOutcome::Advanced,
626 }
627 }
628 EntryPayload::Normal(HighWaterCommand::SetFormatVersion(
629 SetFormatVersionPayload {
630 target,
631 gated_members,
632 },
633 )) => {
634 let target = *target;
635 // Defense-in-depth ahead of the subset check: refuse to
636 // flip the cell to a target outside the LOCAL binary's
637 // readable range. The gate at proposal time
638 // (`target_in_local_readable_range` in `capabilities`)
639 // is the primary safety mechanism, but the apply arm
640 // re-checks here so an older (buggy) binary's committed
641 // out-of-range entry cannot poison a fixed binary's
642 // active write version on replay. Reaching this branch
643 // means a protocol invariant was already broken; the
644 // apply pipeline still must not panic (that would tear
645 // down the leader) — emit a NO-OP outcome and a
646 // distinct counter so the operator sees it.
647 if !(MIN_READABLE_VERSION..=MAX_READABLE_VERSION).contains(&target) {
648 // `debug!` not `error!`: hot path, same rule as the
649 // subset arms below. The operator-visible surface
650 // is the `noop_target_out_of_range` counter — a
651 // non-zero value means an older binary committed
652 // an out-of-range bump (or a peer violated
653 // protocol) and the apply arm's defense-in-depth
654 // contained it. The gate's `warn!` at
655 // `run_activation_gate` is the loud, operator-
656 // facing surface for the in-process path; this
657 // branch only fires on replay of an already-
658 // committed entry, where loud logging on every
659 // restart would just spam.
660 tracing::debug!(
661 target = target,
662 min = MIN_READABLE_VERSION,
663 max = MAX_READABLE_VERSION,
664 "SetFormatVersion apply rejected: target outside local readable range"
665 );
666 crate::observability::record_noop_target_out_of_range();
667 let mut core = self.core.lock();
668 core.last_applied = Some(log_id);
669 HighWaterApplied {
670 value: core.current_value,
671 outcome: ApplyOutcome::FormatActivationTargetOutOfRange { target },
672 }
673 } else {
674 // Evaluate the subset against the membership committed as
675 // of this entry's log position. Apply folds the log in
676 // index order, so `core.last_membership` is exactly that
677 // membership (a `Membership` entry at a lower index has
678 // already updated it; a later one has not). Cover BOTH
679 // voters and learners: openraft replicates/snapshots
680 // learners, so an un-gated learner must force a no-op.
681 let mut core = self.core.lock();
682 let committed_members: std::collections::BTreeSet<u64> = core
683 .last_membership
684 .membership()
685 .nodes()
686 .map(|(node_id, _node)| *node_id)
687 .collect();
688 let outcome = if committed_members.is_subset(gated_members) {
689 // Successful (non-no-op) apply: set the
690 // process-shared active-write-version cell. NO meta
691 // write — durability is the raft log (deterministic
692 // replay re-runs this exact check) plus the snapshot
693 // frame byte.
694 self.active_write_version.set(target);
695 // `debug!` not `info!`: this is on the per-entry apply
696 // hot path (this file carries the
697 // `#[PerformanceCriticalPath]` marker), so
698 // info-or-higher logging is banned. The
699 // operator-visible surface is the `applied` counter
700 // and the `active_write_version` gauge.
701 tracing::debug!(
702 target = target,
703 gated_members = ?gated_members,
704 committed_members = ?committed_members,
705 "SetFormatVersion applied: active write version flipped"
706 );
707 crate::observability::record_applied();
708 crate::observability::record_active_write_version(target);
709 ApplyOutcome::FormatActivated { target }
710 } else {
711 // Membership grew an un-gated member between the gate
712 // and this entry's position. Leave the cell
713 // untouched; the operator re-gates and re-issues. A
714 // no-op writes nothing at `target`, so it cannot
715 // resurrect on restart.
716 //
717 // `debug!` not `warn!`: hot path, same rule as the
718 // success arm above. The operator-visible surface is
719 // the `noop_membership_subset` counter — a non-zero
720 // value during an activation means a membership
721 // change raced the bump.
722 tracing::debug!(
723 target = target,
724 gated_members = ?gated_members,
725 committed_members = ?committed_members,
726 "SetFormatVersion no-op: committed membership is not a subset of the gated set"
727 );
728 crate::observability::record_noop_membership_subset();
729 ApplyOutcome::FormatActivationNoop { target }
730 };
731 core.last_applied = Some(log_id);
732 HighWaterApplied {
733 value: core.current_value,
734 outcome,
735 }
736 }
737 }
738 EntryPayload::Normal(HighWaterCommand::AdvanceDense { key, count }) => {
739 let mut core = self.core.lock();
740 let key_str = key.as_str();
741 let present = core.dense.contains_key(key_str);
742 let outcome = if !present && core.dense.len() as u64 >= core.dense_cap {
743 ApplyOutcome::DenseCardinalityExceeded {
744 cap: core.dense_cap,
745 }
746 } else {
747 let start = core.dense.get(key_str).copied().unwrap_or(0);
748 match start.checked_add(u64::from(*count)) {
749 Some(next) => {
750 core.dense.insert(key_str.to_string(), next);
751 ApplyOutcome::DenseAdvanced { start }
752 }
753 None => ApplyOutcome::DenseOverflow,
754 }
755 };
756 core.last_applied = Some(log_id);
757 HighWaterApplied {
758 value: core.current_value,
759 outcome,
760 }
761 }
762 EntryPayload::Normal(HighWaterCommand::AdvanceDenseBatch { entries }) => {
763 let mut core = self.core.lock();
764 // Phase 1: cardinality. Count distinct keys not already
765 // present; a duplicate entry on the same key counts once
766 // (set semantics), so the new-key count uses a BTreeSet.
767 let new_keys: std::collections::BTreeSet<&str> = entries
768 .iter()
769 .map(|e| e.key.as_str())
770 .filter(|k| !core.dense.contains_key(*k))
771 .collect();
772 let outcome =
773 if core.dense.len() as u64 + new_keys.len() as u64 > core.dense_cap {
774 ApplyOutcome::DenseBatchCardinalityExceeded {
775 cap: core.dense_cap,
776 }
777 } else {
778 // Phase 2: sequential fold into a scratch map seeded
779 // lazily from current counters. Each entry's start is the
780 // running value in the scratch map (or the stored value if
781 // not yet in the scratch map); the running value
782 // accumulates across repeats, so a duplicate key yields
783 // adjacent non-overlapping blocks and the overflow check
784 // is against the ACCUMULATED value — not the shared
785 // pre-batch counter. Any overflow rejects the whole batch.
786 let mut scratch: std::collections::BTreeMap<&str, u64> =
787 std::collections::BTreeMap::new();
788 let mut starts: Vec<u64> = Vec::with_capacity(entries.len());
789 let mut overflow = false;
790 for entry in entries {
791 let key_str = entry.key.as_str();
792 let running = scratch
793 .get(key_str)
794 .copied()
795 .or_else(|| core.dense.get(key_str).copied())
796 .unwrap_or(0);
797 starts.push(running);
798 match running.checked_add(u64::from(entry.count)) {
799 Some(next) => {
800 scratch.insert(key_str, next);
801 }
802 None => {
803 overflow = true;
804 break;
805 }
806 }
807 }
808 if overflow {
809 ApplyOutcome::DenseBatchOverflow
810 } else {
811 // Phase 3: commit. Only now mutate the durable map.
812 for (key_str, next) in scratch {
813 core.dense.insert(key_str.to_string(), next);
814 }
815 ApplyOutcome::DenseBatchAdvanced { starts }
816 }
817 };
818 core.last_applied = Some(log_id);
819 HighWaterApplied {
820 value: core.current_value,
821 outcome,
822 }
823 }
824 EntryPayload::Membership(membership) => {
825 let mut core = self.core.lock();
826 core.last_membership = StoredMembership::new(Some(log_id), membership.clone());
827 core.last_applied = Some(log_id);
828 HighWaterApplied {
829 value: core.current_value,
830 outcome: ApplyOutcome::Advanced,
831 }
832 }
833 };
834
835 if let Some(responder) = responder_opt {
836 responder.send(applied);
837 }
838 }
839 Ok(())
840 }
841
842 async fn get_snapshot_builder(&mut self) -> Self::SnapshotBuilder {
843 self.clone()
844 }
845
846 async fn begin_receiving_snapshot(&mut self) -> Result<SnapData, io::Error> {
847 Ok(Cursor::new(Vec::new()))
848 }
849
850 async fn install_snapshot(
851 &mut self,
852 meta: &SnapMeta,
853 snapshot: SnapData,
854 ) -> Result<(), io::Error> {
855 let bytes = snapshot.into_inner();
856 let payload: HighWaterStateMachineSnapshot =
857 decode_framed(MIN_READABLE_VERSION, MAX_READABLE_VERSION, &bytes)
858 .map_err(|e| codec_io_error("snapshot payload decode", e))?;
859
860 // `meta.last_log_id` is read back by `get_current_snapshot` while
861 // `payload.last_applied` is read back by `applied_state`. `build_snapshot`
862 // derives both from the same `last_applied`, so they always agree for an
863 // honest peer. Reject any disagreement before persisting or mutating: a
864 // mismatch would otherwise desync those two reader methods permanently
865 // and silently. Checked before `store.save` so a rejected install leaves
866 // both the store and in-memory state untouched for openraft to retry.
867 if meta.last_log_id != payload.last_applied {
868 return Err(io::Error::new(
869 io::ErrorKind::InvalidData,
870 format!(
871 "snapshot meta/payload disagree on last_log_id: \
872 meta.last_log_id={:?}, payload.last_applied={:?}",
873 meta.last_log_id, payload.last_applied
874 ),
875 ));
876 }
877
878 // Re-validate every dense key before adopting the snapshot. The dense
879 // map is persisted as raw `String` keys, so — unlike the replicated
880 // `AdvanceDense` log command, whose `SeqKey` re-validates on decode — a
881 // crafted or corrupted snapshot can carry an empty or oversized key that
882 // `decode_framed` above accepts. Enforce `SeqKey`'s invariant here, fail
883 // loud, and do it before the `active_write_version` bump and the persist
884 // below so a rejected install leaves the store and in-memory state
885 // untouched for openraft to retry — exactly like the meta/payload check
886 // above.
887 for key in payload.dense.keys() {
888 if let Err(e) = tsoracle_core::SeqKey::try_new(key.as_str()) {
889 return Err(io::Error::new(
890 io::ErrorKind::InvalidData,
891 format!("snapshot dense key fails the SeqKey invariant: {e}"),
892 ));
893 }
894 }
895
896 // The snapshot's leading version byte is evidence that the sending
897 // leader's active write version was at least this high — which means a
898 // `SetFormatVersion` activation to that version already committed
899 // cluster-wide (the all-members gate required every member, including
900 // this one, to be able to read it). Advance our cell to match, mirroring
901 // `recover_active_write_version`'s max-of-evidence rule. Without this, a
902 // follower that receives a v5 snapshot whose `SetFormatVersion` entry was
903 // purged into the snapshot would keep `active_write_version` at the
904 // baseline while holding a populated dense map, then trip the fail-loud
905 // v4-encode `NotRepresentable` guard on its next `build_snapshot`. The
906 // byte is already within `[MIN, MAX_READABLE_VERSION]` (decode_framed
907 // accepted it above), so this never exceeds what this binary can write.
908 if let Some(&snapshot_version) = bytes.first() {
909 if snapshot_version > self.active_write_version.get() {
910 self.active_write_version.set(snapshot_version);
911 crate::observability::record_active_write_version(snapshot_version);
912 }
913 }
914
915 let persisted = PersistedSnapshot {
916 meta: meta.clone(),
917 data: bytes.clone(),
918 };
919 let envelope = encode_framed(self.active_write_version.get(), &persisted)
920 .map_err(|e| codec_io_error("snapshot envelope serialize", e))?;
921 // Persist, apply the snapshot's state, and publish atomically through
922 // the monotone, serialized commit path — the same one `build_snapshot`
923 // uses, so neither can clobber the other. The store write happens
924 // before the in-memory state is advanced, so a failed write leaves the
925 // SM at its prior state for openraft to retry. A stale install (a
926 // lower `last_log_id` than the already-published snapshot) is dropped
927 // as an accepted no-op: we already cover at least this snapshot, so
928 // adopting it would regress both the durable snapshot and the applied
929 // state.
930 self.commit_snapshot(meta.clone(), bytes, &envelope, |core| {
931 core.current_value = payload.current_value;
932 core.last_applied = payload.last_applied;
933 core.last_membership = payload.last_membership;
934 core.dense = payload.dense;
935 // A v4 snapshot carries dense_cap == 0 (absent sentinel); keep the
936 // constructor-seeded genesis cap. A v5 snapshot carries the real cap.
937 if payload.dense_cap != 0 {
938 core.dense_cap = payload.dense_cap;
939 }
940 })?;
941 Ok(())
942 }
943
944 async fn get_current_snapshot(&mut self) -> Result<Option<SnapOf>, io::Error> {
945 let core = self.core.lock();
946 Ok(core.current_snapshot.as_ref().map(|s| Snapshot {
947 meta: s.meta.clone(),
948 snapshot: Cursor::new(s.data.clone()),
949 }))
950 }
951}
952
953#[cfg(test)]
954mod tests {
955 use super::*;
956
957 use std::collections::BTreeSet;
958
959 use futures::stream;
960 use openraft::EntryPayload;
961 use openraft::entry::RaftEntry;
962 use openraft::storage::EntryResponder;
963 use openraft::type_config::alias::EntryOf;
964
965 use crate::log_entry::HighWaterCommand;
966 use crate::type_config::TypeConfig;
967 use tsoracle_consensus::AdvancePayload;
968
969 // --- Test helpers ---
970
971 /// Install a process-wide `DEBUG`-level subscriber so the `tracing::debug!`
972 /// on the stale-publish reject path actually evaluates its fields (without
973 /// an interested subscriber the macro short-circuits and the argument
974 /// expressions are never executed). Idempotent across tests via `try_init`.
975 fn enable_tracing() {
976 let _ = tracing_subscriber::fmt()
977 .with_max_level(tracing::Level::DEBUG)
978 .with_test_writer()
979 .try_init();
980 }
981
982 /// Build a `LogId` for the toolkit's default leader-id layout
983 /// (`LeaderId<u64, u64>`) with term=1, node_id=1, and the given index.
984 fn log_id(index: u64) -> LogIdOf<TypeConfig> {
985 openraft::testing::log_id::<TypeConfig>(1, 1, index)
986 }
987
988 fn entry(
989 index: u64,
990 payload: EntryPayload<HighWaterCommand, u64, crate::type_config::OpenraftPeer>,
991 ) -> EntryResponder<TypeConfig> {
992 let e: EntryOf<TypeConfig> = match payload {
993 EntryPayload::Blank => EntryOf::<TypeConfig>::new_blank(log_id(index)),
994 EntryPayload::Normal(d) => EntryOf::<TypeConfig>::new_normal(log_id(index), d),
995 EntryPayload::Membership(m) => EntryOf::<TypeConfig>::new_membership(log_id(index), m),
996 };
997 (e, None)
998 }
999
1000 async fn apply_one(
1001 sm: &mut HighWaterStateMachine,
1002 index: u64,
1003 payload: EntryPayload<HighWaterCommand, u64, crate::type_config::OpenraftPeer>,
1004 ) {
1005 sm.apply(stream::iter([Ok(entry(index, payload))]))
1006 .await
1007 .expect("apply");
1008 }
1009
1010 // --- Tests ---
1011
1012 #[tokio::test]
1013 async fn apply_blank_updates_only_log_id() {
1014 let mut sm = HighWaterStateMachine::new();
1015 apply_one(&mut sm, 1, EntryPayload::Blank).await;
1016 assert_eq!(sm.current_value().await, 0);
1017 let (last, _) = sm.applied_state().await.unwrap();
1018 assert_eq!(last.map(|l| l.index), Some(1));
1019 }
1020
1021 #[tokio::test]
1022 async fn apply_normal_advances_value() {
1023 let mut sm = HighWaterStateMachine::new();
1024 apply_one(
1025 &mut sm,
1026 1,
1027 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 100 })),
1028 )
1029 .await;
1030 assert_eq!(sm.current_value().await, 100);
1031 }
1032
1033 #[tokio::test]
1034 async fn apply_normal_holds_monotonic_under_stale_target() {
1035 let mut sm = HighWaterStateMachine::new();
1036 apply_one(
1037 &mut sm,
1038 1,
1039 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 100 })),
1040 )
1041 .await;
1042 apply_one(
1043 &mut sm,
1044 2,
1045 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 50 })),
1046 )
1047 .await;
1048 assert_eq!(sm.current_value().await, 100);
1049 }
1050
1051 #[tokio::test]
1052 async fn apply_normal_equal_target_holds_value() {
1053 let mut sm = HighWaterStateMachine::new();
1054 apply_one(
1055 &mut sm,
1056 1,
1057 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 100 })),
1058 )
1059 .await;
1060 apply_one(
1061 &mut sm,
1062 2,
1063 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 100 })),
1064 )
1065 .await;
1066 assert_eq!(sm.current_value().await, 100);
1067 }
1068
1069 #[tokio::test]
1070 async fn apply_membership_updates_membership_only() {
1071 let mut sm = HighWaterStateMachine::new();
1072 apply_one(
1073 &mut sm,
1074 1,
1075 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 42 })),
1076 )
1077 .await;
1078 let mem = openraft::Membership::new_with_defaults(vec![BTreeSet::from([1u64])], [1u64]);
1079 apply_one(&mut sm, 2, EntryPayload::Membership(mem)).await;
1080 assert_eq!(sm.current_value().await, 42);
1081 let (last, _) = sm.applied_state().await.unwrap();
1082 assert_eq!(last.map(|l| l.index), Some(2));
1083 }
1084
1085 // ---- SetFormatVersion apply tests ----
1086 //
1087 // The apply path keys the flip off a successful (non-no-op) apply: the
1088 // subset check `committed_members ⊆ gated_members` runs against
1089 // `core.last_membership` evaluated at the bump's own log position (apply
1090 // folds the log in index order). On a hit the shared cell is set; on a
1091 // miss the cell is untouched and a no-op outcome is returned. Direct
1092 // cell observation is sufficient since `SetFormatVersionPayload`'s effect
1093 // is exactly "did the cell move?"; we don't need to capture the
1094 // ApplyOutcome through the responder channel for these checks.
1095
1096 /// A membership whose voter config is `voters`, as a `Membership` payload.
1097 /// The concrete `Membership<NID, N>` type is inferred from the call site
1098 /// (e.g. `EntryPayload::Membership(...)` infers it for `TypeConfig`).
1099 fn voters_membership(
1100 voters: &[u64],
1101 ) -> openraft::Membership<u64, crate::type_config::OpenraftPeer> {
1102 openraft::Membership::new_with_defaults(
1103 vec![voters.iter().copied().collect::<BTreeSet<u64>>()],
1104 voters.to_vec(),
1105 )
1106 }
1107
1108 #[tokio::test]
1109 async fn set_format_version_subset_match_sets_cell() {
1110 // The flip is observable because DENSE_WRITE_VERSION (5) is a real
1111 // target distinct from BASELINE (4) and within the readable range
1112 // [MIN, MAX] = [4, 5]. The behavior pinned by this test — "subset OK +
1113 // in-range target ⇒ cell flips" — uses the real dense activation target.
1114 let mut sm = HighWaterStateMachine::new();
1115 let target = tsoracle_openraft_toolkit::DENSE_WRITE_VERSION;
1116 // Establish membership {1, 2} at index 1.
1117 apply_one(
1118 &mut sm,
1119 1,
1120 EntryPayload::Membership(voters_membership(&[1, 2])),
1121 )
1122 .await;
1123 // Bump gated on a superset of the committed membership → subset holds.
1124 apply_one(
1125 &mut sm,
1126 2,
1127 EntryPayload::Normal(HighWaterCommand::SetFormatVersion(
1128 SetFormatVersionPayload {
1129 target,
1130 gated_members: BTreeSet::from([1u64, 2u64, 3u64]),
1131 },
1132 )),
1133 )
1134 .await;
1135 assert_eq!(
1136 sm.active_write_version(),
1137 target,
1138 "cell set on successful subset apply"
1139 );
1140 }
1141
1142 #[tokio::test]
1143 async fn set_format_version_subset_match_out_of_range_target_is_apply_noop() {
1144 // Mirror of the (feature-gated) "subset match" test but for the
1145 // default build: a subset-OK bump whose target is OUT of the
1146 // local readable range must be no-op'd by the apply-arm
1147 // defense-in-depth, regardless of subset success. Pins the
1148 // ordering invariant — the range check runs BEFORE the subset
1149 // check, so an OOR target never sneaks through on a gate
1150 // committed by an older buggy binary.
1151 let mut sm = HighWaterStateMachine::new();
1152 let before = sm.active_write_version();
1153 apply_one(
1154 &mut sm,
1155 1,
1156 EntryPayload::Membership(voters_membership(&[1, 2])),
1157 )
1158 .await;
1159 apply_one(
1160 &mut sm,
1161 2,
1162 EntryPayload::Normal(HighWaterCommand::SetFormatVersion(
1163 SetFormatVersionPayload {
1164 target: 7, // out of range under default features.
1165 gated_members: BTreeSet::from([1u64, 2u64, 3u64]),
1166 },
1167 )),
1168 )
1169 .await;
1170 assert_eq!(
1171 sm.active_write_version(),
1172 before,
1173 "subset-OK + OOR target ⇒ apply-arm defense-in-depth no-op"
1174 );
1175 }
1176
1177 // ---- Apply-arm defense-in-depth: out-of-range target is a no-op ----
1178 //
1179 // The all-members gate at proposal time is the primary safety
1180 // mechanism, but the apply arm independently refuses to flip the
1181 // shared cell to a target outside the LOCAL binary's
1182 // [MIN_READABLE_VERSION, MAX_READABLE_VERSION]. This guards against
1183 // two paths the gate cannot:
1184 //
1185 // 1. Replay of a log committed by an older (buggy) binary that
1186 // let an out-of-range target through. A fixed binary replaying
1187 // that log must NOT propagate the corruption into its in-memory
1188 // cell, since the codec arms reject encode/decode at that
1189 // version and the cluster wedges.
1190 // 2. A protocol-violating peer somehow committing such an entry
1191 // (e.g. a forged proposal). The fail-safe at apply contains
1192 // blast radius to "the flip didn't take effect" instead of
1193 // "all subsequent log appends fail".
1194 //
1195 // The behavior is a NO-OP (cell untouched + `FormatActivationTargetOutOfRange`
1196 // outcome), never a panic. The apply pipeline is on the critical
1197 // path; refusing to apply would tear the leader down rather than
1198 // contain the bad entry.
1199
1200 #[tokio::test]
1201 async fn set_format_version_target_above_max_is_apply_noop() {
1202 let mut sm = HighWaterStateMachine::new();
1203 let before = sm.active_write_version();
1204 // Membership matches the gated set (subset OK) so the only thing
1205 // that can hold the flip back is the range check.
1206 apply_one(
1207 &mut sm,
1208 1,
1209 EntryPayload::Membership(voters_membership(&[1, 2])),
1210 )
1211 .await;
1212 let above_max = tsoracle_openraft_toolkit::MAX_READABLE_VERSION.saturating_add(1);
1213 if above_max == tsoracle_openraft_toolkit::MAX_READABLE_VERSION {
1214 return; // MAX == u8::MAX (impossible in practice).
1215 }
1216 apply_one(
1217 &mut sm,
1218 2,
1219 EntryPayload::Normal(HighWaterCommand::SetFormatVersion(
1220 SetFormatVersionPayload {
1221 target: above_max,
1222 gated_members: BTreeSet::from([1u64, 2u64]),
1223 },
1224 )),
1225 )
1226 .await;
1227 assert_eq!(
1228 sm.active_write_version(),
1229 before,
1230 "apply must refuse to flip to a target above MAX_READABLE_VERSION"
1231 );
1232 }
1233
1234 #[tokio::test]
1235 async fn set_format_version_target_below_min_is_apply_noop() {
1236 let mut sm = HighWaterStateMachine::new();
1237 let before = sm.active_write_version();
1238 apply_one(
1239 &mut sm,
1240 1,
1241 EntryPayload::Membership(voters_membership(&[1, 2])),
1242 )
1243 .await;
1244 // MIN_READABLE_VERSION is at least 1 in any sensible build (the
1245 // codec's compile-time `MIN <= MAX` plus the framing layer's
1246 // version byte both effectively forbid MIN == 0), so MIN - 1 is
1247 // always representable and out of range. This is the bug class
1248 // the finding hit: a target below MIN that the (pre-fix) gate
1249 // accepted because it only checked the upper bound.
1250 let below_min = tsoracle_openraft_toolkit::MIN_READABLE_VERSION - 1;
1251 apply_one(
1252 &mut sm,
1253 2,
1254 EntryPayload::Normal(HighWaterCommand::SetFormatVersion(
1255 SetFormatVersionPayload {
1256 target: below_min,
1257 gated_members: BTreeSet::from([1u64, 2u64]),
1258 },
1259 )),
1260 )
1261 .await;
1262 assert_eq!(
1263 sm.active_write_version(),
1264 before,
1265 "apply must refuse to flip to a target below MIN_READABLE_VERSION"
1266 );
1267 }
1268
1269 #[tokio::test]
1270 async fn set_format_version_target_zero_is_apply_noop() {
1271 // Target 0 is the worst case: 0 is also the legacy-unframed wire
1272 // sentinel, so a 0-stamped record collides with that
1273 // interpretation. The apply arm must refuse the flip rather than
1274 // let it through (the codec would later reject every encode).
1275 let mut sm = HighWaterStateMachine::new();
1276 let before = sm.active_write_version();
1277 apply_one(
1278 &mut sm,
1279 1,
1280 EntryPayload::Membership(voters_membership(&[1, 2])),
1281 )
1282 .await;
1283 apply_one(
1284 &mut sm,
1285 2,
1286 EntryPayload::Normal(HighWaterCommand::SetFormatVersion(
1287 SetFormatVersionPayload {
1288 target: 0,
1289 gated_members: BTreeSet::from([1u64, 2u64]),
1290 },
1291 )),
1292 )
1293 .await;
1294 assert_eq!(
1295 sm.active_write_version(),
1296 before,
1297 "apply must refuse to flip to a 0 target"
1298 );
1299 }
1300
1301 #[tokio::test]
1302 async fn set_format_version_membership_not_subset_is_noop() {
1303 let mut sm = HighWaterStateMachine::new();
1304 let before = sm.active_write_version();
1305 // Membership {1, 2, 9} at index 1; the gate only covered {1, 2}.
1306 apply_one(
1307 &mut sm,
1308 1,
1309 EntryPayload::Membership(voters_membership(&[1, 2, 9])),
1310 )
1311 .await;
1312 apply_one(
1313 &mut sm,
1314 2,
1315 EntryPayload::Normal(HighWaterCommand::SetFormatVersion(
1316 SetFormatVersionPayload {
1317 target: 7,
1318 gated_members: BTreeSet::from([1u64, 2u64]),
1319 },
1320 )),
1321 )
1322 .await;
1323 assert_eq!(
1324 sm.active_write_version(),
1325 before,
1326 "no-op must leave the shared cell untouched"
1327 );
1328 }
1329
1330 #[tokio::test]
1331 async fn set_format_version_covers_learners_not_only_voters() {
1332 let mut sm = HighWaterStateMachine::new();
1333 let before = sm.active_write_version();
1334 // Voter {1}, learner {2}. `new_with_defaults(voter_groups, all_ids)`
1335 // treats node ids in `all_ids` but absent from `voter_groups` as
1336 // learners. A voter-only gate `{1}` must NOT satisfy the subset —
1337 // the learner is a current member and openraft replicates to it,
1338 // so an un-gated learner forces a no-op.
1339 let membership =
1340 openraft::Membership::new_with_defaults(vec![BTreeSet::from([1u64])], [1u64, 2u64]);
1341 apply_one(&mut sm, 1, EntryPayload::Membership(membership)).await;
1342 apply_one(
1343 &mut sm,
1344 2,
1345 EntryPayload::Normal(HighWaterCommand::SetFormatVersion(
1346 SetFormatVersionPayload {
1347 target: 7,
1348 gated_members: BTreeSet::from([1u64]),
1349 },
1350 )),
1351 )
1352 .await;
1353 assert_eq!(
1354 sm.active_write_version(),
1355 before,
1356 "un-gated learner must force a no-op"
1357 );
1358 }
1359
1360 // ---- Replay / recovery confirmation tests ----
1361 //
1362 // These model openraft's deterministic recovery: a fresh state machine
1363 // (fresh BASELINE-seeded cell) re-applies the committed log in index
1364 // order, so the SetFormatVersion subset check re-runs per entry. A
1365 // no-op replays as a no-op (no record was ever written at the higher
1366 // version); a success re-establishes the cell.
1367
1368 #[derive(Clone)]
1369 enum ReplayEntry {
1370 Membership(Vec<u64>),
1371 Bump { target: u8, gated: Vec<u64> },
1372 }
1373
1374 async fn replay(sm: &mut HighWaterStateMachine, log: &[(u64, ReplayEntry)]) {
1375 for (index, kind) in log {
1376 let payload = match kind {
1377 ReplayEntry::Membership(voters) => {
1378 EntryPayload::Membership(voters_membership(voters))
1379 }
1380 ReplayEntry::Bump { target, gated } => EntryPayload::Normal(
1381 HighWaterCommand::SetFormatVersion(SetFormatVersionPayload {
1382 target: *target,
1383 gated_members: gated.iter().copied().collect(),
1384 }),
1385 ),
1386 };
1387 apply_one(sm, *index, payload).await;
1388 }
1389 }
1390
1391 #[tokio::test]
1392 async fn successful_activation_survives_replay() {
1393 let target = tsoracle_openraft_toolkit::DENSE_WRITE_VERSION;
1394 let log = vec![
1395 (1u64, ReplayEntry::Membership(vec![1, 2])),
1396 (
1397 2u64,
1398 ReplayEntry::Bump {
1399 target,
1400 gated: vec![1, 2, 3],
1401 },
1402 ),
1403 ];
1404 // Live apply sets the cell.
1405 let mut live = HighWaterStateMachine::new();
1406 replay(&mut live, &log).await;
1407 assert_eq!(live.active_write_version(), target);
1408
1409 // "Restart": a fresh SM with a fresh BASELINE-seeded cell, replaying
1410 // the same committed log, re-establishes target via the re-run
1411 // subset check. NO meta key consulted.
1412 let mut recovered = HighWaterStateMachine::new();
1413 assert_ne!(
1414 recovered.active_write_version(),
1415 target,
1416 "fresh cell starts at baseline"
1417 );
1418 replay(&mut recovered, &log).await;
1419 assert_eq!(
1420 recovered.active_write_version(),
1421 target,
1422 "replay re-applies the flip"
1423 );
1424 }
1425
1426 #[tokio::test]
1427 async fn out_of_range_bump_replays_as_noop_across_restart() {
1428 // Replay-determinism contract for an out-of-range bump committed
1429 // by an older (buggy) binary: every replay reaches the same
1430 // state (cell stays at baseline) because the apply arm's range
1431 // check uses ONLY the local binary's constants + the entry
1432 // data, never timing or per-replay side state. The fix CANNOT
1433 // resurrect the bad flip on restart.
1434 let log = vec![
1435 (1u64, ReplayEntry::Membership(vec![1, 2])),
1436 (
1437 2u64,
1438 ReplayEntry::Bump {
1439 target: 7, // out of range under default features.
1440 gated: vec![1, 2, 3],
1441 },
1442 ),
1443 ];
1444 let mut live = HighWaterStateMachine::new();
1445 let baseline = live.active_write_version();
1446 replay(&mut live, &log).await;
1447 assert_eq!(
1448 live.active_write_version(),
1449 baseline,
1450 "live apply no-ops on OOR target"
1451 );
1452
1453 let mut recovered = HighWaterStateMachine::new();
1454 replay(&mut recovered, &log).await;
1455 assert_eq!(
1456 recovered.active_write_version(),
1457 baseline,
1458 "replay is deterministic: OOR bump stays a no-op across restart"
1459 );
1460 }
1461
1462 #[tokio::test]
1463 async fn noop_bump_never_advances_across_restart() {
1464 let log = vec![
1465 // Membership has an un-gated member 9; the gate only covered {1, 2}.
1466 (1u64, ReplayEntry::Membership(vec![1, 2, 9])),
1467 (
1468 2u64,
1469 ReplayEntry::Bump {
1470 target: 7,
1471 gated: vec![1, 2],
1472 },
1473 ),
1474 ];
1475 let mut live = HighWaterStateMachine::new();
1476 let baseline = live.active_write_version();
1477 replay(&mut live, &log).await;
1478 assert_eq!(
1479 live.active_write_version(),
1480 baseline,
1481 "no-op leaves the cell"
1482 );
1483
1484 // Restart: replay the same log; the no-op replays as a no-op (the
1485 // subset check fails identically), so the cell stays at baseline.
1486 // The committed entry's mere presence never advances the version.
1487 let mut recovered = HighWaterStateMachine::new();
1488 replay(&mut recovered, &log).await;
1489 assert_eq!(
1490 recovered.active_write_version(),
1491 baseline,
1492 "a no-op'd bump cannot resurrect on restart"
1493 );
1494 }
1495
1496 #[tokio::test]
1497 async fn replay_is_deterministic() {
1498 let in_range = tsoracle_openraft_toolkit::DENSE_WRITE_VERSION;
1499 let log = vec![
1500 (1u64, ReplayEntry::Membership(vec![1, 2])),
1501 (
1502 2u64,
1503 ReplayEntry::Bump {
1504 target: in_range,
1505 gated: vec![1, 2, 3],
1506 },
1507 ),
1508 // A later bump gated on a set that excludes a now-added member.
1509 (3u64, ReplayEntry::Membership(vec![1, 2, 5])),
1510 (
1511 4u64,
1512 ReplayEntry::Bump {
1513 target: in_range,
1514 gated: vec![1, 2],
1515 },
1516 ),
1517 ];
1518 let mut first = HighWaterStateMachine::new();
1519 replay(&mut first, &log).await;
1520 let mut second = HighWaterStateMachine::new();
1521 replay(&mut second, &log).await;
1522 // Index-2 bump succeeded (membership {1,2} ⊆ {1,2,3}); index-4
1523 // bump no-ops (membership {1,2,5} ⊄ {1,2}), so the cell stays at
1524 // `in_range`.
1525 assert_eq!(first.active_write_version(), in_range);
1526 assert_eq!(
1527 first.active_write_version(),
1528 second.active_write_version(),
1529 "two replays of the same committed log yield the same cell"
1530 );
1531 }
1532
1533 #[tokio::test]
1534 async fn replay_is_deterministic_default_features() {
1535 // Default-features replay-determinism check: any mix of
1536 // out-of-range bumps and membership entries produces the same
1537 // cell state across two independent replays, because the apply
1538 // arm folds the log deterministically (range check + subset
1539 // check are both pure functions of entry data + local binary
1540 // constants).
1541 let log = vec![
1542 (1u64, ReplayEntry::Membership(vec![1, 2])),
1543 (
1544 2u64,
1545 ReplayEntry::Bump {
1546 target: 7,
1547 gated: vec![1, 2, 3],
1548 },
1549 ),
1550 (3u64, ReplayEntry::Membership(vec![1, 2, 5])),
1551 (
1552 4u64,
1553 ReplayEntry::Bump {
1554 target: 8,
1555 gated: vec![1, 2],
1556 },
1557 ),
1558 ];
1559 let mut first = HighWaterStateMachine::new();
1560 let baseline = first.active_write_version();
1561 replay(&mut first, &log).await;
1562 let mut second = HighWaterStateMachine::new();
1563 replay(&mut second, &log).await;
1564 assert_eq!(first.active_write_version(), baseline);
1565 assert_eq!(
1566 first.active_write_version(),
1567 second.active_write_version(),
1568 "two replays of the same committed log yield the same cell"
1569 );
1570 }
1571
1572 // ---- Snapshot active-write-version stamping (format migration) ----
1573 //
1574 // These tests pin the contract that the snapshot builder reads the
1575 // node's CURRENT active write version (the shared cell) at build time,
1576 // rather than a hard-coded constant. Today that is BASELINE_WRITE_VERSION
1577 // since no activation has flipped the cell; the assertions are written
1578 // against the accessor (`sm.active_write_version()`) so they remain
1579 // correct after a real activation moves the cell forward and the next
1580 // build emits at the new version.
1581
1582 #[tokio::test]
1583 async fn build_snapshot_stamps_active_write_version() {
1584 // The persisted envelope's leading version byte must equal what the
1585 // accessor reports — not a literal. This proves the plumbing is
1586 // active-version-driven and would catch a regression to a constant.
1587 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
1588 let mut sm = HighWaterStateMachine::with_store(store.clone()).expect("with_store");
1589 apply_one(
1590 &mut sm,
1591 1,
1592 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 321 })),
1593 )
1594 .await;
1595
1596 let active = sm.active_write_version();
1597 sm.build_snapshot().await.expect("build_snapshot");
1598
1599 let envelope_bytes = store.load().expect("load").expect("snapshot present");
1600 assert_eq!(
1601 envelope_bytes[0], active,
1602 "envelope leading version byte must equal the active write version"
1603 );
1604
1605 // The inner payload blob is the on-disk record openraft replays from
1606 // on snapshot install; it must be stamped at the active version too
1607 // so a follower receiving it can decode against its own version
1608 // window.
1609 let persisted: PersistedSnapshot = tsoracle_codec::decode_framed(
1610 tsoracle_openraft_toolkit::MIN_READABLE_VERSION,
1611 tsoracle_openraft_toolkit::MAX_READABLE_VERSION,
1612 &envelope_bytes,
1613 )
1614 .expect("decode envelope");
1615 assert_eq!(
1616 persisted.data[0], active,
1617 "inner snapshot payload leading version byte must equal the active write version"
1618 );
1619 }
1620
1621 #[tokio::test]
1622 async fn old_version_snapshot_is_read_and_reemitted_at_active_version() {
1623 // Migration-on-next-write: a snapshot persisted at version V is
1624 // readable across the multi-version codec and the next
1625 // build_snapshot re-emits at the cluster's ACTIVE write version.
1626 //
1627 // SCOPE: this drives the same-version axis — persist at the active
1628 // write version (BASELINE = 4, since no activation flip is injected
1629 // here) and assert the reopen reads it back and re-emits at the active
1630 // version. The genuine cross-version rewrite (install v4, activate v5,
1631 // assert the rebuild is v5) is the orthogonal axis; the readable range
1632 // is now [4, 5], so the v4 lower version genuinely exists on disk
1633 // pre-activation.
1634 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
1635
1636 // Persist a snapshot via a first SM instance, then drop it.
1637 {
1638 let mut writer = HighWaterStateMachine::with_store(store.clone()).expect("writer SM");
1639 apply_one(
1640 &mut writer,
1641 3,
1642 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 777 })),
1643 )
1644 .await;
1645 writer
1646 .build_snapshot()
1647 .await
1648 .expect("build initial snapshot");
1649 }
1650
1651 // Reopen: the recovery path decodes across the readable range and
1652 // restores state.
1653 let mut sm = HighWaterStateMachine::with_store(store.clone()).expect("reopened SM");
1654 assert_eq!(
1655 sm.current_value().await,
1656 777,
1657 "recovered value across reopen"
1658 );
1659
1660 // The next build re-emits at the active write version.
1661 let active = sm.active_write_version();
1662 apply_one(
1663 &mut sm,
1664 4,
1665 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 888 })),
1666 )
1667 .await;
1668 sm.build_snapshot().await.expect("rebuild snapshot");
1669
1670 let rebuilt = store.load().expect("load").expect("snapshot present");
1671 assert_eq!(
1672 rebuilt[0], active,
1673 "rebuilt snapshot must be emitted at the active write version"
1674 );
1675 let payload: PersistedSnapshot = tsoracle_codec::decode_framed(
1676 tsoracle_openraft_toolkit::MIN_READABLE_VERSION,
1677 tsoracle_openraft_toolkit::MAX_READABLE_VERSION,
1678 &rebuilt,
1679 )
1680 .expect("decode rebuilt envelope");
1681 assert_eq!(payload.meta.last_log_id.map(|l| l.index), Some(4));
1682 }
1683
1684 #[test]
1685 fn snapshot_codec_accepts_full_readable_range() {
1686 // The migration-seam invariant in isolation: the snapshot decoder
1687 // must accept any version in [MIN_READABLE_VERSION,
1688 // MAX_READABLE_VERSION] and reject anything outside it. This is
1689 // what makes an OLD-version on-disk snapshot readable after the
1690 // active version moves forward. Asserted against the codec range
1691 // directly so it covers the full readable range [MIN = 4, MAX = 6] —
1692 // iterating v4 (decode-and-lift), v5 (direct decode), and v6
1693 // (byte-identical to v5).
1694 // Empty dense map + cap 0 so the v4 encode arm (which projects to the
1695 // frozen V4 struct) is lossless; assert_eq!(decoded, payload) holds
1696 // for every version in the readable range because the decode-and-lift
1697 // path restores dense=empty, dense_cap=0 from a v4 blob.
1698 let payload = HighWaterStateMachineSnapshot {
1699 current_value: 5,
1700 last_applied: Some(log_id(2)),
1701 last_membership: StoredMem::default(),
1702 dense: std::collections::BTreeMap::new(),
1703 dense_cap: 0,
1704 };
1705 for version in tsoracle_openraft_toolkit::MIN_READABLE_VERSION
1706 ..=tsoracle_openraft_toolkit::MAX_READABLE_VERSION
1707 {
1708 let framed = tsoracle_codec::encode_framed(version, &payload).expect("encode in range");
1709 let decoded: HighWaterStateMachineSnapshot = tsoracle_codec::decode_framed(
1710 tsoracle_openraft_toolkit::MIN_READABLE_VERSION,
1711 tsoracle_openraft_toolkit::MAX_READABLE_VERSION,
1712 &framed,
1713 )
1714 .expect("decode in range");
1715 assert_eq!(decoded, payload);
1716 }
1717
1718 // One past the readable max must be rejected loudly.
1719 let above = tsoracle_openraft_toolkit::MAX_READABLE_VERSION.saturating_add(1);
1720 let mut framed_above = vec![above];
1721 framed_above.extend_from_slice(&tsoracle_codec::encode_postcard(&payload).expect("body"));
1722 assert!(
1723 tsoracle_codec::decode_framed::<HighWaterStateMachineSnapshot>(
1724 tsoracle_openraft_toolkit::MIN_READABLE_VERSION,
1725 tsoracle_openraft_toolkit::MAX_READABLE_VERSION,
1726 &framed_above,
1727 )
1728 .is_err(),
1729 "a version above the readable max must be rejected"
1730 );
1731 }
1732
1733 #[tokio::test]
1734 async fn active_write_version_survives_reopen_and_drives_writes() {
1735 // The durable active write version is recovered on reopen and is
1736 // the version stamped onto snapshot writes. On `main` the recovered
1737 // value is BASELINE because no activation has flipped it; this
1738 // test pins that (a) the accessor is stable across a reopen and
1739 // (b) the snapshot write uses exactly that recovered value.
1740 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
1741 let recovered_active;
1742 {
1743 let mut sm = HighWaterStateMachine::with_store(store.clone()).expect("first SM");
1744 apply_one(
1745 &mut sm,
1746 1,
1747 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 55 })),
1748 )
1749 .await;
1750 sm.build_snapshot().await.expect("build_snapshot");
1751 recovered_active = sm.active_write_version();
1752 }
1753
1754 let mut reopened = HighWaterStateMachine::with_store(store.clone()).expect("reopened SM");
1755 assert_eq!(
1756 reopened.active_write_version(),
1757 recovered_active,
1758 "active write version must be stable across reopen"
1759 );
1760
1761 apply_one(
1762 &mut reopened,
1763 2,
1764 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 66 })),
1765 )
1766 .await;
1767 reopened.build_snapshot().await.expect("rebuild");
1768 let bytes = store.load().expect("load").expect("snapshot present");
1769 assert_eq!(
1770 bytes[0],
1771 reopened.active_write_version(),
1772 "rebuilt snapshot stamped with the recovered active write version"
1773 );
1774 }
1775
1776 // ---- Snapshot tests ----
1777
1778 #[tokio::test]
1779 async fn build_snapshot_round_trips_payload() {
1780 let mut sm = HighWaterStateMachine::new();
1781 apply_one(
1782 &mut sm,
1783 1,
1784 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 500 })),
1785 )
1786 .await;
1787
1788 let snap = sm.build_snapshot().await.expect("build_snapshot");
1789 let bytes = snap.snapshot.into_inner();
1790 // Use decode_framed (VersionedCodec) so the v4 bytes are decoded via
1791 // the frozen-V4-struct + lift path, not a raw postcard deserialize.
1792 let payload: HighWaterStateMachineSnapshot = tsoracle_codec::decode_framed(
1793 tsoracle_openraft_toolkit::MIN_READABLE_VERSION,
1794 tsoracle_openraft_toolkit::MAX_READABLE_VERSION,
1795 &bytes,
1796 )
1797 .expect("decode snapshot");
1798 assert_eq!(payload.current_value, 500);
1799 assert_eq!(payload.last_applied.map(|l| l.index), Some(1));
1800 }
1801
1802 #[tokio::test]
1803 async fn build_snapshot_uses_fresh_id_each_time() {
1804 let mut sm = HighWaterStateMachine::new();
1805 apply_one(
1806 &mut sm,
1807 1,
1808 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 7 })),
1809 )
1810 .await;
1811
1812 let a = sm.build_snapshot().await.expect("build_snapshot a");
1813 let b = sm.build_snapshot().await.expect("build_snapshot b");
1814 assert_ne!(
1815 a.meta.snapshot_id, b.meta.snapshot_id,
1816 "two snapshots at same last_applied must have distinct ids"
1817 );
1818 }
1819
1820 #[tokio::test]
1821 async fn install_snapshot_replaces_state() {
1822 let mut sm = HighWaterStateMachine::new();
1823 let payload = HighWaterStateMachineSnapshot {
1824 current_value: 999,
1825 last_applied: Some(log_id(5)),
1826 last_membership: StoredMem::default(),
1827 dense: std::collections::BTreeMap::new(),
1828 dense_cap: 0,
1829 };
1830 let bytes = tsoracle_codec::encode_framed(
1831 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
1832 &payload,
1833 )
1834 .expect("serialize payload");
1835
1836 let meta = SnapMeta {
1837 last_log_id: payload.last_applied,
1838 last_membership: payload.last_membership.clone(),
1839 snapshot_id: "test-install-1".to_string(),
1840 };
1841 sm.install_snapshot(&meta, std::io::Cursor::new(bytes))
1842 .await
1843 .expect("install_snapshot");
1844
1845 assert_eq!(sm.current_value().await, 999);
1846 let (last, _) = sm.applied_state().await.unwrap();
1847 assert_eq!(last.map(|l| l.index), Some(5));
1848
1849 let current = sm
1850 .get_current_snapshot()
1851 .await
1852 .expect("get_current_snapshot")
1853 .expect("snapshot present");
1854 assert_eq!(current.meta.snapshot_id, "test-install-1");
1855 }
1856
1857 /// Build a v5 (dense) snapshot whose `dense` map carries one raw-`String`
1858 /// key, paired with a meta whose `last_log_id` agrees with the payload so
1859 /// the install reaches the dense-key validation guard (not the earlier
1860 /// meta/payload-disagreement check). Returns `(meta, framed bytes)`.
1861 fn dense_snapshot_with_key(key: &str) -> (SnapMeta, Vec<u8>) {
1862 let last_applied = Some(log_id(9));
1863 let payload = HighWaterStateMachineSnapshot {
1864 current_value: 999,
1865 last_applied,
1866 last_membership: StoredMem::default(),
1867 dense: std::collections::BTreeMap::from([(key.to_string(), 7u64)]),
1868 dense_cap: 64,
1869 };
1870 let bytes =
1871 tsoracle_codec::encode_framed(tsoracle_openraft_toolkit::DENSE_WRITE_VERSION, &payload)
1872 .expect("serialize v5 dense payload");
1873 let meta = SnapMeta {
1874 last_log_id: last_applied,
1875 last_membership: StoredMem::default(),
1876 snapshot_id: "test-invalid-dense-key".to_string(),
1877 };
1878 (meta, bytes)
1879 }
1880
1881 #[tokio::test]
1882 async fn install_snapshot_rejects_empty_dense_key() {
1883 let mut sm = HighWaterStateMachine::new();
1884 let before = sm.current_value().await;
1885 let (meta, bytes) = dense_snapshot_with_key("");
1886
1887 let err = sm
1888 .install_snapshot(&meta, std::io::Cursor::new(bytes))
1889 .await
1890 .expect_err("install of a snapshot with an empty dense key must be rejected");
1891 assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
1892
1893 // Rejected before any mutation: the wholesale state replacement (which
1894 // would have set current_value to 999) did not happen, and the invalid
1895 // key is absent from the dense map.
1896 assert_eq!(sm.current_value().await, before);
1897 assert_eq!(sm.dense_value(""), 0);
1898 }
1899
1900 #[tokio::test]
1901 async fn install_snapshot_rejects_oversized_dense_key() {
1902 let mut sm = HighWaterStateMachine::new();
1903 let before = sm.current_value().await;
1904 let oversized = "a".repeat(tsoracle_core::MAX_SEQ_KEY_LEN + 1);
1905 let (meta, bytes) = dense_snapshot_with_key(&oversized);
1906
1907 let err = sm
1908 .install_snapshot(&meta, std::io::Cursor::new(bytes))
1909 .await
1910 .expect_err("install of a snapshot with an oversized dense key must be rejected");
1911 assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
1912
1913 assert_eq!(sm.current_value().await, before);
1914 assert_eq!(sm.dense_value(&oversized), 0);
1915 }
1916
1917 #[tokio::test]
1918 async fn install_snapshot_accepts_valid_dense_key() {
1919 // The guard must not reject an honest v5 snapshot: a valid dense key
1920 // installs and is observable afterwards.
1921 let mut sm = HighWaterStateMachine::new();
1922 let (meta, bytes) = dense_snapshot_with_key("orders");
1923
1924 sm.install_snapshot(&meta, std::io::Cursor::new(bytes))
1925 .await
1926 .expect("install of a snapshot with a valid dense key must succeed");
1927
1928 assert_eq!(sm.current_value().await, 999);
1929 assert_eq!(sm.dense_value("orders"), 7);
1930 }
1931
1932 #[tokio::test]
1933 async fn get_current_snapshot_initially_none() {
1934 let mut sm = HighWaterStateMachine::new();
1935 let s = sm
1936 .get_current_snapshot()
1937 .await
1938 .expect("get_current_snapshot");
1939 assert!(s.is_none());
1940 }
1941
1942 // ---- SnapshotStore integration ----
1943
1944 use crate::snapshot_store::{InMemorySnapshotStore, SnapshotStore};
1945
1946 #[tokio::test]
1947 async fn build_snapshot_writes_through_store() {
1948 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
1949 let mut sm = HighWaterStateMachine::with_store(store.clone()).expect("with_store");
1950 apply_one(
1951 &mut sm,
1952 1,
1953 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 42 })),
1954 )
1955 .await;
1956 sm.build_snapshot().await.expect("build_snapshot");
1957 assert!(
1958 store.load().expect("load").is_some(),
1959 "store must contain a persisted snapshot after build_snapshot"
1960 );
1961 }
1962
1963 #[tokio::test]
1964 async fn with_store_recovers_prior_snapshot_on_construction() {
1965 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
1966 {
1967 let mut sm = HighWaterStateMachine::with_store(store.clone()).expect("first SM");
1968 apply_one(
1969 &mut sm,
1970 1,
1971 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: 99 })),
1972 )
1973 .await;
1974 sm.build_snapshot().await.expect("build_snapshot");
1975 }
1976 let mut sm = HighWaterStateMachine::with_store(store).expect("reopened SM");
1977 assert_eq!(sm.current_value().await, 99);
1978 // `applied_state` must report the snapshot's last_log_id after reopen —
1979 // without this, openraft re-applies from index 0 and panics on missing
1980 // log entries that the snapshot already covered.
1981 let (last, _) = sm.applied_state().await.unwrap();
1982 assert_eq!(last.map(|l| l.index), Some(1));
1983 let snap = sm
1984 .get_current_snapshot()
1985 .await
1986 .expect("get_current_snapshot")
1987 .expect("snapshot present after reopen");
1988 assert_eq!(snap.meta.last_log_id.map(|l| l.index), Some(1));
1989 }
1990
1991 #[tokio::test]
1992 async fn default_constructor_matches_new() {
1993 // `Default` is the canonical "no-arg" entry point for embedders that
1994 // build the SM via `..Default::default()`; pinning behavior here keeps
1995 // the in-memory snapshot store as the unsurprising default.
1996 let mut sm = HighWaterStateMachine::default();
1997 assert_eq!(sm.current_value().await, 0);
1998 let snap = sm
1999 .get_current_snapshot()
2000 .await
2001 .expect("get_current_snapshot");
2002 assert!(snap.is_none());
2003 }
2004
2005 #[tokio::test]
2006 async fn state_machine_defaults_to_baseline_active_write_version() {
2007 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
2008 let sm = HighWaterStateMachine::with_store(store).expect("with_store");
2009 assert_eq!(
2010 sm.active_write_version(),
2011 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION
2012 );
2013 }
2014
2015 #[tokio::test]
2016 async fn state_machine_reads_the_same_shared_cell() {
2017 // The SM and the log store hold clones of ONE cell. A set() on the
2018 // cell (a later phase's activation apply will do this) is observed
2019 // identically by every clone.
2020 let cell = tsoracle_openraft_toolkit::ActiveWriteVersion::default();
2021 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
2022 let sm = HighWaterStateMachine::with_store_and_active_version(store, cell.clone())
2023 .expect("with_store_and_active_version");
2024 assert_eq!(
2025 sm.active_write_version(),
2026 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION
2027 );
2028 cell.set(7);
2029 assert_eq!(sm.active_write_version(), 7);
2030 }
2031
2032 #[tokio::test]
2033 async fn begin_receiving_snapshot_returns_empty_cursor() {
2034 // openraft hands the returned cursor to the snapshot-receiving network
2035 // path; the contract is "empty, writable buffer." Anything non-empty
2036 // would corrupt the install on the receiving side.
2037 let mut sm = HighWaterStateMachine::new();
2038 let cursor = sm
2039 .begin_receiving_snapshot()
2040 .await
2041 .expect("begin_receiving_snapshot");
2042 assert!(cursor.into_inner().is_empty());
2043 }
2044
2045 #[tokio::test]
2046 async fn with_store_errors_on_malformed_persisted_envelope() {
2047 // Hardens the recovery path against on-disk corruption: a snapshot
2048 // blob that doesn't decode as `PersistedSnapshot` must surface as a
2049 // structured `io::Error` from `with_store`, not a silent reset to
2050 // the default state (which would lose the value across restart).
2051 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
2052 store.save(b"not a postcard envelope").unwrap();
2053 let Err(err) = HighWaterStateMachine::with_store(store) else {
2054 panic!("with_store must reject malformed envelope");
2055 };
2056 let msg = err.to_string();
2057 assert!(msg.contains("envelope decode"));
2058 }
2059
2060 #[tokio::test]
2061 async fn with_store_errors_on_malformed_inner_payload() {
2062 // Envelope decodes but the inner `HighWaterStateMachineSnapshot` does
2063 // not — also surfaces as `io::Error` rather than silent state reset.
2064 let envelope = PersistedSnapshot {
2065 meta: SnapMeta::default(),
2066 data: b"not a postcard payload".to_vec(),
2067 };
2068 let bytes = tsoracle_openraft_toolkit::encode(
2069 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2070 &envelope,
2071 )
2072 .unwrap();
2073 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
2074 store.save(&bytes).unwrap();
2075 let Err(err) = HighWaterStateMachine::with_store(store) else {
2076 panic!("with_store must reject malformed inner payload");
2077 };
2078 let msg = err.to_string();
2079 assert!(msg.contains("payload decode"));
2080 }
2081
2082 #[tokio::test]
2083 async fn install_snapshot_writes_through_store() {
2084 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
2085 let mut sm = HighWaterStateMachine::with_store(store.clone()).expect("with_store");
2086 let payload = HighWaterStateMachineSnapshot {
2087 current_value: 700,
2088 last_applied: Some(log_id(10)),
2089 last_membership: StoredMem::default(),
2090 dense: std::collections::BTreeMap::new(),
2091 dense_cap: 0,
2092 };
2093 let bytes = tsoracle_codec::encode_framed(
2094 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2095 &payload,
2096 )
2097 .unwrap();
2098 let meta = SnapMeta {
2099 last_log_id: payload.last_applied,
2100 last_membership: payload.last_membership.clone(),
2101 snapshot_id: "install-1".into(),
2102 };
2103 sm.install_snapshot(&meta, std::io::Cursor::new(bytes))
2104 .await
2105 .expect("install_snapshot");
2106 // Reopen with the same store: install must have written through.
2107 let mut sm2 = HighWaterStateMachine::with_store(store).expect("reopened SM");
2108 assert_eq!(sm2.current_value().await, 700);
2109 let (last, _) = sm2.applied_state().await.unwrap();
2110 assert_eq!(last.map(|l| l.index), Some(10));
2111 }
2112
2113 /// Encode a `(meta, framed-payload-bytes)` pair for an install at
2114 /// `(value, last_applied)`, with `meta.last_log_id` matching the payload so
2115 /// the meta/payload consistency check passes and the install reaches the
2116 /// publish path.
2117 fn install_payload(value: u64, last_applied: LogId) -> (SnapMeta, Vec<u8>) {
2118 let payload = HighWaterStateMachineSnapshot {
2119 current_value: value,
2120 last_applied: Some(last_applied),
2121 last_membership: StoredMem::default(),
2122 dense: std::collections::BTreeMap::new(),
2123 dense_cap: 0,
2124 };
2125 let bytes = tsoracle_codec::encode_framed(
2126 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2127 &payload,
2128 )
2129 .expect("serialize payload");
2130 let meta = SnapMeta {
2131 last_log_id: payload.last_applied,
2132 last_membership: payload.last_membership.clone(),
2133 snapshot_id: format!("snap-{}", last_applied.index),
2134 };
2135 (meta, bytes)
2136 }
2137
2138 #[tokio::test]
2139 async fn install_snapshot_does_not_regress_published_snapshot() {
2140 // `build_snapshot` publishes both the durable store (:256) and the
2141 // in-memory `current_snapshot` (:258) AFTER releasing the core lock,
2142 // and openraft can drive a stale build on a clone concurrently with a
2143 // newer install. The publish must therefore be monotone by
2144 // `last_log_id`: a later-arriving, lower-indexed snapshot must not roll
2145 // the durable + in-memory snapshot back to a stale state (which a
2146 // subsequent restart would then recover, after openraft has already
2147 // purged its log past the newer snapshot). We exercise that invariant
2148 // through the install path here — install and build share the same
2149 // monotone publish, so a regressing install stands in for a stale build
2150 // that finishes its write last.
2151 enable_tracing();
2152 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
2153 let mut sm = HighWaterStateMachine::with_store(store.clone()).expect("with_store");
2154
2155 let (newer_meta, newer_bytes) = install_payload(80, log_id(8));
2156 sm.install_snapshot(&newer_meta, Cursor::new(newer_bytes))
2157 .await
2158 .expect("install newer");
2159
2160 // A stale snapshot (lower last_log_id) arriving afterwards is a no-op:
2161 // we are already at least this far, so it must neither error nor
2162 // regress state.
2163 let (older_meta, older_bytes) = install_payload(50, log_id(5));
2164 sm.install_snapshot(&older_meta, Cursor::new(older_bytes))
2165 .await
2166 .expect("stale install must be an accepted no-op, not an error");
2167
2168 assert_eq!(
2169 sm.current_value().await,
2170 80,
2171 "value must not regress to the stale snapshot"
2172 );
2173 let (last, _) = sm.applied_state().await.unwrap();
2174 assert_eq!(
2175 last.map(|l| l.index),
2176 Some(8),
2177 "last_applied must not regress"
2178 );
2179 let current = sm
2180 .get_current_snapshot()
2181 .await
2182 .expect("get_current_snapshot")
2183 .expect("snapshot present");
2184 assert_eq!(
2185 current.meta.last_log_id.map(|l| l.index),
2186 Some(8),
2187 "in-memory current_snapshot must not regress"
2188 );
2189
2190 // The durable store must also still hold the newer snapshot — this is
2191 // the recovery-critical half: reopening must not resurrect the stale
2192 // state.
2193 let mut reopened = HighWaterStateMachine::with_store(store).expect("reopen");
2194 assert_eq!(
2195 reopened.current_value().await,
2196 80,
2197 "durable store must not have been rolled back to the stale snapshot"
2198 );
2199 let (last, _) = reopened.applied_state().await.unwrap();
2200 assert_eq!(last.map(|l| l.index), Some(8));
2201 }
2202
2203 #[test]
2204 fn supersedes_published_is_monotone_by_last_log_id() {
2205 // The exact rule the publish gate enforces. Equal or greater
2206 // last_log_id supersedes; a lower one — or `None` against a `Some` —
2207 // does not. `None` is the minimum: a fresh/pre-genesis snapshot never
2208 // displaces an established one, but anything fills an empty slot.
2209 assert!(
2210 supersedes_published(Some(log_id(5)), None),
2211 "any Some fills an empty slot"
2212 );
2213 assert!(
2214 supersedes_published(None, None),
2215 "first publish into an empty slot is allowed"
2216 );
2217 assert!(
2218 supersedes_published(Some(log_id(8)), Some(log_id(5))),
2219 "higher index supersedes"
2220 );
2221 assert!(
2222 supersedes_published(Some(log_id(5)), Some(log_id(5))),
2223 "equal index may republish (same state, fresh snapshot_id)"
2224 );
2225 assert!(
2226 !supersedes_published(Some(log_id(5)), Some(log_id(8))),
2227 "lower index must not regress"
2228 );
2229 assert!(
2230 !supersedes_published(None, Some(log_id(1))),
2231 "None must not regress an established snapshot"
2232 );
2233 }
2234
2235 #[tokio::test]
2236 async fn commit_snapshot_drops_a_stale_publish() {
2237 // Exercise the build-side of the TOCTOU directly at the seam
2238 // build_snapshot shares with install. Once a newer snapshot is durable
2239 // and published, a commit carrying a lower last_log_id — a build whose
2240 // state was captured before the newer install but whose write lands
2241 // last — must be dropped: it reports `Ok(false)`, runs no adopt
2242 // callback, and leaves both the store and current_snapshot byte-for-byte
2243 // untouched. (The full build path can only reach this state under a
2244 // real scheduling race, so we drive the shared commit seam directly.)
2245 enable_tracing();
2246 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
2247 let mut sm = HighWaterStateMachine::with_store(store.clone()).expect("with_store");
2248
2249 let (newer_meta, newer_bytes) = install_payload(80, log_id(8));
2250 sm.install_snapshot(&newer_meta, Cursor::new(newer_bytes))
2251 .await
2252 .expect("install newer");
2253 let durable_after_newer = store
2254 .load()
2255 .expect("load")
2256 .expect("durable snapshot present");
2257
2258 let (stale_meta, stale_data) = install_payload(50, log_id(5));
2259 let envelope = tsoracle_openraft_toolkit::encode(
2260 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2261 &PersistedSnapshot {
2262 meta: stale_meta.clone(),
2263 data: stale_data.clone(),
2264 },
2265 )
2266 .expect("encode envelope");
2267
2268 let adopt_ran = std::cell::Cell::new(false);
2269 let adopted = sm
2270 .commit_snapshot(stale_meta, stale_data, &envelope, |_| adopt_ran.set(true))
2271 .expect("commit_snapshot");
2272
2273 assert!(!adopted, "stale commit must report it was not adopted");
2274 assert!(
2275 !adopt_ran.get(),
2276 "adopt callback must not run for a dropped commit"
2277 );
2278 assert_eq!(
2279 store.load().expect("load").as_deref(),
2280 Some(durable_after_newer.as_slice()),
2281 "durable snapshot must be unchanged — a stale commit must not write"
2282 );
2283 let current = sm
2284 .get_current_snapshot()
2285 .await
2286 .expect("get_current_snapshot")
2287 .expect("snapshot present");
2288 assert_eq!(
2289 current.meta.last_log_id.map(|l| l.index),
2290 Some(8),
2291 "in-memory current_snapshot must be unchanged"
2292 );
2293 }
2294
2295 #[tokio::test]
2296 async fn install_snapshot_rejects_meta_payload_log_id_mismatch() {
2297 // `meta.last_log_id` (read back by `get_current_snapshot`) and
2298 // `payload.last_applied` (read back by `applied_state`) must agree —
2299 // honest peers always build them from the same `last_applied`, so a
2300 // disagreement signals corruption or a peer bug. Installing it anyway
2301 // would silently desync the two reader methods, so the install must be
2302 // rejected, leave state untouched, and persist nothing.
2303 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
2304 let mut sm = HighWaterStateMachine::with_store(store.clone()).expect("with_store");
2305
2306 let payload = HighWaterStateMachineSnapshot {
2307 current_value: 123,
2308 last_applied: Some(log_id(5)),
2309 last_membership: StoredMem::default(),
2310 dense: std::collections::BTreeMap::new(),
2311 dense_cap: 0,
2312 };
2313 let bytes = tsoracle_codec::encode_framed(
2314 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2315 &payload,
2316 )
2317 .expect("serialize payload");
2318
2319 let meta = SnapMeta {
2320 last_log_id: Some(log_id(6)),
2321 last_membership: payload.last_membership.clone(),
2322 snapshot_id: "mismatch-1".into(),
2323 };
2324
2325 let Err(err) = sm
2326 .install_snapshot(&meta, std::io::Cursor::new(bytes))
2327 .await
2328 else {
2329 panic!("install_snapshot must reject meta/payload last_log_id mismatch");
2330 };
2331 let msg = err.to_string();
2332 assert!(
2333 msg.contains("last_log_id"),
2334 "error should name the mismatched field: {msg}"
2335 );
2336
2337 assert_eq!(sm.current_value().await, 0);
2338 let (last, _) = sm.applied_state().await.unwrap();
2339 assert_eq!(last, None);
2340 assert!(
2341 sm.get_current_snapshot().await.unwrap().is_none(),
2342 "rejected install must not publish a current snapshot"
2343 );
2344 assert!(
2345 store.load().expect("load").is_none(),
2346 "rejected install must not persist an envelope"
2347 );
2348 }
2349
2350 // ---- Property tests ----
2351
2352 use proptest::prelude::*;
2353
2354 fn rt() -> tokio::runtime::Runtime {
2355 tokio::runtime::Builder::new_current_thread()
2356 .enable_all()
2357 .build()
2358 .unwrap()
2359 }
2360
2361 proptest! {
2362 /// Monotonicity invariant: applying an arbitrary sequence of `Advance`
2363 /// values leaves the state machine at `max(0, max(at_leasts))`, and
2364 /// `current_value` is non-decreasing at every intermediate step.
2365 #[test]
2366 fn p1_monotonicity_under_arbitrary_bumps(
2367 targets in prop::collection::vec(any::<u64>(), 0..=64)
2368 ) {
2369 rt().block_on(async {
2370 let mut sm = HighWaterStateMachine::new();
2371 let mut prev = 0u64;
2372 let mut idx = 0u64;
2373 for t in &targets {
2374 idx += 1;
2375 apply_one(
2376 &mut sm,
2377 idx,
2378 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: *t })),
2379 )
2380 .await;
2381 let now = sm.current_value().await;
2382 prop_assert!(now >= prev, "value went backwards: prev={prev} now={now}");
2383 prev = now;
2384 }
2385 let expected = targets.iter().copied().max().unwrap_or(0);
2386 prop_assert_eq!(sm.current_value().await, expected);
2387 Ok(())
2388 })?;
2389 }
2390
2391 /// Snapshot payload round-trip: build_snapshot -> install_snapshot
2392 /// preserves (current_value, last_applied, last_membership) across
2393 /// arbitrary apply sequences.
2394 #[test]
2395 fn p2_snapshot_payload_round_trip(
2396 bumps in prop::collection::vec(any::<u64>(), 0..=32)
2397 ) {
2398 rt().block_on(async {
2399 let mut sm = HighWaterStateMachine::new();
2400 let mut idx = 0u64;
2401 for t in &bumps {
2402 idx += 1;
2403 apply_one(
2404 &mut sm,
2405 idx,
2406 EntryPayload::Normal(HighWaterCommand::Advance(AdvancePayload { at_least: *t })),
2407 )
2408 .await;
2409 }
2410
2411 let snap = sm.build_snapshot().await.expect("build_snapshot");
2412 let meta = snap.meta.clone();
2413 let bytes = snap.snapshot.into_inner();
2414
2415 let mut sm2 = HighWaterStateMachine::new();
2416 sm2.install_snapshot(&meta, std::io::Cursor::new(bytes))
2417 .await
2418 .expect("install_snapshot");
2419
2420 prop_assert_eq!(sm2.current_value().await, sm.current_value().await);
2421 let (a_last, a_mem) = sm.applied_state().await.unwrap();
2422 let (b_last, b_mem) = sm2.applied_state().await.unwrap();
2423 prop_assert_eq!(a_last, b_last);
2424 prop_assert_eq!(a_mem, b_mem);
2425 Ok(())
2426 })?;
2427 }
2428 }
2429
2430 #[test]
2431 fn snapshot_payload_versioned_codec_matches_legacy_frame() {
2432 use tsoracle_codec::{decode_framed, encode_framed};
2433 // The v4 encode arm projects to the frozen `HighWaterStateMachineSnapshotV4`
2434 // struct (3 fields, no dense/cap). This test proves:
2435 //
2436 // 1. The v4 framed bytes are `[4, 7, 0, 0, 0, 0]` — the on-disk v4
2437 // layout did not change from the pre-dense era.
2438 // 2. Those bytes equal what the frozen V4 struct produces when encoded
2439 // directly — a real pre-dense reader can still decode the output.
2440 // 3. A v4-decode-and-lift roundtrip restores the original value with
2441 // empty dense map + cap 0 ("absent" sentinel).
2442 let payload = HighWaterStateMachineSnapshot {
2443 current_value: 7,
2444 last_applied: None,
2445 last_membership: StoredMembership::default(),
2446 dense: std::collections::BTreeMap::new(),
2447 dense_cap: 0,
2448 };
2449 let via_seam = encode_framed(tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION, &payload)
2450 .expect("encode_framed");
2451
2452 // Prove the v4 body equals encoding the frozen V4 struct directly.
2453 let v4_direct = tsoracle_openraft_toolkit::encode(
2454 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2455 &HighWaterStateMachineSnapshotV4 {
2456 current_value: 7,
2457 last_applied: None,
2458 last_membership: StoredMembership::default(),
2459 },
2460 )
2461 .expect("v4 direct encode");
2462 assert_eq!(
2463 via_seam, v4_direct,
2464 "v4 encode must project to the frozen V4 struct"
2465 );
2466 assert_eq!(via_seam, vec![4, 7, 0, 0, 0, 0]);
2467
2468 // decode-and-lift: v4 bytes decode into the current layout with empty
2469 // dense map and cap 0 (the "absent" sentinel), which matches `payload`.
2470 let back: HighWaterStateMachineSnapshot = decode_framed(
2471 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2472 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2473 &via_seam,
2474 )
2475 .expect("decode_framed");
2476 assert_eq!(back, payload);
2477 }
2478
2479 #[test]
2480 fn snapshot_payload_versioned_codec_rejects_foreign_version() {
2481 use tsoracle_codec::{CodecError, decode_framed};
2482 let framed = vec![0xFFu8, 7, 0, 0, 0, 0];
2483 let err = decode_framed::<HighWaterStateMachineSnapshot>(
2484 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2485 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2486 &framed,
2487 )
2488 .expect_err("must reject");
2489 assert!(matches!(
2490 err,
2491 CodecError::VersionUnsupported { actual: 0xFF, .. }
2492 ));
2493 }
2494
2495 #[test]
2496 fn snapshot_payload_pins_v4_layout() {
2497 use tsoracle_codec::encode_framed;
2498 // Hand-built v4 frame: [BASELINE_WRITE_VERSION | postcard(V4_payload)],
2499 // produced through the VersionedCodec seam that build_snapshot uses
2500 // when active_write_version == BASELINE_WRITE_VERSION. The v4 encode
2501 // arm projects to the frozen `HighWaterStateMachineSnapshotV4` struct
2502 // (3 fields, no dense/cap), so the body bytes are unchanged from the
2503 // pre-dense era. Leading byte advanced 3 -> 4 when OpenraftPeer gained
2504 // the admin_endpoint field (a breaking on-disk change for membership);
2505 // the postcard body is unchanged. Body [7, 0, 0, 0, 0] =
2506 // current_value 7, last_applied None, then default StoredMembership
2507 // (None log id + empty configs + empty nodes). Reordering or inserting
2508 // a field in the V4 struct changes these bytes and trips this test,
2509 // forcing a deliberate review.
2510 let payload = HighWaterStateMachineSnapshot {
2511 current_value: 7,
2512 last_applied: None,
2513 last_membership: StoredMembership::default(),
2514 dense: std::collections::BTreeMap::new(),
2515 dense_cap: 0,
2516 };
2517 let framed = encode_framed(tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION, &payload)
2518 .expect("encode");
2519 // v4 encode projects to the frozen V4 struct: body is still
2520 // [7, 0, 0, 0, 0] — the dense/cap bytes are NOT in the output.
2521 assert_eq!(framed, vec![4, 7, 0, 0, 0, 0]);
2522 }
2523
2524 #[test]
2525 fn log_entry_pins_v4_layout() {
2526 use crate::log_codec::OpenraftLogCodec;
2527 use tsoracle_openraft_toolkit::LogStoreCodec;
2528 // The bytes RocksdbLogStore<TypeConfig> persists per entry: the
2529 // active write version byte (prepended by the store) +
2530 // OpenraftLogCodec entry body — the v4 frame around a Normal entry
2531 // carrying Advance(AdvancePayload { at_least: 5 })
2532 // at log id (term 1, node 1, index 1). Body [1,1,1,1,0,5] = leader
2533 // (term 1, node 1), index 1, EntryPayload::Normal tag (1), Advance
2534 // variant (0), at_least 5 — byte-identical to the pre-seam layout.
2535 let lid = openraft::testing::log_id::<TypeConfig>(1, 1, 1);
2536 let entry: EntryOf<TypeConfig> = EntryOf::<TypeConfig>::new_normal(
2537 lid,
2538 HighWaterCommand::Advance(AdvancePayload { at_least: 5 }),
2539 );
2540 let body = <OpenraftLogCodec as LogStoreCodec<TypeConfig>>::encode_entry(
2541 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2542 &entry,
2543 )
2544 .expect("encode entry body");
2545 let mut framed = vec![tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION];
2546 framed.extend_from_slice(&body);
2547 assert_eq!(framed, vec![4, 1, 1, 1, 1, 0, 5]);
2548 }
2549
2550 #[test]
2551 fn meta_vote_pins_v4_layout() {
2552 use crate::log_codec::OpenraftLogCodec;
2553 use tsoracle_openraft_toolkit::LogStoreCodec;
2554 // The bytes RocksdbLogStore<TypeConfig> persists in the meta column for a
2555 // Vote: active write version byte + OpenraftLogCodec vote body. Body [7, 3, 1] =
2556 // leader (term 7, node 3), committed flag true. A layout change to Vote
2557 // trips this test. This is the recovery-critical field the framing
2558 // protects: a foreign version loud-rejects instead of misdecoding.
2559 let vote: openraft::type_config::alias::VoteOf<TypeConfig> =
2560 openraft::Vote::new_committed(7, 3);
2561 let body = <OpenraftLogCodec as LogStoreCodec<TypeConfig>>::encode_vote(
2562 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2563 &vote,
2564 )
2565 .expect("encode vote body");
2566 let mut framed = vec![tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION];
2567 framed.extend_from_slice(&body);
2568 assert_eq!(framed, vec![4, 7, 3, 1]);
2569 }
2570
2571 #[tokio::test]
2572 async fn install_snapshot_rejects_foreign_schema_version() {
2573 // A streamed snapshot framed with a foreign version must be rejected
2574 // as InvalidData before any store write or state mutation.
2575 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
2576 let mut sm = HighWaterStateMachine::with_store(store.clone()).expect("with_store");
2577 let meta = SnapMeta {
2578 last_log_id: None,
2579 last_membership: StoredMembership::default(),
2580 snapshot_id: "test".to_string(),
2581 };
2582 let foreign = vec![0xFF, 7, 0, 0, 0, 0];
2583 let err = sm
2584 .install_snapshot(&meta, Cursor::new(foreign))
2585 .await
2586 .expect_err("must reject foreign version");
2587 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
2588
2589 assert_eq!(sm.current_value().await, 0);
2590 let (last, _) = sm.applied_state().await.unwrap();
2591 assert_eq!(last, None);
2592 assert!(
2593 sm.get_current_snapshot().await.unwrap().is_none(),
2594 "rejected install must not publish a current snapshot"
2595 );
2596 assert!(
2597 store.load().expect("load").is_none(),
2598 "rejected install must not persist an envelope"
2599 );
2600 }
2601
2602 // ---- Snapshot codec cross-version tests (Task 4) ----
2603 //
2604 // These tests pin the postcard correctness guarantees added by Task 4:
2605 // - v5 round-trip: encode a dense snapshot, decode back equal.
2606 // - v4 forward-compat: encode an empty-dense snapshot at v4; the bytes
2607 // equal what the frozen V4 struct produces; decode lifts to empty dense.
2608 // - v4 old-bytes: hand-encode a V4 struct, decode as current → empty dense.
2609
2610 #[test]
2611 fn snapshot_codec_v5_roundtrip_with_nonempty_dense_map() {
2612 use tsoracle_codec::{decode_framed, encode_framed};
2613 use tsoracle_openraft_toolkit::DENSE_WRITE_VERSION;
2614
2615 let mut dense = std::collections::BTreeMap::new();
2616 dense.insert("orders".to_string(), 42u64);
2617 dense.insert("invoices".to_string(), 1000u64);
2618
2619 let payload = HighWaterStateMachineSnapshot {
2620 current_value: 777,
2621 last_applied: Some(log_id(9)),
2622 last_membership: StoredMem::default(),
2623 dense: dense.clone(),
2624 dense_cap: 500,
2625 };
2626
2627 let framed = encode_framed(DENSE_WRITE_VERSION, &payload).expect("v5 encode");
2628 assert_eq!(framed[0], DENSE_WRITE_VERSION, "leading byte must be v5");
2629
2630 let decoded: HighWaterStateMachineSnapshot = decode_framed(
2631 tsoracle_openraft_toolkit::MIN_READABLE_VERSION,
2632 tsoracle_openraft_toolkit::MAX_READABLE_VERSION,
2633 &framed,
2634 )
2635 .expect("v5 decode");
2636 assert_eq!(
2637 decoded, payload,
2638 "v5 roundtrip must produce the original payload"
2639 );
2640 assert_eq!(decoded.dense, dense, "dense map must survive roundtrip");
2641 assert_eq!(decoded.dense_cap, 500, "dense_cap must survive roundtrip");
2642 }
2643
2644 #[test]
2645 fn snapshot_payload_round_trips_at_v6() {
2646 use tsoracle_codec::{decode_framed, encode_framed};
2647 use tsoracle_openraft_toolkit::BATCH_WRITE_VERSION;
2648
2649 let mut dense = std::collections::BTreeMap::new();
2650 dense.insert("orders".to_string(), 9u64);
2651
2652 let payload = HighWaterStateMachineSnapshot {
2653 current_value: 777,
2654 last_applied: Some(log_id(9)),
2655 last_membership: StoredMem::default(),
2656 dense: dense.clone(),
2657 dense_cap: 32,
2658 };
2659
2660 let framed = encode_framed(BATCH_WRITE_VERSION, &payload).expect("v6 encode");
2661 assert_eq!(framed[0], BATCH_WRITE_VERSION, "leading byte must be v6");
2662
2663 let decoded: HighWaterStateMachineSnapshot = decode_framed(
2664 tsoracle_openraft_toolkit::MIN_READABLE_VERSION,
2665 tsoracle_openraft_toolkit::MAX_READABLE_VERSION,
2666 &framed,
2667 )
2668 .expect("v6 decode");
2669 assert_eq!(
2670 decoded, payload,
2671 "v6 roundtrip must produce the original payload"
2672 );
2673 assert_eq!(decoded.dense, dense, "dense map must survive v6 roundtrip");
2674 assert_eq!(decoded.dense_cap, 32, "dense_cap must survive v6 roundtrip");
2675 }
2676
2677 #[test]
2678 fn snapshot_codec_v4_forward_compat_empty_dense() {
2679 use tsoracle_codec::{decode_framed, encode_framed};
2680
2681 // A pre-dense snapshot: dense is empty, cap is 0.
2682 let payload = HighWaterStateMachineSnapshot {
2683 current_value: 42,
2684 last_applied: Some(log_id(3)),
2685 last_membership: StoredMem::default(),
2686 dense: std::collections::BTreeMap::new(),
2687 dense_cap: 0,
2688 };
2689
2690 let framed = encode_framed(tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION, &payload)
2691 .expect("v4 encode");
2692
2693 // The v4 bytes must equal encoding the frozen V4 struct directly,
2694 // so a pre-dense reader can still decode them.
2695 let v4_struct_bytes = tsoracle_openraft_toolkit::encode(
2696 tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION,
2697 &HighWaterStateMachineSnapshotV4 {
2698 current_value: 42,
2699 last_applied: Some(log_id(3)),
2700 last_membership: StoredMem::default(),
2701 },
2702 )
2703 .expect("v4 struct direct encode");
2704 assert_eq!(
2705 framed, v4_struct_bytes,
2706 "v4 encode of empty-dense payload must be byte-identical to frozen V4 struct"
2707 );
2708
2709 // Decode: the v4 bytes lift into the current layout with empty dense, cap 0.
2710 let decoded: HighWaterStateMachineSnapshot = decode_framed(
2711 tsoracle_openraft_toolkit::MIN_READABLE_VERSION,
2712 tsoracle_openraft_toolkit::MAX_READABLE_VERSION,
2713 &framed,
2714 )
2715 .expect("v4 decode");
2716 assert_eq!(
2717 decoded, payload,
2718 "v4 forward-compat decode must produce empty-dense payload"
2719 );
2720 assert!(decoded.dense.is_empty(), "decoded v4 dense must be empty");
2721 assert_eq!(
2722 decoded.dense_cap, 0,
2723 "decoded v4 cap must be 0 (absent sentinel)"
2724 );
2725 }
2726
2727 #[test]
2728 fn snapshot_codec_v4_old_bytes_decode_lifts_to_empty_dense() {
2729 use tsoracle_codec::{decode_framed, encode_postcard};
2730
2731 // Hand-encode a V4 struct (the 3-field pre-dense layout) as postcard,
2732 // then frame it as a v4 record. This simulates bytes written by an old
2733 // (pre-dense) node.
2734 let v4 = HighWaterStateMachineSnapshotV4 {
2735 current_value: 99,
2736 last_applied: Some(log_id(7)),
2737 last_membership: StoredMem::default(),
2738 };
2739 let v4_body = encode_postcard(&v4).expect("postcard encode V4");
2740 let mut framed = vec![tsoracle_openraft_toolkit::BASELINE_WRITE_VERSION];
2741 framed.extend_from_slice(&v4_body);
2742
2743 // Decode through the VersionedCodec path.
2744 let decoded: HighWaterStateMachineSnapshot = decode_framed(
2745 tsoracle_openraft_toolkit::MIN_READABLE_VERSION,
2746 tsoracle_openraft_toolkit::MAX_READABLE_VERSION,
2747 &framed,
2748 )
2749 .expect("decode old v4 bytes");
2750
2751 assert_eq!(decoded.current_value, 99, "current_value lifts correctly");
2752 assert_eq!(
2753 decoded.last_applied.map(|l| l.index),
2754 Some(7),
2755 "last_applied lifts correctly"
2756 );
2757 assert!(
2758 decoded.dense.is_empty(),
2759 "old v4 bytes must lift to empty dense map"
2760 );
2761 assert_eq!(
2762 decoded.dense_cap, 0,
2763 "old v4 bytes must lift to cap 0 (absent)"
2764 );
2765 }
2766
2767 // ---- AdvanceDense apply tests ----
2768 //
2769 // Drive the `AdvanceDense` apply arm directly via `apply_one` and verify
2770 // the observable side effects through `dense_value()` (the pre/post
2771 // counter), the `last_applied` log position, and the fact that the
2772 // high-water is NOT touched by a dense advance.
2773 //
2774 // The four cases the plan mandates:
2775 // 1. First advance of "orders" by 5 → start=0, dense["orders"]=5.
2776 // 2. Second advance of "orders" by 3 → start=5, dense["orders"]=8.
2777 // 3. A new key when the cap is full → DenseCardinalityExceeded: counter
2778 // stays 0 and the map size does not grow.
2779 // 4. A counter at u64::MAX with count=1 → DenseOverflow: counter
2780 // stays at u64::MAX.
2781 //
2782 // We verify start and outcome by observing `dense_value` BEFORE the apply
2783 // (which gives `start`) and AFTER (which gives `start + count` on success,
2784 // unchanged on rejection). The high-water is pinned at 0 throughout (no
2785 // `Advance` entries), which confirms the apply arm does NOT touch it.
2786
2787 fn dense_key(s: &str) -> tsoracle_core::SeqKey {
2788 tsoracle_core::SeqKey::try_new(s).expect("valid key")
2789 }
2790
2791 #[tokio::test]
2792 async fn dense_advance_first_key_starts_at_zero() {
2793 let mut sm = HighWaterStateMachine::new();
2794 let key = dense_key("orders");
2795
2796 // Pre-condition: key is absent → start would be 0.
2797 assert_eq!(sm.dense_value("orders"), 0);
2798
2799 apply_one(
2800 &mut sm,
2801 1,
2802 EntryPayload::Normal(HighWaterCommand::AdvanceDense {
2803 key: key.clone(),
2804 count: 5,
2805 }),
2806 )
2807 .await;
2808
2809 // Post-condition: counter advanced by 5; start was 0.
2810 assert_eq!(
2811 sm.dense_value("orders"),
2812 5,
2813 "first advance by 5: dense[orders] = 5 (start=0)"
2814 );
2815 // High-water must be untouched.
2816 assert_eq!(
2817 sm.current_value().await,
2818 0,
2819 "AdvanceDense must not touch the high-water"
2820 );
2821 let (last, _) = sm.applied_state().await.unwrap();
2822 assert_eq!(last.map(|l| l.index), Some(1));
2823 }
2824
2825 #[tokio::test]
2826 async fn dense_advance_second_apply_starts_at_previous_value() {
2827 let mut sm = HighWaterStateMachine::new();
2828 let key = dense_key("orders");
2829
2830 // First advance: start=0, next=5.
2831 apply_one(
2832 &mut sm,
2833 1,
2834 EntryPayload::Normal(HighWaterCommand::AdvanceDense {
2835 key: key.clone(),
2836 count: 5,
2837 }),
2838 )
2839 .await;
2840 assert_eq!(sm.dense_value("orders"), 5);
2841
2842 // Capture start before second apply.
2843 let start_before = sm.dense_value("orders"); // == 5
2844
2845 apply_one(
2846 &mut sm,
2847 2,
2848 EntryPayload::Normal(HighWaterCommand::AdvanceDense {
2849 key: key.clone(),
2850 count: 3,
2851 }),
2852 )
2853 .await;
2854
2855 // The second block starts at 5 (the pre-advance value).
2856 assert_eq!(
2857 start_before, 5,
2858 "second advance: start must be the prior counter value (5)"
2859 );
2860 assert_eq!(
2861 sm.dense_value("orders"),
2862 8,
2863 "second advance by 3: dense[orders] = 8"
2864 );
2865 }
2866
2867 #[tokio::test]
2868 async fn dense_advance_cardinality_cap_rejects_new_key() {
2869 // Cap set to 1: one key is allowed; a second new key must be rejected.
2870 let mut sm = HighWaterStateMachine::new_with_dense_cap(1);
2871 let k1 = dense_key("orders");
2872 let k2 = dense_key("users");
2873
2874 // First key is accepted (map was empty, cap=1 → len 0 < 1 → OK).
2875 apply_one(
2876 &mut sm,
2877 1,
2878 EntryPayload::Normal(HighWaterCommand::AdvanceDense {
2879 key: k1.clone(),
2880 count: 1,
2881 }),
2882 )
2883 .await;
2884 assert_eq!(sm.dense_value("orders"), 1, "first key accepted");
2885
2886 // Second NEW key: map len==1 >= cap=1 → DenseCardinalityExceeded.
2887 // The counter for the new key must stay at 0 (rejection is a no-op on
2888 // the dense map).
2889 apply_one(
2890 &mut sm,
2891 2,
2892 EntryPayload::Normal(HighWaterCommand::AdvanceDense {
2893 key: k2.clone(),
2894 count: 1,
2895 }),
2896 )
2897 .await;
2898 assert_eq!(
2899 sm.dense_value("users"),
2900 0,
2901 "new key rejected by cardinality cap: dense[users] must be 0"
2902 );
2903 // Existing key still advances (already present → cap check skipped).
2904 apply_one(
2905 &mut sm,
2906 3,
2907 EntryPayload::Normal(HighWaterCommand::AdvanceDense {
2908 key: k1.clone(),
2909 count: 4,
2910 }),
2911 )
2912 .await;
2913 assert_eq!(
2914 sm.dense_value("orders"),
2915 5,
2916 "existing key still advances despite cap being full"
2917 );
2918 }
2919
2920 #[tokio::test]
2921 async fn dense_advance_overflow_leaves_counter_unchanged() {
2922 // Seed the counter at u64::MAX by using a large initial count and then
2923 // bridging up. Since count is u32 (max 2^32-1) we need multiple applies
2924 // to reach u64::MAX. It is simpler to directly manipulate the core in
2925 // the test module (same file, so private access is allowed).
2926 let sm = HighWaterStateMachine::new();
2927 {
2928 let mut core = sm.core.lock();
2929 core.dense.insert("x".to_string(), u64::MAX);
2930 }
2931 // Apply AdvanceDense with count=1: checked_add(u64::MAX, 1) overflows.
2932 let mut sm_mut = sm;
2933 apply_one(
2934 &mut sm_mut,
2935 1,
2936 EntryPayload::Normal(HighWaterCommand::AdvanceDense {
2937 key: dense_key("x"),
2938 count: 1,
2939 }),
2940 )
2941 .await;
2942 // Overflow: counter must remain at u64::MAX (rejection is a no-op).
2943 assert_eq!(
2944 sm_mut.dense_value("x"),
2945 u64::MAX,
2946 "overflow must not modify the counter"
2947 );
2948 }
2949
2950 // ---- AdvanceDenseBatch apply tests ----
2951 //
2952 // These tests pin the all-or-nothing semantics of the batch apply arm:
2953 // 1. A successful batch advances ALL counters atomically; starts are
2954 // the pre-advance values in request order (including duplicate-key
2955 // blocks which are adjacent and non-overlapping within a single batch).
2956 // 2. A cardinality violation rejects the ENTIRE batch — no counter moves,
2957 // not even keys already in the map.
2958 // 3. An overflow caused by ACCUMULATED advances across duplicate entries
2959 // rejects the ENTIRE batch — a per-entry-against-prestate check would
2960 // pass both entries individually, so the sequential fold is the only
2961 // correct approach.
2962 //
2963 // Outcome is verified by observing `dense_value` before and after each
2964 // apply (the same mechanism the single-key AdvanceDense tests use), since
2965 // `ApplyOutcome` is only delivered via the openraft responder channel,
2966 // which the test harness does not wire.
2967
2968 use crate::log_entry::DenseAdvance;
2969
2970 #[tokio::test]
2971 async fn dense_batch_applies_atomically_and_gaplessly() {
2972 // A two-key batch: both keys start absent (counter = 0) and the batch
2973 // advances "orders" by 5 and "users" by 2. After the batch both counters
2974 // must reflect exactly those advances — gaplessly starting from 0.
2975 let mut sm = HighWaterStateMachine::new_with_dense_cap(8);
2976
2977 // Pre-condition: both keys absent.
2978 assert_eq!(sm.dense_value("orders"), 0);
2979 assert_eq!(sm.dense_value("users"), 0);
2980
2981 apply_one(
2982 &mut sm,
2983 1,
2984 EntryPayload::Normal(HighWaterCommand::AdvanceDenseBatch {
2985 entries: vec![
2986 DenseAdvance {
2987 key: dense_key("orders"),
2988 count: 5,
2989 },
2990 DenseAdvance {
2991 key: dense_key("users"),
2992 count: 2,
2993 },
2994 ],
2995 }),
2996 )
2997 .await;
2998
2999 // Post-condition: both counters advanced from 0.
3000 assert_eq!(
3001 sm.dense_value("orders"),
3002 5,
3003 "first batch: orders must advance from 0 to 5"
3004 );
3005 assert_eq!(
3006 sm.dense_value("users"),
3007 2,
3008 "first batch: users must advance from 0 to 2"
3009 );
3010 // High-water must be untouched.
3011 assert_eq!(
3012 sm.current_value().await,
3013 0,
3014 "AdvanceDenseBatch must not touch the high-water"
3015 );
3016
3017 // Second batch: advance "orders" by 3 — start must be 5 (current value).
3018 apply_one(
3019 &mut sm,
3020 2,
3021 EntryPayload::Normal(HighWaterCommand::AdvanceDenseBatch {
3022 entries: vec![DenseAdvance {
3023 key: dense_key("orders"),
3024 count: 3,
3025 }],
3026 }),
3027 )
3028 .await;
3029
3030 assert_eq!(
3031 sm.dense_value("orders"),
3032 8,
3033 "second batch: orders must advance from 5 to 8"
3034 );
3035 assert_eq!(
3036 sm.dense_value("users"),
3037 2,
3038 "users must be unchanged by a batch that does not mention it"
3039 );
3040 }
3041
3042 #[tokio::test]
3043 async fn dense_batch_cardinality_is_atomic() {
3044 // Cap = 1: only one distinct key is permitted. Seed "a" with a first
3045 // batch (1 key, accepted), then submit a batch of ["a", "b"] — "b" is
3046 // new and would push the count to 2 > cap. The ENTIRE batch must be
3047 // rejected atomically: "a" must NOT advance and "b" must remain absent.
3048 let mut sm = HighWaterStateMachine::new_with_dense_cap(1);
3049
3050 // Seed "a" at 1.
3051 apply_one(
3052 &mut sm,
3053 1,
3054 EntryPayload::Normal(HighWaterCommand::AdvanceDenseBatch {
3055 entries: vec![DenseAdvance {
3056 key: dense_key("a"),
3057 count: 1,
3058 }],
3059 }),
3060 )
3061 .await;
3062 assert_eq!(sm.dense_value("a"), 1, "seed apply must succeed");
3063
3064 // The cardinality-exceeded batch: "a" already exists (no new key there)
3065 // but "b" is new and would push the distinct-key count to 2 > cap=1.
3066 apply_one(
3067 &mut sm,
3068 2,
3069 EntryPayload::Normal(HighWaterCommand::AdvanceDenseBatch {
3070 entries: vec![
3071 DenseAdvance {
3072 key: dense_key("a"),
3073 count: 1,
3074 },
3075 DenseAdvance {
3076 key: dense_key("b"),
3077 count: 1,
3078 },
3079 ],
3080 }),
3081 )
3082 .await;
3083
3084 // Atomic rejection: neither "a" nor "b" must move.
3085 assert_eq!(
3086 sm.dense_value("a"),
3087 1,
3088 "cardinality rejection must not advance existing key a"
3089 );
3090 assert_eq!(
3091 sm.dense_value("b"),
3092 0,
3093 "cardinality rejection must leave new key b absent"
3094 );
3095 }
3096
3097 #[tokio::test]
3098 async fn dense_batch_overflow_is_atomic_and_accumulates_duplicates() {
3099 // Seed "k" near the ceiling so that two entries on "k" within a single
3100 // batch — each advancing by 4 — ACCUMULATE to an overflow: the running
3101 // value after the first entry is (u64::MAX - 5 + 4) = u64::MAX - 1,
3102 // and after the second entry checked_add(u64::MAX - 1, 4) overflows.
3103 // A naive per-entry check against the pre-batch stored value (u64::MAX - 5)
3104 // would pass both (each count=4 fits), so the sequential fold over a
3105 // scratch map is the only correct implementation.
3106 //
3107 // The batch must be rejected entirely: "k" must remain at u64::MAX - 5.
3108 let mut sm = HighWaterStateMachine::new_with_dense_cap(8);
3109 {
3110 let mut core = sm.core.lock();
3111 core.dense.insert("k".to_string(), u64::MAX - 5);
3112 }
3113
3114 apply_one(
3115 &mut sm,
3116 1,
3117 EntryPayload::Normal(HighWaterCommand::AdvanceDenseBatch {
3118 entries: vec![
3119 DenseAdvance {
3120 key: dense_key("k"),
3121 count: 4,
3122 },
3123 DenseAdvance {
3124 key: dense_key("k"),
3125 count: 4,
3126 },
3127 ],
3128 }),
3129 )
3130 .await;
3131
3132 // Entire batch must be rejected: k stays at u64::MAX - 5.
3133 assert_eq!(
3134 sm.dense_value("k"),
3135 u64::MAX - 5,
3136 "overflow from accumulated duplicate advances must not modify any counter"
3137 );
3138 }
3139
3140 // ---- Regression test: install_snapshot advances active_write_version ----
3141 //
3142 // Bug: a follower that received a v5 snapshot whose `SetFormatVersion` entry
3143 // had already been purged into the snapshot would keep `active_write_version`
3144 // at `BASELINE_WRITE_VERSION` (4) while now holding a populated dense map.
3145 // The next `build_snapshot` call would then hit the fail-loud
3146 // `NotRepresentable` guard in `encode_version`'s v4 arm and return an error.
3147 //
3148 // Fix: after decoding the incoming snapshot, `install_snapshot` now advances
3149 // the `active_write_version` cell to the snapshot's leading version byte if
3150 // it exceeds the current cell value.
3151
3152 #[tokio::test]
3153 async fn install_snapshot_advances_active_write_version_to_snapshot_byte() {
3154 // ---- Step 1: build a v5-framed snapshot on a source SM ----
3155 //
3156 // Create a source SM whose active_write_version is DENSE_WRITE_VERSION
3157 // (5) and whose dense map is non-empty — the combination that would have
3158 // triggered the NotRepresentable guard on the follower.
3159 let v5_cell = ActiveWriteVersion::default();
3160 v5_cell.set(DENSE_WRITE_VERSION);
3161
3162 let store_src: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
3163 let mut src =
3164 HighWaterStateMachine::with_store_and_active_version(store_src, v5_cell.clone())
3165 .expect("source SM");
3166
3167 // Populate the dense map and advance the high-water so the snapshot is
3168 // non-trivial. `with_store_and_active_version` starts at the genesis
3169 // cardinality cap; apply an AdvanceDense to ensure dense is non-empty.
3170 apply_one(
3171 &mut src,
3172 1,
3173 EntryPayload::Normal(HighWaterCommand::AdvanceDense {
3174 key: dense_key("orders"),
3175 count: 42,
3176 }),
3177 )
3178 .await;
3179
3180 assert_eq!(src.active_write_version(), DENSE_WRITE_VERSION);
3181 assert_eq!(src.dense_value("orders"), 42, "source dense populated");
3182
3183 // Build the v5 snapshot: the bytes handed back by openraft are what
3184 // gets streamed to the follower via `begin_receiving_snapshot` +
3185 // `install_snapshot`.
3186 let snap = src
3187 .build_snapshot()
3188 .await
3189 .expect("build_snapshot on source");
3190 let snap_meta = snap.meta.clone();
3191 // `snap.snapshot` is a `Cursor<Vec<u8>>`; the inner Vec is the framed
3192 // payload bytes that `install_snapshot` receives.
3193 let snap_bytes = snap.snapshot.into_inner();
3194
3195 // Confirm the snapshot is v5-framed (the first byte is the version).
3196 assert_eq!(
3197 snap_bytes[0], DENSE_WRITE_VERSION,
3198 "source snapshot must be v5-framed"
3199 );
3200
3201 // ---- Step 2: build a fresh target SM at BASELINE_WRITE_VERSION ----
3202 //
3203 // This models a newly-joining follower that never applied the
3204 // SetFormatVersion entry (it was already purged into the snapshot it is
3205 // about to receive).
3206 let mut target = HighWaterStateMachine::new();
3207 assert_eq!(
3208 target.active_write_version(),
3209 BASELINE_WRITE_VERSION,
3210 "target must start at baseline (the pre-bug state)"
3211 );
3212 assert_eq!(target.dense_value("orders"), 0, "target dense starts empty");
3213
3214 // ---- Step 3: install the v5 snapshot onto the target ----
3215 target
3216 .install_snapshot(&snap_meta, std::io::Cursor::new(snap_bytes))
3217 .await
3218 .expect("install_snapshot must succeed");
3219
3220 // ---- Step 4: regression assertions ----
3221
3222 // PRIMARY: the cell must advance to DENSE_WRITE_VERSION after install.
3223 // Before the fix this stayed at BASELINE_WRITE_VERSION (4) — the exact
3224 // bug being pinned by this test.
3225 assert_eq!(
3226 target.active_write_version(),
3227 DENSE_WRITE_VERSION,
3228 "install_snapshot must advance active_write_version to the snapshot's \
3229 leading version byte (regression: was left at BASELINE before fix)"
3230 );
3231
3232 // Dense state must be restored from the snapshot.
3233 assert_eq!(
3234 target.dense_value("orders"),
3235 42,
3236 "dense map must be restored from the installed snapshot"
3237 );
3238
3239 // CRITICAL: build_snapshot on the target must now SUCCEED. Before the
3240 // fix, this would error with NotRepresentable because the v4 encode arm
3241 // would be selected (cell=4) while dense is non-empty.
3242 target
3243 .build_snapshot()
3244 .await
3245 .expect("build_snapshot on target after install must NOT return NotRepresentable");
3246 }
3247
3248 // ---- Regression test: with_store recovers active_write_version ----
3249 //
3250 // Sibling of `install_snapshot_advances_active_write_version_to_snapshot_byte`.
3251 // Bug: the public `with_store` / `with_store_and_active_version` load path
3252 // rehydrated `core.dense` (and `dense_cap`) from a persisted v5 snapshot but
3253 // left the freshly-defaulted `active_write_version` cell at
3254 // `BASELINE_WRITE_VERSION` (4). The next `build_snapshot` would then select
3255 // the v4 encode arm with a non-empty dense map and hit the fail-loud
3256 // `NotRepresentable` guard; `advance_dense` would also stay gated.
3257 //
3258 // Production dodges this because standalone bootstrap recovers the cell from
3259 // log + snapshot evidence *before* constructing the SM, but the public
3260 // `with_store` API documents durable-snapshot rehydration, so the bare path
3261 // must recover the version byte too (a safe floor: the snapshot byte proves
3262 // an activation to at least that version committed cluster-wide).
3263 //
3264 // Fix: after `store.load()`, advance the cell to the persisted envelope's
3265 // leading version byte when it exceeds the cell's current value.
3266
3267 #[tokio::test]
3268 async fn with_store_recovers_active_write_version_from_v5_snapshot() {
3269 // ---- Step 1: persist a v5/dense snapshot through a source SM ----
3270 let v5_cell = ActiveWriteVersion::default();
3271 v5_cell.set(DENSE_WRITE_VERSION);
3272
3273 let store: Arc<dyn SnapshotStore> = Arc::new(InMemorySnapshotStore::new());
3274 {
3275 let mut src = HighWaterStateMachine::with_store_and_active_version(
3276 store.clone(),
3277 v5_cell.clone(),
3278 )
3279 .expect("source SM");
3280 apply_one(
3281 &mut src,
3282 1,
3283 EntryPayload::Normal(HighWaterCommand::AdvanceDense {
3284 key: dense_key("orders"),
3285 count: 42,
3286 }),
3287 )
3288 .await;
3289 src.build_snapshot().await.expect("build v5 snapshot");
3290 }
3291
3292 // Confirm the persisted envelope is v5-framed.
3293 let persisted_bytes = store.load().expect("load").expect("snapshot present");
3294 assert_eq!(
3295 persisted_bytes[0], DENSE_WRITE_VERSION,
3296 "persisted envelope must be v5-framed"
3297 );
3298
3299 // ---- Step 2: reopen via the bare public API (fresh BASELINE cell) ----
3300 // This models an embedder that uses `with_store` with a durable backend
3301 // holding v5 data — the pre-bug state left the cell at baseline.
3302 let mut reopened =
3303 HighWaterStateMachine::with_store(store.clone()).expect("reopen via with_store");
3304
3305 // ---- Step 3: regression assertions ----
3306
3307 // PRIMARY: the cell must recover to DENSE_WRITE_VERSION from the
3308 // snapshot's leading byte. Before the fix this stayed at BASELINE (4).
3309 assert_eq!(
3310 reopened.active_write_version(),
3311 DENSE_WRITE_VERSION,
3312 "with_store must recover active_write_version from the persisted \
3313 snapshot's leading version byte (regression: was left at BASELINE)"
3314 );
3315
3316 // Dense state must be rehydrated from the snapshot.
3317 assert_eq!(
3318 reopened.dense_value("orders"),
3319 42,
3320 "dense map must be rehydrated from the persisted snapshot"
3321 );
3322
3323 // CRITICAL: build_snapshot must now SUCCEED. Before the fix, the v4
3324 // encode arm (cell=4) with a non-empty dense map hit NotRepresentable.
3325 reopened
3326 .build_snapshot()
3327 .await
3328 .expect("build_snapshot after reopen must NOT return NotRepresentable");
3329 }
3330}