Single state machine across main and object stores

libs/cluster/Server/ClusterConfig.cs

@@ -1301,6 +1301,7 @@ public ClusterConfig BumpLocalNodeConfigEpoch()
         /// Check if sender has same local worker epoch as the receiver node and resolve collision.
         /// </summary>
         /// <param name="senderConfig">Incoming configuration object.</param>
+        /// <param name="logger"></param>
         /// <returns>ClusterConfig object with updates.</returns>
         public ClusterConfig HandleConfigEpochCollision(ClusterConfig senderConfig, ILogger logger = null)
         {

libs/cluster/Server/ClusterProvider.cs

@@ -201,7 +201,13 @@ public void OnCheckpointInitiated(out long CheckpointCoveredAofAddress)
         {
             Debug.Assert(serverOptions.EnableCluster);
             if (serverOptions.EnableAOF && clusterManager.CurrentConfig.LocalNodeRole == NodeRole.REPLICA)
-                CheckpointCoveredAofAddress = replicationManager.ReplicationOffset;
+            {
+                // When the replica takes a checkpoint on encountering the checkpoint end marker, it needs to truncate the AOF only
+                // until the checkpoint start marker. Otherwise, we will be left with an AOF that starts at the checkpoint end marker.
+                // ReplicationCheckpointStartOffset is set by { ReplicaReplayTask.Consume -> AofProcessor.ProcessAofRecordInternal } when
+                // it encounters the checkpoint start marker.
+                CheckpointCoveredAofAddress = replicationManager.ReplicationCheckpointStartOffset;
+            }
             else
                 CheckpointCoveredAofAddress = storeWrapper.appendOnlyFile.TailAddress;
 

libs/cluster/Server/Replication/PrimaryOps/DisklessReplication/ReplicationSnapshotIterator.cs

@@ -194,10 +194,6 @@ public void OnStop(bool completed, long numberOfRecords, bool isMainStore, long
             // Wait for flush and response to complete
             replicationSyncManager.WaitForFlush().GetAwaiter().GetResult();
 
-            // Enqueue commit end marker
-            var entryType = isMainStore ? AofEntryType.MainStoreStreamingCheckpointEndCommit : AofEntryType.ObjectStoreStreamingCheckpointEndCommit;
-            replicationSyncManager.ClusterProvider.storeWrapper.EnqueueCommit(entryType, targetVersion);
-
             logger?.LogTrace("{OnStop} {store} {numberOfRecords} {targetVersion}",
                 nameof(OnStop), isMainStore ? "MAIN STORE" : "OBJECT STORE", numberOfRecords, targetVersion);
 

libs/cluster/Server/Replication/ReplicaOps/ReplicaReplayTask.cs

@@ -60,7 +60,11 @@ public unsafe void Consume(byte* record, int recordLength, long currentAddress,
                 var payloadLength = storeWrapper.appendOnlyFile.UnsafeGetLength(ptr);
                 if (payloadLength > 0)
                 {
-                    aofProcessor.ProcessAofRecordInternal(ptr + entryLength, payloadLength, true);
+                    aofProcessor.ProcessAofRecordInternal(ptr + entryLength, payloadLength, true, out var isCheckpointStart);
+                    // Encountered checkpoint start marker, log the ReplicationCheckpointStartOffset so we know the correct AOF truncation
+                    // point when we take a checkpoint at the checkpoint end marker
+                    if (isCheckpointStart)
+                        ReplicationCheckpointStartOffset = ReplicationOffset;
                     entryLength += TsavoriteLog.UnsafeAlign(payloadLength);
                 }
                 else if (payloadLength < 0)

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, bool> checkpointVersionShiftStart;
+        public Action<bool, long, long, bool> 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, bool isStreaming)
+            => checkpointVersionShiftStart?.Invoke(isMainStore, oldVersion, newVersion, isStreaming);
+
+        public override void CheckpointVersionShiftEnd(long oldVersion, long newVersion, bool isStreaming)
+            => checkpointVersionShiftEnd?.Invoke(isMainStore, oldVersion, newVersion, isStreaming);
 
         public void DeleteLogCheckpoint(Guid logToken)
             => deviceFactory.Delete(checkpointNamingScheme.LogCheckpointBase(logToken));

libs/cluster/Server/Replication/ReplicationManager.cs

@@ -53,6 +53,12 @@ public long ReplicationOffset
             set { replicationOffset = value; }
         }
 
+        /// <summary>
+        /// Replication offset corresponding to the checkpoint start marker. We will truncate only to this point after taking a checkpoint (the checkpoint
+        /// is taken only when we encounter a checkpoint end marker).
+        /// </summary>
+        public long ReplicationCheckpointStartOffset;
+
         /// <summary>
         /// Replication offset until which AOF address is valid for old primary if failover has occurred
         /// </summary>
@@ -112,9 +118,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(RecoveryStatus.InitializeRecover))
@@ -157,14 +167,44 @@ 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, bool isStreaming)
+        {
+            if (clusterProvider.clusterManager.CurrentConfig.LocalNodeRole == NodeRole.REPLICA)
+                return;
+
+            if (isStreaming)
+            {
+                if (isMainStore)
+                    storeWrapper.EnqueueCommit(AofEntryType.MainStoreStreamingCheckpointStartCommit, newVersion);
+                else
+                    storeWrapper.EnqueueCommit(AofEntryType.ObjectStoreStreamingCheckpointStartCommit, newVersion);
+            }
+            else
+            {
+                // We enqueue a single checkpoint start marker, since we have unified checkpointing
+                if (isMainStore)
+                    storeWrapper.EnqueueCommit(AofEntryType.CheckpointStartCommit, newVersion);
+            }
+        }
+
+        void CheckpointVersionShiftEnd(bool isMainStore, long oldVersion, long newVersion, bool isStreaming)
         {
             if (clusterProvider.clusterManager.CurrentConfig.LocalNodeRole == NodeRole.REPLICA)
                 return;
-            var entryType = clusterProvider.serverOptions.ReplicaDisklessSync ?
-                (isMainStore ? AofEntryType.MainStoreStreamingCheckpointStartCommit : AofEntryType.ObjectStoreStreamingCheckpointStartCommit) :
-                (isMainStore ? AofEntryType.MainStoreCheckpointStartCommit : AofEntryType.ObjectStoreCheckpointStartCommit);
-            storeWrapper.EnqueueCommit(entryType, newVersion);
+
+            if (isStreaming)
+            {
+                if (isMainStore)
+                    storeWrapper.EnqueueCommit(AofEntryType.MainStoreStreamingCheckpointEndCommit, newVersion);
+                else
+                    storeWrapper.EnqueueCommit(AofEntryType.ObjectStoreStreamingCheckpointEndCommit, newVersion);
+            }
+            else
+            {
+                // We enqueue a single checkpoint end marker, since we have unified checkpointing
+                if (isMainStore)
+                    storeWrapper.EnqueueCommit(AofEntryType.CheckpointEndCommit, newVersion);
+            }
         }
 
         /// <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/AofEntryType.cs

@@ -42,41 +42,34 @@ public enum AofEntryType : byte
         /// </summary>
         TxnAbort = 0x22,
         /// <summary>
-        /// Checkpoint for main store start
+        /// Checkpoint start marker for unified checkpoint
         /// </summary>
-        MainStoreCheckpointStartCommit = 0x30,
+        CheckpointStartCommit = 0x30,
         /// <summary>
-        /// Checkpoint for object store start
+        /// Checkpoint end marker for unified checkpoint
         /// </summary>
-        ObjectStoreCheckpointStartCommit = 0x31,
-        /// <summary>
-        /// Checkpoint for main store end
-        /// </summary>
-        MainStoreCheckpointEndCommit = 0x32,
+        CheckpointEndCommit = 0x32,
         /// <summary>
-        /// Checkpoint for object store end
-        /// </summary>
-        ObjectStoreCheckpointEndCommit = 0x33,
-        /// <summary>
-        /// Streaming checkpoint for main store start
+        /// Streaming checkpoint start marker for main store
         /// </summary>
         MainStoreStreamingCheckpointStartCommit = 0x40,
         /// <summary>
-        /// Streaming checkpoint for object store start
+        /// Streaming checkpoint start marker for object store
         /// </summary>
         ObjectStoreStreamingCheckpointStartCommit = 0x41,
         /// <summary>
-        /// Streaming checkpoint for main store end
+        /// Streaming checkpoint end marker for main store
         /// </summary>
         MainStoreStreamingCheckpointEndCommit = 0x42,
         /// <summary>
-        /// Streaming checkpoint for object store end
+        /// Streaming checkpoint end marker for object store
         /// </summary>
         ObjectStoreStreamingCheckpointEndCommit = 0x43,
         /// <summary>
         /// StoredProcedure
         /// </summary>
         StoredProcedure = 0x50,
+
         /// <summary>
         /// Flush all
         /// </summary>
@@ -85,6 +78,17 @@ public enum AofEntryType : byte
         /// Flush db
         /// </summary>
         FlushDb = 0x61,
+
+        #region Deprecated markers
+        /// <summary>
+        /// Deprecated with unified checkpointing: Checkpoint for object store start
+        /// </summary>
+        ObjectStoreCheckpointStartCommit = 0x31,
+        /// <summary>
+        /// Deprecated with unified checkpointing: Checkpoint for object store end
+        /// </summary>
+        ObjectStoreCheckpointEndCommit = 0x33,
+        #endregion
     }
 
     internal enum AofStoreType : byte

libs/server/AOF/AofHeader.cs

@@ -12,7 +12,7 @@ struct AofHeader
         // * Layout, size, contents of this struct
         // * Any of the AofEntryType or AofStoreType enums' existing value mappings
         // * SpanByte format or header
-        const byte AofHeaderVersion = 1;
+        const byte AofHeaderVersion = 2;
 
         /// <summary>
         /// Version of AOF