fix(streaming): release ADO.NET messages during queue handoff

src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs

@@ -1,21 +1,60 @@
 namespace Orleans.Streaming.AdoNet;
 
+internal interface IStreamMessageQueries
+{
+    Task<IList<AdoNetStreamMessage>> GetStreamMessagesAsync(
+        string serviceId,
+        string providerId,
+        string queueId,
+        int maxCount,
+        int maxAttempts,
+        int visibilityTimeout,
+        int removalTimeout,
+        int evictionInterval,
+        int evictionBatchSize);
+
+    Task<IList<AdoNetStreamConfirmationAck>> ConfirmStreamMessagesAsync(
+        string serviceId,
+        string providerId,
+        string queueId,
+        IList<AdoNetStreamConfirmation> messages);
+
+    Task<IList<AdoNetStreamConfirmationAck>> ReleaseStreamMessagesAsync(
+        string serviceId,
+        string providerId,
+        string queueId,
+        IList<AdoNetStreamConfirmation> messages);
+}
+
 /// <summary>
 /// Receives message batches from an individual queue of an ADO.NET provider.
 /// </summary>
-internal partial class AdoNetQueueAdapterReceiver(string providerId, string queueId, AdoNetStreamOptions streamOptions, ClusterOptions clusterOptions, SimpleQueueCacheOptions cacheOptions, RelationalOrleansQueries queries, Serializer<AdoNetBatchContainer> serializer, ILogger<AdoNetQueueAdapterReceiver> logger) : IQueueAdapterReceiver
+internal partial class AdoNetQueueAdapterReceiver(string providerId, string queueId, AdoNetStreamOptions streamOptions, ClusterOptions clusterOptions, SimpleQueueCacheOptions cacheOptions, IStreamMessageQueries queries, Serializer<AdoNetBatchContainer> serializer, ILogger<AdoNetQueueAdapterReceiver> logger) : IQueueAdapterReceiver
 {
     private readonly ILogger<AdoNetQueueAdapterReceiver> _logger = logger;
+    private readonly object _lock = new();
+    private readonly Dictionary<long, PendingMessage> _pendingMessages = [];
 
     /// <summary>
     /// Flags that no further work should be attempted.
     /// </summary>
     private bool _shutdown;
 
-    /// <summary>
-    /// Helps shutdown wait for any outstanding storage operation.
-    /// </summary>
-    private Task? _outstandingTask;
+    private int _activeOperations;
+    private TaskCompletionSource? _operationsCompleted;
+
+    public AdoNetQueueAdapterReceiver(
+        string providerId,
+        string queueId,
+        AdoNetStreamOptions streamOptions,
+        ClusterOptions clusterOptions,
+        SimpleQueueCacheOptions cacheOptions,
+        RelationalOrleansQueries queries,
+        Serializer<AdoNetBatchContainer> serializer,
+        ILogger<AdoNetQueueAdapterReceiver> logger)
+        : this(providerId, queueId, streamOptions, clusterOptions, cacheOptions, new RelationalStreamMessageQueries(queries), serializer, logger)
+    {
+    }
 
     /// <summary>
     /// This receiver does not require initialization.
@@ -27,29 +66,61 @@ internal partial class AdoNetQueueAdapterReceiver(string providerId, string queu
     /// </summary>
     public async Task Shutdown(TimeSpan timeout)
     {
-        // disable any further attempts to access storage
-        _shutdown = true;
+        Task? operationsCompleted;
+        lock (_lock)
+        {
+            _shutdown = true;
+            operationsCompleted = _operationsCompleted?.Task;
+        }
 
-        // wait for any outstanding storage operation to complete.
-        var outstandingTask = _outstandingTask;
-        if (outstandingTask is not null)
+        if (operationsCompleted is not null)
         {
             try
             {
-                await outstandingTask.WaitAsync(timeout);
+                await operationsCompleted.WaitAsync(timeout);
             }
             catch (Exception ex)
             {
                 LogShutdownFault(ex, clusterOptions.ServiceId, providerId, queueId);
+                return;
             }
         }
+
+        List<AdoNetStreamConfirmation> pending;
+        lock (_lock)
+        {
+            RemoveExpiredPendingMessages();
+            if (_pendingMessages.Count == 0)
+            {
+                return;
+            }
+
+            pending = _pendingMessages
+                .Select(static item => new AdoNetStreamConfirmation(item.Key, item.Value.Dequeued))
+                .ToList();
+        }
+
+        try
+        {
+            var released = await queries.ReleaseStreamMessagesAsync(clusterOptions.ServiceId, providerId, queueId, pending).WaitAsync(timeout);
+            lock (_lock)
+            {
+                foreach (var message in released)
+                {
+                    _pendingMessages.Remove(message.MessageId);
+                }
+            }
+        }
+        catch (Exception ex)
+        {
+            LogReleaseFailed(ex, clusterOptions.ServiceId, providerId, queueId, pending);
+        }
     }
 
     /// <inheritdoc />
     public async Task<IList<IBatchContainer>> GetQueueMessagesAsync(int maxCount)
     {
-        // if shutdown has been called then we refuse further requests gracefully
-        if (_shutdown)
+        if (!TryBeginOperation())
         {
             return [];
         }
@@ -59,8 +130,7 @@ public async Task<IList<IBatchContainer>> GetQueueMessagesAsync(int maxCount)
 
         try
         {
-            // grab a message batch from storage while pinning the task so shutdown can wait for it
-            var task = queries.GetStreamMessagesAsync(
+            var messages = await queries.GetStreamMessagesAsync(
                 clusterOptions.ServiceId,
                 providerId,
                 queueId,
@@ -71,9 +141,14 @@ public async Task<IList<IBatchContainer>> GetQueueMessagesAsync(int maxCount)
                 streamOptions.EvictionInterval.TotalSecondsCeiling(),
                 streamOptions.EvictionBatchSize);
 
-            _outstandingTask = task;
-
-            var messages = await task;
+            lock (_lock)
+            {
+                RemoveExpiredPendingMessages();
+                foreach (var message in messages)
+                {
+                    _pendingMessages[message.MessageId] = new(message.Dequeued, message.ExpiresOn);
+                }
+            }
 
             // convert the messages into standard batch containers
             return messages.Select(x => AdoNetBatchContainer.FromMessage(serializer, x)).Cast<IBatchContainer>().ToList();
@@ -85,7 +160,7 @@ public async Task<IList<IBatchContainer>> GetQueueMessagesAsync(int maxCount)
         }
         finally
         {
-            _outstandingTask = null;
+            EndOperation();
         }
     }
 
@@ -98,31 +173,122 @@ public async Task MessagesDeliveredAsync(IList<IBatchContainer> messages)
             return;
         }
 
+        if (!TryBeginOperation())
+        {
+            return;
+        }
+
         // get the identifiers for the messages to confirm
         var items = messages.Cast<AdoNetBatchContainer>().Select(x => new AdoNetStreamConfirmation(x.SequenceToken.SequenceNumber, x.Dequeued)).ToList();
 
         try
         {
-            // execute the confirmation while pinning the task so shutdown can wait for it
-            var task = queries.ConfirmStreamMessagesAsync(clusterOptions.ServiceId, providerId, queueId, items);
-            _outstandingTask = task;
-
             try
             {
-                await task;
+                var confirmed = await queries.ConfirmStreamMessagesAsync(clusterOptions.ServiceId, providerId, queueId, items);
+                var receipts = items.ToDictionary(static item => item.MessageId, static item => item.Dequeued);
+                lock (_lock)
+                {
+                    foreach (var message in confirmed)
+                    {
+                        if (receipts.TryGetValue(message.MessageId, out var receipt)
+                            && _pendingMessages.TryGetValue(message.MessageId, out var pending)
+                            && receipt == pending.Dequeued)
+                        {
+                            _pendingMessages.Remove(message.MessageId);
+                        }
+                    }
+                }
             }
             catch (Exception ex)
             {
-                LogConfirmationFailed(ex, clusterOptions.ClusterId, providerId, queueId, items);
+                LogConfirmationFailed(ex, clusterOptions.ServiceId, providerId, queueId, items);
                 throw;
             }
         }
         finally
         {
-            _outstandingTask = null;
+            EndOperation();
         }
     }
 
+    private bool TryBeginOperation()
+    {
+        lock (_lock)
+        {
+            if (_shutdown)
+            {
+                return false;
+            }
+
+            if (_activeOperations++ == 0)
+            {
+                _operationsCompleted = new(TaskCreationOptions.RunContinuationsAsynchronously);
+            }
+
+            return true;
+        }
+    }
+
+    private void EndOperation()
+    {
+        TaskCompletionSource? operationsCompleted = null;
+        lock (_lock)
+        {
+            if (--_activeOperations == 0)
+            {
+                operationsCompleted = _operationsCompleted;
+                _operationsCompleted = null;
+            }
+        }
+
+        operationsCompleted?.TrySetResult();
+    }
+
+    private void RemoveExpiredPendingMessages()
+    {
+        var now = DateTime.UtcNow;
+        var expired = _pendingMessages
+            .Where(message => message.Value.ExpiresOn <= now)
+            .Select(static message => message.Key)
+            .ToList();
+        foreach (var messageId in expired)
+        {
+            _pendingMessages.Remove(messageId);
+        }
+    }
+
+    private readonly record struct PendingMessage(int Dequeued, DateTime ExpiresOn);
+
+    private sealed class RelationalStreamMessageQueries(RelationalOrleansQueries queries) : IStreamMessageQueries
+    {
+        public Task<IList<AdoNetStreamMessage>> GetStreamMessagesAsync(
+            string serviceId,
+            string providerId,
+            string queueId,
+            int maxCount,
+            int maxAttempts,
+            int visibilityTimeout,
+            int removalTimeout,
+            int evictionInterval,
+            int evictionBatchSize) =>
+            queries.GetStreamMessagesAsync(serviceId, providerId, queueId, maxCount, maxAttempts, visibilityTimeout, removalTimeout, evictionInterval, evictionBatchSize);
+
+        public Task<IList<AdoNetStreamConfirmationAck>> ConfirmStreamMessagesAsync(
+            string serviceId,
+            string providerId,
+            string queueId,
+            IList<AdoNetStreamConfirmation> messages) =>
+            queries.ConfirmStreamMessagesAsync(serviceId, providerId, queueId, messages);
+
+        public Task<IList<AdoNetStreamConfirmationAck>> ReleaseStreamMessagesAsync(
+            string serviceId,
+            string providerId,
+            string queueId,
+            IList<AdoNetStreamConfirmation> messages) =>
+            queries.ReleaseStreamMessagesAsync(serviceId, providerId, queueId, messages);
+    }
+
     #region Logging
 
     [LoggerMessage(1, LogLevel.Error, "Failed to get messages from ({ServiceId}, {ProviderId}, {QueueId})")]
@@ -134,5 +300,8 @@ public async Task MessagesDeliveredAsync(IList<IBatchContainer> messages)
     [LoggerMessage(3, LogLevel.Warning, "Handled fault while shutting down receiver for ({ServiceId}, {ProviderId}, {QueueId})")]
     private partial void LogShutdownFault(Exception exception, string serviceId, string providerId, string queueId);
 
+    [LoggerMessage(4, LogLevel.Warning, "Failed to release messages while shutting down receiver for ({ServiceId}, {ProviderId}, {QueueId}, {@Items})")]
+    private partial void LogReleaseFailed(Exception exception, string serviceId, string providerId, string queueId, List<AdoNetStreamConfirmation> items);
+
     #endregion Logging
 }

src/AdoNet/Orleans.Streaming.AdoNet/MySQL-Streaming.sql

@@ -375,8 +375,8 @@ WHILE LOCATE(_Delimiter1, _Items) > 0 DO
 
     SET _Value = SUBSTRING_INDEX(_Items, _Delimiter1, 1);
     SET _MessageId = CAST(SUBSTRING_INDEX(_Value, _Delimiter2, 1) AS UNSIGNED);
-    SET _Dequeued = CAST(SUBSTRING_INDEX(_Value, _Delimiter2, -1) AS UNSIGNED);
-    
+    SET _Dequeued = CAST(SUBSTRING_INDEX(_Value, _Delimiter2, -1) AS SIGNED);
+
     INSERT INTO _ItemsTable
     (
         ServiceId,
@@ -414,22 +414,35 @@ FROM
         AND M.ProviderId = I.ProviderId
         AND M.QueueId = I.QueueId
         AND M.MessageId = I.MessageId
-        AND M.Dequeued = I.Dequeued
+        AND M.Dequeued = ABS(I.Dequeued)
 ORDER BY
     M.ServiceId,
     M.ProviderId,
     M.QueueId,
 	M.MessageId
 FOR UPDATE;
 
-/* delete the elected batch */
-DELETE M
-FROM OrleansStreamMessage AS M
-INNER JOIN _Batch AS B
-    ON M.ServiceId = B.ServiceId
-    AND M.ProviderId = B.ProviderId
-    AND M.QueueId = B.QueueId
-    AND M.MessageId = B.MessageId;
+IF EXISTS (SELECT 1 FROM _ItemsTable WHERE Dequeued < 0) THEN
+    /* negative dequeue receipts release messages for immediate redelivery */
+    UPDATE OrleansStreamMessage AS M
+    INNER JOIN _Batch AS B
+        ON M.ServiceId = B.ServiceId
+        AND M.ProviderId = B.ProviderId
+        AND M.QueueId = B.QueueId
+        AND M.MessageId = B.MessageId
+    SET
+        M.VisibleOn = UTC_TIMESTAMP(6),
+        M.ModifiedOn = UTC_TIMESTAMP(6);
+ELSE
+    /* delete the elected batch */
+    DELETE M
+    FROM OrleansStreamMessage AS M
+    INNER JOIN _Batch AS B
+        ON M.ServiceId = B.ServiceId
+        AND M.ProviderId = B.ProviderId
+        AND M.QueueId = B.QueueId
+        AND M.MessageId = B.MessageId;
+END IF;
 
 /* return the ack */
 SELECT