Reset state machine on exception during operation (#1497)

* Fix exception paths in state machine driver

* update version

* improvement

* make last version stuff private fields

* refactor

Author: badrishc

Version.props

@@ -1,6 +1,6 @@
 <Project>
 	<!-- VersionPrefix property for builds and packages -->
 	<PropertyGroup>
-		<VersionPrefix>1.0.91</VersionPrefix>
+		<VersionPrefix>1.0.92</VersionPrefix>
 	</PropertyGroup>
 </Project>

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

@@ -47,8 +47,7 @@ public virtual void GlobalBeforeEnteringState(SystemState next, StateMachineDriv
                     store.CheckpointVersionShiftEnd(lastVersion, next.Version, isStreaming);
 
                     Debug.Assert(stateMachineDriver.GetNumActiveTransactions(lastVersion) == 0, $"Active transactions in last version: {stateMachineDriver.GetNumActiveTransactions(lastVersion)}");
-                    stateMachineDriver.lastVersionTransactionsDone = null;
-                    stateMachineDriver.lastVersion = 0;
+                    stateMachineDriver.ResetLastVersion();
                     // Grab final logical address (end of fuzzy region)
                     store._hybridLogCheckpoint.info.finalLogicalAddress = store.hlogBase.GetTailAddress();
 
@@ -93,13 +92,7 @@ public virtual void GlobalAfterEnteringState(SystemState next, StateMachineDrive
                 case Phase.IN_PROGRESS:
                     // State machine should wait for active transactions in the last version to complete (drain out).
                     // Note that we allow new transactions to process in parallel.
-                    if (stateMachineDriver.GetNumActiveTransactions(lastVersion) > 0)
-                    {
-                        stateMachineDriver.lastVersion = lastVersion;
-                        stateMachineDriver.lastVersionTransactionsDone = new(0);
-                    }
-                    if (stateMachineDriver.GetNumActiveTransactions(lastVersion) > 0)
-                        stateMachineDriver.AddToWaitingList(stateMachineDriver.lastVersionTransactionsDone);
+                    stateMachineDriver.TrackLastVersion(lastVersion);
                     break;
             }
         }

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

@@ -33,8 +33,7 @@ public void GlobalBeforeEnteringState(SystemState next, StateMachineDriver state
                     // Verify full transaction barrier
                     Debug.Assert(stateMachineDriver.GetNumActiveTransactions(lastVersion) == 0);
                     Debug.Assert(stateMachineDriver.GetNumActiveTransactions(next.Version) == 0);
-                    stateMachineDriver.lastVersionTransactionsDone = null;
-                    stateMachineDriver.lastVersion = 0;
+                    stateMachineDriver.ResetLastVersion();
 
                     // Set up the transition to new version of HT
                     var numChunks = (int)(store.state[store.resizeInfo.version].size / Constants.kSizeofChunk);
@@ -70,13 +69,7 @@ public void GlobalAfterEnteringState(SystemState next, StateMachineDriver stateM
                 case Phase.PREPARE_GROW:
                     // State machine should wait for active transactions in the last version to complete (drain out).
                     // Note that we DO NOT allow new transactions to start in PREPARE_GROW (i.e., this is a full barrier)
-                    if (stateMachineDriver.GetNumActiveTransactions(lastVersion) > 0)
-                    {
-                        stateMachineDriver.lastVersion = lastVersion;
-                        stateMachineDriver.lastVersionTransactionsDone = new(0);
-                    }
-                    if (stateMachineDriver.GetNumActiveTransactions(lastVersion) > 0)
-                        stateMachineDriver.AddToWaitingList(stateMachineDriver.lastVersionTransactionsDone);
+                    stateMachineDriver.TrackLastVersion(lastVersion);
                     break;
 
                 case Phase.IN_PROGRESS_GROW:

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

@@ -21,11 +21,12 @@ public class StateMachineDriver
         TaskCompletionSource<bool> stateMachineCompleted;
         // All threads have entered the given state
         SemaphoreSlim waitForTransitionIn;
+        Exception waitForTransitionInException;
         // All threads have exited the given state
         SemaphoreSlim waitForTransitionOut;
         // Transactions drained in last version
-        public long lastVersion;
-        public SemaphoreSlim lastVersionTransactionsDone;
+        long lastVersion;
+        SemaphoreSlim lastVersionTransactionsDone;
         List<IStateMachineCallback> callbacks;
         readonly LightEpoch epoch;
         readonly ILogger logger;
@@ -55,13 +56,35 @@ void DecrementActiveTransactions(long txnVersion)
         {
             if (Interlocked.Decrement(ref NumActiveTransactions[txnVersion & 0x1]) == 0)
             {
-                if (lastVersionTransactionsDone != null && txnVersion == lastVersion)
+                var _lastVersionTransactionsDone = lastVersionTransactionsDone;
+                if (_lastVersionTransactionsDone != null && txnVersion == lastVersion)
                 {
-                    lastVersionTransactionsDone.Release();
+                    _lastVersionTransactionsDone.Release();
                 }
             }
         }
 
+        internal void TrackLastVersion(long version)
+        {
+            if (GetNumActiveTransactions(version) > 0)
+            {
+                // Set version number first, then create semaphore
+                lastVersion = version;
+                lastVersionTransactionsDone = new(0);
+            }
+
+            // We have to re-check the number of active transactions after assigning lastVersion and lastVersionTransactionsDone
+            if (GetNumActiveTransactions(version) > 0)
+                AddToWaitingList(lastVersionTransactionsDone);
+        }
+
+        internal void ResetLastVersion()
+        {
+            // First null semaphore, then reset version number
+            lastVersionTransactionsDone = null;
+            lastVersion = 0;
+        }
+
         /// <summary>
         /// Acquire a transaction version - this should be called before
         /// BeginLockable is called for all sessions in the transaction.
@@ -211,7 +234,7 @@ void GlobalStateMachineStep(SystemState expectedState)
             // Release waiters for new phase
             _ = waitForTransitionOut?.Release(int.MaxValue);
 
-            // Write new semaphore
+            // Write new semaphores
             waitForTransitionOut = new SemaphoreSlim(0);
             waitForTransitionIn = new SemaphoreSlim(0);
 
@@ -261,13 +284,31 @@ public async Task WaitForCompletion(SystemState currentState)
 
         void MakeTransitionWorker(SystemState nextState)
         {
-            stateMachine.GlobalAfterEnteringState(nextState, this);
-            waitForTransitionIn.Release(int.MaxValue);
+            try
+            {
+                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
+                waitForTransitionInException = e;
+
+                logger?.LogError(e, "Exception in state machine transition worker");
+            }
+            finally
+            {
+                waitForTransitionIn.Release(int.MaxValue);
+            }
         }
 
         async Task ProcessWaitingListAsync(CancellationToken token = default)
         {
             await waitForTransitionIn.WaitAsync(token);
+            if (waitForTransitionInException != null)
+            {
+                throw waitForTransitionInException;
+            }
             foreach (var waiter in waitingList)
             {
                 await waiter.WaitAsync(token);
@@ -288,6 +329,7 @@ async Task RunStateMachine(CancellationToken token = default)
             }
             catch (Exception e)
             {
+                FastForwardStateMachineToRest();
                 logger?.LogError(e, "Exception in state machine");
                 ex = e;
                 throw;
@@ -311,5 +353,31 @@ async Task RunStateMachine(CancellationToken token = default)
                 }
             }
         }
+
+        void FastForwardStateMachineToRest()
+        {
+            // Move system state to the next REST phase
+            while (systemState.Phase != Phase.REST)
+            {
+                systemState.Word = stateMachine.NextState(systemState).Word;
+            }
+
+            // Reset last version
+            ResetLastVersion();
+
+            // Release any waiters on existing transition-out semaphore
+            if (waitForTransitionOut?.CurrentCount == 0)
+                _ = waitForTransitionOut?.Release(int.MaxValue);
+
+            // Clear semaphores
+            waitForTransitionOut = null;
+            waitForTransitionIn = null;
+
+            // Clear exception if any
+            waitForTransitionInException = null;
+
+            // Clear waiting list
+            waitingList.Clear();
+        }
     }
 }