Skip to main content

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}