Simplify the Tsavorite checkpoint state machine (#1059)

* Simpler state machine for checkpointing

* cleanup

* updates

* remove dead code

* updates

* updates

* updates

* update

* kill code

* updates

* simplify LightEpoch

* move epvs to test

* nits

* updates

* updates

* formatting

* fix garnet

* nit

* comments

* remove manualLockingActive

* update the barrier condition and remove checkpoint version switch barrier flag.

* remove INTERMEDIATE state

* Remove CPR_SHIFT_DETECTED and LartchDestination.Retry

* add black box test for checkpointing version switch state machine
remove deprecated white box test

* add transaction test

* clean the test

* cleanup

* Refactor the phases of various machines

* format

* remove sessionName

* update LightEpoch based on PR comment

* fix break

* Use session-local isAcquiredLockable as signal for threads to decide whether to spin or operate in PREPARE phase. Moved isAcquiredLockable to Ctx. This commit also removes a race that will re-establish the invariant that no threads are operating in PREPARE while any thread is operating in IN_PROGRESS phase.

* address review comments

* nit

libs/host/GarnetServer.cs

@@ -302,7 +302,6 @@ private void CreateMainStore(IClusterFactory clusterFactory, out string checkpoi
 
             // Run checkpoint on its own thread to control p99
             kvSettings.ThrottleCheckpointFlushDelayMs = opts.CheckpointThrottleFlushDelayMs;
-            kvSettings.CheckpointVersionSwitchBarrier = opts.EnableCluster;
 
             if (opts.EnableCluster)
             {
@@ -330,7 +329,6 @@ private void CreateObjectStore(IClusterFactory clusterFactory, CustomCommandMana
 
                 // Run checkpoint on its own thread to control p99
                 objKvSettings.ThrottleCheckpointFlushDelayMs = opts.CheckpointThrottleFlushDelayMs;
-                objKvSettings.CheckpointVersionSwitchBarrier = opts.EnableCluster;
 
                 if (opts.EnableCluster)
                     objKvSettings.CheckpointManager = clusterFactory.CreateCheckpointManager(

libs/server/Metrics/Info/GarnetInfoMetrics.cs

@@ -207,7 +207,6 @@ private void PopulateStoreStats(StoreWrapper storeWrapper)
                 [
                     new("CurrentVersion", storeWrapper.store.CurrentVersion.ToString()),
                     new("LastCheckpointedVersion", storeWrapper.store.LastCheckpointedVersion.ToString()),
-                    new("RecoveredVersion", storeWrapper.store.RecoveredVersion.ToString()),
                     new("SystemState", storeWrapper.store.SystemState.ToString()),
                     new("IndexSize", storeWrapper.store.IndexSize.ToString()),
                     new("LogDir", storeWrapper.serverOptions.LogDir),
@@ -234,7 +233,6 @@ private void PopulateObjectStoreStats(StoreWrapper storeWrapper)
                 [
                     new("CurrentVersion", storeWrapper.objectStore.CurrentVersion.ToString()),
                     new("LastCheckpointedVersion", storeWrapper.objectStore.LastCheckpointedVersion.ToString()),
-                    new("RecoveredVersion", storeWrapper.objectStore.RecoveredVersion.ToString()),
                     new("SystemState", storeWrapper.objectStore.SystemState.ToString()),
                     new("IndexSize", storeWrapper.objectStore.IndexSize.ToString()),
                     new("LogDir", storeWrapper.serverOptions.LogDir),

libs/server/StoreWrapper.cs

@@ -668,11 +668,11 @@ private async void IndexAutoGrowTask(CancellationToken token)
 
                     if (!indexMaxedOut)
                         indexMaxedOut = GrowIndexIfNeeded(StoreType.Main, serverOptions.AdjustedIndexMaxCacheLines, store.OverflowBucketAllocations,
-                            () => store.IndexSize, () => store.GrowIndex());
+                            () => store.IndexSize, async () => await store.GrowIndexAsync());
 
                     if (!objectStoreIndexMaxedOut)
                         objectStoreIndexMaxedOut = GrowIndexIfNeeded(StoreType.Object, serverOptions.AdjustedObjectStoreIndexMaxCacheLines, objectStore.OverflowBucketAllocations,
-                            () => objectStore.IndexSize, () => objectStore.GrowIndex());
+                            () => objectStore.IndexSize, async () => await objectStore.GrowIndexAsync());
                 }
             }
             catch (Exception ex)

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

@@ -312,6 +312,8 @@ public async ValueTask IsCheckpointCompletedAsync(CancellationToken token = defa
             s.Release();
         }
 
+        public SemaphoreSlim GetCheckpointSemaphore() => checkpointSemaphore;
+
         /// <summary>
         /// Public facing persistence API
         /// </summary>

libs/storage/Tsavorite/cs/src/core/ClientSession/ClientSession.cs

@@ -40,14 +40,12 @@ public sealed class ClientSession<TKey, TValue, TInput, TOutput, TContext, TFunc
         internal ulong sharedLockCount;
         internal ulong exclusiveLockCount;
 
-        bool isAcquiredLockable;
-
         ScanCursorState<TKey, TValue> scanCursorState;
 
         internal void AcquireLockable<TSessionFunctions>(TSessionFunctions sessionFunctions)
             where TSessionFunctions : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {
-            CheckIsNotAcquiredLockable();
+            CheckIsNotAcquiredLockable(sessionFunctions);
 
             while (true)
             {
@@ -60,39 +58,43 @@ internal void AcquireLockable<TSessionFunctions>(TSessionFunctions sessionFuncti
                 }
 
                 store.IncrementNumLockingSessions();
-                isAcquiredLockable = true;
+                sessionFunctions.Ctx.isAcquiredLockable = true;
 
                 if (!IsInPreparePhase())
                     break;
-                InternalReleaseLockable();
+                InternalReleaseLockable(sessionFunctions);
                 _ = Thread.Yield();
             }
         }
 
-        internal void ReleaseLockable()
+        internal void ReleaseLockable<TSessionFunctions>(TSessionFunctions sessionFunctions)
+            where TSessionFunctions : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {
-            CheckIsAcquiredLockable();
+            CheckIsAcquiredLockable(sessionFunctions);
             if (TotalLockCount > 0)
                 throw new TsavoriteException($"EndLockable called with locks held: {sharedLockCount} shared locks, {exclusiveLockCount} exclusive locks");
-            InternalReleaseLockable();
+            InternalReleaseLockable(sessionFunctions);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private void InternalReleaseLockable()
+        private void InternalReleaseLockable<TSessionFunctions>(TSessionFunctions sessionFunctions)
+            where TSessionFunctions : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {
-            isAcquiredLockable = false;
+            sessionFunctions.Ctx.isAcquiredLockable = false;
             store.DecrementNumLockingSessions();
         }
 
-        internal void CheckIsAcquiredLockable()
+        internal void CheckIsAcquiredLockable<TSessionFunctions>(TSessionFunctions sessionFunctions)
+            where TSessionFunctions : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {
-            if (!isAcquiredLockable)
+            if (!sessionFunctions.Ctx.isAcquiredLockable)
                 throw new TsavoriteException("Lockable method call when BeginLockable has not been called");
         }
 
-        void CheckIsNotAcquiredLockable()
+        void CheckIsNotAcquiredLockable<TSessionFunctions>(TSessionFunctions sessionFunctions)
+            where TSessionFunctions : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {
-            if (isAcquiredLockable)
+            if (sessionFunctions.Ctx.isAcquiredLockable)
                 throw new TsavoriteException("BeginLockable cannot be called twice (call EndLockable first)");
         }
 
@@ -133,7 +135,9 @@ public void Dispose()
 
             // By the time Dispose is called, we should have no outstanding locks, so can use the BasicContext's sessionFunctions.
             _ = CompletePending(bContext.sessionFunctions, true);
-            store.DisposeClientSession(ID, ctx.phase);
+
+            if (store.RevivificationManager.IsEnabled)
+                MergeRevivificationStatsTo(ref store.RevivificationManager.stats, reset: true);
         }
 
         /// <summary>
@@ -388,7 +392,7 @@ private async ValueTask WaitForCommitAsync<TSessionFunctionsWrapper>(TSessionFun
         {
             token.ThrowIfCancellationRequested();
 
-            if (!ctx.prevCtx.pendingReads.IsEmpty || !ctx.pendingReads.IsEmpty)
+            if (!ctx.pendingReads.IsEmpty)
                 throw new TsavoriteException("Make sure all async operations issued on this session are awaited and completed first");
 
             // Complete all pending sync operations on session
@@ -536,11 +540,6 @@ internal void UnsafeSuspendThread()
             store.epoch.Suspend();
         }
 
-        void IClientSession.AtomicSwitch(long version)
-        {
-            _ = TsavoriteKV<TKey, TValue, TStoreFunctions, TAllocator>.AtomicSwitch(ctx, ctx.prevCtx, version);
-        }
-
         /// <inheritdoc/>
         public void MergeRevivificationStatsTo(ref RevivificationStats to, bool reset) => ctx.MergeRevivificationStatsTo(ref to, reset);
 
@@ -552,7 +551,8 @@ void IClientSession.AtomicSwitch(long version)
         /// </summary>
         internal bool IsInPreparePhase()
         {
-            return store.SystemState.Phase == Phase.PREPARE || store.SystemState.Phase == Phase.PREPARE_GROW;
+            var storeState = store.stateMachineDriver.SystemState;
+            return storeState.Phase == Phase.PREPARE || storeState.Phase == Phase.PREPARE_GROW;
         }
 
         #endregion Other Operations

libs/storage/Tsavorite/cs/src/core/ClientSession/IClientSession.cs

@@ -5,15 +5,12 @@ namespace Tsavorite.core
 {
     class SessionInfo
     {
-        public string sessionName;
         public bool isActive;
         public IClientSession session;
     }
 
     internal interface IClientSession
     {
-        void AtomicSwitch(long version);
-
         void MergeRevivificationStatsTo(ref RevivificationStats globalStats, bool reset);
 
         void ResetRevivificationStats();

libs/storage/Tsavorite/cs/src/core/ClientSession/LockableContext.cs

@@ -37,7 +37,7 @@ internal LockableContext(ClientSession<TKey, TValue, TInput, TOutput, TContext,
         public void BeginLockable() => clientSession.AcquireLockable(sessionFunctions);
 
         /// <inheritdoc/>
-        public void EndLockable() => clientSession.ReleaseLockable();
+        public void EndLockable() => clientSession.ReleaseLockable(sessionFunctions);
 
         #endregion Begin/EndLockable
 
@@ -236,7 +236,7 @@ internal static void DoManualUnlock<TLockableKey>(ClientSession<TKey, TValue, TI
         public void Lock<TLockableKey>(TLockableKey[] keys, int start, int count)
             where TLockableKey : ILockableKey
         {
-            clientSession.CheckIsAcquiredLockable();
+            clientSession.CheckIsAcquiredLockable(sessionFunctions);
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected(), "Trying to protect an already-protected epoch for LockableUnsafeContext.Lock()");
             bool lockAquired = false;
             while (!lockAquired)
@@ -287,7 +287,7 @@ public bool TryLock<TLockableKey>(TLockableKey[] keys, TimeSpan timeout, Cancell
         public bool TryLock<TLockableKey>(TLockableKey[] keys, int start, int count, TimeSpan timeout, CancellationToken cancellationToken)
             where TLockableKey : ILockableKey
         {
-            clientSession.CheckIsAcquiredLockable();
+            clientSession.CheckIsAcquiredLockable(sessionFunctions);
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected(), "Trying to protect an already-protected epoch for LockableUnsafeContext.Lock()");
 
             clientSession.UnsafeResumeThread(sessionFunctions);
@@ -320,7 +320,7 @@ public bool TryPromoteLock<TLockableKey>(TLockableKey key, CancellationToken can
         public bool TryPromoteLock<TLockableKey>(TLockableKey key, TimeSpan timeout, CancellationToken cancellationToken)
             where TLockableKey : ILockableKey
         {
-            clientSession.CheckIsAcquiredLockable();
+            clientSession.CheckIsAcquiredLockable(sessionFunctions);
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected(), "Trying to protect an already-protected epoch for LockableUnsafeContext.Lock()");
 
             clientSession.UnsafeResumeThread(sessionFunctions);
@@ -341,7 +341,7 @@ public bool TryPromoteLock<TLockableKey>(TLockableKey key, TimeSpan timeout, Can
         public void Unlock<TLockableKey>(TLockableKey[] keys, int start, int count)
             where TLockableKey : ILockableKey
         {
-            clientSession.CheckIsAcquiredLockable();
+            clientSession.CheckIsAcquiredLockable(sessionFunctions);
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected(), "Trying to protect an already-protected epoch for LockableUnsafeContext.Unlock()");
 
             clientSession.UnsafeResumeThread(sessionFunctions);

libs/storage/Tsavorite/cs/src/core/ClientSession/LockableUnsafeContext.cs

@@ -48,7 +48,7 @@ internal LockableUnsafeContext(ClientSession<TKey, TValue, TInput, TOutput, TCon
         public void BeginLockable() => clientSession.AcquireLockable(sessionFunctions);
 
         /// <inheritdoc/>
-        public void EndLockable() => clientSession.ReleaseLockable();
+        public void EndLockable() => clientSession.ReleaseLockable(sessionFunctions);
         #endregion Begin/EndLockable
 
         #region Key Locking
@@ -72,7 +72,7 @@ internal LockableUnsafeContext(ClientSession<TKey, TValue, TInput, TOutput, TCon
         public void Lock<TLockableKey>(TLockableKey[] keys, int start, int count)
             where TLockableKey : ILockableKey
         {
-            clientSession.CheckIsAcquiredLockable();
+            clientSession.CheckIsAcquiredLockable(sessionFunctions);
             Debug.Assert(clientSession.store.epoch.ThisInstanceProtected(), "Epoch protection required for LockableUnsafeContext.Lock()");
             while (true)
             {
@@ -120,7 +120,7 @@ public bool TryLock<TLockableKey>(TLockableKey[] keys, TimeSpan timeout, Cancell
         public bool TryLock<TLockableKey>(TLockableKey[] keys, int start, int count, TimeSpan timeout, CancellationToken cancellationToken)
             where TLockableKey : ILockableKey
         {
-            clientSession.CheckIsAcquiredLockable();
+            clientSession.CheckIsAcquiredLockable(sessionFunctions);
             Debug.Assert(clientSession.store.epoch.ThisInstanceProtected(), "Epoch protection required for LockableUnsafeContext.Lock()");
 
             return LockableContext<TKey, TValue, TInput, TOutput, TContext, TFunctions, TStoreFunctions, TAllocator>.DoManualTryLock(sessionFunctions, clientSession, keys, start, count, timeout, cancellationToken);
@@ -145,7 +145,7 @@ public bool TryPromoteLock<TLockableKey>(TLockableKey key, CancellationToken can
         public bool TryPromoteLock<TLockableKey>(TLockableKey key, TimeSpan timeout, CancellationToken cancellationToken)
             where TLockableKey : ILockableKey
         {
-            clientSession.CheckIsAcquiredLockable();
+            clientSession.CheckIsAcquiredLockable(sessionFunctions);
             Debug.Assert(clientSession.store.epoch.ThisInstanceProtected(), "Epoch protection required for LockableUnsafeContext.Lock()");
 
             return LockableContext<TKey, TValue, TInput, TOutput, TContext, TFunctions, TStoreFunctions, TAllocator>.DoManualTryPromoteLock(sessionFunctions, clientSession, key, timeout, cancellationToken);
@@ -158,7 +158,7 @@ public bool TryPromoteLock<TLockableKey>(TLockableKey key, TimeSpan timeout, Can
         public void Unlock<TLockableKey>(TLockableKey[] keys, int start, int count)
             where TLockableKey : ILockableKey
         {
-            clientSession.CheckIsAcquiredLockable();
+            clientSession.CheckIsAcquiredLockable(sessionFunctions);
             Debug.Assert(clientSession.store.epoch.ThisInstanceProtected(), "Epoch protection required for LockableUnsafeContext.Unlock()");
 
             LockableContext<TKey, TValue, TInput, TOutput, TContext, TFunctions, TStoreFunctions, TAllocator>.DoManualUnlock(clientSession, keys, start, start + count - 1);