[Event Hubs] Change trigger checkpointing behavior (Azure#42299)

jsquire · web-flow · commit 7cdc94d333f8 · 2024-03-01T16:20:20.000-05:00
* [Event Hubs] Change trigger checkpointing behavior

The focus of these changes is to change the behavior
of the Event Hubs trigger to not checkpoint when the
listener is shutting down. This is necessary to
prevent potential data loss from occurring when
shutting down Function retries.  Because the trigger
cannot know if the Function host would have retried
a failure if it were not shutting down, we cannot
assume that it is safe to checkpoint.  This change
ensures that the batch of events being processed
when shut down was requested will be retried by
another instance or the next time the Function app
is run.  Though this may cause duplicate processing,
it is consistent with the Event Hubs at-least-once
guarantee.
diff --git a/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/CHANGELOG.md b/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/CHANGELOG.md
@@ -10,6 +10,8 @@
 
 ### Other Changes
 
+- Adjusted checkpointing logic to no longer write a checkpoint when the listener is shutting down.  This was necessary to prevent potential data loss from occurring when shutting down Function retries.  Because the trigger cannot know if the Function host would have retried a failure if it were not shutting down, we cannot assume that it is safe to checkpoint. This change ensures that the batch of events being processed when shut down was requested will be retried by another instance or the next time the Function app is run.
+
 - Updated the `Azure.Messaging.EventHubs`, which includes a new build of the AMQP transport library.  The notable bug fix addresses an obscure race condition when a cancellation token is signaled while service operations are being invoked concurrently which caused those operations to hang.
 
 ## 6.1.0 (2024-02-13)
diff --git a/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/src/Listeners/EventHubListener.PartitionProcessor.cs b/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/src/Listeners/EventHubListener.PartitionProcessor.cs
@@ -42,6 +42,7 @@ internal class PartitionProcessor : IEventProcessor, IDisposable
             private Task _cachedEventsBackgroundTask;
             private CancellationTokenSource _cachedEventsBackgroundTaskCts;
             private SemaphoreSlim _cachedEventsGuard;
+            private readonly CancellationToken _listenerCancellationToken;
             private readonly CancellationToken _functionExecutionToken;
             private readonly CancellationTokenSource _ownershipLostTokenSource;
 
@@ -50,7 +51,7 @@ internal class PartitionProcessor : IEventProcessor, IDisposable
             /// </summary>
             internal PartitionProcessorEventsManager CachedEventsManager { get; }
 
-            public PartitionProcessor(EventHubOptions options, ITriggeredFunctionExecutor executor, ILogger logger, bool singleDispatch, CancellationToken functionExecutionToken)
+            public PartitionProcessor(EventHubOptions options, ITriggeredFunctionExecutor executor, ILogger logger, bool singleDispatch, CancellationToken listenerCancellationToken, CancellationToken functionExecutionToken)
             {
                 _executor = executor;
                 _singleDispatch = singleDispatch;
@@ -59,6 +60,7 @@ public PartitionProcessor(EventHubOptions options, ITriggeredFunctionExecutor ex
                 _firstFunctionInvocation = true;
                 _maxWaitTime = options.MaxWaitTime;
                 _minimumBatchesEnabled = options.MinEventBatchSize > 1; // 1 is the default
+                _listenerCancellationToken = listenerCancellationToken;
                 _functionExecutionToken = functionExecutionToken;
                 _ownershipLostTokenSource = new CancellationTokenSource();
 
@@ -215,8 +217,10 @@ public async Task ProcessEventsAsync(EventProcessorHostPartition context, IEnume
                     // and wait to send until we receive enough events or total max wait time has passed.
                 }
 
-                // Checkpoint if we processed any events and cancellation has not been signaled.
-                // Don't checkpoint if no events. This can reset the sequence counter to 0.
+                // Checkpoint if we processed any events, the listener is not stopping, and
+                // cancellation has not been signaled.  Don't checkpoint if no events. This
+                // can reset the sequence counter to 0.
+                //
                 // Note: we intentionally checkpoint the batch regardless of function
                 // success/failure. EventHub doesn't support any sort "poison event" model,
                 // so that is the responsibility of the user's function currently. E.g.
@@ -225,10 +229,12 @@ public async Task ProcessEventsAsync(EventProcessorHostPartition context, IEnume
                 // Don't checkpoint if cancellation has been requested as this can lead to data loss,
                 // since the user may not actually process the event.
 
-                if (eventToCheckpoint != null &&
+                if (eventToCheckpoint != null
                     // IMPORTANT - explicitly check each token to avoid data loss as the linkedCts is not canceled atomically when each of the
                     // sources are canceled.
-                    !_functionExecutionToken.IsCancellationRequested && !_ownershipLostTokenSource.IsCancellationRequested)
+                    && !_listenerCancellationToken.IsCancellationRequested
+                    && !_functionExecutionToken.IsCancellationRequested
+                    && !_ownershipLostTokenSource.IsCancellationRequested)
                 {
                     await CheckpointAsync(eventToCheckpoint, context).ConfigureAwait(false);
                 }
diff --git a/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/src/Listeners/EventHubListener.cs b/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/src/Listeners/EventHubListener.cs
@@ -28,6 +28,7 @@ internal sealed partial class EventHubListener : IListener, IEventProcessorFacto
         private readonly ILoggerFactory _loggerFactory;
         private readonly ILogger _logger;
         private string _details;
+        private CancellationTokenSource _listenerCancellationTokenSource;
         private CancellationTokenSource _functionExecutionCancellationTokenSource;
         private readonly IDrainModeManager _drainModeManager;
         private volatile bool _disposed;
@@ -50,6 +51,7 @@ public EventHubListener(
             _checkpointStore = checkpointStore;
             _options = options;
             _logger = _loggerFactory.CreateLogger<EventHubListener>();
+            _listenerCancellationTokenSource = new CancellationTokenSource();
             _functionExecutionCancellationTokenSource = new CancellationTokenSource();
             _drainModeManager = drainModeManager;
 
@@ -88,6 +90,7 @@ void IListener.Cancel()
 
         void IDisposable.Dispose()
         {
+            _listenerCancellationTokenSource.Cancel();
             _functionExecutionCancellationTokenSource.Cancel();
 
 #pragma warning disable AZC0102
@@ -110,6 +113,8 @@ public async Task StartAsync(CancellationToken cancellationToken)
 
         public async Task StopAsync(CancellationToken cancellationToken)
         {
+            _listenerCancellationTokenSource.Cancel();
+
             if (!_drainModeManager.IsDrainModeEnabled)
             {
                 _functionExecutionCancellationTokenSource.Cancel();
@@ -122,7 +127,13 @@ public async Task StopAsync(CancellationToken cancellationToken)
 
         IEventProcessor IEventProcessorFactory.CreatePartitionProcessor()
         {
-            return new PartitionProcessor(_options, _executor, _loggerFactory.CreateLogger<PartitionProcessor>(), _singleDispatch, _functionExecutionCancellationTokenSource.Token);
+            return new PartitionProcessor(
+                _options,
+                _executor,
+                _loggerFactory.CreateLogger<PartitionProcessor>(),
+                _singleDispatch,
+                _listenerCancellationTokenSource.Token,
+                _functionExecutionCancellationTokenSource.Token);
         }
 
         public IScaleMonitor GetMonitor()
diff --git a/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/tests/EventHubEndToEndTests.cs b/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/tests/EventHubEndToEndTests.cs
@@ -8,6 +8,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
+using Azure.Core.Diagnostics;
 using Azure.Core.TestFramework;
 using Azure.Messaging.EventHubs;
 using Azure.Messaging.EventHubs.Consumer;
@@ -93,14 +94,18 @@ public async Task EventHub_StringBinding()
         [Test]
         public async Task EventHub_SingleDispatch()
         {
+            var watch = ValueStopwatch.StartNew();
             var (jobHost, host) = BuildHost<EventHubTestSingleDispatchJobs>();
+
             using (jobHost)
             {
                 await jobHost.CallAsync(nameof(EventHubTestSingleDispatchJobs.SendEvent_TestHub), new { input = "data" });
 
-                bool result = _eventWait1.WaitOne(Timeout);
+                bool result = _eventWait1.WaitOne(GetRemainingTimeoutMilliseconds(watch.GetElapsedTime()));
                 Assert.True(result);
 
+                using var cancellationSource = new CancellationTokenSource(GetRemainingTimeoutMilliseconds(watch.GetElapsedTime()));
+                await AwaitCheckpointing(host, cancellationSource.Token);
                 await StopWithDrainAsync(host);
             }
 
@@ -158,14 +163,18 @@ public async Task EventHub_SingleDispatch_ConsumerGroup()
         [Test]
         public async Task EventHub_SingleDispatch_BinaryData()
         {
+            var watch = ValueStopwatch.StartNew();
             var (jobHost, host) = BuildHost<EventHubTestSingleDispatchJobsBinaryData>();
+
             using (jobHost)
             {
                 await jobHost.CallAsync(nameof(EventHubTestSingleDispatchJobsBinaryData.SendEvent_TestHub), new { input = "data" });
 
-                bool result = _eventWait1.WaitOne(Timeout);
+                bool result = _eventWait1.WaitOne(GetRemainingTimeoutMilliseconds(watch.GetElapsedTime()));
                 Assert.True(result);
 
+                using var cancellationSource = new CancellationTokenSource(GetRemainingTimeoutMilliseconds(watch.GetElapsedTime()));
+                await AwaitCheckpointing(host, cancellationSource.Token);
                 await StopWithDrainAsync(host);
             }
 
@@ -211,30 +220,6 @@ public async Task EventHub_CollectorPartitionKey()
             }
         }
 
-        private static void AssertSingleDispatchLogs(IHost host)
-        {
-            IEnumerable<LogMessage> logMessages = host.GetTestLoggerProvider()
-                .GetAllLogMessages();
-
-            Assert.AreEqual(logMessages.Where(x => !string.IsNullOrEmpty(x.FormattedMessage)
-                && x.FormattedMessage.Contains("Trigger Details:")
-                && x.FormattedMessage.Contains("Offset:")).Count(), 1);
-
-            Assert.True(logMessages.Where(x => !string.IsNullOrEmpty(x.FormattedMessage)
-                && x.FormattedMessage.Contains("OpenAsync")).Any());
-
-            Assert.True(logMessages.Where(x => !string.IsNullOrEmpty(x.FormattedMessage)
-                && x.FormattedMessage.Contains("CheckpointAsync")
-                && x.FormattedMessage.Contains("lease")
-                && x.FormattedMessage.Contains("offset")
-                && x.FormattedMessage.Contains("sequenceNumber")).Any());
-
-            Assert.True(logMessages.Where(x => !string.IsNullOrEmpty(x.FormattedMessage)
-                && x.FormattedMessage.Contains("Sending events to EventHub")).Any());
-
-            AssertAzureSdkLogs(logMessages);
-        }
-
         [Test]
         public async Task CanSendAndReceive_ConnectionStringInConfiguration()
         {
@@ -332,15 +317,19 @@ public async Task AssertCanSendReceiveMessage(Action<IHostBuilder> hostConfigura
         [Test]
         public async Task EventHub_MultipleDispatch()
         {
+            var watch = ValueStopwatch.StartNew();
             var (jobHost, host) = BuildHost<EventHubTestMultipleDispatchJobs>();
+
             using (jobHost)
             {
                 int numEvents = 5;
                 await jobHost.CallAsync(nameof(EventHubTestMultipleDispatchJobs.SendEvents_TestHub), new { numEvents = numEvents, input = "data" });
 
-                bool result = _eventWait1.WaitOne(Timeout);
+                bool result = _eventWait1.WaitOne(GetRemainingTimeoutMilliseconds(watch.GetElapsedTime()));
                 Assert.True(result);
 
+                using var cancellationSource = new CancellationTokenSource(GetRemainingTimeoutMilliseconds(watch.GetElapsedTime()));
+                await AwaitCheckpointing(host, cancellationSource.Token);
                 await StopWithDrainAsync(host);
             }
 
@@ -350,15 +339,19 @@ public async Task EventHub_MultipleDispatch()
         [Test]
         public async Task EventHub_MultipleDispatch_BinaryData()
         {
+            var watch = ValueStopwatch.StartNew();
             var (jobHost, host) = BuildHost<EventHubTestMultipleDispatchJobsBinaryData>();
+
             using (jobHost)
             {
                 int numEvents = 5;
                 await jobHost.CallAsync(nameof(EventHubTestMultipleDispatchJobsBinaryData.SendEvents_TestHub), new { numEvents = numEvents, input = "data" });
 
-                bool result = _eventWait1.WaitOne(Timeout);
+                bool result = _eventWait1.WaitOne(GetRemainingTimeoutMilliseconds(watch.GetElapsedTime()));
                 Assert.True(result);
 
+                using var cancellationSource = new CancellationTokenSource(GetRemainingTimeoutMilliseconds(watch.GetElapsedTime()));
+                await AwaitCheckpointing(host, cancellationSource.Token);
                 await StopWithDrainAsync(host);
             }
 
@@ -370,6 +363,7 @@ public async Task EventHub_MultipleDispatch_MinBatchSize()
         {
             const int minEventBatchSize = 5;
 
+            var watch = ValueStopwatch.StartNew();
             var (jobHost, host) = BuildHost<EventHubTestMultipleDispatchMinBatchSizeJobs>(
                 builder =>
                 {
@@ -393,9 +387,11 @@ public async Task EventHub_MultipleDispatch_MinBatchSize()
                 int numEvents = 5;
                 await jobHost.CallAsync(nameof(EventHubTestMultipleDispatchMinBatchSizeJobs.SendEvents_TestHub), new { numEvents = numEvents, input = "data" });
 
-                bool result = _eventWait1.WaitOne(Timeout);
+                bool result = _eventWait1.WaitOne(GetRemainingTimeoutMilliseconds(watch.GetElapsedTime()));
                 Assert.True(result);
 
+                using var cancellationSource = new CancellationTokenSource(GetRemainingTimeoutMilliseconds(watch.GetElapsedTime()));
+                await AwaitCheckpointing(host, cancellationSource.Token);
                 await StopWithDrainAsync(host);
             }
 
@@ -558,6 +554,30 @@ public async Task EventHub_InitialOffsetFromEnqueuedTime()
             }
         }
 
+        private static void AssertSingleDispatchLogs(IHost host)
+        {
+            IEnumerable<LogMessage> logMessages = host.GetTestLoggerProvider()
+                .GetAllLogMessages();
+
+            Assert.AreEqual(logMessages.Where(x => !string.IsNullOrEmpty(x.FormattedMessage)
+                && x.FormattedMessage.Contains("Trigger Details:")
+                && x.FormattedMessage.Contains("Offset:")).Count(), 1);
+
+            Assert.True(logMessages.Where(x => !string.IsNullOrEmpty(x.FormattedMessage)
+                && x.FormattedMessage.Contains("OpenAsync")).Any());
+
+            Assert.True(logMessages.Where(x => !string.IsNullOrEmpty(x.FormattedMessage)
+                && x.FormattedMessage.Contains("CheckpointAsync")
+                && x.FormattedMessage.Contains("lease")
+                && x.FormattedMessage.Contains("offset")
+                && x.FormattedMessage.Contains("sequenceNumber")).Any());
+
+            Assert.True(logMessages.Where(x => !string.IsNullOrEmpty(x.FormattedMessage)
+                && x.FormattedMessage.Contains("Sending events to EventHub")).Any());
+
+            AssertAzureSdkLogs(logMessages);
+        }
+
         private static void AssertMultipleDispatchLogsMinBatch(IHost host)
         {
             IEnumerable<LogMessage> logMessages = host.GetTestLoggerProvider()
@@ -606,17 +626,36 @@ private static void AssertMultipleDispatchLogs(IHost host)
             AssertAzureSdkLogs(logMessages);
         }
 
+        private static void AssertAzureSdkLogs(IEnumerable<LogMessage> logMessages)
+        {
+            Assert.True(logMessages.Any(x => x.Category.StartsWith("Azure.")));
+        }
+
         private static async Task StopWithDrainAsync(IHost host)
         {
-            // Enable drain mode so checkpointing occurs when stopping
             var drainModeManager = host.Services.GetService<IDrainModeManager>();
             await drainModeManager.EnableDrainModeAsync(CancellationToken.None);
             await host.StopAsync();
         }
 
-        private static void AssertAzureSdkLogs(IEnumerable<LogMessage> logMessages)
+         private static async Task AwaitCheckpointing(IHost host, CancellationToken cancellationToken)
         {
-            Assert.True(logMessages.Any(x => x.Category.StartsWith("Azure.")));
+            var logMessages = host.GetTestLoggerProvider()
+                .GetAllLogMessages();
+
+            while (true)
+            {
+                if (logMessages.Where(x => !string.IsNullOrEmpty(x.FormattedMessage)
+                    && x.FormattedMessage.Contains("CheckpointAsync")
+                    && x.FormattedMessage.Contains("sequenceNumber")).Any())
+                {
+                    return;
+                }
+
+                // No need to explicitly check the cancellation token;
+                // the delay will throw if the token is signaled.
+                await Task.Delay(100, cancellationToken).ConfigureAwait(false);
+            }
         }
 
         public class EventHubTestSingleDispatchJobs
diff --git a/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/tests/EventHubListenerTests.cs b/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/tests/EventHubListenerTests.cs
diff --git a/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/tests/WebJobsEventHubTestBase.cs b/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/tests/WebJobsEventHubTestBase.cs
