Implement TaskManager (#1498)

* add task manager implementation

* Add test for replica to stop background collect task

* add calls to resume suspend maintenance tasks

* separate primary only from generic background tasks

* misc fixes

* fix formatting

* loggerFactory null check

* simplify task manager implementation

* addressing comments

* no allocation lookup placement category

* add verbose log message

* update RegisterAndRun

* pass global cts in when waiting for task manager tasks to complete

* check if TaskManager disposed before using global cts

* add TaskManager tests

* fix test

* addressing comments and adding more tests

* add test cleanup with exception

* bump version

* add Allure tests for task manager

Author: vazois

Version.props

@@ -1,6 +1,6 @@
 <Project>
 	<!-- VersionPrefix property for builds and packages -->
 	<PropertyGroup>
-		<VersionPrefix>1.0.92</VersionPrefix>
+		<VersionPrefix>1.0.93</VersionPrefix>
 	</PropertyGroup>
 </Project>

libs/cluster/Server/Failover/ReplicaFailoverSession.cs

@@ -145,6 +145,9 @@ private bool TakeOverAsPrimary()
                 if (!clusterProvider.replicationManager.InitializeCheckpointStore())
                     logger?.LogWarning("Failed acquiring latest memory checkpoint metadata at {method}", nameof(TakeOverAsPrimary));
                 _ = clusterProvider.BumpAndWaitForEpochTransition();
+
+                // Resume all background maintenance that were possibly shutdown when this node became a replica
+                clusterProvider.storeWrapper.StartPrimaryTasks();
             }
             finally
             {

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

@@ -80,6 +80,9 @@ async Task<string> TryBeginReplicaSync(bool downgradeLock)
                     if (!disklessSync)
                         storeWrapper.Reset();
 
+                    // Suspend background tasks that may interfere with AOF
+                    await storeWrapper.SuspendPrimaryOnlyTasks();
+
                     // Send request to primary
                     //      Primary will initiate background task and start sending checkpoint data
                     //

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

@@ -122,6 +122,9 @@ async Task<string> ReplicaSyncAttachTask(bool downgradeLock)
                     // Reset the database in preparation for connecting to primary
                     storeWrapper.Reset();
 
+                    // Suspend background tasks that may interfere with AOF
+                    await storeWrapper.SuspendPrimaryOnlyTasks();
+
                     // Send request to primary
                     //      Primary will initiate background task and start sending checkpoint data
                     //

libs/cluster/Session/ReplicaOfCommand.cs

@@ -45,6 +45,7 @@ private bool TryREPLICAOF(out bool invalidParameters)
                     clusterProvider.replicationManager.TryUpdateForFailover();
                     clusterProvider.replicationManager.ResetReplayIterator();
                     UnsafeBumpAndWaitForEpochTransition();
+                    clusterProvider.storeWrapper.StartPrimaryTasks();
                 }
                 finally
                 {

libs/server/Databases/MultiDatabaseManager.cs

@@ -524,7 +524,6 @@ public override bool GrowIndexesIfNeeded(CancellationToken token = default)
             {
                 var activeDbIdsMapSize = activeDbIds.ActualSize;
                 var activeDbIdsMapSnapshot = activeDbIds.Map;
-
                 var databasesMapSnapshot = databases.Map;
 
                 for (var i = 0; i < activeDbIdsMapSize; i++)

libs/server/Resp/RespServerSessionSlotVerify.cs

@@ -21,6 +21,11 @@ internal sealed unsafe partial class RespServerSession : ServerSessionBase
         bool NetworkKeyArraySlotVerify(Span<ArgSlice> keys, bool readOnly, int count = -1)
             => clusterSession != null && clusterSession.NetworkKeyArraySlotVerify(keys, readOnly, SessionAsking, ref dcurr, ref dend, count);
 
+        /// <summary>
+        /// Validate if this command can be served based on the current slot assignment
+        /// </summary>
+        /// <param name="cmd"></param>
+        /// <returns></returns>
         bool CanServeSlot(RespCommand cmd)
         {
             Debug.Assert(clusterSession != null);

libs/server/StoreWrapper.cs

@@ -130,6 +130,11 @@ public sealed class StoreWrapper
         /// </summary>
         internal readonly LuaTimeoutManager luaTimeoutManager;
 
+        /// <summary>
+        /// Background task manager instance
+        /// </summary>
+        internal readonly TaskManager taskManager;
+
         private IDatabaseManager databaseManager;
         SingleWriterMultiReaderLock databaseManagerLock;
 
@@ -151,7 +156,9 @@ public sealed class StoreWrapper
 
         internal readonly CancellationTokenSource ctsCommit;
 
-        // True if StoreWrapper instance is disposed
+        /// <summary>
+        /// True if StoreWrapper instance is disposed
+        /// </summary>
         bool disposed;
 
         /// <summary>
@@ -197,6 +204,7 @@ public StoreWrapper(
             this.sessionLogger = loggerFactory?.CreateLogger("Session");
             this.accessControlList = accessControlList;
             this.GarnetObjectSerializer = new GarnetObjectSerializer(this.customCommandManager);
+            this.taskManager = new TaskManager(loggerFactory?.CreateLogger("TaskManager"));
             this.loggingFrequency = TimeSpan.FromSeconds(serverOptions.LoggingFrequency);
 
             logger?.LogTrace("StoreWrapper logging frequency: {loggingFrequency} seconds.", this.loggingFrequency);
@@ -654,7 +662,7 @@ async Task AutoCheckpointBasedOnAofSizeLimit(long aofSizeLimit, CancellationToke
             }
         }
 
-        async Task CommitTask(int commitFrequencyMs, ILogger logger = null, CancellationToken token = default)
+        async Task CommitTask(int commitFrequencyMs, CancellationToken token = default, ILogger logger = null)
         {
             try
             {
@@ -768,9 +776,12 @@ async Task ExpiredKeyDeletionScanTask(int expiredKeyDeletionScanFrequencySecs, C
         public (long numExpiredKeysFound, long totalRecordsScanned) ExpiredKeyDeletionScan(int dbId)
             => databaseManager.ExpiredKeyDeletionScan(dbId);
 
-        /// <summary>Grows indexes of both main store and object store if current size is too small.</summary>
+        /// <summary>
+        /// Grows indexes of both main store and object store if current size is too small.
+        /// </summary>
         /// <param name="token"></param>
-        private async void IndexAutoGrowTask(CancellationToken token)
+        /// <returns></returns>
+        private async Task IndexAutoGrowTask(CancellationToken token)
         {
             try
             {
@@ -799,37 +810,16 @@ internal void Start()
             clusterProvider?.Start();
             luaTimeoutManager?.Start();
 
-            if (serverOptions.AofSizeLimit.Length > 0)
-            {
-                var aofSizeLimitBytes = 1L << serverOptions.AofSizeLimitSizeBits();
-                Task.Run(async () => await AutoCheckpointBasedOnAofSizeLimit(aofSizeLimitBytes, ctsCommit.Token, logger));
-            }
-
-            if (serverOptions.CommitFrequencyMs > 0 && serverOptions.EnableAOF)
+            // Start background maintenance tasks that should run only on the primary
+            if (clusterProvider == null || clusterProvider.IsPrimary())
             {
-                Task.Run(async () => await CommitTask(serverOptions.CommitFrequencyMs, logger, ctsCommit.Token));
+                StartPrimaryTasks();
             }
 
-            if (serverOptions.CompactionFrequencySecs > 0 && serverOptions.CompactionType != LogCompactionType.None)
-            {
-                Task.Run(async () => await CompactionTask(serverOptions.CompactionFrequencySecs, ctsCommit.Token));
-            }
-
-            if (serverOptions.ExpiredObjectCollectionFrequencySecs > 0)
-            {
-                Task.Run(async () => await ObjectCollectTask(serverOptions.ExpiredObjectCollectionFrequencySecs, ctsCommit.Token));
-            }
-
-            if (serverOptions.ExpiredKeyDeletionScanFrequencySecs > 0)
-            {
-                Task.Run(async () => await ExpiredKeyDeletionScanTask(serverOptions.ExpiredKeyDeletionScanFrequencySecs, ctsCommit.Token));
-            }
-
-            if (serverOptions.AdjustedIndexMaxCacheLines > 0 || serverOptions.AdjustedObjectStoreIndexMaxCacheLines > 0)
-            {
-                Task.Run(() => IndexAutoGrowTask(ctsCommit.Token));
-            }
+            // Start generic node tasks
+            StartGenericNodeTasks();
 
+            // Start object size trackers
             databaseManager.StartObjectSizeTrackers(ctsCommit.Token);
         }
 
@@ -914,9 +904,63 @@ public void Dispose()
             monitor?.Dispose();
             luaTimeoutManager?.Dispose();
             ctsCommit?.Cancel();
+            taskManager.Dispose();
             databaseManager.Dispose();
 
             ctsCommit?.Dispose();
         }
+
+        /// <summary>
+        /// Suspend background task that may interfere with the replicas AOF
+        /// </summary>
+        /// <returns></returns>
+        public async Task SuspendPrimaryOnlyTasks()
+        {
+            await taskManager.Cancel(TaskPlacementCategory.Primary);
+        }
+
+        /// <summary>
+        /// Start background maintenance tasks that should only run when this node is a primary
+        /// </summary>
+        /// <returns></returns>
+        public void StartPrimaryTasks()
+        {
+            if (serverOptions.AofSizeLimit.Length > 0)
+            {
+                var aofSizeLimitBytes = 1L << serverOptions.AofSizeLimitSizeBits();
+                taskManager.RegisterAndRun(TaskType.AofSizeLimitTask, (token) => AutoCheckpointBasedOnAofSizeLimit(aofSizeLimitBytes, token, logger));
+            }
+
+            if (serverOptions.CommitFrequencyMs > 0 && serverOptions.EnableAOF)
+            {
+                taskManager.RegisterAndRun(TaskType.CommitTask, (token) => CommitTask(serverOptions.CommitFrequencyMs, token, logger));
+            }
+
+            if (serverOptions.CompactionFrequencySecs > 0 && serverOptions.CompactionType != LogCompactionType.None)
+            {
+                taskManager.RegisterAndRun(TaskType.CompactionTask, (token) => CompactionTask(serverOptions.CompactionFrequencySecs, token));
+            }
+
+            if (serverOptions.ExpiredObjectCollectionFrequencySecs > 0)
+            {
+                taskManager.RegisterAndRun(TaskType.ObjectCollectTask, (token) => ObjectCollectTask(serverOptions.ExpiredObjectCollectionFrequencySecs, token));
+            }
+
+            if (serverOptions.ExpiredKeyDeletionScanFrequencySecs > 0)
+            {
+                taskManager.RegisterAndRun(TaskType.ExpiredKeyDeletionTask, (token) => ExpiredKeyDeletionScanTask(serverOptions.ExpiredKeyDeletionScanFrequencySecs, token));
+            }
+        }
+
+        /// <summary>
+        /// Start background maintenance generic tasks
+        /// </summary>
+        public void StartGenericNodeTasks()
+        {
+            if (serverOptions.AdjustedIndexMaxCacheLines > 0 || serverOptions.AdjustedObjectStoreIndexMaxCacheLines > 0)
+            {
+                taskManager.RegisterAndRun(TaskType.IndexAutoGrowTask, (token) => IndexAutoGrowTask(token));
+            }
+        }
     }
 }