[Tsavorite] Fix checkpoint abort abandoning in-flight flushes and stale checkpoint selection on recovery - #2172

Two independent defects in Tsavorite's checkpoint/recovery paths, both found while triaging intermittent CI failures. They are unrelated in mechanism but adjacent in code, and both are silent — each produces a wrong result rather than an error.

1. An aborted checkpoint releases the devices its flushes are still writing to, so orphaned completions land on the next checkpoint's state.
2. Recovering to "the latest checkpoint" could silently recover an older one, losing every update made after it.

Part 1 — Aborted checkpoints abandon their in-flight flushes

Root Cause

An aborted checkpoint state machine releases the devices its flushes are still writing to without waiting for those flushes to complete. The completion callbacks then report through flush state that is shared across checkpoints — mainIndexCheckpointCallbackCount, mainIndexCheckpointError and mainIndexCheckpointTcs in TsavoriteBase, and the checkpointCallbackCount / checkpointError / checkpointTcs trio in MallocFixedPageSize. Both BeginMainIndexCheckpoint and BeginCheckpoint reset that state at the start of every checkpoint, so a completion that outlives its own checkpoint lands on the next one's state.

Three windows leave a flush unawaited:

1. Index flush. IndexCheckpointSMTask issues TakeIndexFuzzyCheckpoint() at Phase.PREPARE but only registers the resulting tasks on the driver's waiting list at Phase.WAIT_INDEX_CHECKPOINT. An abort in between reaches OnAbort, whose store._indexCheckpoint.Reset() disposes main_ht_device under the live writes.

2. Snapshot / foldover flush. StateMachineDriver.ProcessWaitingListAsync rethrows waitForTransitionInException before it iterates the waiting list, so an abort around Phase.WAIT_FLUSH reaches HybridLogCheckpointSMTask.OnAbort with the snapshot flush still running, and store._hybridLogCheckpoint.Dispose() releases the snapshot devices and object-log flush buffers it is writing through.

3. Issuance failure. BeginMainIndexCheckpoint completed its task with TrySetException when a chunk could not be submitted, while the chunks it had already issued were still in flight — unlike MallocFixedPageSize.BeginCheckpoint, which correctly retires only the levels that were never issued. This also meant the waits added for (1) and (2) could return early.

Depending on when the orphaned completion lands it either records an error into the next checkpoint's state (the ERR checkpoint failed, check server logs seen on CI), completes that checkpoint's flush before its data is on disk, or drives the outstanding count negative so the flush task never completes at all. Only the first is loud.

Making the flush task mean "every write this checkpoint issued has called back" exposed two further ordering defects in the issue loops themselves, both of which had to be fixed for the waits above to be trustworthy:

• A device may invoke a write's completion callback synchronously and then throw out of the same WriteAsync. If that write was the last one outstanding, the callback drove the count to zero and resolved the task with success before the catch recorded the error — a failed index checkpoint reporting success.
• In MallocFixedPageSize.BeginCheckpoint, the error path advanced the loop variable (for (i++; ...)) before retiring the remaining levels. A throw from the sector-aligned staging code, which runs outside the inner try, therefore skipped the failing level's own retire, the count never reached zero, and the checkpoint hung.

Description of Change

Both OnAbort paths now await the flush they left outstanding before releasing what it is writing to, discarding its outcome — the checkpoint has already failed, and the exception that aborted it is the actionable one. The two index-checkpoint issue loops are restructured so the flush task always means "every write this checkpoint issued has called back", and can never be resolved while issuance is still in progress.

Key changes:

• IndexCheckpointSMTask.OnAbort calls the new TsavoriteBase.WaitForIndexCheckpointFlushCompletion() before _indexCheckpoint.Reset().
• HybridLogCheckpointSMTask.OnAbort calls the new TsavoriteBase.WaitForCheckpointFlush(...) on _hybridLogCheckpoint.flushedTask before _hybridLogCheckpoint.Dispose().
• Both BeginMainIndexCheckpoint and MallocFixedPageSize.BeginCheckpoint hold an issuance sentinel: the outstanding count starts at numChunks + 1 (numLevels + 1 for the overflow buckets) and the sentinel is retired in a finally, after issuance has ended and any error has been recorded. A synchronous callback can no longer resolve the task ahead of the catch that records the failure.
• BeginMainIndexCheckpoint tracks how many chunks are accounted for (a callback is pending, or the submit failed and the catch retired it in place), records the error before retiring, and retires only the chunks that were counted but never issued — letting the last real completion finish the task.
• MallocFixedPageSize.BeginCheckpoint tracks accounted levels the same way instead of skipping the failing level with for (i++; ...), so a throw from the staging code outside the inner try still retires its own level.
• New RetireMainIndexCheckpointChunk() extracts the decrement-and-complete logic previously inline in AsyncPageFlushCallback, so issuance and completion retire through one path.
• MallocFixedPageSize.GetCheckpointTask() returns checkpointTcs?.Task; it previously threw NullReferenceException before the first checkpoint, which the new wait would hit.

Blocking in OnAbort is safe: it runs on the state machine's own thread with no epoch held — the same thread that would have awaited those tasks on the normal path — and IStateMachineTask.OnAbort is a default interface method, so it stays synchronous and the public API is unchanged.

Affected types:

• TsavoriteBase (Index/Recovery/IndexCheckpoint.cs) — owns the shared index-flush state; gains WaitForIndexCheckpointFlushCompletion(), static WaitForCheckpointFlush(Task), GetMainIndexCheckpointTask() and RetireMainIndexCheckpointChunk().
• IndexCheckpointSMTask, HybridLogCheckpointSMTask — the two OnAbort implementations.
• MallocFixedPageSize<T> — overflow-bucket allocator with the same shared-state pattern.

Edge Cases

Behavior change

A submission rejected synchronously by the device now surfaces as TsavoriteIOException wrapping the rejection, rather than the raw exception — consistent with how the identical failure already arrives from a flush completion, and the cause is preserved as the inner exception. RecoveryTestFailOnSectorSize used Assert.ThrowsAsync (exact type match) and is updated to Assert.CatchAsync<TsavoriteException>, which keeps its intent and is robust to which path reports the failure.

Part 2 — Recovery could silently restore a stale checkpoint

Root Cause

RecoverAsync() with no version argument means "recover to the latest checkpoint", and reaches FindRecoveryInfo(-1) → GetClosestHybridLogCheckpointInfo. That scan is written to examine every token and keep the highest version, precisely because — as its own comment says — the "file system is not guaranteed to return tokens in order of freshness".

But the loop also carried an early exit for the case where an exact target version has been found:

if (current.info.nextVersion > requestedVersion)
    break;   // no closer version can exist

nextVersion is always >= 1, so with requestedVersion == -1 that test is true for every checkpoint. The scan broke on the first token and returned whatever the checkpoint manager enumerated first.

That enumeration is not ordered by checkpoint age. LocalStorageNamedDeviceFactory.ListContents sorts the checkpoint directories by LastWriteTime, which comes from the system clock; its file-time granularity is coarse (~15.6 ms on Windows), so two checkpoints taken in quick succession can carry the same timestamp. OrderByDescending is a stable sort, so ties fall back to the order the filesystem returned — on NTFS, by name, effectively the GUID, which is unrelated to when the checkpoint was taken. The older checkpoint then wins roughly half the time, and every update made after it is silently lost.

This is what made NamespaceTests.RecoveryAsync(*, MLSD) fail intermittently on Windows CI with a post-recovery value mismatch.

Description of Change

Restrict the early exit to a real requested version, so the -1 case runs the full scan. No new selection logic was needed: distanceToTarget is computed as long.MaxValue - version when requestedVersion == -1, and closestVersion starts at long.MaxValue, so the existing accumulate logic below already keeps the smallest distance and therefore the highest version.

if (requestedVersion != -1 && current.info.nextVersion > requestedVersion)

The same branch also overwrote closest without disposing it, leaking the snapshot file device of any candidate the scan had already accepted — the accumulate branch below it disposes correctly. closest.Dispose() is added there; HybridLogCheckpointInfo.Dispose() is null-safe and ends with this = default, so it is safe on a not-yet-assigned candidate.

Blast radius. GetLatestCheckpointTokens and GetLatestCheckpointVersion route through the same helper, so both were affected. The cluster replication path in CheckpointStore.GetLatestCheckpointEntryFromDisk — documented as returning "the latest checkpoint entry from disk by scanning all available tokens" — could likewise hand a replica a stale checkpoint.

Not a bug: GetClosestIndexCheckpointInfo picks the first compatible index checkpoint (IsCompatible: indexInfo.finalLogicalAddress <= recoveryInfo.finalLogicalAddress). An older compatible index checkpoint is safe by design, because recovery replays the log forward from it.

What NOT to Do (for future agents)

• ❌ Don't complete a checkpoint's flush task to "unblock" a failure while writes it issued are still in flight. That is precisely the bug: it releases the waiters that dispose the device those writes target, and the resulting completions are charged to whatever checkpoint owns the shared state by then.
• ❌ Don't assume the driver's waiting list covers every flush. The index flush is issued one phase before it is registered, and ProcessWaitingListAsync rethrows the transition exception before it ever reaches the list.
• ❌ Don't treat a synchronous throw out of WriteAsync as "the write never happened." A device may invoke the completion callback and then throw; the TryClaimIoUnitRelease / TryClaimRetirement guards are what distinguish the two, and retiring without them double-counts.
• ❌ Don't let the outstanding count reach zero while the issue loop is still running. Without the issuance sentinel, a synchronous completion resolves the task before the catch records the error, and the checkpoint reports success for a write that failed.
• ❌ Don't make OnAbort async. It is a default interface method on a public interface; changing its signature is a breaking API change, and blocking there is already correct.
• ❌ Don't rely on checkpoint enumeration order to mean "newest first." ListContents sorts by a coarse filesystem timestamp with a stable sort, so equal timestamps fall back to name order. Checkpoint version is the only ordering that means anything.
• ❌ Don't "fix" this in ListContents by sorting on something else. The selection scan already knows the versions; ordering the directory listing would only make the defect harder to hit, not impossible.

Tests

Seven new tests, each verified to fail when its corresponding product change is reverted:

Part 1

• CheckpointFailureTests.AbortedCheckpointWaitsForTheIndexCheckpointFlushItIssued — holds back ht.dat write completions, aborts at Phase.IN_PROGRESS, and asserts nothing is in flight when SAVE reports the failure. Without the fix: 2 writes outstanding.
• CheckpointFailureTests.AbortedCheckpointWaitsForTheSnapshotFlushItIssued — same for snapshot.dat, aborting at Phase.WAIT_FLUSH. Without the fix: 1 write outstanding.
• IndexCheckpointIoTests.IndexCheckpointDoesNotCompleteWhileAnIssuedChunkIsStillWriting — fails chunk 1's submission with chunk 0 still writing, and asserts the checkpoint neither completes early nor hangs on the chunks that were never issued.
• IndexCheckpointIoTests.IndexCheckpointFailsWhenTheFinalChunkCompletesThenThrows — the last chunk's write calls back and then throws; asserts the checkpoint fails. Without the sentinel it reports success.
• IndexCheckpointIoTests.OverflowBucketCheckpointFailsWhenTheFinalLevelCompletesThenThrows — the same for MallocFixedPageSize.BeginCheckpoint.

Both Garnet tests synchronize on observed state rather than a timer: they spin until the abort has fired with a write verifiably in flight, then assert SAVE stays blocked for a grace period before the deferred completions are released. A timer started before anything is deferred could drain the held writes ahead of the abort on a slow run and pass against unfixed code.

New harness DeferringCheckpointDeviceFactoryCreator withholds DeviceIOCompletionCallback invocations for a named checkpoint file, so a test can abort with writes verifiably in flight. ThrowOnNthWriteDevice gains the equivalent DeferWriteCompletions and CompleteBeforeThrowing knobs.

Part 2 — new fixture RecoverLatestVersionTests, which removes the timing dependence entirely by backdating the newer checkpoint's cpr-checkpoints and index-checkpoints directories, so enumeration is guaranteed to hand the older checkpoint over first:

• RecoverLatestIgnoresDirTimeOrder(Snapshot) and RecoverLatestIgnoresDirTimeOrder(FoldOver) — take two checkpoints writing different values, backdate the second, recover with -1, and assert the second value is read back. Without the fix both fail with Expected: 222, But was: 111 — on Linux, deterministically.
• RecoverToVersionPicksThatVersion — asserts the targeted-version path is unchanged, covering the early exit that is still reachable for a real requestedVersion.

Validation

• CheckpointFailureTests 10/10 across 10 consecutive runs.
• Full Garnet.test: 1116 passed, 0 failed.
• Tsavorite: recovery 242, core 290, hlog 580 — all green.
• Garnet.test.cluster.replication: 110 passed, 0 failed — covers CheckpointStore.GetLatestCheckpointEntryFromDisk, the cluster consumer of Part 2.
• NamespaceTests (the Part 2 CI failure) 12× locally plus 6× under taskset -c 0,1 — all pass.
• Both solutions build on net8.0 and net10.0; dotnet format --verify-no-changes clean on both.

The Part 1 CI flake did not reproduce locally on Linux across 1600 abort+retry iterations, consistent with the report; that fix rests on the deterministic tests above rather than on reproducing the original timing. Part 2 does reproduce deterministically once the enumeration order is pinned.

Issues Fixed

Fixes #2171