Single state machine across main and object stores

libs/cluster/Server/Replication/ReplicationLogCheckpointManager.cs

@@ -25,14 +25,16 @@ internal sealed class ReplicationLogCheckpointManager(
         public string RecoveredReplicationId = string.Empty;
 
         readonly bool isMainStore = isMainStore;
-        public Action<bool, long, long> checkpointVersionShift;
+        public Action<bool, long, long> checkpointVersionShiftStart;
+        public Action<bool, long, long> checkpointVersionShiftEnd;
 
         readonly bool safelyRemoveOutdated = removeOutdated;
 
-        public override void CheckpointVersionShift(long oldVersion, long newVersion)
-        {
-            checkpointVersionShift?.Invoke(isMainStore, oldVersion, newVersion);
-        }
+        public override void CheckpointVersionShiftStart(long oldVersion, long newVersion)
+            => checkpointVersionShiftStart?.Invoke(isMainStore, oldVersion, newVersion);
+
+        public override void CheckpointVersionShiftEnd(long oldVersion, long newVersion)
+            => checkpointVersionShiftEnd?.Invoke(isMainStore, oldVersion, newVersion);
 
         public void DeleteLogCheckpoint(Guid logToken)
             => deviceFactory.Delete(checkpointNamingScheme.LogCheckpointBase(logToken));

libs/cluster/Server/Replication/ReplicationManager.cs

@@ -107,9 +107,13 @@ public ReplicationManager(ClusterProvider clusterProvider, ILogger logger = null
             ReplicationOffset = 0;
 
             // Set the appendOnlyFile field for all stores
-            clusterProvider.GetReplicationLogCheckpointManager(StoreType.Main).checkpointVersionShift = CheckpointVersionShift;
+            clusterProvider.GetReplicationLogCheckpointManager(StoreType.Main).checkpointVersionShiftStart = CheckpointVersionShiftStart;
+            clusterProvider.GetReplicationLogCheckpointManager(StoreType.Main).checkpointVersionShiftEnd = CheckpointVersionShiftEnd;
             if (storeWrapper.objectStore != null)
-                clusterProvider.GetReplicationLogCheckpointManager(StoreType.Object).checkpointVersionShift = CheckpointVersionShift;
+            {
+                clusterProvider.GetReplicationLogCheckpointManager(StoreType.Object).checkpointVersionShiftStart = CheckpointVersionShiftStart;
+                clusterProvider.GetReplicationLogCheckpointManager(StoreType.Object).checkpointVersionShiftEnd = CheckpointVersionShiftEnd;
+            }
 
             // If this node starts as replica, it cannot serve requests until it is connected to primary
             if (clusterProvider.clusterManager.CurrentConfig.LocalNodeRole == NodeRole.REPLICA && clusterProvider.serverOptions.Recover && !StartRecovery())
@@ -152,7 +156,7 @@ public ReplicationManager(ClusterProvider clusterProvider, ILogger logger = null
 
         public string GetBufferPoolStats() => networkPool.GetStats();
 
-        void CheckpointVersionShift(bool isMainStore, long oldVersion, long newVersion)
+        void CheckpointVersionShiftStart(bool isMainStore, long oldVersion, long newVersion)
         {
             if (clusterProvider.clusterManager.CurrentConfig.LocalNodeRole == NodeRole.REPLICA)
                 return;
@@ -162,6 +166,16 @@ void CheckpointVersionShift(bool isMainStore, long oldVersion, long newVersion)
             storeWrapper.EnqueueCommit(entryType, newVersion);
         }
 
+        void CheckpointVersionShiftEnd(bool isMainStore, long oldVersion, long newVersion)
+        {
+            if (clusterProvider.clusterManager.CurrentConfig.LocalNodeRole == NodeRole.REPLICA)
+                return;
+            var entryType = clusterProvider.serverOptions.ReplicaDisklessSync ?
+                (isMainStore ? AofEntryType.MainStoreStreamingCheckpointEndCommit : AofEntryType.ObjectStoreStreamingCheckpointEndCommit) :
+                (isMainStore ? AofEntryType.MainStoreCheckpointEndCommit : AofEntryType.ObjectStoreCheckpointEndCommit);
+            storeWrapper.EnqueueCommit(entryType, newVersion);
+        }
+
         /// <summary>
         /// Acquire recovery and checkpoint locks to prevent checkpoints and parallel recovery tasks
         /// </summary>

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/Lua/LuaRunner.cs

@@ -1988,6 +1988,7 @@ internal int UnsafeRunPreambleForRunner()
         private void RunInTransaction<TResponse>(ref TResponse response)
             where TResponse : struct, IResponseAdapter
         {
+            var txnVersion = respServerSession.storageSession.stateMachineDriver.AcquireTransactionVersion();
             try
             {
                 respServerSession.storageSession.lockableContext.BeginLockable();
@@ -1996,6 +1997,10 @@ private void RunInTransaction<TResponse>(ref TResponse response)
                 respServerSession.SetTransactionMode(true);
                 txnKeyEntries.LockAllKeys();
 
+                txnVersion = respServerSession.storageSession.stateMachineDriver.VerifyTransactionVersion(txnVersion);
+                respServerSession.storageSession.lockableContext.LocksAcquired(txnVersion);
+                if (!respServerSession.storageSession.objectStoreLockableContext.IsNull)
+                    respServerSession.storageSession.objectStoreLockableContext.LocksAcquired(txnVersion);
                 RunCommon(ref response);
             }
             finally
@@ -2005,6 +2010,7 @@ private void RunInTransaction<TResponse>(ref TResponse response)
                 respServerSession.storageSession.lockableContext.EndLockable();
                 if (!respServerSession.storageSession.objectStoreLockableContext.IsNull)
                     respServerSession.storageSession.objectStoreLockableContext.EndLockable();
+                respServerSession.storageSession.stateMachineDriver.EndTransaction(txnVersion);
             }
         }
 

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/Resp/RespServerSession.cs

@@ -229,7 +229,7 @@ public RespServerSession(
             // Associate new session with default user and automatically authenticate, if possible
             this.AuthenticateUser(Encoding.ASCII.GetBytes(this.storeWrapper.accessControlList.GetDefaultUserHandle().User.Name));
 
-            txnManager = new TransactionManager(this, storageSession, scratchBufferManager, storeWrapper.serverOptions.EnableCluster, logger);
+            txnManager = new TransactionManager(this, storageSession, scratchBufferManager, storeWrapper.serverOptions.StateMachineDriver, storeWrapper.serverOptions.EnableCluster, logger);
             storageSession.txnManager = txnManager;
 
             clusterSession = storeWrapper.clusterProvider?.CreateClusterSession(txnManager, this._authenticator, this._userHandle, sessionMetrics, basicGarnetApi, networkSender, logger);

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/Storage/Session/StorageSession.cs

@@ -45,6 +45,7 @@ sealed partial class StorageSession : IDisposable
         public readonly FunctionsState functionsState;
 
         public TransactionManager txnManager;
+        public StateMachineDriver stateMachineDriver;
         readonly ILogger logger;
         private readonly CollectionItemBroker itemBroker;
 
@@ -64,7 +65,7 @@ public StorageSession(StoreWrapper storeWrapper,
             this.scratchBufferManager = scratchBufferManager;
             this.logger = logger;
             this.itemBroker = storeWrapper.itemBroker;
-
+            this.stateMachineDriver = storeWrapper.serverOptions.StateMachineDriver;
             parseState.Initialize();
 
             functionsState = storeWrapper.CreateFunctionsState();