Fix Checkpoint Version Shift for Diskless Replication (#1068)

* correctly add commit marker for streaming checkpoint in the AOF

* enqueue only on checkpoint version shift

* remove default value

* rename checkpoint start marker

* add end marker and change enqueue

* guard checkpoint end marker commit using EnableCluster flag

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

@@ -26,7 +26,6 @@ internal sealed unsafe class SnapshotIteratorManager
         long currentFlushEventCount = 0;
         long lastFlushEventCount = 0;
 
-
         public long CheckpointCoveredAddress { get; private set; }
 
         public SnapshotIteratorManager(ReplicationSyncManager replicationSyncManager, CancellationToken cancellationToken, ILogger logger = null)
@@ -195,8 +194,9 @@ public void OnStop(bool completed, long numberOfRecords, bool isMainStore, long
             // Wait for flush and response to complete
             replicationSyncManager.WaitForFlush().GetAwaiter().GetResult();
 
-            // Enqueue version change commit
-            replicationSyncManager.ClusterProvider.storeWrapper.EnqueueCommit(isMainStore, targetVersion);
+            // 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/ReplicationManager.cs

@@ -156,7 +156,10 @@ void CheckpointVersionShift(bool isMainStore, long oldVersion, long newVersion)
         {
             if (clusterProvider.clusterManager.CurrentConfig.LocalNodeRole == NodeRole.REPLICA)
                 return;
-            storeWrapper.EnqueueCommit(isMainStore, newVersion, streaming: true);
+            var entryType = clusterProvider.serverOptions.ReplicaDisklessSync ?
+                (isMainStore ? AofEntryType.MainStoreStreamingCheckpointStartCommit : AofEntryType.ObjectStoreStreamingCheckpointStartCommit) :
+                (isMainStore ? AofEntryType.MainStoreCheckpointStartCommit : AofEntryType.ObjectStoreCheckpointStartCommit);
+            storeWrapper.EnqueueCommit(entryType, newVersion);
         }
 
         /// <summary>

libs/server/AOF/AofEntryType.cs

@@ -3,21 +3,79 @@
 
 namespace Garnet.server
 {
-    enum AofEntryType : byte
+    public enum AofEntryType : byte
     {
+        /// <summary>
+        /// Store upsert
+        /// </summary>
         StoreUpsert = 0x00,
+        /// <summary>
+        /// Store RMW
+        /// </summary>
         StoreRMW = 0x01,
+        /// <summary>
+        /// Store delete
+        /// </summary>
         StoreDelete = 0x02,
+        /// <summary>
+        /// Object store upsert
+        /// </summary>
         ObjectStoreUpsert = 0x10,
+        /// <summary>
+        /// Object store RMW
+        /// </summary>
         ObjectStoreRMW = 0x11,
+        /// <summary>
+        /// Object store delete
+        /// </summary>
         ObjectStoreDelete = 0x12,
+        /// <summary>
+        /// Transaction start
+        /// </summary>
         TxnStart = 0x20,
+        /// <summary>
+        /// Transaction commit
+        /// </summary>
         TxnCommit = 0x21,
+        /// <summary>
+        /// Transaction abort
+        /// </summary>
         TxnAbort = 0x22,
-        MainStoreCheckpointCommit = 0x30,
-        ObjectStoreCheckpointCommit = 0x31,
-        MainStoreStreamingCheckpointCommit = 0x40,
-        ObjectStoreStreamingCheckpointCommit = 0x41,
+        /// <summary>
+        /// Checkpoint for main store start
+        /// </summary>
+        MainStoreCheckpointStartCommit = 0x30,
+        /// <summary>
+        /// Checkpoint for object store start
+        /// </summary>
+        ObjectStoreCheckpointStartCommit = 0x31,
+        /// <summary>
+        /// Checkpoint for main store end
+        /// </summary>
+        MainStoreCheckpointEndCommit = 0x32,
+        /// <summary>
+        /// Checkpoint for object store end
+        /// </summary>
+        ObjectStoreCheckpointEndCommit = 0x33,
+        /// <summary>
+        /// Streaming checkpoint for main store start
+        /// </summary>
+        MainStoreStreamingCheckpointStartCommit = 0x40,
+        /// <summary>
+        /// Streaming checkpoint for object store start
+        /// </summary>
+        ObjectStoreStreamingCheckpointStartCommit = 0x41,
+        /// <summary>
+        /// Streaming checkpoint for main store end
+        /// </summary>
+        MainStoreStreamingCheckpointEndCommit = 0x42,
+        /// <summary>
+        /// Streaming checkpoint for object store end
+        /// </summary>
+        ObjectStoreStreamingCheckpointEndCommit = 0x43,
+        /// <summary>
+        /// StoredProcedure
+        /// </summary>
         StoredProcedure = 0x50,
     }
 

libs/server/AOF/AofProcessor.cs

@@ -200,42 +200,41 @@ public unsafe void ProcessAofRecordInternal(byte* ptr, int length, bool asReplic
             switch (header.opType)
             {
                 case AofEntryType.TxnStart:
-                    inflightTxns[header.sessionID] = new List<byte[]>();
+                    inflightTxns[header.sessionID] = [];
                     break;
                 case AofEntryType.TxnAbort:
                 case AofEntryType.TxnCommit:
                     // We encountered a transaction end without start - this could happen because we truncated the AOF
                     // after a checkpoint, and the transaction belonged to the previous version. It can safely
                     // be ignored.
                     break;
-                case AofEntryType.MainStoreCheckpointCommit:
-                    if (asReplica)
-                    {
-                        if (header.storeVersion > storeWrapper.store.CurrentVersion)
-                        {
-                            storeWrapper.TakeCheckpoint(false, StoreType.Main, logger);
-                        }
-                    }
+                case AofEntryType.MainStoreCheckpointStartCommit:
+                    if (asReplica && header.storeVersion > storeWrapper.store.CurrentVersion)
+                        _ = storeWrapper.TakeCheckpoint(false, StoreType.Main, logger);
                     break;
-                case AofEntryType.ObjectStoreCheckpointCommit:
-                    if (asReplica)
-                    {
-                        if (header.storeVersion > storeWrapper.objectStore.CurrentVersion)
-                        {
-                            storeWrapper.TakeCheckpoint(false, StoreType.Object, logger);
-                        }
-                    }
+                case AofEntryType.ObjectStoreCheckpointStartCommit:
+                    if (asReplica && header.storeVersion > storeWrapper.objectStore.CurrentVersion)
+                        _ = storeWrapper.TakeCheckpoint(false, StoreType.Object, logger);
                     break;
-                case AofEntryType.MainStoreStreamingCheckpointCommit:
+                case AofEntryType.MainStoreCheckpointEndCommit:
+                case AofEntryType.ObjectStoreCheckpointEndCommit:
+                    break;
+                case AofEntryType.MainStoreStreamingCheckpointStartCommit:
                     Debug.Assert(storeWrapper.serverOptions.ReplicaDisklessSync);
-                    if (header.storeVersion > storeWrapper.store.CurrentVersion)
+                    if (asReplica && header.storeVersion > storeWrapper.store.CurrentVersion)
                         storeWrapper.store.SetVersion(header.storeVersion);
                     break;
-                case AofEntryType.ObjectStoreStreamingCheckpointCommit:
+                case AofEntryType.MainStoreStreamingCheckpointEndCommit:
+                    Debug.Assert(storeWrapper.serverOptions.ReplicaDisklessSync);
+                    break;
+                case AofEntryType.ObjectStoreStreamingCheckpointStartCommit:
                     Debug.Assert(storeWrapper.serverOptions.ReplicaDisklessSync);
-                    if (header.storeVersion > storeWrapper.store.CurrentVersion)
+                    if (asReplica && header.storeVersion > storeWrapper.store.CurrentVersion)
                         storeWrapper.objectStore.SetVersion(header.storeVersion);
                     break;
+                case AofEntryType.ObjectStoreStreamingCheckpointEndCommit:
+                    Debug.Assert(storeWrapper.serverOptions.ReplicaDisklessSync);
+                    break;
                 default:
                     ReplayOp(ptr);
                     break;
@@ -423,7 +422,8 @@ static AofStoreType ToAofStoreType(AofEntryType type)
                 AofEntryType.StoreUpsert or AofEntryType.StoreRMW or AofEntryType.StoreDelete => AofStoreType.MainStoreType,
                 AofEntryType.ObjectStoreUpsert or AofEntryType.ObjectStoreRMW or AofEntryType.ObjectStoreDelete => AofStoreType.ObjectStoreType,
                 AofEntryType.TxnStart or AofEntryType.TxnCommit or AofEntryType.TxnAbort or AofEntryType.StoredProcedure => AofStoreType.TxnType,
-                AofEntryType.MainStoreCheckpointCommit or AofEntryType.ObjectStoreCheckpointCommit => AofStoreType.CheckpointType,
+                AofEntryType.MainStoreCheckpointStartCommit or AofEntryType.ObjectStoreCheckpointStartCommit or AofEntryType.MainStoreStreamingCheckpointStartCommit or AofEntryType.ObjectStoreStreamingCheckpointStartCommit => AofStoreType.CheckpointType,
+                AofEntryType.MainStoreCheckpointEndCommit or AofEntryType.ObjectStoreCheckpointEndCommit or AofEntryType.MainStoreStreamingCheckpointEndCommit or AofEntryType.ObjectStoreStreamingCheckpointEndCommit => AofStoreType.CheckpointType,
                 _ => throw new GarnetException($"Conversion to AofStoreType not possible for {type}"),
             };
         }

libs/server/StoreWrapper.cs

@@ -495,18 +495,13 @@ void DoCompaction()
         /// <summary>
         /// Append a checkpoint commit to the AOF
         /// </summary>
-        /// <param name="isMainStore"></param>
+        /// <param name="entryType"></param>
         /// <param name="version"></param>
-        /// <param name="streaming"></param>
-        public void EnqueueCommit(bool isMainStore, long version, bool streaming = false)
+        public void EnqueueCommit(AofEntryType entryType, long version)
         {
-            var opType = streaming ?
-                isMainStore ? AofEntryType.MainStoreStreamingCheckpointCommit : AofEntryType.ObjectStoreStreamingCheckpointCommit :
-                isMainStore ? AofEntryType.MainStoreCheckpointCommit : AofEntryType.ObjectStoreCheckpointCommit;
-
             AofHeader header = new()
             {
-                opType = isMainStore ? AofEntryType.MainStoreCheckpointCommit : AofEntryType.ObjectStoreCheckpointCommit,
+                opType = entryType,
                 storeVersion = version,
                 sessionID = -1
             };
@@ -849,18 +844,30 @@ private async Task InitiateCheckpoint(bool full, CheckpointType checkpointType,
             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);
+                }
             }
             else
             {
                 if (storeType is StoreType.Main or StoreType.All)
+                {
                     storeCheckpointResult = await store.TakeHybridLogCheckpointAsync(checkpointType, tryIncremental);
+                    if (serverOptions.EnableCluster && clusterProvider.IsPrimary()) EnqueueCommit(AofEntryType.MainStoreCheckpointEndCommit, store.CurrentVersion);
+                }
 
                 if (objectStore != null && (storeType == StoreType.Object || storeType == StoreType.All))
+                {
                     objectStoreCheckpointResult = await objectStore.TakeHybridLogCheckpointAsync(checkpointType, tryIncremental);
+                    if (serverOptions.EnableCluster && clusterProvider.IsPrimary()) EnqueueCommit(AofEntryType.ObjectStoreCheckpointEndCommit, objectStore.CurrentVersion);
+                }
             }
 
             // If cluster is enabled the replication manager is responsible for truncating AOF