Fix checkpoint failure when serializing a concurrently-mutated object (#2101)

libs/server/Objects/Hash/HashObjectImpl.cs

@@ -456,8 +456,15 @@ private void HashExpire(ref ObjectInput input, ref ObjectOutput output, byte res
 
         private void HashTimeToLive(ref ObjectInput input, ref ObjectOutput output, byte respProtocolVersion)
         {
-            DeleteExpiredItems();
-
+            // This is a read operation, so it must not mutate the object: it runs under a shared lock and may execute
+            // while the flush path is serializing this same instance. Expired fields are reported as absent by
+            // ContainsKey (via GetExpiration) without being removed; the mutating paths purge them from the live object.
+            //
+            // Retention is bounded, not unbounded: a field can only become expirable through HEXPIRE/HPEXPIRE, which is
+            // an RMW that itself calls DeleteExpiredItems() first. So the expired-but-unpurged set can never exceed the
+            // fields that were carrying a TTL at the last mutating operation, and a read-only workload cannot grow it.
+            // That memory is reclaimed by the next mutating operation, by HCOLLECT, or by the background collector
+            // (ExpiredObjectCollectionFrequencySecs). This matches SortedSetTimeToLive (ZTTL), which has never purged.
             var isMilliseconds = input.arg1 == 1;
             var isTimestamp = input.arg2 == 1;
             var numFields = input.parseState.Count;

libs/storage/Tsavorite/cs/src/core/Allocator/HeapObjectBase.cs

@@ -62,10 +62,18 @@ public void Serialize(BinaryWriter writer)
                 // CopyUpdater does that, as it must ensure the object's (v1) data is not changed during the checkpoint.
                 if (SerializationPhase == SerializationPhase.REST && MakeTransition(SerializationPhase.REST, SerializationPhase.SERIALIZING))
                 {
-                    // Directly serialize to wire, do not cache serialized state
-                    WriteType(writer, isNull: false);
-                    DoSerialize(writer);
-                    SerializationPhase = SerializationPhase.REST;
+                    // Directly serialize to wire, do not cache serialized state. Restore REST even if DoSerialize
+                    // throws: leaving the object in SERIALIZING would make every later Serialize() and
+                    // CacheSerializedObjectData() spin on it forever, wedging the checkpoint pipeline.
+                    try
+                    {
+                        WriteType(writer, isNull: false);
+                        DoSerialize(writer);
+                    }
+                    finally
+                    {
+                        SerializationPhase = SerializationPhase.REST;
+                    }
                     return;
                 }
 
@@ -144,10 +152,20 @@ public void CacheSerializedObjectData(ref LogRecord dstLogRecord, ref RMWInfo rm
             {
                 if (SerializationPhase == (int)SerializationPhase.REST && MakeTransition(SerializationPhase.REST, SerializationPhase.SERIALIZING))
                 {
-                    using var ms = new MemoryStream();
-                    using var writer = new BinaryWriter(ms, Encoding.UTF8);
-                    DoSerialize(writer);
-                    serializedBytes = ms.ToArray();
+                    try
+                    {
+                        using var ms = new MemoryStream();
+                        using var writer = new BinaryWriter(ms, Encoding.UTF8);
+                        DoSerialize(writer);
+                        serializedBytes = ms.ToArray();
+                    }
+                    catch
+                    {
+                        // Publish no partial capture, and do not strand the object in SERIALIZING.
+                        serializedBytes = null;
+                        SerializationPhase = SerializationPhase.REST;
+                        throw;
+                    }
 
                     SerializationPhase = SerializationPhase.SERIALIZED;    // This is the only place .SERIALIZED is set
                     break;
@@ -167,9 +185,18 @@ public void CacheSerializedObjectData(ref LogRecord dstLogRecord, ref RMWInfo rm
         /// <inheritdoc />
         public void ClearSerializedObjectData()
         {
-            // Do not disturb an unrelated object that is currently serializing.
-            if (Interlocked.Exchange(ref serializedBytes, null) is not null)
-                SerializationPhase = SerializationPhase.REST;
+            // Release the cached (v) bytes so they can be GC'd, but deliberately leave the phase terminal.
+            //
+            // Bytes are only ever cached by CacheSerializedObjectData for a source object that a CopyUpdate
+            // has just superseded, and Clone() is a shallow copy: the (v+1) record that replaced this one
+            // shares this object's internal collections and keeps mutating them. Returning to REST would let
+            // a later checkpoint take the direct-serialize path and enumerate those still-shared, still-live
+            // collections with no synchronization against the writer.
+            //
+            // SERIALIZED with null bytes is the state Serialize() already documents as "superseded after
+            // checkpoint completion": it writes a null indicator. That is safe because the superseding record
+            // always sits at a higher address and carries the live data.
+            _ = Interlocked.Exchange(ref serializedBytes, null);
         }
     }
 }

libs/storage/Tsavorite/cs/src/core/Index/Checkpointing/StateMachineDriver.cs

@@ -15,21 +15,41 @@ namespace Tsavorite.core
     /// </summary>
     public class StateMachineDriver
     {
+        // Globally published phase and version.
         SystemState systemState;
+
+        // The single state machine currently owning this driver; null while idle.
         IStateMachine stateMachine;
+
+        // Already-started tasks that must complete before the driver can leave the current phase.
+        // ProcessWaitingListAsync awaits these only after the transition-in epoch barrier completes.
         readonly List<(Task task, StateMachineTaskType type)> waitingList;
+
+        // Completion source for the entire state-machine run, not an individual phase transition.
         TaskCompletionSource<bool> stateMachineCompleted;
-        // All threads have entered the given state
+
+        // Semaphore associated with the currently published state. MakeTransitionWorker releases it after
+        // prior-epoch participants have advanced or suspended and GlobalAfterEnteringState has completed.
+        // This does not include completion of tasks in waitingList.
         SemaphoreSlim waitForTransitionIn;
+
+        // GlobalAfterEnteringState may run on an arbitrary epoch-drain thread, so its exception is captured
+        // here and rethrown by ProcessWaitingListAsync on the state-machine driver path.
         Exception waitForTransitionInException;
-        // All threads have exited the given state
+
+        // Semaphore associated with the currently published state. GlobalStateMachineStep releases it as
+        // soon as the next state is published; it does not wait for the transition-in epoch barrier.
         SemaphoreSlim waitForTransitionOut;
-        // Transactions drained in last version
+
+        // Version whose active transactions must drain before the state machine can advance.
         long lastVersion;
         TaskCompletionSource<bool> lastVersionTransactionsDone;
+
         List<IStateMachineCallback> callbacks;
         readonly LightEpoch epoch;
         readonly ILogger logger;
+
+        // Active transaction counts are indexed by version parity; only two adjacent versions can be active.
         readonly long[] NumActiveTransactions;
 
         public SystemState SystemState => SystemState.Copy(ref systemState);
@@ -157,6 +177,7 @@ public void EndTransaction(long txnVersion)
 
         internal void AddToWaitingList(Task waiter, StateMachineTaskType type)
         {
+            // Callers start the operation before registering it. The driver awaits it after transition-in.
             if (waiter != null)
                 waitingList.Add((waiter, type));
         }
@@ -219,22 +240,25 @@ void GlobalStateMachineStep(SystemState expectedState)
 
             var nextState = stateMachine.NextState(systemState);
 
+            // Run task-specific work while systemState still identifies the previous phase.
             stateMachine.GlobalBeforeEnteringState(nextState, this);
 
-            // Execute any additional registered callbacks
+            // External callbacks have the same before-publication ordering as the state-machine task hooks.
             if (callbacks != null)
             {
                 foreach (var callback in callbacks)
                     callback.BeforeEnteringState(nextState);
             }
 
-            // Write new phase
+            // Publish the new phase and version so subsequent session refreshes observe nextState.
             systemState.Word = nextState.Word;
 
-            // Release waiters for new phase
+            // Release the semaphore associated with the phase just exited.
             _ = waitForTransitionOut?.Release(int.MaxValue);
 
-            // Write new semaphores
+            // Install semaphores for the newly published phase. Its transition-out semaphore is released
+            // when the following phase is published; its transition-in semaphore is released by the
+            // epoch-drain callback below. These assignments occur after systemState is published.
             waitForTransitionOut = new SemaphoreSlim(0);
             waitForTransitionIn = new SemaphoreSlim(0);
 
@@ -244,6 +268,9 @@ void GlobalStateMachineStep(SystemState expectedState)
             try
             {
                 epoch.Resume();
+
+                // Associate MakeTransitionWorker with the prior epoch. It becomes eligible only after
+                // participants still announcing that epoch have advanced through ProtectAndDrain or suspended.
                 epoch.BumpCurrentEpoch(() => MakeTransitionWorker(nextState));
             }
             finally
@@ -259,6 +286,7 @@ void GlobalStateMachineStep(SystemState expectedState)
         /// <returns></returns>
         public async Task WaitForStateChange(SystemState currentState)
         {
+            // Capture before rechecking state so a racing transition that releases this semaphore is observed.
             var _waitForTransitionOut = waitForTransitionOut;
             if (SystemState.Equal(currentState, systemState))
             {
@@ -273,7 +301,11 @@ public async Task WaitForStateChange(SystemState currentState)
         /// <returns></returns>
         public async Task WaitForCompletion(SystemState currentState)
         {
+            // First wait until currentState is no longer published.
             await WaitForStateChange(currentState).ConfigureAwait(false);
+
+            // Then capture the newly published state and wait until its epoch transition and
+            // GlobalAfterEnteringState hooks complete. Phase waiting-list tasks are not included.
             currentState = systemState;
             var _waitForTransitionIn = waitForTransitionIn;
             if (SystemState.Equal(currentState, systemState))
@@ -286,29 +318,35 @@ void MakeTransitionWorker(SystemState nextState)
         {
             try
             {
+                // This is an epoch-drain action and may execute synchronously from BumpCurrentEpoch
+                // or later on any thread that advances or suspends epoch protection.
                 stateMachine.GlobalAfterEnteringState(nextState, this);
             }
             catch (Exception e)
             {
-                // Store the exception to be thrown by state machine driver
-                // We do not throw here as this epoch action may be executed in a different thread context
+                // Propagate on the driver path rather than throwing on an arbitrary epoch-drain thread.
                 waitForTransitionInException = e;
 
                 logger?.LogError(e, "Exception in state machine transition worker");
             }
             finally
             {
+                // Signal that the epoch transition and all after-transition hooks have finished.
                 waitForTransitionIn.Release(int.MaxValue);
             }
         }
 
         async Task ProcessWaitingListAsync(CancellationToken token = default)
         {
+            // Do not process phase tasks until the prior epoch has drained and after-transition hooks finish.
             await waitForTransitionIn.WaitAsync(token).ConfigureAwait(false);
             if (waitForTransitionInException != null)
             {
                 throw waitForTransitionInException;
             }
+
+            // These tasks were started by state-machine hooks and may have progressed concurrently with
+            // the epoch transition. Awaiting them here prevents the driver from publishing the next phase.
             foreach (var (task, type) in waitingList)
             {
                 try
@@ -331,6 +369,7 @@ async Task RunStateMachine(CancellationToken token = default)
             {
                 do
                 {
+                    // Publish one transition, then wait for both transition-in and its registered phase work.
                     GlobalStateMachineStep(systemState);
                     await ProcessWaitingListAsync(token).ConfigureAwait(false);
                 } while (systemState.Phase != Phase.REST);
@@ -377,14 +416,14 @@ void FastForwardStateMachineToRest()
             if (waitForTransitionOut?.CurrentCount == 0)
                 _ = waitForTransitionOut?.Release(int.MaxValue);
 
-            // Clear semaphores
+            // Failure recovery does not execute skipped transition hooks. Discard their synchronization state.
             waitForTransitionOut = null;
             waitForTransitionIn = null;
 
-            // Clear exception if any
+            // Clear any exception captured from an after-transition hook.
             waitForTransitionInException = null;
 
-            // Clear waiting list
+            // The failed run no longer waits for phase-specific asynchronous work.
             waitingList.Clear();
         }
     }

libs/storage/Tsavorite/cs/src/core/Index/Tsavorite/Implementation/BlockAllocate.cs

@@ -60,6 +60,17 @@ bool TryAllocateRecord<TInput, TOutput, TContext, TSessionFunctionsWrapper>(TSes
             if ((minRevivAddress <= stackCtx.hei.Address) && (!options.elideSourceRecord || stackCtx.hei.Address != stackCtx.recSrc.LogicalAddress))
                 minRevivAddress = stackCtx.hei.Address;
 
+            // A (v+1) record must never land below the checkpoint's fuzzy-region start: the snapshot includes every record
+            // below that address regardless of its version bit, so reusing a lower address would leak (v+1) data into the
+            // checkpoint image. This applies to both reuse paths below, including an allocation that was saved before the
+            // version shift and is only now being retried as (v+1).
+            if (sessionFunctions.Ctx.IsInV1)
+            {
+                var fuzzyStartAddress = _hybridLogCheckpoint.info.startLogicalAddress;
+                if (fuzzyStartAddress > minRevivAddress)
+                    minRevivAddress = fuzzyStartAddress;
+            }
+
             if (options.recycle && operationState.retryNewLogicalAddress != kInvalidAddress
                     && GetAllocationForRetry(sessionFunctions, ref operationState, minRevivAddress, in sizeInfo, out newLogicalAddress, out newPhysicalAddress))
             {
@@ -68,12 +79,6 @@ bool TryAllocateRecord<TInput, TOutput, TContext, TSessionFunctionsWrapper>(TSes
             }
             if (RevivificationManager.UseFreeRecordPool)
             {
-                if (sessionFunctions.Ctx.IsInV1)
-                {
-                    var fuzzyStartAddress = _hybridLogCheckpoint.info.startLogicalAddress;
-                    if (fuzzyStartAddress > minRevivAddress)
-                        minRevivAddress = fuzzyStartAddress;
-                }
                 if (TryTakeFreeRecord<TInput, TOutput, TContext, TSessionFunctionsWrapper>(sessionFunctions, in sizeInfo, minRevivAddress, out newLogicalAddress, out newPhysicalAddress))
                 {
                     new LogRecord(newPhysicalAddress).PrepareForRevivification(ref sizeInfo);

libs/storage/Tsavorite/cs/src/core/Index/Tsavorite/Implementation/Helpers.cs

@@ -106,6 +106,27 @@ private bool IsFrozen<TInput, TOutput, TContext, TSessionFunctionsWrapper>(TSess
         internal long GetMinRevivifiableAddress()
             => RevivificationManager.GetMinRevivifiableAddress(hlogBase.GetTailAddress(), hlogBase.ReadOnlyAddress);
 
+        /// <summary>
+        /// Dispose the resources of an in-memory source record that a newly-CAS'd record has just superseded, unless an
+        /// ongoing checkpoint has frozen it.
+        /// </summary>
+        /// <remarks>
+        /// Disposal clears the record's heap fields, which returns the value's <see cref="ObjectIdMap"/> slot to that page's
+        /// free list for reuse by another record. The snapshot flush reads object ids from its page copy but resolves them
+        /// against the live map, so disposing a frozen record lets the flush serialize a freed - or recycled, and therefore
+        /// unrelated - object in its place. A frozen record must keep its value until the checkpoint has captured it; the
+        /// value is then accounted for and released when the page is evicted.
+        /// </remarks>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private void OnDisposeSupersededSource<TInput, TOutput, TContext, TSessionFunctionsWrapper>(TSessionFunctionsWrapper sessionFunctions,
+                ref OperationStackContext<TStoreFunctions, TAllocator> stackCtx, ref LogRecord logRecord)
+            where TSessionFunctionsWrapper : ISessionFunctionsWrapper<TInput, TOutput, TContext, TStoreFunctions, TAllocator>
+        {
+            if (IsFrozen<TInput, TOutput, TContext, TSessionFunctionsWrapper>(sessionFunctions, ref stackCtx, logRecord.Info))
+                return;
+            OnDispose(ref logRecord, DisposeReason.Deleted);
+        }
+
         [MethodImpl(MethodImplOptions.NoInlining)]
         private (bool elided, bool added) TryElideAndTransferToFreeList<TInput, TOutput, TContext, TSessionFunctionsWrapper>(TSessionFunctionsWrapper sessionFunctions,
                 ref OperationStackContext<TStoreFunctions, TAllocator> stackCtx, ref LogRecord logRecord)