Fix snapshot flush and checkpoint exception propagation [dev]

[badrishc]

Fix snapshot flush ObjectIdMap use-after-free and checkpoint exception propagation

Root cause: During snapshot checkpoint flush, ObjectAllocatorImpl.WriteAsync accessed the page's ObjectIdMap after releasing epoch protection. Under low memory with small segments, HeadAddress could advance during epoch gaps, causing page eviction which clears the ObjectIdMap (tail = -1). Subsequent objectIdMap.Get() hit a Debug.Assert failure, throwing an unhandled exception in FlushRunner. Since FlushRunner ran via Task.Run, the exception was unobserved, the completion semaphore was never released, and the checkpoint hung indefinitely.

Fixes:

• ObjectAllocatorImpl: Capture objectIdMap under epoch protection with a HeadAddress guard. Skip objectIdMap dereference (KeyOverflow, ValueOverflow, ValueObject) when HeadAddress advances past the current record, marking the record invalid instead. All objectIdMap accesses now occur strictly under epoch protection.
• AllocatorBase: Change AsyncFlushPagesForSnapshot and AsyncFlushDeltaToDevice to return Task instead of SemaphoreSlim, so exceptions in FlushRunner propagate naturally through the Task to the state machine.
• StateMachineDriver: Unify waiting list as List<(StateMachineTaskType, Task)>. Add StateMachineTaskType enum for identifying faulting tasks in error logs. ProcessWaitingListAsync logs the task type before re-throwing exceptions.
• AsyncFlushPageForSnapshotCallback: Move result.Release() into finally block to ensure release even if dirty-bit clearing throws.

Test fixes:

• ClusterTestUtils.Checkpoint: Retry on 'checkpoint already in progress' (race with on-demand checkpoint from replication).
• ClusterTestContext.AttachAndWaitForSync: Fix loop bounds bug (was i < replica_count, should be i < replicaStartIndex + replicaCount). Improve parameter names for clarity.
• ClusterReplicationBaseTests: Use Task.WhenAll for parallel task waiting with diagnostic status on timeout.

[Copilot]

Pull request overview

This PR addresses a snapshot-checkpoint flush correctness issue (ObjectIdMap use-after-free under epoch gaps) and improves checkpoint failure behavior by propagating flush exceptions through the checkpoint state machine rather than leaving them unobserved/hanging.

Changes:

• Adjust snapshot flush logic in ObjectAllocatorImpl to only access ObjectIdMap under epoch protection and to handle HeadAddress advancement by skipping object serialization and invalidating affected records.
• Refactor Tsavorite checkpoint waiting from SemaphoreSlim-based completion to Task-based waiting, and enhance state-machine diagnostics via StateMachineTaskType.
• Stabilize cluster replication tests by fixing replica attach loop bounds, adding checkpoint retry on “already in progress”, and improving timeout diagnostics.

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
test/Garnet.test.cluster/ReplicationTests/ClusterReplicationBaseTests.cs	Increases test timeout, fixes shard query index, and improves parallel task waiting/diagnostics.
test/Garnet.test.cluster/ClusterTestUtils.cs	Retries checkpoint when a checkpoint is already in progress to reduce replication race flakiness.
test/Garnet.test.cluster/ClusterTestContext.cs	Fixes replica attach loop bounds and parameter clarity; aligns primary index usage in sync waits.
libs/storage/Tsavorite/cs/src/core/Utilities/PageAsyncResultTypes.cs	Updates flush completion tracking structure in support of Task-based waiting changes.
libs/storage/Tsavorite/cs/src/core/Index/Recovery/IndexCheckpoint.cs	Tags checkpoint waiters with task types for improved fault logging.
libs/storage/Tsavorite/cs/src/core/Index/CheckpointManagement/RecoveryInfo.cs	Converts checkpoint flush completion tracking field from semaphore to task.
libs/storage/Tsavorite/cs/src/core/Index/Checkpointing/StateMachineTaskType.cs	Adds enum to identify faulting checkpoint wait tasks in logs.
libs/storage/Tsavorite/cs/src/core/Index/Checkpointing/StateMachineDriver.cs	Switches waiting list to typed tasks and logs task type on failure before rethrowing.
libs/storage/Tsavorite/cs/src/core/Index/Checkpointing/SnapshotCheckpointSMTask.cs	Waits on Task-based snapshot flush completion and tags it for diagnostics.
libs/storage/Tsavorite/cs/src/core/Index/Checkpointing/IncrementalSnapshotCheckpointSMTask.cs	Waits on Task-based delta flush completion and tags it for diagnostics.
libs/storage/Tsavorite/cs/src/core/Index/Checkpointing/FoldOverSMTask.cs	Converts foldover flush completion semaphore to a Task waiter and tags it for diagnostics.
libs/storage/Tsavorite/cs/src/core/Allocator/ObjectAllocatorImpl.cs	Fixes epoch scoping around ObjectIdMap usage during snapshot flush to avoid use-after-free.
libs/storage/Tsavorite/cs/src/core/Allocator/AllocatorBase.cs	Refactors snapshot/delta flush completion to Task-based signaling and adjusts callback release behavior.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 14 out of 14 changed files in this pull request and generated 2 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 3 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (1)

libs/storage/Tsavorite/cs/src/core/Index/Recovery/IndexCheckpoint.cs:86

• BeginMainIndexCheckpoint() uses a fire-and-forget Task.Run(FlushRunner) when throttling is enabled, but FlushRunner exceptions are not observed and mainIndexCheckpointTcs is never faulted/completed in that case. That can leave WAIT_INDEX_CHECKPOINT waiting forever. Consider wrapping FlushRunner in try/catch that calls mainIndexCheckpointTcs.TrySetException(ex), and/or storing/awaiting the Task so faults propagate (similar to the other flush paths that now return Tasks).

            mainIndexCheckpointTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);

            if (throttleCheckpointFlushDelayMs >= 0)
                Task.Run(FlushRunner);
            else
                FlushRunner();

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[vazois]

LGTM!