Ensure Streaming Checkpoint Cancellation on Timeout (#1385)

* cancel streaming checkpoint if timeout occurred

* propagate cancellation token to stateMachineDriver when initiating a checkpoint

Author: vazois

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

@@ -280,47 +280,60 @@ async Task TakeStreamingCheckpoint()
 
                 // Iterate through main store
                 var mainStoreCheckpointTask = ClusterProvider.storeWrapper.store.
-                    TakeFullCheckpointAsync(CheckpointType.StreamingSnapshot, streamingSnapshotIteratorFunctions: manager.mainStoreSnapshotIterator);
+                    TakeFullCheckpointAsync(CheckpointType.StreamingSnapshot, cancellationToken: cts.Token, streamingSnapshotIteratorFunctions: manager.mainStoreSnapshotIterator);
 
                 var result = await WaitOrDie(checkpointTask: mainStoreCheckpointTask, iteratorManager: manager);
                 if (!result.success)
-                    throw new InvalidOperationException("Main store checkpoint stream failed!");
+                    throw new GarnetException("Main store checkpoint stream failed!");
 
                 if (!ClusterProvider.serverOptions.DisableObjects)
                 {
                     // Iterate through object store
                     var objectStoreCheckpointTask = ClusterProvider.storeWrapper.objectStore.
-                        TakeFullCheckpointAsync(CheckpointType.StreamingSnapshot, streamingSnapshotIteratorFunctions: manager.objectStoreSnapshotIterator);
+                        TakeFullCheckpointAsync(CheckpointType.StreamingSnapshot, cancellationToken: cts.Token, streamingSnapshotIteratorFunctions: manager.objectStoreSnapshotIterator);
                     result = await WaitOrDie(checkpointTask: objectStoreCheckpointTask, iteratorManager: manager);
                     if (!result.success)
-                        throw new InvalidOperationException("Object store checkpoint stream failed!");
+                        throw new GarnetException("Object store checkpoint stream failed!");
                 }
 
                 // Note: We do not truncate the AOF here as this was just a "virtual" checkpoint
-
+                // WaitOrDie is needed here to check if streaming checkpoint is making progress.
+                // We cannot use a timeout on the cancellationToken because we don't know in total how long the streaming checkpoint will take
                 async ValueTask<(bool success, Guid token)> WaitOrDie(ValueTask<(bool success, Guid token)> checkpointTask, SnapshotIteratorManager iteratorManager)
                 {
-                    var timeout = replicaSyncTimeout;
-                    var delay = TimeSpan.FromSeconds(1);
-                    while (true)
+                    try
                     {
-                        // Check if cancellation requested
-                        cts.Token.ThrowIfCancellationRequested();
+                        var timeout = replicaSyncTimeout;
+                        var delay = TimeSpan.FromSeconds(1);
+                        while (true)
+                        {
+                            // Check if cancellation requested
+                            cts.Token.ThrowIfCancellationRequested();
 
-                        // Wait for stream sync to make some progress
-                        await Task.Delay(delay);
+                            // Wait for stream sync to make some progress
+                            await Task.Delay(delay);
 
-                        // Check if checkpoint has completed
-                        if (checkpointTask.IsCompleted)
-                            return await checkpointTask;
+                            // Check if checkpoint has completed
+                            if (checkpointTask.IsCompleted)
+                                return await checkpointTask;
 
-                        // Check if we made some progress
-                        timeout = !manager.IsProgressing() ? timeout.Subtract(delay) : replicaSyncTimeout;
+                            // Check if we made some progress
+                            timeout = !manager.IsProgressing() ? timeout.Subtract(delay) : replicaSyncTimeout;
 
-                        // Throw timeout equals to zero
-                        if (timeout.TotalSeconds <= 0)
-                            throw new TimeoutException("Streaming snapshot checkpoint timed out");
+                            // Throw timeout equals to zero
+                            if (timeout.TotalSeconds <= 0)
+                            {
+                                cts.Cancel();
+                                throw new TimeoutException("Streaming snapshot checkpoint timed out");
+                            }
+                        }
+                    }
+                    catch (Exception ex)
+                    {
+                        logger?.LogError(ex, "{method} faulted", nameof(WaitOrDie));
+                        cts.Cancel();
                     }
+                    return (false, default);
                 }
             }
         }

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

@@ -299,7 +299,11 @@ async Task RunStateMachine(CancellationToken token = default)
                 _ = Interlocked.Exchange(ref stateMachine, null);
                 if (ex != null)
                 {
-                    _ = _stateMachineCompleted.TrySetException(ex);
+                    // If the state machine stopped due to cancellation, propagate cancellation to the completion TCS
+                    if (ex is OperationCanceledException || ex is TaskCanceledException)
+                        _ = _stateMachineCompleted.TrySetCanceled();
+                    else
+                        _ = _stateMachineCompleted.TrySetException(ex);
                 }
                 else
                 {

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

@@ -188,12 +188,13 @@ public TsavoriteKV(KVSettings<TKey, TValue> kvSettings, TStoreFunctions storeFun
         /// <param name="token">Checkpoint token</param>
         /// <param name="checkpointType">Checkpoint type</param>
         /// <param name="streamingSnapshotIteratorFunctions">Iterator for streaming snapshot records</param>
+        /// <param name="cancellationToken">Caller's cancellation token</param>
         /// <returns>
         /// Whether we successfully initiated the checkpoint (initiation may
         /// fail if we are already taking a checkpoint or performing some other
         /// operation such as growing the index). Use CompleteCheckpointAsync to wait completion.
         /// </returns>
-        public bool TryInitiateFullCheckpoint(out Guid token, CheckpointType checkpointType, IStreamingSnapshotIteratorFunctions<TKey, TValue> streamingSnapshotIteratorFunctions = null)
+        public bool TryInitiateFullCheckpoint(out Guid token, CheckpointType checkpointType, IStreamingSnapshotIteratorFunctions<TKey, TValue> streamingSnapshotIteratorFunctions = null, CancellationToken cancellationToken = default)
         {
             IStateMachine stateMachine;
 
@@ -208,7 +209,7 @@ public bool TryInitiateFullCheckpoint(out Guid token, CheckpointType checkpointT
             {
                 stateMachine = Checkpoint.Full(this, checkpointType, out token);
             }
-            return stateMachineDriver.Register(stateMachine);
+            return stateMachineDriver.Register(stateMachine, cancellationToken);
         }
 
         /// <summary>
@@ -228,7 +229,7 @@ public bool TryInitiateFullCheckpoint(out Guid token, CheckpointType checkpointT
         public async ValueTask<(bool success, Guid token)> TakeFullCheckpointAsync(CheckpointType checkpointType,
             CancellationToken cancellationToken = default, IStreamingSnapshotIteratorFunctions<TKey, TValue> streamingSnapshotIteratorFunctions = null)
         {
-            var success = TryInitiateFullCheckpoint(out Guid token, checkpointType, streamingSnapshotIteratorFunctions);
+            var success = TryInitiateFullCheckpoint(out Guid token, checkpointType, streamingSnapshotIteratorFunctions, cancellationToken);
 
             if (success)
                 await CompleteCheckpointAsync(cancellationToken).ConfigureAwait(false);
@@ -241,10 +242,10 @@ public bool TryInitiateFullCheckpoint(out Guid token, CheckpointType checkpointT
         /// </summary>
         /// <param name="token">Checkpoint token</param>
         /// <returns>Whether we could initiate the checkpoint. Use CompleteCheckpointAsync to wait completion.</returns>
-        public bool TryInitiateIndexCheckpoint(out Guid token)
+        public bool TryInitiateIndexCheckpoint(out Guid token, CancellationToken cancellationToken = default)
         {
             var stateMachine = Checkpoint.IndexOnly(this, out token);
-            return stateMachineDriver.Register(stateMachine);
+            return stateMachineDriver.Register(stateMachine, cancellationToken);
         }
 
         /// <summary>
@@ -261,7 +262,7 @@ public bool TryInitiateIndexCheckpoint(out Guid token)
         /// </returns>
         public async ValueTask<(bool success, Guid token)> TakeIndexCheckpointAsync(CancellationToken cancellationToken = default)
         {
-            var success = TryInitiateIndexCheckpoint(out Guid token);
+            var success = TryInitiateIndexCheckpoint(out Guid token, cancellationToken);
 
             if (success)
                 await CompleteCheckpointAsync(cancellationToken).ConfigureAwait(false);
@@ -277,7 +278,7 @@ public bool TryInitiateIndexCheckpoint(out Guid token)
         /// <param name="tryIncremental">For snapshot, try to store as incremental delta over last snapshot</param>
         /// <returns>Whether we could initiate the checkpoint. Use CompleteCheckpointAsync to wait completion.</returns>
         public bool TryInitiateHybridLogCheckpoint(out Guid token, CheckpointType checkpointType, bool tryIncremental = false,
-            IStreamingSnapshotIteratorFunctions<TKey, TValue> streamingSnapshotIteratorFunctions = null)
+            IStreamingSnapshotIteratorFunctions<TKey, TValue> streamingSnapshotIteratorFunctions = null, CancellationToken cancellationToken = default)
         {
             IStateMachine stateMachine;
 
@@ -305,7 +306,7 @@ public bool TryInitiateHybridLogCheckpoint(out Guid token, CheckpointType checkp
                     stateMachine = Checkpoint.HybridLogOnly(this, checkpointType, out token);
                 }
             }
-            return stateMachineDriver.Register(stateMachine);
+            return stateMachineDriver.Register(stateMachine, cancellationToken);
         }
 
         /// <summary>
@@ -340,7 +341,7 @@ public bool CanTakeIncrementalCheckpoint(CheckpointType checkpointType, out Guid
         public async ValueTask<(bool success, Guid token)> TakeHybridLogCheckpointAsync(CheckpointType checkpointType,
             bool tryIncremental = false, CancellationToken cancellationToken = default)
         {
-            var success = TryInitiateHybridLogCheckpoint(out Guid token, checkpointType, tryIncremental);
+            var success = TryInitiateHybridLogCheckpoint(out Guid token, checkpointType, tryIncremental, cancellationToken: cancellationToken);
 
             if (success)
                 await CompleteCheckpointAsync(cancellationToken).ConfigureAwait(false);