Kill AOF BulkConsume When Connection is Dead

libs/cluster/Server/Replication/PrimaryOps/AofSyncTaskInfo.cs

@@ -90,6 +90,9 @@ public unsafe void Consume(byte* payloadPtr, int payloadLength, long currentAddr
 
         public void Throttle()
         {
+            if (!garnetClient.IsConnected)
+                ExceptionUtils.ThrowException(new GarnetException("AOF stream client disconnected!"));
+
             // Trigger flush while we are out of epoch protection
             garnetClient.CompletePending(false);
             garnetClient.Throttle();
@@ -108,11 +111,11 @@ public async Task ReplicaSyncTask()
 
                 iter = clusterProvider.storeWrapper.appendOnlyFile.ScanSingle(startAddress, long.MaxValue, scanUncommitted: true, recover: false, logger: logger);
 
-                while (true)
-                {
-                    if (cts.Token.IsCancellationRequested) break;
-                    await iter.BulkConsumeAllAsync(this, clusterProvider.serverOptions.ReplicaSyncDelayMs, maxChunkSize: 1 << 20, cts.Token);
-                }
+                await iter.BulkConsumeAllAsync(
+                    this,
+                    clusterProvider.serverOptions.ReplicaSyncDelayMs,
+                    maxChunkSize: 1 << 20,
+                    cts.Token).ConfigureAwait(false);
             }
             catch (Exception ex)
             {

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

@@ -95,15 +95,12 @@ public async void ReplicaReplayTask()
             try
             {
                 activeReplay.ReadLock();
-                while (true)
-                {
-                    replicaReplayTaskCts.Token.ThrowIfCancellationRequested();
-                    await replayIterator.BulkConsumeAllAsync(
-                        this,
-                        clusterProvider.serverOptions.ReplicaSyncDelayMs,
-                        maxChunkSize: 1 << 20,
-                        replicaReplayTaskCts.Token);
-                }
+
+                await replayIterator.BulkConsumeAllAsync(
+                    this,
+                    clusterProvider.serverOptions.ReplicaSyncDelayMs,
+                    maxChunkSize: 1 << 20,
+                    replicaReplayTaskCts.Token);
             }
             catch (Exception ex)
             {

libs/common/ExceptionUtils.cs

@@ -0,0 +1,14 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+using System;
+using System.Runtime.CompilerServices;
+
+namespace Garnet.common
+{
+    public static class ExceptionUtils
+    {
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        public static void ThrowException(Exception e) => throw e;
+    }
+}

libs/storage/Tsavorite/cs/src/core/TsavoriteLog/TsavoriteLogIterator.cs

@@ -475,6 +475,9 @@ public unsafe bool TryConsumeNext<T>(T consumer) where T : ILogEntryConsumer
         /// <returns>whether a next entry is present</returns>
         public unsafe bool TryBulkConsumeNext<T>(T consumer, int maxChunkSize = 0) where T : IBulkLogEntryConsumer
         {
+            // Throttle and implicitly check for consumer liveness
+            consumer.Throttle();
+
             if (maxChunkSize == 0) maxChunkSize = allocator.PageSize;
 
             if (disposed)
@@ -500,7 +503,7 @@ public unsafe bool TryBulkConsumeNext<T>(T consumer, int maxChunkSize = 0) where
                         epoch.ProtectAndDrain();
                     }
 
-                    var hasNext = GetNextInternal(out long startPhysicalAddress, out int newEntryLength, out long startLogicalAddress, out long endLogicalAddress, out bool isCommitRecord, out bool onFrame);
+                    var hasNext = GetNextInternal(out long startPhysicalAddress, out var newEntryLength, out var startLogicalAddress, out var endLogicalAddress, out bool isCommitRecord, out bool onFrame);
 
                     if (!hasNext)
                     {
@@ -509,7 +512,7 @@ public unsafe bool TryBulkConsumeNext<T>(T consumer, int maxChunkSize = 0) where
                     }
 
                     // GetNextInternal returns only the payload length, so adjust the totalLength
-                    int totalLength = headerSize + Align(newEntryLength);
+                    var totalLength = headerSize + Align(newEntryLength);
 
                     // Expand the records in iteration, as long as as they are on the same physical page
                     while (ExpandGetNextInternal(startPhysicalAddress, ref totalLength, out _, out endLogicalAddress, out isCommitRecord))