expose create cookie delegate

Author: vazois

libs/cluster/Server/Replication/ReplicationLogCheckpointManager.cs

@@ -11,27 +11,55 @@
 
 namespace Garnet.cluster
 {
-    internal sealed class ReplicationLogCheckpointManager(
-        INamedDeviceFactoryCreator deviceFactoryCreator,
-        ICheckpointNamingScheme checkpointNamingScheme,
-        bool isMainStore,
-        bool removeOutdated = false,
-        int fastCommitThrottleFreq = 0,
-        ILogger logger = null) : DeviceLogCommitCheckpointManager(deviceFactoryCreator, checkpointNamingScheme, removeOutdated: false, fastCommitThrottleFreq, logger), IDisposable
+    internal sealed class ReplicationLogCheckpointManager : DeviceLogCommitCheckpointManager, IDisposable
     {
         public long CurrentSafeAofAddress = 0;
         public long RecoveredSafeAofAddress = 0;
 
         public string CurrentReplicationId = string.Empty;
         public string RecoveredReplicationId = string.Empty;
 
-        readonly bool isMainStore = isMainStore;
+        readonly bool isMainStore;
         public Action<bool, long, long, bool> checkpointVersionShiftStart;
         public Action<bool, long, long, bool> checkpointVersionShiftEnd;
 
-        readonly bool safelyRemoveOutdated = removeOutdated;
+        readonly bool safelyRemoveOutdated;
 
-        readonly ILogger logger = logger;
+        readonly ILogger logger;
+
+        public ReplicationLogCheckpointManager(
+            INamedDeviceFactoryCreator deviceFactoryCreator,
+            ICheckpointNamingScheme checkpointNamingScheme,
+            bool isMainStore,
+            bool safelyRemoveOutdated = false,
+            int fastCommitThrottleFreq = 0,
+            ILogger logger = null)
+            : base(deviceFactoryCreator, checkpointNamingScheme, removeOutdated: false, fastCommitThrottleFreq, logger)
+        {
+            this.isMainStore = isMainStore;
+            this.safelyRemoveOutdated = safelyRemoveOutdated;
+            this.createCookieDelegate = CreateCookie;
+            this.logger = logger;
+
+            // Pre-append cookie in commitMetadata.
+            // cookieMetadata 52 bytes
+            // 1. 4 bytes to track size of cookie
+            // 2. 8 bytes for checkpointCoveredAddress
+            // 3. 40 bytes for primaryReplicationId
+            unsafe byte[] CreateCookie()
+            {
+                var cookie = new byte[sizeof(int) + sizeof(long) + CurrentReplicationId.Length];
+                var primaryReplIdBytes = Encoding.ASCII.GetBytes(CurrentReplicationId);
+                fixed (byte* ptr = cookie)
+                fixed (byte* pridPtr = primaryReplIdBytes)
+                {
+                    *(int*)ptr = sizeof(long) + CurrentReplicationId.Length;
+                    *(long*)(ptr + 4) = CurrentSafeAofAddress;
+                    Buffer.MemoryCopy(pridPtr, ptr + 12, primaryReplIdBytes.Length, primaryReplIdBytes.Length);
+                }
+                return cookie;
+            }
+        }
 
         public override void CheckpointVersionShiftStart(long oldVersion, long newVersion, bool isStreaming)
             => checkpointVersionShiftStart?.Invoke(isMainStore, oldVersion, newVersion, isStreaming);
@@ -63,31 +91,6 @@ public IDevice GetDevice(CheckpointFileType retStateType, Guid fileToken)
 
         #region ICheckpointManager
 
-        /// <summary>
-        /// Pre-append cookie in commitMetadata.
-        /// cookieMetadata 52 bytes
-        /// 1. 4 bytes to track size of cookie
-        /// 2. 8 bytes for checkpointCoveredAddress
-        /// 3. 40 bytes for primaryReplicationId
-        /// </summary>
-        /// <returns></returns>
-        private unsafe byte[] CreateCookie()
-        {
-            var cookie = new byte[sizeof(int) + sizeof(long) + CurrentReplicationId.Length];
-            var primaryReplIdBytes = Encoding.ASCII.GetBytes(CurrentReplicationId);
-            fixed (byte* ptr = cookie)
-            fixed (byte* pridPtr = primaryReplIdBytes)
-            {
-                *(int*)ptr = sizeof(long) + CurrentReplicationId.Length;
-                *(long*)(ptr + 4) = CurrentSafeAofAddress;
-                Buffer.MemoryCopy(pridPtr, ptr + 12, primaryReplIdBytes.Length, primaryReplIdBytes.Length);
-            }
-            return cookie;
-        }
-
-        /// <inheritdoc />
-        public override byte[] GetCookie() => CreateCookie();
-
         private HybridLogRecoveryInfo ConverMetadata(byte[] checkpointMetadata)
         {
             var success = true;

libs/storage/Tsavorite/cs/src/core/Index/CheckpointManagement/DeviceLogCommitCheckpointManager.cs

@@ -31,6 +31,8 @@ public class DeviceLogCommitCheckpointManager : ILogCommitManager, ICheckpointMa
         protected readonly ICheckpointNamingScheme checkpointNamingScheme;
         private readonly SemaphoreSlim semaphore;
 
+        protected Func<byte[]> createCookieDelegate;
+
         private readonly bool removeOutdated;
         private SectorAlignedBufferPool bufferPool;
 
@@ -173,7 +175,10 @@ public byte[] GetCommitMetadata(long commitNum)
         #region ICheckpointManager
 
         /// <inheritdoc />
-        public virtual byte[] GetCookie() => null;
+        public void AddCookieDelegate(Func<byte[]> createCookieAction) => this.createCookieDelegate = createCookieAction;
+
+        /// <inheritdoc />
+        public byte[] GetCookie() => createCookieDelegate == null ? null : createCookieDelegate();
 
         /// <inheritdoc />
         public unsafe void CommitIndexCheckpoint(Guid indexToken, byte[] commitMetadata)

libs/storage/Tsavorite/cs/src/core/Index/Checkpointing/TsavoriteStateMachineProperties.cs

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
+using System;
+
 namespace Tsavorite.core
 {
     public partial class TsavoriteKV<TKey, TValue, TStoreFunctions, TAllocator> : TsavoriteBase
@@ -9,6 +11,12 @@ public partial class TsavoriteKV<TKey, TValue, TStoreFunctions, TAllocator> : Ts
     {
         internal long lastVersion;
 
+        /// <summary>
+        /// Add cookie delegate
+        /// </summary>
+        /// <param name="cookieCreationDelegate"></param>
+        public void AddCookieDelegate(Func<byte[]> cookieCreationDelegate) => checkpointManager.AddCookieDelegate(cookieCreationDelegate);
+
         private byte[] recoveredCommitCookie;
         /// <summary>
         /// User-specified commit cookie persisted with last recovered commit

libs/storage/Tsavorite/cs/src/core/Index/Recovery/ICheckpointManager.cs

@@ -30,6 +30,12 @@ namespace Tsavorite.core
     /// </summary>
     public interface ICheckpointManager : IDisposable
     {
+        /// <summary>
+        /// Add create cookie action
+        /// </summary>
+        /// <param name="createCookieAction"></param>
+        void AddCookieDelegate(Func<byte[]> createCookieAction);
+
         /// <summary>
         /// Get current cookie
         /// </summary>

libs/storage/Tsavorite/cs/test/SimpleRecoveryTest.cs

@@ -149,7 +149,7 @@ private async ValueTask SimpleRecoveryTest1_Worker(CheckpointType checkpointType
             }
 
             if (testCommitCookie)
-                store1.AddCookie(commitCookie);
+                store1.AddCookieDelegate(() => commitCookie);
             _ = store1.TryInitiateFullCheckpoint(out Guid token, checkpointType);
             if (completionSyncMode == CompletionSyncMode.Sync)
                 store1.CompleteCheckpointAsync().AsTask().GetAwaiter().GetResult();