Some fixes to database checkpointing (#1370)

* wip

* wip

* wip

* fix

* Fixed comments

Author: TalZaccai

libs/server/Databases/IDatabaseManager.cs

@@ -97,7 +97,7 @@ public interface IDatabaseManager : IDisposable
         public void RecoverCheckpoint(bool replicaRecover = false, bool recoverMainStoreFromToken = false, bool recoverObjectStoreFromToken = false, CheckpointMetadata metadata = null);
 
         /// <summary>
-        /// Take checkpoint of all active databases
+        /// Take checkpoint of all active databases if checkpointing is not in progress
         /// </summary>
         /// <param name="background">True if method can return before checkpoint is taken</param>
         /// <param name="logger">Logger</param>
@@ -106,7 +106,7 @@ public interface IDatabaseManager : IDisposable
         public bool TakeCheckpoint(bool background, ILogger logger = null, CancellationToken token = default);
 
         /// <summary>
-        /// Take checkpoint of specified database ID
+        /// Take checkpoint of specified database ID if checkpointing is not in progress
         /// </summary>
         /// <param name="background">True if method can return before checkpoint is taken</param>
         /// <param name="dbId">ID of database to checkpoint</param>

libs/server/Databases/MultiDatabaseManager.cs

@@ -36,6 +36,9 @@ internal class MultiDatabaseManager : DatabaseManagerBase
         // The swap-db operation should take a write lock and any operation that should be swap-db-safe should take a read lock.
         SingleWriterMultiReaderLock databasesContentLock;
 
+        // Lock for synchronizing checkpointing of all active DBs (if more than one)
+        SingleWriterMultiReaderLock multiDbCheckpointingLock;
+
         // Reusable task array for tracking checkpointing of multiple DBs
         // Used by recurring checkpointing task if multiple DBs exist
         Task[] checkpointTasks;
@@ -152,20 +155,40 @@ public override bool TakeCheckpoint(bool background, ILogger logger = null, Canc
             var lockAcquired = TryGetDatabasesContentReadLock(token);
             if (!lockAcquired) return false;
 
-            try
+            var checkpointTask = Task.Run(async () =>
             {
-                var activeDbIdsMapSize = activeDbIds.ActualSize;
-                var activeDbIdsMapSnapshot = activeDbIds.Map;
-                Array.Copy(activeDbIdsMapSnapshot, dbIdsToCheckpoint, activeDbIdsMapSize);
+                var checkpointLockTaken = false;
 
-                TakeDatabasesCheckpointAsync(activeDbIdsMapSize, logger: logger, token: token).GetAwaiter().GetResult();
-            }
-            finally
-            {
-                databasesContentLock.ReadUnlock();
-            }
+                try
+                {
+                    var activeDbIdsMapSize = activeDbIds.ActualSize;
 
-            return true;
+                    if (activeDbIdsMapSize > 1)
+                    {
+                        if (!multiDbCheckpointingLock.TryWriteLock())
+                            return false;
+
+                        checkpointLockTaken = true;
+                    }
+
+                    var activeDbIdsMapSnapshot = activeDbIds.Map;
+                    Array.Copy(activeDbIdsMapSnapshot, dbIdsToCheckpoint, activeDbIdsMapSize);
+
+                    return await TakeDatabasesCheckpointAsync(activeDbIdsMapSize, logger: logger, token: token);
+                }
+                finally
+                {
+                    if (checkpointLockTaken)
+                        multiDbCheckpointingLock.WriteUnlock();
+
+                    databasesContentLock.ReadUnlock();
+                }
+            }, token).GetAwaiter();
+
+            if (background)
+                return true;
+
+            return checkpointTask.GetResult();
         }
 
         /// <inheritdoc/>
@@ -175,7 +198,8 @@ public override bool TakeCheckpoint(bool background, int dbId, ILogger logger =
             var databasesMapSnapshot = databases.Map;
             Debug.Assert(dbId < databasesMapSize && databasesMapSnapshot[dbId] != null);
 
-            if (!TryPauseCheckpointsContinuousAsync(dbId, token).GetAwaiter().GetResult())
+            // Check if checkpoint already in progress
+            if (!TryPauseCheckpoints(dbId))
                 return false;
 
             var checkpointTask = TakeCheckpointAsync(databasesMapSnapshot[dbId], logger: logger, token: token).ContinueWith(
@@ -212,12 +236,12 @@ public override async Task TakeOnDemandCheckpointAsync(DateTimeOffset entryTime,
 
             var db = databasesMapSnapshot[dbId];
 
-            // Take lock to ensure no other task will be taking a checkpoint
+            // Check if checkpoint already in progress
             var checkpointsPaused = TryPauseCheckpoints(dbId);
 
             try
             {
-                // If an external task has taken a checkpoint beyond the provided entryTime return
+                // If another checkpoint is in progress or a checkpoint was taken beyond the provided entryTime - return
                 if (!checkpointsPaused || db.LastSaveTime > entryTime)
                     return;
 
@@ -241,10 +265,21 @@ public override async Task TaskCheckpointBasedOnAofSizeLimitAsync(long aofSizeLi
             var lockAcquired = TryGetDatabasesContentReadLock(token);
             if (!lockAcquired) return;
 
+            var activeDbIdsMapSize = activeDbIds.ActualSize;
+
+            var checkpointLockTaken = false;
+
             try
             {
+                if (activeDbIdsMapSize > 1)
+                {
+                    if (!multiDbCheckpointingLock.TryWriteLock())
+                        return;
+
+                    checkpointLockTaken = true;
+                }
+
                 var databasesMapSnapshot = databases.Map;
-                var activeDbIdsMapSize = activeDbIds.ActualSize;
                 var activeDbIdsMapSnapshot = activeDbIds.Map;
 
                 var dbIdsIdx = 0;
@@ -270,6 +305,9 @@ public override async Task TaskCheckpointBasedOnAofSizeLimitAsync(long aofSizeLi
             }
             finally
             {
+                if (checkpointLockTaken)
+                    multiDbCheckpointingLock.WriteUnlock();
+
                 databasesContentLock.ReadUnlock();
             }
         }
@@ -948,7 +986,7 @@ private void CopyDatabases(IDatabaseManager src, bool enableAof)
         /// <param name="logger">Logger</param>
         /// <param name="token">Cancellation token</param>
         /// <returns>False if checkpointing already in progress</returns>
-        private async Task TakeDatabasesCheckpointAsync(int dbIdsCount, ILogger logger = null,
+        private async Task<bool> TakeDatabasesCheckpointAsync(int dbIdsCount, ILogger logger = null,
             CancellationToken token = default)
         {
             Debug.Assert(checkpointTasks != null);
@@ -958,40 +996,32 @@ private async Task TakeDatabasesCheckpointAsync(int dbIdsCount, ILogger logger =
                 checkpointTasks[i] = Task.CompletedTask;
 
             var lockAcquired = TryGetDatabasesContentReadLock(token);
-            if (!lockAcquired) return;
+            if (!lockAcquired) return false;
 
             try
             {
                 var databaseMapSnapshot = databases.Map;
 
-                var currIdx = 0;
-                while (currIdx < dbIdsCount)
+                for (var currIdx = 0; currIdx < dbIdsCount; currIdx++)
                 {
                     var dbId = dbIdsToCheckpoint[currIdx];
 
-                    // Prevent parallel checkpoint
-                    if (!await TryPauseCheckpointsContinuousAsync(dbId, token))
+                    // If a checkpoint is already in progress for this database, skip it
+                    if (!TryPauseCheckpoints(dbId))
                         continue;
 
                     checkpointTasks[currIdx] = TakeCheckpointAsync(databaseMapSnapshot[dbId], logger: logger, token: token).ContinueWith(
                         t =>
                         {
-                            try
-                            {
-                                if (t.IsCompletedSuccessfully)
-                                {
-                                    var storeTailAddress = t.Result.Item1;
-                                    var objectStoreTailAddress = t.Result.Item2;
-                                    UpdateLastSaveData(dbId, storeTailAddress, objectStoreTailAddress);
-                                }
-                            }
-                            finally
-                            {
-                                ResumeCheckpoints(dbId);
-                            }
-                        }, TaskContinuationOptions.ExecuteSynchronously);
+                            ResumeCheckpoints(dbId);
 
-                    currIdx++;
+                            if (!t.IsCompletedSuccessfully)
+                                return;
+
+                            var storeTailAddress = t.Result.Item1;
+                            var objectStoreTailAddress = t.Result.Item2;
+                            UpdateLastSaveData(dbId, storeTailAddress, objectStoreTailAddress);
+                        }, TaskContinuationOptions.ExecuteSynchronously);
                 }
 
                 await Task.WhenAll(checkpointTasks);
@@ -1004,6 +1034,8 @@ private async Task TakeDatabasesCheckpointAsync(int dbIdsCount, ILogger logger =
             {
                 databasesContentLock.ReadUnlock();
             }
+
+            return true;
         }
 
         private void UpdateLastSaveData(int dbId, long? storeTailAddress, long? objectStoreTailAddress)
@@ -1022,24 +1054,6 @@ private void UpdateLastSaveData(int dbId, long? storeTailAddress, long? objectSt
             }
         }
 
-        private async Task<bool> TryPauseCheckpointsContinuousAsync(int dbId, CancellationToken token = default)
-        {
-            var databasesMapSize = databases.ActualSize;
-            var databasesMapSnapshot = databases.Map;
-            Debug.Assert(dbId < databasesMapSize && databasesMapSnapshot[dbId] != null);
-
-            var db = databasesMapSnapshot[dbId];
-            var checkpointsPaused = TryPauseCheckpoints(db);
-
-            while (!checkpointsPaused && !token.IsCancellationRequested && !Disposed)
-            {
-                await Task.Yield();
-                checkpointsPaused = TryPauseCheckpoints(db);
-            }
-
-            return checkpointsPaused;
-        }
-
         public override void Dispose()
         {
             if (Disposed) return;

libs/server/Databases/SingleDatabaseManager.cs

@@ -127,7 +127,8 @@ public override void ResumeCheckpoints(int dbId)
         /// <inheritdoc/>
         public override bool TakeCheckpoint(bool background, ILogger logger = null, CancellationToken token = default)
         {
-            if (!TryPauseCheckpointsContinuousAsync(defaultDatabase.Id, token: token).GetAwaiter().GetResult())
+            // Check if checkpoint already in progress
+            if (!TryPauseCheckpoints(defaultDatabase.Id))
                 return false;
 
             var checkpointTask = TakeCheckpointAsync(defaultDatabase, logger: logger, token: token).ContinueWith(

libs/server/Resp/AdminCommands.cs

@@ -880,7 +880,7 @@ private bool NetworkSAVE()
 
             if (!storeWrapper.TakeCheckpoint(false, dbId: dbId, logger: logger))
             {
-                while (!RespWriteUtils.TryWriteError("ERR checkpoint already in progress"u8, ref dcurr, dend))
+                while (!RespWriteUtils.TryWriteError(CmdStrings.RESP_ERR_CHECKPOINT_ALREADY_IN_PROGRESS, ref dcurr, dend))
                     SendAndReset();
             }
             else
@@ -1003,7 +1003,7 @@ private bool NetworkBGSAVE()
             }
             else
             {
-                while (!RespWriteUtils.TryWriteError("ERR checkpoint already in progress"u8, ref dcurr, dend))
+                while (!RespWriteUtils.TryWriteError(CmdStrings.RESP_ERR_CHECKPOINT_ALREADY_IN_PROGRESS, ref dcurr, dend))
                     SendAndReset();
             }
 

libs/server/Resp/CmdStrings.cs

@@ -303,6 +303,7 @@ static partial class CmdStrings
         public static ReadOnlySpan<byte> RESP_ERR_FLUSHALL_READONLY_REPLICA => "ERR You can't write against a read only replica."u8;
         public static ReadOnlySpan<byte> RESP_ERR_ZSET_MEMBER => "ERR could not decode requested zset member"u8;
         public static ReadOnlySpan<byte> RESP_ERR_EXPDELSCAN_INVALID => "ERR Cannot execute EXPDELSCAN with background expired key deletion scan enabled"u8;
+        public static ReadOnlySpan<byte> RESP_ERR_CHECKPOINT_ALREADY_IN_PROGRESS => "ERR checkpoint already in progress"u8;
 
         /// <summary>
         /// Response string templates

test/Garnet.test/MultiDatabaseTests.cs

@@ -1278,6 +1278,71 @@ public void MultiDatabaseAofRecoverObjectTest()
             }
         }
 
+        [Test]
+        public void MultiDatabaseSaveInProgressTest()
+        {
+            using (var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig()))
+            {
+                var db1 = redis.GetDatabase(0);
+                var db2 = redis.GetDatabase(1);
+
+                // Check no saves present
+                var lastsave = (int)db1.Execute("LASTSAVE");
+                var lastsave1 = (int)db1.Execute("LASTSAVE", "0");
+                var lastsave2 = (int)db2.Execute("LASTSAVE", "1");
+                ClassicAssert.AreEqual(0, lastsave);
+                ClassicAssert.AreEqual(0, lastsave1);
+                ClassicAssert.AreEqual(0, lastsave2);
+
+                // Issue background save to DB 0
+                var res1 = db1.Execute("BGSAVE", "0");
+                ClassicAssert.AreEqual("Background saving started", res1.ToString());
+
+                // Issue background save to DB 1 while DB 0 save is in progress - legal
+                var res2 = db1.Execute("BGSAVE", "1");
+                ClassicAssert.AreEqual("Background saving started", res2.ToString());
+
+                // Issue general background save while DB 0 save is in progress - legal
+                var res = db1.Execute("BGSAVE");
+                ClassicAssert.AreEqual("Background saving started", res.ToString());
+
+                // Wait for saves to complete
+                do
+                {
+                    Thread.Sleep(10);
+                    lastsave = (int)db1.Execute("LASTSAVE");
+                    lastsave1 = (int)db1.Execute("LASTSAVE", "0");
+                    lastsave2 = (int)db2.Execute("LASTSAVE", "1");
+                }
+                while (lastsave == 0 || lastsave1 == 0 || lastsave2 == 0);
+
+                // Add some data
+                for (var i = 0; i < 1024; i++)
+                {
+                    db1.StringSet($"k{i}", new string('x', 256));
+                    db2.StringSet($"k{i}", new string('x', 256));
+                    db1.ListLeftPush($"k{i}o", new string('x', 256));
+                    db2.ListLeftPush($"k{i}o", new string('x', 256));
+                }
+
+                // Issue general background save
+                res = db1.Execute("BGSAVE");
+                ClassicAssert.AreEqual("Background saving started", res.ToString());
+
+                // Issue background save to DB 0 while general save is in progress - illegal
+                Assert.Throws<RedisServerException>(() => db1.Execute("BGSAVE", "0"),
+                    Encoding.ASCII.GetString(CmdStrings.RESP_ERR_CHECKPOINT_ALREADY_IN_PROGRESS));
+
+                // Wait for save to complete
+                do
+                {
+                    Thread.Sleep(10);
+                    lastsave = (int)db1.Execute("LASTSAVE");
+                }
+                while (lastsave == 0);
+            }
+        }
+
         [Test]
         [TestCase(false)]
         [TestCase(true)]

test/Garnet.test/RespAdminCommandsTests.cs

@@ -211,6 +211,40 @@ public void SeSaveRecoverObjectTest()
             }
         }
 
+        [Test]
+        [TestCase(SaveType.BackgroundSave)]
+#pragma warning disable CS0618 // Type or member is obsolete
+        [TestCase(SaveType.ForegroundSave)]
+#pragma warning restore CS0618 // Type or member is obsolete
+        public void SeSaveInProgressTest(SaveType saveType)
+        {
+            using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig(allowAdmin: true));
+            var server = redis.GetServer(TestUtils.EndPoint);
+            var db = redis.GetDatabase(0);
+
+            var lastSave = server.LastSave();
+
+            // Check no saves present
+            ClassicAssert.AreEqual(DateTimeOffset.FromUnixTimeSeconds(0).Ticks, lastSave.Ticks);
+
+            // Add some data
+            for (var i = 0; i < 1024; i++)
+            {
+                db.StringSet($"k{i}", new string('x', 256));
+                db.ListLeftPush($"k{i}o", new string('x', 256));
+            }
+
+            // Issue background save
+            server.Save(SaveType.BackgroundSave);
+
+            // Issue another save while one is in progress
+            Assert.Throws<RedisServerException>(() => server.Save(saveType),
+                Encoding.ASCII.GetString(CmdStrings.RESP_ERR_CHECKPOINT_ALREADY_IN_PROGRESS));
+
+            // Wait for save to complete
+            while (server.LastSave() == lastSave) Thread.Sleep(10);
+        }
+
         [Test]
         public void SeSaveRecoverCustomObjectTest()
         {