Single state machine across main and object stores

libs/host/GarnetServer.cs

@@ -227,6 +227,7 @@ private void InitializeServer()
             if (!setMax && !ThreadPool.SetMaxThreads(maxThreads, maxCPThreads))
                 throw new Exception($"Unable to call ThreadPool.SetMaxThreads with {maxThreads}, {maxCPThreads}");
 
+            opts.Initialize(loggerFactory);
             CreateMainStore(clusterFactory, out var checkpointDir);
             CreateObjectStore(clusterFactory, customCommandManager, checkpointDir, out var objectStoreSizeTracker);
 
@@ -324,7 +325,7 @@ private void CreateObjectStore(IClusterFactory clusterFactory, CustomCommandMana
             objectStoreSizeTracker = null;
             if (!opts.DisableObjects)
             {
-                objKvSettings = opts.GetObjectStoreSettings(this.loggerFactory?.CreateLogger("TsavoriteKV  [obj]"),
+                objKvSettings = opts.GetObjectStoreSettings(loggerFactory,
                     out var objHeapMemorySize, out var objReadCacheHeapMemorySize);
 
                 // Run checkpoint on its own thread to control p99

libs/server/AOF/AofProcessor.cs

@@ -210,11 +210,10 @@ public unsafe void ProcessAofRecordInternal(byte* ptr, int length, bool asReplic
                     break;
                 case AofEntryType.MainStoreCheckpointStartCommit:
                     if (asReplica && header.storeVersion > storeWrapper.store.CurrentVersion)
-                        _ = storeWrapper.TakeCheckpoint(false, StoreType.Main, logger);
+                        _ = storeWrapper.TakeCheckpoint(false, logger);
                     break;
                 case AofEntryType.ObjectStoreCheckpointStartCommit:
-                    if (asReplica && header.storeVersion > storeWrapper.objectStore.CurrentVersion)
-                        _ = storeWrapper.TakeCheckpoint(false, StoreType.Object, logger);
+                    // With unified checkpoint, we do not need to take object store checkpoint separately
                     break;
                 case AofEntryType.MainStoreCheckpointEndCommit:
                 case AofEntryType.ObjectStoreCheckpointEndCommit:

libs/server/Resp/AdminCommands.cs

@@ -818,7 +818,7 @@ private bool NetworkSAVE()
                 return AbortWithWrongNumberOfArguments(nameof(RespCommand.SAVE));
             }
 
-            if (!storeWrapper.TakeCheckpoint(false, StoreType.All, logger))
+            if (!storeWrapper.TakeCheckpoint(false, logger))
             {
                 while (!RespWriteUtils.TryWriteError("ERR checkpoint already in progress"u8, ref dcurr, dend))
                     SendAndReset();
@@ -853,7 +853,7 @@ private bool NetworkBGSAVE()
                 return AbortWithWrongNumberOfArguments(nameof(RespCommand.BGSAVE));
             }
 
-            var success = storeWrapper.TakeCheckpoint(true, StoreType.All, logger);
+            var success = storeWrapper.TakeCheckpoint(true, logger);
             if (success)
             {
                 while (!RespWriteUtils.TryWriteSimpleString("Background saving started"u8, ref dcurr, dend))

libs/server/Servers/GarnetServerOptions.cs

@@ -446,6 +446,16 @@ public class GarnetServerOptions : ServerOptions
         /// </summary>
         public UnixFileMode UnixSocketPermission { get; set; }
 
+        /// <summary>
+        /// Epoch instance used by server
+        /// </summary>
+        public LightEpoch Epoch;
+
+        /// <summary>
+        /// Common state machine driver used by Garnet
+        /// </summary>
+        public StateMachineDriver StateMachineDriver;
+
         /// <summary>
         /// Constructor
         /// </summary>
@@ -454,6 +464,16 @@ public GarnetServerOptions(ILogger logger = null) : base(logger)
             this.logger = logger;
         }
 
+        /// <summary>
+        /// Initialize Garnet server options
+        /// </summary>
+        /// <param name="loggerFactory"></param>
+        public void Initialize(ILoggerFactory loggerFactory = null)
+        {
+            Epoch = new LightEpoch();
+            StateMachineDriver = new(Epoch, loggerFactory?.CreateLogger($"StateMachineDriver"));
+        }
+
         /// <summary>
         /// Get main store settings
         /// </summary>
@@ -466,15 +486,16 @@ public KVSettings<SpanByte, SpanByte> GetSettings(ILoggerFactory loggerFactory,
             if (MutablePercent is < 10 or > 95)
                 throw new Exception("MutablePercent must be between 10 and 95");
 
-            KVSettings<SpanByte, SpanByte> kvSettings = new(baseDir: null, logger: logger);
-
             var indexCacheLines = IndexSizeCachelines("hash index size", IndexSize);
-            kvSettings = new()
+
+            KVSettings<SpanByte, SpanByte> kvSettings = new()
             {
                 IndexSize = indexCacheLines * 64L,
                 PreallocateLog = false,
                 MutableFraction = MutablePercent / 100.0,
                 PageSize = 1L << PageSizeBits(),
+                Epoch = Epoch,
+                StateMachineDriver = StateMachineDriver,
                 loggerFactory = loggerFactory,
                 logger = loggerFactory?.CreateLogger("TsavoriteKV [main]")
             };
@@ -618,23 +639,26 @@ public static int MemorySizeBits(string memorySize, string storePageSize, out in
         /// <summary>
         /// Get KVSettings for the object store log
         /// </summary>
-        public KVSettings<byte[], IGarnetObject> GetObjectStoreSettings(ILogger logger, out long objHeapMemorySize, out long objReadCacheHeapMemorySize)
+        public KVSettings<byte[], IGarnetObject> GetObjectStoreSettings(ILoggerFactory loggerFactory, out long objHeapMemorySize, out long objReadCacheHeapMemorySize)
         {
             objReadCacheHeapMemorySize = default;
 
             if (ObjectStoreMutablePercent is < 10 or > 95)
                 throw new Exception("ObjectStoreMutablePercent must be between 10 and 95");
 
-            KVSettings<byte[], IGarnetObject> kvSettings = new(baseDir: null, logger: logger);
-
             var indexCacheLines = IndexSizeCachelines("object store hash index size", ObjectStoreIndexSize);
-            kvSettings = new()
+            KVSettings<byte[], IGarnetObject> kvSettings = new()
             {
                 IndexSize = indexCacheLines * 64L,
                 PreallocateLog = false,
                 MutableFraction = ObjectStoreMutablePercent / 100.0,
-                PageSize = 1L << ObjectStorePageSizeBits()
+                PageSize = 1L << ObjectStorePageSizeBits(),
+                Epoch = Epoch,
+                StateMachineDriver = StateMachineDriver,
+                loggerFactory = loggerFactory,
+                logger = loggerFactory?.CreateLogger("TsavoriteKV  [obj]")
             };
+
             logger?.LogInformation("[Object Store] Using page size of {PageSize}", PrettySize(kvSettings.PageSize));
             logger?.LogInformation("[Object Store] Each page can hold ~{PageSize} key-value pairs of objects", kvSettings.PageSize / 24);
 

libs/server/StoreWrapper.cs

@@ -761,28 +761,27 @@ public async Task TakeOnDemandCheckpoint(DateTimeOffset entryTime)
             }
 
             // Necessary to take a checkpoint because the latest checkpoint is before entryTime
-            await CheckpointTask(StoreType.All, logger: logger);
+            await CheckpointTask(logger);
         }
 
         /// <summary>
         /// Take checkpoint
         /// </summary>
         /// <param name="background"></param>
-        /// <param name="storeType"></param>
         /// <param name="logger"></param>
         /// <returns></returns>
-        public bool TakeCheckpoint(bool background, StoreType storeType = StoreType.All, ILogger logger = null)
+        public bool TakeCheckpoint(bool background, ILogger logger = null)
         {
             // Prevent parallel checkpoint
             if (!TryPauseCheckpoints()) return false;
             if (background)
-                Task.Run(async () => await CheckpointTask(storeType, logger));
+                Task.Run(async () => await CheckpointTask(logger));
             else
-                CheckpointTask(storeType, logger).ConfigureAwait(false).GetAwaiter().GetResult();
+                CheckpointTask(logger).ConfigureAwait(false).GetAwaiter().GetResult();
             return true;
         }
 
-        private async Task CheckpointTask(StoreType storeType, ILogger logger = null)
+        private async Task CheckpointTask(ILogger logger = null)
         {
             try
             {
@@ -803,12 +802,11 @@ private async Task CheckpointTask(StoreType storeType, ILogger logger = null)
                     tryIncremental = false;
 
                 var checkpointType = serverOptions.UseFoldOverCheckpoints ? CheckpointType.FoldOver : CheckpointType.Snapshot;
-                await InitiateCheckpoint(full, checkpointType, tryIncremental, storeType, logger);
+                await InitiateCheckpoint(full, checkpointType, tryIncremental, logger);
                 if (full)
                 {
-                    if (storeType is StoreType.Main or StoreType.All)
-                        this.lastSaveStoreTailAddress = lastSaveStoreTailAddress;
-                    if (storeType is StoreType.Object or StoreType.All)
+                    this.lastSaveStoreTailAddress = lastSaveStoreTailAddress;
+                    if (objectStore != null)
                         this.lastSaveObjectStoreTailAddress = lastSaveObjectStoreTailAddress;
                 }
                 lastSaveTime = DateTimeOffset.UtcNow;
@@ -823,9 +821,9 @@ private async Task CheckpointTask(StoreType storeType, ILogger logger = null)
             }
         }
 
-        private async Task InitiateCheckpoint(bool full, CheckpointType checkpointType, bool tryIncremental, StoreType storeType, ILogger logger = null)
+        private async Task InitiateCheckpoint(bool full, CheckpointType checkpointType, bool tryIncremental, ILogger logger = null)
         {
-            logger?.LogInformation("Initiating checkpoint; full = {full}, type = {checkpointType}, tryIncremental = {tryIncremental}, storeType = {storeType}", full, checkpointType, tryIncremental, storeType);
+            logger?.LogInformation("Initiating checkpoint; full = {full}, type = {checkpointType}, tryIncremental = {tryIncremental}", full, checkpointType, tryIncremental);
 
             long CheckpointCoveredAofAddress = 0;
             if (appendOnlyFile != null)
@@ -839,41 +837,47 @@ private async Task InitiateCheckpoint(bool full, CheckpointType checkpointType,
                     logger?.LogInformation("Will truncate AOF to {tailAddress} after checkpoint (files deleted after next commit)", CheckpointCoveredAofAddress);
             }
 
-            (bool success, Guid token) storeCheckpointResult = default;
-            (bool success, Guid token) objectStoreCheckpointResult = default;
+            (bool success, Guid token) checkpointResult = default;
+            IStateMachine sm;
             if (full)
             {
-                if (storeType is StoreType.Main or StoreType.All)
-                {
-                    storeCheckpointResult = await store.TakeFullCheckpointAsync(checkpointType);
-                    if (serverOptions.EnableCluster && clusterProvider.IsPrimary()) EnqueueCommit(AofEntryType.MainStoreCheckpointEndCommit, store.CurrentVersion);
-                }
-
-                if (objectStore != null && (storeType == StoreType.Object || storeType == StoreType.All))
-                {
-                    objectStoreCheckpointResult = await objectStore.TakeFullCheckpointAsync(checkpointType);
-                    if (serverOptions.EnableCluster && clusterProvider.IsPrimary()) EnqueueCommit(AofEntryType.ObjectStoreCheckpointEndCommit, objectStore.CurrentVersion);
-                }
+                sm = objectStore == null ?
+                        Checkpoint.Full(store, checkpointType, out checkpointResult.token) :
+                        Checkpoint.Full(store, objectStore, checkpointType, out checkpointResult.token);
             }
             else
             {
-                if (storeType is StoreType.Main or StoreType.All)
+                tryIncremental &= store.CanTakeIncrementalCheckpoint(checkpointType, out checkpointResult.token);
+                if (objectStore != null)
+                    tryIncremental &= objectStore.CanTakeIncrementalCheckpoint(checkpointType, out var guid2) & checkpointResult.token == guid2;
+
+                if (tryIncremental)
                 {
-                    storeCheckpointResult = await store.TakeHybridLogCheckpointAsync(checkpointType, tryIncremental);
-                    if (serverOptions.EnableCluster && clusterProvider.IsPrimary()) EnqueueCommit(AofEntryType.MainStoreCheckpointEndCommit, store.CurrentVersion);
+                    sm = objectStore == null ?
+                        Checkpoint.IncrementalHybridLogOnly(store, checkpointResult.token) :
+                        Checkpoint.IncrementalHybridLogOnly(store, objectStore, checkpointResult.token);
                 }
-
-                if (objectStore != null && (storeType == StoreType.Object || storeType == StoreType.All))
+                else
                 {
-                    objectStoreCheckpointResult = await objectStore.TakeHybridLogCheckpointAsync(checkpointType, tryIncremental);
-                    if (serverOptions.EnableCluster && clusterProvider.IsPrimary()) EnqueueCommit(AofEntryType.ObjectStoreCheckpointEndCommit, objectStore.CurrentVersion);
+                    sm = objectStore == null ?
+                        Checkpoint.HybridLogOnly(store, checkpointType, out checkpointResult.token) :
+                        Checkpoint.HybridLogOnly(store, objectStore, checkpointType, out checkpointResult.token);
                 }
             }
 
+            checkpointResult.success = await serverOptions.StateMachineDriver.RunAsync(sm);
+
+            if (serverOptions.EnableCluster && clusterProvider.IsPrimary())
+            {
+                EnqueueCommit(AofEntryType.MainStoreCheckpointEndCommit, store.CurrentVersion);
+                if (objectStore != null)
+                    EnqueueCommit(AofEntryType.ObjectStoreCheckpointEndCommit, objectStore.CurrentVersion);
+            }
+
             // If cluster is enabled the replication manager is responsible for truncating AOF
             if (serverOptions.EnableCluster && serverOptions.EnableAOF)
             {
-                clusterProvider.SafeTruncateAOF(storeType, full, CheckpointCoveredAofAddress, storeCheckpointResult.token, objectStoreCheckpointResult.token);
+                clusterProvider.SafeTruncateAOF(StoreType.All, full, CheckpointCoveredAofAddress, checkpointResult.token, checkpointResult.token);
             }
             else
             {

libs/server/Transaction/TransactionManager.cs

@@ -292,23 +292,43 @@ internal void GetKeysForValidation(byte* recvBufferPtr, out ArgSlice[] keys, out
             readOnly = keyEntries.IsReadOnly;
         }
 
+        void BeginLockable(StoreType transactionStoreType)
+        {
+            while (true)
+            {
+                if (transactionStoreType is StoreType.All or StoreType.Main)
+                {
+                    lockableContext.BeginLockable();
+                }
+                if (transactionStoreType is StoreType.All or StoreType.Object)
+                {
+                    if (objectStoreBasicContext.IsNull)
+                        throw new Exception("Trying to perform object store transaction with object store disabled");
+                    if (objectStoreLockableContext.TryBeginLockable())
+                    {
+                        // If we managed to begin lockable for the object store, we MUST be in the same version as the main store
+                        break;
+                    }
+                    objectStoreLockableContext.Refresh();
+                    if (transactionStoreType is StoreType.All or StoreType.Main)
+                    {
+                        lockableContext.EndLockable();
+                        lockableContext.Refresh();
+                    }
+                    continue;
+                }
+                break;
+            }
+        }
+
         internal bool Run(bool internal_txn = false, bool fail_fast_on_lock = false, TimeSpan lock_timeout = default)
         {
             // Save watch keys to lock list
             if (!internal_txn)
                 watchContainer.SaveKeysToLock(this);
 
             // Acquire lock sessions
-            if (transactionStoreType == StoreType.All || transactionStoreType == StoreType.Main)
-            {
-                lockableContext.BeginLockable();
-            }
-            if (transactionStoreType == StoreType.All || transactionStoreType == StoreType.Object)
-            {
-                if (objectStoreBasicContext.IsNull)
-                    throw new Exception("Trying to perform object store transaction with object store disabled");
-                objectStoreLockableContext.BeginLockable();
-            }
+            BeginLockable(transactionStoreType);
 
             bool lockSuccess;
             if (fail_fast_on_lock)

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

@@ -13,9 +13,9 @@ public interface IStreamingSnapshotIteratorFunctions<TKey, TValue>
         /// <summary>Iteration is starting.</summary>
         /// <param name="checkpointToken">Checkpoint token</param>
         /// <param name="currentVersion">Current version of database</param>
-        /// <param name="targetVersion">Target version of database</param>
+        /// <param name="nextVersion">Next version of database</param>
         /// <returns>True to continue iteration, else false</returns>
-        bool OnStart(Guid checkpointToken, long currentVersion, long targetVersion);
+        bool OnStart(Guid checkpointToken, long currentVersion, long nextVersion);
 
         /// <summary>Next record in the streaming snapshot.</summary>
         /// <param name="key">Reference to the current record's key</param>

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

@@ -67,6 +67,27 @@ internal void AcquireLockable<TSessionFunctions>(TSessionFunctions sessionFuncti
             }
         }
 
+        internal bool TryAcquireLockable<TSessionFunctions>(TSessionFunctions sessionFunctions)
+            where TSessionFunctions : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
+        {
+            CheckIsNotAcquiredLockable(sessionFunctions);
+
+            // Checkpoints cannot complete while we have active locking sessions.
+            if (IsInPreparePhase())
+            {
+                return false;
+            }
+
+            store.IncrementNumLockingSessions();
+            sessionFunctions.Ctx.isAcquiredLockable = true;
+
+            if (!IsInPreparePhase())
+                return true;
+
+            InternalReleaseLockable(sessionFunctions);
+            return false;
+        }
+
         internal void ReleaseLockable<TSessionFunctions>(TSessionFunctions sessionFunctions)
             where TSessionFunctions : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {

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

@@ -19,6 +19,11 @@ public interface ILockableContext<TKey>
         /// </summary>
         void BeginLockable();
 
+        /// <summary>
+        /// Try to begin a series of lock operations on possibly multiple keys; call before any locks are taken.
+        /// </summary>
+        bool TryBeginLockable();
+
         /// <summary>
         /// Ends a series of lock operations on possibly multiple keys; call after all locks are released.
         /// </summary>

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

@@ -36,6 +36,9 @@ internal LockableContext(ClientSession<TKey, TValue, TInput, TOutput, TContext,
         /// <inheritdoc/>
         public void BeginLockable() => clientSession.AcquireLockable(sessionFunctions);
 
+        /// <inheritdoc/>
+        public bool TryBeginLockable() => clientSession.TryAcquireLockable(sessionFunctions);
+
         /// <inheritdoc/>
         public void EndLockable() => clientSession.ReleaseLockable(sessionFunctions);
 

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

@@ -47,6 +47,9 @@ internal LockableUnsafeContext(ClientSession<TKey, TValue, TInput, TOutput, TCon
         /// <inheritdoc/>
         public void BeginLockable() => clientSession.AcquireLockable(sessionFunctions);
 
+        /// <inheritdoc/>
+        public bool TryBeginLockable() => clientSession.TryAcquireLockable(sessionFunctions);
+
         /// <inheritdoc/>
         public void EndLockable() => clientSession.ReleaseLockable(sessionFunctions);
         #endregion Begin/EndLockable