From e1bc42689cd93651b66a15b004fed52e8e095feb Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Wed, 12 Aug 2026 21:23:09 -0700 Subject: [PATCH 1/6] fix(streaming): release ADO.NET messages on handoff Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../AdoNetQueueAdapterReceiver.cs | 65 ++++++++++++++++++- .../MySQL-Streaming.sql | 35 ++++++---- .../PostgreSQL-Streaming.sql | 43 ++++++++++++ .../SQLServer-Streaming.sql | 41 ++++++++++++ .../Storage/RelationalOrleansQueries.cs | 57 ++++++++++++++-- .../AdoNetQueueAdapterReceiverTests.cs | 38 +++++++++++ .../RelationalOrleansQueriesTests.cs | 37 +++++++++++ .../StreamingDiagnosticObserver.cs | 52 +++++++++++---- 8 files changed, 337 insertions(+), 31 deletions(-) diff --git a/src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs b/src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs index 77f44eff836..2c5c117d1f1 100644 --- a/src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs +++ b/src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs @@ -6,6 +6,7 @@ namespace Orleans.Streaming.AdoNet; internal partial class AdoNetQueueAdapterReceiver(string providerId, string queueId, AdoNetStreamOptions streamOptions, ClusterOptions clusterOptions, SimpleQueueCacheOptions cacheOptions, RelationalOrleansQueries queries, Serializer serializer, ILogger logger) : IQueueAdapterReceiver { private readonly ILogger _logger = logger; + private readonly Dictionary _pendingMessages = []; /// /// Flags that no further work should be attempted. @@ -43,6 +44,35 @@ public async Task Shutdown(TimeSpan timeout) LogShutdownFault(ex, clusterOptions.ServiceId, providerId, queueId); } } + + RemoveExpiredPendingMessages(); + if (_outstandingTask is not null || _pendingMessages.Count == 0) + { + return; + } + + var pending = _pendingMessages + .Select(static item => new AdoNetStreamConfirmation(item.Key, item.Value.Dequeued)) + .ToList(); + + try + { + var task = queries.ReleaseStreamMessagesAsync(clusterOptions.ServiceId, providerId, queueId, pending); + _outstandingTask = task; + var released = await task.WaitAsync(timeout); + foreach (var message in released) + { + _pendingMessages.Remove(message.MessageId); + } + } + catch (Exception ex) + { + LogReleaseFailed(ex, clusterOptions.ServiceId, providerId, queueId, pending); + } + finally + { + _outstandingTask = null; + } } /// @@ -74,6 +104,11 @@ public async Task> GetQueueMessagesAsync(int maxCount) _outstandingTask = task; var messages = await task; + 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().ToList(); @@ -109,7 +144,17 @@ public async Task MessagesDeliveredAsync(IList messages) try { - await task; + var confirmed = await task; + var receipts = items.ToDictionary(static item => item.MessageId, static item => item.Dequeued); + 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) { @@ -123,6 +168,21 @@ public async Task MessagesDeliveredAsync(IList messages) } } + 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); + #region Logging [LoggerMessage(1, LogLevel.Error, "Failed to get messages from ({ServiceId}, {ProviderId}, {QueueId})")] @@ -134,5 +194,8 @@ public async Task MessagesDeliveredAsync(IList 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 items); + #endregion Logging } diff --git a/src/AdoNet/Orleans.Streaming.AdoNet/MySQL-Streaming.sql b/src/AdoNet/Orleans.Streaming.AdoNet/MySQL-Streaming.sql index ed709933a54..7caa8b2e4e0 100644 --- a/src/AdoNet/Orleans.Streaming.AdoNet/MySQL-Streaming.sql +++ b/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,7 +414,7 @@ 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, @@ -422,14 +422,27 @@ ORDER BY 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 diff --git a/src/AdoNet/Orleans.Streaming.AdoNet/PostgreSQL-Streaming.sql b/src/AdoNet/Orleans.Streaming.AdoNet/PostgreSQL-Streaming.sql index 96d752ef816..817742dab8f 100644 --- a/src/AdoNet/Orleans.Streaming.AdoNet/PostgreSQL-Streaming.sql +++ b/src/AdoNet/Orleans.Streaming.AdoNet/PostgreSQL-Streaming.sql @@ -316,6 +316,7 @@ AS $$ #VARIABLE_CONFLICT USE_COLUMN DECLARE _Count INT; + _Now TIMESTAMP(6) WITHOUT TIME ZONE := CURRENT_TIMESTAMP AT TIME ZONE 'UTC'; BEGIN CREATE TEMP TABLE _ItemsTable @@ -335,6 +336,48 @@ SELECT FROM UNNEST(string_to_array(_Items, '|')) AS Value; +/* negative dequeue receipts release messages for immediate redelivery */ +IF EXISTS (SELECT 1 FROM _ItemsTable WHERE Dequeued < 0) THEN + RETURN QUERY + WITH Batch AS + ( + SELECT + M.* + FROM + OrleansStreamMessage AS M + INNER JOIN _ItemsTable AS I + ON I.MessageId = M.MessageId + AND -I.Dequeued = M.Dequeued + WHERE + ServiceId = _ServiceId + AND ProviderId = _ProviderId + AND QueueId = _QueueId + ORDER BY + ServiceId, + ProviderId, + QueueId, + MessageId + FOR UPDATE + ) + UPDATE OrleansStreamMessage AS M + SET + VisibleOn = _Now, + ModifiedOn = _Now + FROM + Batch AS B + WHERE + M.ServiceId = B.ServiceId + AND M.ProviderId = B.ProviderId + AND M.QueueId = B.QueueId + AND M.MessageId = B.MessageId + RETURNING + M.ServiceId, + M.ProviderId, + M.QueueId, + M.MessageId; + RETURN; +END IF; + RETURN QUERY WITH Batch AS ( diff --git a/src/AdoNet/Orleans.Streaming.AdoNet/SQLServer-Streaming.sql b/src/AdoNet/Orleans.Streaming.AdoNet/SQLServer-Streaming.sql index 5e193dfea45..742792f7102 100644 --- a/src/AdoNet/Orleans.Streaming.AdoNet/SQLServer-Streaming.sql +++ b/src/AdoNet/Orleans.Streaming.AdoNet/SQLServer-Streaming.sql @@ -415,6 +415,47 @@ FROM /* count the number of messages to delete so we can use order by in the next query */ DECLARE @Count INT = (SELECT COUNT(*) FROM @ItemsTable); +/* negative dequeue receipts release messages for immediate redelivery */ +IF EXISTS (SELECT 1 FROM @ItemsTable WHERE Dequeued < 0) +BEGIN + DECLARE @Now DATETIME2(7) = SYSUTCDATETIME(); + + WITH Batch AS + ( + SELECT TOP (@Count) + M.* + FROM + OrleansStreamMessage AS M WITH (UPDLOCK, READPAST, READCOMMITTEDLOCK, ROWLOCK) + WHERE + ServiceId = @ServiceId + AND ProviderId = @ProviderId + AND QueueId = @QueueId + AND EXISTS + ( + SELECT * + FROM @ItemsTable AS I + WHERE I.MessageId = M.MessageId + AND -I.Dequeued = M.Dequeued + ) + ORDER BY + ServiceId, + ProviderId, + QueueId, + MessageId + ) + UPDATE Batch + SET + VisibleOn = @Now, + ModifiedOn = @Now + OUTPUT + Inserted.ServiceId, + Inserted.ProviderId, + Inserted.QueueId, + Inserted.MessageId; + + RETURN; +END; + /* delete messages in the exact same order as the clustered index to avoid deadlocks with other queries */ /* skip messages being changed concurrently since their dequeue receipt may no longer match */ WITH Batch AS diff --git a/src/AdoNet/Shared/Storage/RelationalOrleansQueries.cs b/src/AdoNet/Shared/Storage/RelationalOrleansQueries.cs index 86086c5e77e..a00c90978bb 100644 --- a/src/AdoNet/Shared/Storage/RelationalOrleansQueries.cs +++ b/src/AdoNet/Shared/Storage/RelationalOrleansQueries.cs @@ -495,10 +495,46 @@ internal Task> ConfirmStreamMessagesAsync(str return Task.FromResult>([]); } - // this builds a string in the form "1:2|3:4|5:6" where the first number is the message id and the second is the dequeue counter which acts as a receipt - // while we have more efficient ways of passing this data per RDMS, we use a string here to ensure call compatibility across ADONET providers - // it is the responsibility of the RDMS implementation to parse this string and apply it correctly - var items = messages.Aggregate(new StringBuilder(), (b, m) => b.Append(b.Length > 0 ? "|" : "").Append(m.MessageId).Append(':').Append(m.Dequeued), b => b.ToString()); + return ReadAsync>( + dbStoredQueries.ConfirmStreamMessagesKey, + record => new AdoNetStreamConfirmationAck( + (string)record[nameof(AdoNetStreamConfirmationAck.ServiceId)], + (string)record[nameof(AdoNetStreamConfirmationAck.ProviderId)], + (string)record[nameof(AdoNetStreamConfirmationAck.QueueId)], + (long)record[nameof(AdoNetStreamConfirmationAck.MessageId)]), + command => new DbStoredQueries.Columns(command) + { + ServiceId = serviceId, + ProviderId = providerId, + QueueId = queueId, + Items = FormatStreamConfirmations(messages, release: false) + }, + result => result.ToList()); + } + + /// + /// Makes unconfirmed stream messages immediately available for redelivery. + /// + /// The service identifier. + /// The provider identifier. + /// The queue identifier. + /// The messages to release. + /// A list of released messages. + /// + /// The dequeue counter acts as a receipt so that a stale receiver cannot release a message + /// which has already been dequeued by a new receiver. + /// + internal Task> ReleaseStreamMessagesAsync(string serviceId, string providerId, string queueId, IList messages) + { + ArgumentNullException.ThrowIfNull(serviceId); + ArgumentNullException.ThrowIfNull(providerId); + ArgumentNullException.ThrowIfNull(queueId); + ArgumentNullException.ThrowIfNull(messages); + + if (messages.Count == 0) + { + return Task.FromResult>([]); + } return ReadAsync>( dbStoredQueries.ConfirmStreamMessagesKey, @@ -512,11 +548,22 @@ internal Task> ConfirmStreamMessagesAsync(str ServiceId = serviceId, ProviderId = providerId, QueueId = queueId, - Items = items + Items = FormatStreamConfirmations(messages, release: true) }, result => result.ToList()); } + // Builds a provider-neutral receipt list in the form "1:2|3:4|5:6". + private static string FormatStreamConfirmations(IList messages, bool release) => + messages.Aggregate( + new StringBuilder(), + (builder, message) => builder + .Append(builder.Length > 0 ? "|" : "") + .Append(message.MessageId) + .Append(':') + .Append(release ? -message.Dequeued : message.Dequeued), + static builder => builder.ToString()); + /// /// Applies delivery failure logic to a stream message, such as making the message visible again or moving it to dead letters. /// diff --git a/test/Extensions/Orleans.AdoNet.Tests/Streaming/AdoNetQueueAdapterReceiverTests.cs b/test/Extensions/Orleans.AdoNet.Tests/Streaming/AdoNetQueueAdapterReceiverTests.cs index 6dd266e7006..3424f21d50c 100644 --- a/test/Extensions/Orleans.AdoNet.Tests/Streaming/AdoNetQueueAdapterReceiverTests.cs +++ b/test/Extensions/Orleans.AdoNet.Tests/Streaming/AdoNetQueueAdapterReceiverTests.cs @@ -147,6 +147,44 @@ public async Task AdoNetQueueAdapterReceiver_GetsMessages_ConfirmsMessages() Assert.False(storedConfirmed.ContainsKey(ackValid.MessageId)); } + /// + /// Tests that shutting down a receiver immediately releases its unconfirmed messages. + /// + [SkippableFact] + public async Task AdoNetQueueAdapterReceiver_Shutdown_ReleasesUnconfirmedMessages() + { + var serviceId = $"Service-{Guid.NewGuid()}"; + var providerId = $"Provider-{Guid.NewGuid()}"; + var queueId = $"Queue-{Guid.NewGuid()}"; + var clusterOptions = new ClusterOptions { ServiceId = serviceId }; + var streamOptions = new AdoNetStreamOptions + { + Invariant = invariant, + ConnectionString = _storage.ConnectionString, + VisibilityTimeout = TimeSpan.FromMinutes(5), + EvictionBatchSize = 0 + }; + var cacheOptions = new SimpleQueueCacheOptions(); + var serializer = _fixture.Serializer.GetSerializer(); + var logger = NullLogger.Instance; + var streamId = StreamId.Create("MyNamespace", "MyKey"); + var payload = serializer.SerializeToArray(new AdoNetBatchContainer(streamId, [new TestModel(1)], null!)); + var ack = await _queries.QueueStreamMessageAsync(serviceId, providerId, queueId, payload, 100); + + var receiver = new AdoNetQueueAdapterReceiver(providerId, queueId, streamOptions, clusterOptions, cacheOptions, _queries, serializer, logger); + var first = Assert.IsType(Assert.Single(await receiver.GetQueueMessagesAsync(1))); + Assert.Equal(1, first.Dequeued); + + await receiver.Shutdown(TimeSpan.FromSeconds(10)); + + var replacement = new AdoNetQueueAdapterReceiver(providerId, queueId, streamOptions, clusterOptions, cacheOptions, _queries, serializer, logger); + var redelivered = Assert.IsType(Assert.Single(await replacement.GetQueueMessagesAsync(1))); + Assert.Equal(ack.MessageId, redelivered.SequenceToken.SequenceNumber); + Assert.Equal(2, redelivered.Dequeued); + await replacement.MessagesDeliveredAsync([redelivered]); + await replacement.Shutdown(TimeSpan.FromSeconds(10)); + } + /// /// Tests that waits for the outstanding task. /// diff --git a/test/Extensions/Orleans.AdoNet.Tests/Streaming/RelationalOrleansQueriesTests.cs b/test/Extensions/Orleans.AdoNet.Tests/Streaming/RelationalOrleansQueriesTests.cs index 29b5d851700..42dcd99f1d3 100644 --- a/test/Extensions/Orleans.AdoNet.Tests/Streaming/RelationalOrleansQueriesTests.cs +++ b/test/Extensions/Orleans.AdoNet.Tests/Streaming/RelationalOrleansQueriesTests.cs @@ -619,6 +619,43 @@ public async Task RelationalOrleansQueries_DoesNotDequeueSingleMessageBeforeVisi Assert.Empty(second); } + /// + /// Tests that a message can be released for immediate redelivery using its dequeue receipt. + /// + [SkippableFact] + public async Task RelationalOrleansQueries_ReleasesMessageUsingDequeueReceipt() + { + var serviceId = $"Service-{Guid.NewGuid()}"; + var providerId = $"Provider-{Guid.NewGuid()}"; + var queueId = $"Queue-{Guid.NewGuid()}"; + var payload = RandomPayload(); + var expiryTimeout = 100; + var maxCount = 1; + var maxAttempts = 5; + var visibilityTimeout = 100; + var removalTimeout = 100; + var evictionInterval = 100; + var evictionBatchSize = 0; + + var ack = await _queries.QueueStreamMessageAsync(serviceId, providerId, queueId, payload, expiryTimeout); + var first = Assert.Single(await _queries.GetStreamMessagesAsync(serviceId, providerId, queueId, maxCount, maxAttempts, visibilityTimeout, removalTimeout, evictionInterval, evictionBatchSize)); + var firstReceipt = new AdoNetStreamConfirmation(first.MessageId, first.Dequeued); + + var released = Assert.Single(await _queries.ReleaseStreamMessagesAsync(serviceId, providerId, queueId, [firstReceipt])); + Assert.Equal(ack.MessageId, released.MessageId); + + var second = Assert.Single(await _queries.GetStreamMessagesAsync(serviceId, providerId, queueId, maxCount, maxAttempts, visibilityTimeout, removalTimeout, evictionInterval, evictionBatchSize)); + Assert.Equal(first.Dequeued + 1, second.Dequeued); + + Assert.Empty(await _queries.ReleaseStreamMessagesAsync(serviceId, providerId, queueId, [firstReceipt])); + Assert.Empty(await _queries.GetStreamMessagesAsync(serviceId, providerId, queueId, maxCount, maxAttempts, visibilityTimeout, removalTimeout, evictionInterval, evictionBatchSize)); + + var secondReceipt = new AdoNetStreamConfirmation(second.MessageId, second.Dequeued); + Assert.Single(await _queries.ReleaseStreamMessagesAsync(serviceId, providerId, queueId, [secondReceipt])); + var third = Assert.Single(await _queries.GetStreamMessagesAsync(serviceId, providerId, queueId, maxCount, maxAttempts, visibilityTimeout, removalTimeout, evictionInterval, evictionBatchSize)); + await _queries.ConfirmStreamMessagesAsync(serviceId, providerId, queueId, [new AdoNetStreamConfirmation(third.MessageId, third.Dequeued)]); + } + /// /// Tests that a single message is not dequeued again after expiry /// diff --git a/test/TestInfrastructure/TestExtensions/Diagnostics/StreamingDiagnosticObserver.cs b/test/TestInfrastructure/TestExtensions/Diagnostics/StreamingDiagnosticObserver.cs index 91ef437da7c..10588faa25b 100644 --- a/test/TestInfrastructure/TestExtensions/Diagnostics/StreamingDiagnosticObserver.cs +++ b/test/TestInfrastructure/TestExtensions/Diagnostics/StreamingDiagnosticObserver.cs @@ -178,13 +178,25 @@ await _events /// public async Task WaitForItemDeliveryCountAsync(StreamId streamId, int expectedCount, string? streamProvider, CancellationToken cancellationToken) { - await _events - .OfType() - .Where(e => MatchesStream(e.StreamId, e.StreamProvider, streamId, streamProvider)) - .Take(expectedCount) - .LastOrDefaultAsync() - .ToTask(cancellationToken) - .ConfigureAwait(false); + var actualCount = 0; + try + { + await _events + .OfType() + .Where(e => MatchesStream(e.StreamId, e.StreamProvider, streamId, streamProvider)) + .Do(_ => actualCount++) + .Take(expectedCount) + .LastOrDefaultAsync() + .ToTask(cancellationToken) + .ConfigureAwait(false); + } + catch (OperationCanceledException exception) when (cancellationToken.IsCancellationRequested) + { + throw new OperationCanceledException( + $"Canceled while waiting for {expectedCount} item deliveries on stream {streamId} from provider '{streamProvider ?? ""}'; observed {actualCount}.", + exception, + cancellationToken); + } } /// @@ -192,13 +204,25 @@ await _events /// public async Task WaitForItemDeliveryCountAsync(StreamId streamId, Guid subscriptionId, int expectedCount, string? streamProvider, CancellationToken cancellationToken) { - await _events - .OfType() - .Where(e => MatchesSubscription(e.StreamId, e.SubscriptionId, e.StreamProvider, streamId, subscriptionId, streamProvider)) - .Take(expectedCount) - .LastOrDefaultAsync() - .ToTask(cancellationToken) - .ConfigureAwait(false); + var actualCount = 0; + try + { + await _events + .OfType() + .Where(e => MatchesSubscription(e.StreamId, e.SubscriptionId, e.StreamProvider, streamId, subscriptionId, streamProvider)) + .Do(_ => actualCount++) + .Take(expectedCount) + .LastOrDefaultAsync() + .ToTask(cancellationToken) + .ConfigureAwait(false); + } + catch (OperationCanceledException exception) when (cancellationToken.IsCancellationRequested) + { + throw new OperationCanceledException( + $"Canceled while waiting for {expectedCount} item deliveries on stream {streamId}, subscription {subscriptionId}, from provider '{streamProvider ?? ""}'; observed {actualCount}.", + exception, + cancellationToken); + } } /// From 8df8955b0cdc9c8b11fc5a46da6105d73a8dcbf4 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 13 Aug 2026 02:06:02 -0700 Subject: [PATCH 2/6] test(streaming): count diagnostic events atomically Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../Diagnostics/StreamingDiagnosticObserver.cs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/TestInfrastructure/TestExtensions/Diagnostics/StreamingDiagnosticObserver.cs b/test/TestInfrastructure/TestExtensions/Diagnostics/StreamingDiagnosticObserver.cs index 10588faa25b..5e398f5e9d9 100644 --- a/test/TestInfrastructure/TestExtensions/Diagnostics/StreamingDiagnosticObserver.cs +++ b/test/TestInfrastructure/TestExtensions/Diagnostics/StreamingDiagnosticObserver.cs @@ -184,7 +184,7 @@ public async Task WaitForItemDeliveryCountAsync(StreamId streamId, int expectedC await _events .OfType() .Where(e => MatchesStream(e.StreamId, e.StreamProvider, streamId, streamProvider)) - .Do(_ => actualCount++) + .Do(_ => Interlocked.Increment(ref actualCount)) .Take(expectedCount) .LastOrDefaultAsync() .ToTask(cancellationToken) @@ -193,7 +193,7 @@ await _events catch (OperationCanceledException exception) when (cancellationToken.IsCancellationRequested) { throw new OperationCanceledException( - $"Canceled while waiting for {expectedCount} item deliveries on stream {streamId} from provider '{streamProvider ?? ""}'; observed {actualCount}.", + $"Canceled while waiting for {expectedCount} item deliveries on stream {streamId} from provider '{streamProvider ?? ""}'; observed {Volatile.Read(ref actualCount)}.", exception, cancellationToken); } @@ -210,7 +210,7 @@ public async Task WaitForItemDeliveryCountAsync(StreamId streamId, Guid subscrip await _events .OfType() .Where(e => MatchesSubscription(e.StreamId, e.SubscriptionId, e.StreamProvider, streamId, subscriptionId, streamProvider)) - .Do(_ => actualCount++) + .Do(_ => Interlocked.Increment(ref actualCount)) .Take(expectedCount) .LastOrDefaultAsync() .ToTask(cancellationToken) @@ -219,7 +219,7 @@ await _events catch (OperationCanceledException exception) when (cancellationToken.IsCancellationRequested) { throw new OperationCanceledException( - $"Canceled while waiting for {expectedCount} item deliveries on stream {streamId}, subscription {subscriptionId}, from provider '{streamProvider ?? ""}'; observed {actualCount}.", + $"Canceled while waiting for {expectedCount} item deliveries on stream {streamId}, subscription {subscriptionId}, from provider '{streamProvider ?? ""}'; observed {Volatile.Read(ref actualCount)}.", exception, cancellationToken); } From 7ea1b0d977fe4000c65e5f78475aabed086c1c38 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Fri, 14 Aug 2026 09:23:09 -0700 Subject: [PATCH 3/6] fix(streaming): await ADO.NET receiver handoff Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../AdoNetQueueAdapterReceiver.cs | 206 +++++++++++++----- .../AdoNetQueueAdapterReceiverTests.cs | 134 ++++++++++-- 2 files changed, 271 insertions(+), 69 deletions(-) diff --git a/src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs b/src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs index 2c5c117d1f1..a88c34e57fa 100644 --- a/src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs +++ b/src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs @@ -1,11 +1,38 @@ namespace Orleans.Streaming.AdoNet; +internal interface IStreamMessageQueries +{ + Task> GetStreamMessagesAsync( + string serviceId, + string providerId, + string queueId, + int maxCount, + int maxAttempts, + int visibilityTimeout, + int removalTimeout, + int evictionInterval, + int evictionBatchSize); + + Task> ConfirmStreamMessagesAsync( + string serviceId, + string providerId, + string queueId, + IList messages); + + Task> ReleaseStreamMessagesAsync( + string serviceId, + string providerId, + string queueId, + IList messages); +} + /// /// Receives message batches from an individual queue of an ADO.NET provider. /// -internal partial class AdoNetQueueAdapterReceiver(string providerId, string queueId, AdoNetStreamOptions streamOptions, ClusterOptions clusterOptions, SimpleQueueCacheOptions cacheOptions, RelationalOrleansQueries queries, Serializer serializer, ILogger logger) : IQueueAdapterReceiver +internal partial class AdoNetQueueAdapterReceiver(string providerId, string queueId, AdoNetStreamOptions streamOptions, ClusterOptions clusterOptions, SimpleQueueCacheOptions cacheOptions, IStreamMessageQueries queries, Serializer serializer, ILogger logger) : IQueueAdapterReceiver { private readonly ILogger _logger = logger; + private readonly object _lock = new(); private readonly Dictionary _pendingMessages = []; /// @@ -13,10 +40,21 @@ internal partial class AdoNetQueueAdapterReceiver(string providerId, string queu /// private bool _shutdown; - /// - /// Helps shutdown wait for any outstanding storage operation. - /// - private Task? _outstandingTask; + private int _activeOperations; + private TaskCompletionSource? _operationsCompleted; + + internal AdoNetQueueAdapterReceiver( + string providerId, + string queueId, + AdoNetStreamOptions streamOptions, + ClusterOptions clusterOptions, + SimpleQueueCacheOptions cacheOptions, + RelationalOrleansQueries queries, + Serializer serializer, + ILogger logger) + : this(providerId, queueId, streamOptions, clusterOptions, cacheOptions, new RelationalStreamMessageQueries(queries), serializer, logger) + { + } /// /// This receiver does not require initialization. @@ -28,58 +66,61 @@ internal partial class AdoNetQueueAdapterReceiver(string providerId, string queu /// 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; } } - RemoveExpiredPendingMessages(); - if (_outstandingTask is not null || _pendingMessages.Count == 0) + List pending; + lock (_lock) { - return; - } + RemoveExpiredPendingMessages(); + if (_pendingMessages.Count == 0) + { + return; + } - var pending = _pendingMessages - .Select(static item => new AdoNetStreamConfirmation(item.Key, item.Value.Dequeued)) - .ToList(); + pending = _pendingMessages + .Select(static item => new AdoNetStreamConfirmation(item.Key, item.Value.Dequeued)) + .ToList(); + } try { - var task = queries.ReleaseStreamMessagesAsync(clusterOptions.ServiceId, providerId, queueId, pending); - _outstandingTask = task; - var released = await task.WaitAsync(timeout); - foreach (var message in released) + var released = await queries.ReleaseStreamMessagesAsync(clusterOptions.ServiceId, providerId, queueId, pending).WaitAsync(timeout); + lock (_lock) { - _pendingMessages.Remove(message.MessageId); + foreach (var message in released) + { + _pendingMessages.Remove(message.MessageId); + } } } catch (Exception ex) { LogReleaseFailed(ex, clusterOptions.ServiceId, providerId, queueId, pending); } - finally - { - _outstandingTask = null; - } } /// public async Task> GetQueueMessagesAsync(int maxCount) { - // if shutdown has been called then we refuse further requests gracefully - if (_shutdown) + if (!TryBeginOperation()) { return []; } @@ -89,8 +130,7 @@ public async Task> 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, @@ -101,13 +141,13 @@ public async Task> GetQueueMessagesAsync(int maxCount) streamOptions.EvictionInterval.TotalSecondsCeiling(), streamOptions.EvictionBatchSize); - _outstandingTask = task; - - var messages = await task; - RemoveExpiredPendingMessages(); - foreach (var message in messages) + lock (_lock) { - _pendingMessages[message.MessageId] = new(message.Dequeued, message.ExpiresOn); + RemoveExpiredPendingMessages(); + foreach (var message in messages) + { + _pendingMessages[message.MessageId] = new(message.Dequeued, message.ExpiresOn); + } } // convert the messages into standard batch containers @@ -120,7 +160,7 @@ public async Task> GetQueueMessagesAsync(int maxCount) } finally { - _outstandingTask = null; + EndOperation(); } } @@ -133,39 +173,76 @@ public async Task MessagesDeliveredAsync(IList messages) return; } + if (!TryBeginOperation()) + { + return; + } + // get the identifiers for the messages to confirm var items = messages.Cast().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 { - var confirmed = await task; + var confirmed = await queries.ConfirmStreamMessagesAsync(clusterOptions.ServiceId, providerId, queueId, items); var receipts = items.ToDictionary(static item => item.MessageId, static item => item.Dequeued); - foreach (var message in confirmed) + lock (_lock) { - if (receipts.TryGetValue(message.MessageId, out var receipt) - && _pendingMessages.TryGetValue(message.MessageId, out var pending) - && receipt == pending.Dequeued) + foreach (var message in confirmed) { - _pendingMessages.Remove(message.MessageId); + 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() @@ -183,6 +260,35 @@ private void RemoveExpiredPendingMessages() private readonly record struct PendingMessage(int Dequeued, DateTime ExpiresOn); + private sealed class RelationalStreamMessageQueries(RelationalOrleansQueries queries) : IStreamMessageQueries + { + public Task> 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> ConfirmStreamMessagesAsync( + string serviceId, + string providerId, + string queueId, + IList messages) => + queries.ConfirmStreamMessagesAsync(serviceId, providerId, queueId, messages); + + public Task> ReleaseStreamMessagesAsync( + string serviceId, + string providerId, + string queueId, + IList messages) => + queries.ReleaseStreamMessagesAsync(serviceId, providerId, queueId, messages); + } + #region Logging [LoggerMessage(1, LogLevel.Error, "Failed to get messages from ({ServiceId}, {ProviderId}, {QueueId})")] diff --git a/test/Extensions/Orleans.AdoNet.Tests/Streaming/AdoNetQueueAdapterReceiverTests.cs b/test/Extensions/Orleans.AdoNet.Tests/Streaming/AdoNetQueueAdapterReceiverTests.cs index 3424f21d50c..3ede05a1dae 100644 --- a/test/Extensions/Orleans.AdoNet.Tests/Streaming/AdoNetQueueAdapterReceiverTests.cs +++ b/test/Extensions/Orleans.AdoNet.Tests/Streaming/AdoNetQueueAdapterReceiverTests.cs @@ -10,6 +10,103 @@ namespace Tester.AdoNet.Streaming; +/// +/// Provider-independent lifecycle tests for . +/// +[Collection(TestEnvironmentFixture.DefaultCollection)] +[TestCategory("BVT"), TestCategory("AdoNet"), TestCategory("Streaming")] +[TestProvider("None")] +[TestSuite("BVT")] +[TestArea("Streaming")] +public class AdoNetQueueAdapterReceiverLifecycleTests(TestEnvironmentFixture fixture) +{ + [Fact] + public async Task AdoNetQueueAdapterReceiver_Shutdown_WaitsForDequeueBookkeepingBeforeRelease() + { + var serviceId = $"Service-{Guid.NewGuid()}"; + var providerId = $"Provider-{Guid.NewGuid()}"; + var queueId = $"Queue-{Guid.NewGuid()}"; + var clusterOptions = new ClusterOptions { ServiceId = serviceId }; + var streamOptions = new AdoNetStreamOptions + { + VisibilityTimeout = TimeSpan.FromMinutes(5), + EvictionBatchSize = 0 + }; + var cacheOptions = new SimpleQueueCacheOptions(); + var serializer = fixture.Serializer.GetSerializer(); + var logger = NullLogger.Instance; + var payload = serializer.SerializeToArray(new AdoNetBatchContainer(StreamId.Create("MyNamespace", "MyKey"), [new TestModel(1)], null!)); + var now = DateTime.UtcNow; + var message = new AdoNetStreamMessage(serviceId, providerId, queueId, 42, 1, now.AddMinutes(5), now.AddHours(1), now, now, payload); + var dequeueStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var continueDequeue = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var queries = new BlockingStreamMessageQueries(message, dequeueStarted, continueDequeue); + + var receiver = new AdoNetQueueAdapterReceiver(providerId, queueId, streamOptions, clusterOptions, cacheOptions, queries, serializer, logger); + var getTask = receiver.GetQueueMessagesAsync(1); + await dequeueStarted.Task; + + var shutdownTask = receiver.Shutdown(TimeSpan.FromSeconds(10)); + Assert.False(shutdownTask.IsCompleted); + + continueDequeue.SetResult(); + var dequeued = Assert.IsType(Assert.Single(await getTask)); + await shutdownTask; + + Assert.Equal(message.MessageId, dequeued.SequenceToken.SequenceNumber); + var released = Assert.Single(queries.Released); + Assert.Equal(message.MessageId, released.MessageId); + Assert.Equal(message.Dequeued, released.Dequeued); + } + + [GenerateSerializer] + [Alias("Tester.AdoNet.Streaming.AdoNetQueueAdapterReceiverLifecycleTests.TestModel")] + public record TestModel( + [property: Id(0)] int Value); + + private sealed class BlockingStreamMessageQueries( + AdoNetStreamMessage message, + TaskCompletionSource dequeueStarted, + TaskCompletionSource continueDequeue) : IStreamMessageQueries + { + public IList Released { get; private set; } = []; + + public async Task> GetStreamMessagesAsync( + string serviceId, + string providerId, + string queueId, + int maxCount, + int maxAttempts, + int visibilityTimeout, + int removalTimeout, + int evictionInterval, + int evictionBatchSize) + { + dequeueStarted.SetResult(); + await continueDequeue.Task; + return [message]; + } + + public Task> ConfirmStreamMessagesAsync( + string serviceId, + string providerId, + string queueId, + IList messages) => + throw new NotSupportedException(); + + public Task> ReleaseStreamMessagesAsync( + string serviceId, + string providerId, + string queueId, + IList messages) + { + Released = messages.ToList(); + return Task.FromResult>( + [new(serviceId, providerId, queueId, message.MessageId)]); + } + } +} + /// /// Tests for against SQL Server. /// @@ -188,41 +285,40 @@ public async Task AdoNetQueueAdapterReceiver_Shutdown_ReleasesUnconfirmedMessage /// /// Tests that waits for the outstanding task. /// - /// [SkippableFact] public async Task AdoNetQueueAdapterReceiver_Shutdown_WaitsForOutstandingTask() { - // arrange - receiver - var serviceId = "MyServiceId"; - var clusterOptions = new ClusterOptions - { - ServiceId = serviceId - }; - var providerId = "MyProviderId"; - var queueId = "MyQueueId"; + var serviceId = $"Service-{Guid.NewGuid()}"; + var providerId = $"Provider-{Guid.NewGuid()}"; + var queueId = $"Queue-{Guid.NewGuid()}"; + var clusterOptions = new ClusterOptions { ServiceId = serviceId }; var streamOptions = new AdoNetStreamOptions { Invariant = invariant, - ConnectionString = _storage.ConnectionString + ConnectionString = _storage.ConnectionString, + VisibilityTimeout = TimeSpan.FromMinutes(5), + EvictionBatchSize = 0 }; var cacheOptions = new SimpleQueueCacheOptions(); var serializer = _fixture.Serializer.GetSerializer(); var logger = NullLogger.Instance; var receiver = new AdoNetQueueAdapterReceiver(providerId, queueId, streamOptions, clusterOptions, cacheOptions, _queries, serializer, logger); - await receiver.Initialize(TimeSpan.FromSeconds(10)); - - // arrange - enqueue a message var payload = serializer.SerializeToArray(new AdoNetBatchContainer(StreamId.Create("MyNamespace", "MyKey"), [new TestModel(1)], null!)); - await _queries.QueueStreamMessageAsync(serviceId, providerId, queueId, payload, 100); - - // act - start getting messages from the receiver - var getTask = receiver.GetQueueMessagesAsync(10); + var ack = await _queries.QueueStreamMessageAsync(serviceId, providerId, queueId, payload, 100); - // act - shutdown the receiver + var getTask = receiver.GetQueueMessagesAsync(1); await receiver.Shutdown(TimeSpan.FromSeconds(10)); - // assert - the outstanding task completes before the shutdown task Assert.True(getTask.IsCompleted); + var first = Assert.IsType(Assert.Single(await getTask)); + Assert.Equal(1, first.Dequeued); + + var replacement = new AdoNetQueueAdapterReceiver(providerId, queueId, streamOptions, clusterOptions, cacheOptions, _queries, serializer, logger); + var redelivered = Assert.IsType(Assert.Single(await replacement.GetQueueMessagesAsync(1))); + Assert.Equal(ack.MessageId, redelivered.SequenceToken.SequenceNumber); + Assert.Equal(2, redelivered.Dequeued); + await replacement.MessagesDeliveredAsync([redelivered]); + await replacement.Shutdown(TimeSpan.FromSeconds(10)); } public Task DisposeAsync() => Task.CompletedTask; From 2cbafa4082120cf9b709a4df6a4f1a5081e155da Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Sat, 15 Aug 2026 10:18:50 -0700 Subject: [PATCH 4/6] fix(streaming): preserve cold stream registration handoff Restore the ADO.NET receiver constructor used by ActivatorUtilities and prevent later queue reads from overtaking cold-stream registration. Retain the first cache pin until subscriber cursors attach, with cleanup on removal and cache pressure. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../AdoNetQueueAdapterReceiver.cs | 2 +- .../PersistentStreamPullingAgent.cs | 47 +++++++++++++++++-- .../StreamConsumerCollection.cs | 21 ++++++++- .../AdoNetQueueAdapterReceiverTests.cs | 5 ++ .../PersistentStreamPullingAgentTests.cs | 45 ++++++++++++++++++ 5 files changed, 113 insertions(+), 7 deletions(-) diff --git a/src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs b/src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs index a88c34e57fa..3a6bcdca5af 100644 --- a/src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs +++ b/src/AdoNet/Orleans.Streaming.AdoNet/AdoNetQueueAdapterReceiver.cs @@ -43,7 +43,7 @@ internal partial class AdoNetQueueAdapterReceiver(string providerId, string queu private int _activeOperations; private TaskCompletionSource? _operationsCompleted; - internal AdoNetQueueAdapterReceiver( + public AdoNetQueueAdapterReceiver( string providerId, string queueId, AdoNetStreamOptions streamOptions, diff --git a/src/Orleans.Streaming/PersistentStreams/PersistentStreamPullingAgent.cs b/src/Orleans.Streaming/PersistentStreams/PersistentStreamPullingAgent.cs index 444d99b1efd..524d6cc5c5f 100644 --- a/src/Orleans.Streaming/PersistentStreams/PersistentStreamPullingAgent.cs +++ b/src/Orleans.Streaming/PersistentStreams/PersistentStreamPullingAgent.cs @@ -329,6 +329,11 @@ private async Task AddSubscriber_Impl( data.LastProcessedToken = startToken; data.PendingStartToken = null; data.IsRegistered = true; + if (streamDataCollection.AllConsumers().All(static consumer => consumer.IsRegistered)) + { + streamDataCollection.ReleaseRegistrationPin(); + } + StreamingEvents.EmitSubscriptionAttached(streamProviderName, streamId.StreamId, subscriptionId.Guid, streamConsumer, Silo); if (data.State == StreamConsumerDataState.Inactive) RunConsumerCursor(data).Ignore(); // Start delivering events if not actively doing so @@ -421,7 +426,10 @@ public void RemoveSubscriber_Impl(GuidId subscriptionId, QualifiedStreamId strea } if (streamData.Count == 0) + { + streamData.DisposeAll(logger); pubSubCache.Remove(streamId); + } } private Task RunQueuePump(QueueId queueId, CancellationToken cancellationToken) @@ -507,6 +515,11 @@ private async Task ReadFromQueue(QueueId myQueueId, IQueueAdapterReceiver? return false; } + if (pubSubCache.Values.Any(static stream => stream.RegistrationTask is { IsCompleted: false })) + { + return false; + } + var now = _timeProvider.GetUtcNow().UtcDateTime; TagList? tags = null; @@ -548,6 +561,14 @@ private async Task ReadFromQueue(QueueId myQueueId, IQueueAdapterReceiver? if (queueCache is not null && queueCache.IsUnderPressure()) { + foreach (var stream in pubSubCache.Values) + { + if (stream.Count == 0 && stream.RegistrationTask is null) + { + stream.ReleaseRegistrationPin(); + } + } + // Under back pressure. Exit the loop. Will attempt again in the next timer callback. LogInfoStreamCacheUnderPressure(); return false; @@ -589,6 +610,18 @@ private async Task ReadFromQueue(QueueId myQueueId, IQueueAdapterReceiver? if (pubSubCache.TryGetValue(streamId, out var streamData)) { streamData.RefreshActivity(now); + if (streamData.Count == 0 || streamData.AllConsumers().Any(static consumer => !consumer.IsRegistered)) + { + if (streamData.PendingStartToken is null || startToken.Older(streamData.PendingStartToken)) + { + streamData.PendingStartToken = startToken; + } + } + else + { + streamData.PendingStartToken = null; + } + StartInactiveCursors(streamData, startToken); } else @@ -690,13 +723,16 @@ private void RegisterStream(QualifiedStreamId streamId, StreamSequenceToken firs return; } - var streamData = new StreamConsumerCollection(now); + var streamData = new StreamConsumerCollection(now) + { + PendingStartToken = firstToken + }; // Create a fake cursor to point into a cache. // That way we will not purge the event from the cache, until we talk to pub sub. // This will help ensure the "casual consistency" between pre-existing subscripton (of a potentially new already subscribed consumer) // and later production. - var pinCursor = queueCache?.GetCacheCursor(streamId, firstToken); + streamData.RegistrationPin = queueCache?.GetCacheCursor(streamId, firstToken); streamData.RegistrationTask = RegisterStreamAsync(); pubSubCache.Add(streamId, streamData); @@ -747,9 +783,10 @@ async Task RegisterStreamAsync() { streamData.RegistrationTask = null; - // Disposed after all initial subscriber handshakes complete so the first - // batch stays pinned until each subscriber has its own cache cursor. - pinCursor?.Dispose(); + if (streamData.Count > 0 && streamData.AllConsumers().All(static consumer => consumer.IsRegistered)) + { + streamData.ReleaseRegistrationPin(); + } } } diff --git a/src/Orleans.Streaming/PersistentStreams/StreamConsumerCollection.cs b/src/Orleans.Streaming/PersistentStreams/StreamConsumerCollection.cs index fb197158024..a1eb4bcffa3 100644 --- a/src/Orleans.Streaming/PersistentStreams/StreamConsumerCollection.cs +++ b/src/Orleans.Streaming/PersistentStreams/StreamConsumerCollection.cs @@ -23,6 +23,14 @@ internal sealed class StreamConsumerCollection [NonSerialized] public Task? RegistrationTask; + // Retains the first cached token until subscriptions discovered after producer registration + // have attached their cursors. + [NonSerialized] + public StreamSequenceToken? PendingStartToken; + + [NonSerialized] + public IQueueCacheCursor? RegistrationPin; + public StreamConsumerCollection(DateTime now) { queueData = new Dictionary(); @@ -31,7 +39,10 @@ public StreamConsumerCollection(DateTime now) public StreamConsumerData AddConsumer(GuidId subscriptionId, QualifiedStreamId streamId, IStreamConsumerExtension streamConsumer, string? filterData, DateTime now) { - var consumerData = new StreamConsumerData(subscriptionId, streamId, streamConsumer, filterData); + var consumerData = new StreamConsumerData(subscriptionId, streamId, streamConsumer, filterData) + { + PendingStartToken = PendingStartToken + }; queueData.Add(subscriptionId, consumerData); lastActivityTime = now; return consumerData; @@ -62,6 +73,7 @@ public IEnumerable AllConsumers() public void DisposeAll(ILogger logger) { + ReleaseRegistrationPin(); foreach (StreamConsumerData consumer in queueData.Values) { consumer.SafeDisposeCursor(logger); @@ -69,6 +81,13 @@ public void DisposeAll(ILogger logger) queueData.Clear(); } + public void ReleaseRegistrationPin() + { + RegistrationPin?.Dispose(); + RegistrationPin = null; + PendingStartToken = null; + } + public int Count { diff --git a/test/Extensions/Orleans.AdoNet.Tests/Streaming/AdoNetQueueAdapterReceiverTests.cs b/test/Extensions/Orleans.AdoNet.Tests/Streaming/AdoNetQueueAdapterReceiverTests.cs index 3ede05a1dae..008abe3eb4f 100644 --- a/test/Extensions/Orleans.AdoNet.Tests/Streaming/AdoNetQueueAdapterReceiverTests.cs +++ b/test/Extensions/Orleans.AdoNet.Tests/Streaming/AdoNetQueueAdapterReceiverTests.cs @@ -3,6 +3,7 @@ using Orleans.Configuration; using Orleans.Streaming.AdoNet; using Orleans.Tests.SqlUtils; +using System.Runtime.CompilerServices; using TestExtensions; using UnitTests.General; using static System.String; @@ -20,6 +21,10 @@ namespace Tester.AdoNet.Streaming; [TestArea("Streaming")] public class AdoNetQueueAdapterReceiverLifecycleTests(TestEnvironmentFixture fixture) { + [Fact] + public void AdoNetQueueAdapterReceiver_CanBeCreatedByAdapterFactory() => + RuntimeHelpers.RunClassConstructor(typeof(AdoNetQueueAdapter).TypeHandle); + [Fact] public async Task AdoNetQueueAdapterReceiver_Shutdown_WaitsForDequeueBookkeepingBeforeRelease() { diff --git a/test/Orleans.Streaming.Tests/StreamingTests/PersistentStreamPullingAgentTests.cs b/test/Orleans.Streaming.Tests/StreamingTests/PersistentStreamPullingAgentTests.cs index 14a3f2d790d..ac778615f72 100644 --- a/test/Orleans.Streaming.Tests/StreamingTests/PersistentStreamPullingAgentTests.cs +++ b/test/Orleans.Streaming.Tests/StreamingTests/PersistentStreamPullingAgentTests.cs @@ -60,10 +60,15 @@ public async Task ReadFromQueue_DoesNotWaitForColdStreamRegistration() Assert.True(await readTask, "ReadFromQueue should return true indicating data was read"); + Assert.False(await testAccessor.ReadFromQueue(queueId, receiver, 1)); + await receiver.Received(1).GetQueueMessagesAsync(Arg.Any(), Arg.Any()); + // Completing registration should resolve the tracked task and clear it. registration.SetResult(new HashSet()); await registrationTask; + Assert.True(await testAccessor.ReadFromQueue(queueId, receiver, 1)); Assert.Null(streamData.RegistrationTask); + await receiver.Received(2).GetQueueMessagesAsync(Arg.Any(), Arg.Any()); } [TestSuite("BVT")] @@ -106,6 +111,46 @@ public async Task ReadFromQueue_ClearsRegistrationTaskWhenColdStreamRegistration Assert.True(streamData.StreamRegistered); } + [TestSuite("BVT")] + [TestProvider("None")] + [TestArea("Streaming")] + [Fact, TestCategory("BVT"), TestCategory("Streaming")] + public async Task ReadFromQueue_PreservesColdStreamStartTokenForLateSubscriber() + { + var pubSub = Substitute.For(); + pubSub.RegisterProducer(default, default) + .ReturnsForAnyArgs(Task.FromResult>(new HashSet())); + + var queueId = QueueId.GetQueueId("queue", 0u, 0u); + var streamId = StreamId.Create("namespace", Guid.NewGuid()); + var qualifiedStreamId = new QualifiedStreamId("provider", streamId); + var firstToken = new EventSequenceTokenV2(1); + var receiver = Substitute.For(); + receiver.GetQueueMessagesAsync(Arg.Any(), Arg.Any()) + .Returns(Task.FromResult>( + [ + new GeneratedBatchContainer(streamId, 1, firstToken), + ])); + + var agent = CreateAgent(pubSub, queueId); + var testAccessor = (PersistentStreamPullingAgent.ITestAccessor)agent; + await InitializeAgent(agent); + + Assert.True(await testAccessor.ReadFromQueue(queueId, receiver, 1)); + + var streamData = (await testAccessor.GetPubSubCache()).Single().Value; + Assert.Equal(firstToken, streamData.PendingStartToken); + + var consumerData = streamData.AddConsumer( + GuidId.GetGuidId(Guid.NewGuid()), + qualifiedStreamId, + new RecordingConsumer(), + filterData: null, + now: DateTime.UtcNow); + + Assert.Equal(firstToken, consumerData.PendingStartToken); + } + [TestSuite("BVT")] [TestProvider("None")] [TestArea("Streaming")] From 04056b78e0fd573a1c7b9539efcebbc153ce40fe Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Sun, 16 Aug 2026 02:12:37 -0700 Subject: [PATCH 5/6] docs(streaming): explain registration read gate Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../PersistentStreams/PersistentStreamPullingAgent.cs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/Orleans.Streaming/PersistentStreams/PersistentStreamPullingAgent.cs b/src/Orleans.Streaming/PersistentStreams/PersistentStreamPullingAgent.cs index 524d6cc5c5f..318586ae94c 100644 --- a/src/Orleans.Streaming/PersistentStreams/PersistentStreamPullingAgent.cs +++ b/src/Orleans.Streaming/PersistentStreams/PersistentStreamPullingAgent.cs @@ -515,6 +515,7 @@ private async Task ReadFromQueue(QueueId myQueueId, IQueueAdapterReceiver? return false; } + // Pause all queue reads so a cold stream's first batch stays pinned until registration completes. if (pubSubCache.Values.Any(static stream => stream.RegistrationTask is { IsCompleted: false })) { return false; From c8b9a67f2a6a9f92937249adedf59837ad04eb51 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Mon, 17 Aug 2026 02:37:29 -0700 Subject: [PATCH 6/6] fix(streaming): resume from first available batch Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../PersistentStreamPullingAgent.cs | 70 +++++++------------ .../StreamConsumerCollection.cs | 22 +----- .../PersistentStreamPullingAgentTests.cs | 40 ----------- 3 files changed, 28 insertions(+), 104 deletions(-) diff --git a/src/Orleans.Streaming/PersistentStreams/PersistentStreamPullingAgent.cs b/src/Orleans.Streaming/PersistentStreams/PersistentStreamPullingAgent.cs index 318586ae94c..f6da1b22d42 100644 --- a/src/Orleans.Streaming/PersistentStreams/PersistentStreamPullingAgent.cs +++ b/src/Orleans.Streaming/PersistentStreams/PersistentStreamPullingAgent.cs @@ -329,11 +329,6 @@ private async Task AddSubscriber_Impl( data.LastProcessedToken = startToken; data.PendingStartToken = null; data.IsRegistered = true; - if (streamDataCollection.AllConsumers().All(static consumer => consumer.IsRegistered)) - { - streamDataCollection.ReleaseRegistrationPin(); - } - StreamingEvents.EmitSubscriptionAttached(streamProviderName, streamId.StreamId, subscriptionId.Guid, streamConsumer, Silo); if (data.State == StreamConsumerDataState.Inactive) RunConsumerCursor(data).Ignore(); // Start delivering events if not actively doing so @@ -365,7 +360,16 @@ private async Task DoHandshakeWithConsumer( if (requestedToken != null) { consumerData.SafeDisposeCursor(logger); - consumerData.Cursor = queueCache.GetCacheCursor(consumerData.StreamId, requestedToken); + try + { + consumerData.Cursor = queueCache.GetCacheCursor(consumerData.StreamId, requestedToken); + } + catch (QueueCacheMissException) when (cacheToken is not null) + { + // A cold stream's triggering batch is the receiver's first available + // message, so resume there if the consumer's prior token was evicted. + consumerData.Cursor = queueCache.GetCacheCursor(consumerData.StreamId, cacheToken); + } } else { @@ -426,10 +430,7 @@ public void RemoveSubscriber_Impl(GuidId subscriptionId, QualifiedStreamId strea } if (streamData.Count == 0) - { - streamData.DisposeAll(logger); pubSubCache.Remove(streamId); - } } private Task RunQueuePump(QueueId queueId, CancellationToken cancellationToken) @@ -562,14 +563,6 @@ private async Task ReadFromQueue(QueueId myQueueId, IQueueAdapterReceiver? if (queueCache is not null && queueCache.IsUnderPressure()) { - foreach (var stream in pubSubCache.Values) - { - if (stream.Count == 0 && stream.RegistrationTask is null) - { - stream.ReleaseRegistrationPin(); - } - } - // Under back pressure. Exit the loop. Will attempt again in the next timer callback. LogInfoStreamCacheUnderPressure(); return false; @@ -611,18 +604,6 @@ private async Task ReadFromQueue(QueueId myQueueId, IQueueAdapterReceiver? if (pubSubCache.TryGetValue(streamId, out var streamData)) { streamData.RefreshActivity(now); - if (streamData.Count == 0 || streamData.AllConsumers().Any(static consumer => !consumer.IsRegistered)) - { - if (streamData.PendingStartToken is null || startToken.Older(streamData.PendingStartToken)) - { - streamData.PendingStartToken = startToken; - } - } - else - { - streamData.PendingStartToken = null; - } - StartInactiveCursors(streamData, startToken); } else @@ -724,16 +705,13 @@ private void RegisterStream(QualifiedStreamId streamId, StreamSequenceToken firs return; } - var streamData = new StreamConsumerCollection(now) - { - PendingStartToken = firstToken - }; + var streamData = new StreamConsumerCollection(now); // Create a fake cursor to point into a cache. // That way we will not purge the event from the cache, until we talk to pub sub. // This will help ensure the "casual consistency" between pre-existing subscripton (of a potentially new already subscribed consumer) // and later production. - streamData.RegistrationPin = queueCache?.GetCacheCursor(streamId, firstToken); + var pinCursor = queueCache?.GetCacheCursor(streamId, firstToken); streamData.RegistrationTask = RegisterStreamAsync(); pubSubCache.Add(streamId, streamData); @@ -783,11 +761,7 @@ async Task RegisterStreamAsync() finally { streamData.RegistrationTask = null; - - if (streamData.Count > 0 && streamData.AllConsumers().All(static consumer => consumer.IsRegistered)) - { - streamData.ReleaseRegistrationPin(); - } + pinCursor?.Dispose(); } } @@ -927,10 +901,20 @@ private async Task RunConsumerCursor(StreamConsumerData consumerData) { consumerData.LastProcessedToken = newToken.Token; consumerData.LastToken = newToken; - IQueueCacheCursor newCursor = queueCache!.GetCacheCursor(consumerData.StreamId, newToken.Token); // queueCache must be non-null here: consumerData.Cursor was only ever populated via queueCache.GetCacheCursor. - // The handshake token points to an already processed event, we need to advance the cursor to - // the next event. - newCursor.MoveNext(); + IQueueCacheCursor newCursor; + try + { + newCursor = queueCache!.GetCacheCursor(consumerData.StreamId, newToken.Token); // queueCache must be non-null here: consumerData.Cursor was only ever populated via queueCache.GetCacheCursor. + // The handshake token points to an already processed event, so advance past it. + newCursor.MoveNext(); + } + catch (QueueCacheMissException) + { + // The current batch is the receiver's first available message. + // Keep it pending when the consumer resumes from an evicted token. + newCursor = queueCache!.GetCacheCursor(consumerData.StreamId, batch.SequenceToken); + } + consumerData.SafeDisposeCursor(logger); consumerData.Cursor = newCursor; } diff --git a/src/Orleans.Streaming/PersistentStreams/StreamConsumerCollection.cs b/src/Orleans.Streaming/PersistentStreams/StreamConsumerCollection.cs index a1eb4bcffa3..e77c8f21a3a 100644 --- a/src/Orleans.Streaming/PersistentStreams/StreamConsumerCollection.cs +++ b/src/Orleans.Streaming/PersistentStreams/StreamConsumerCollection.cs @@ -23,14 +23,6 @@ internal sealed class StreamConsumerCollection [NonSerialized] public Task? RegistrationTask; - // Retains the first cached token until subscriptions discovered after producer registration - // have attached their cursors. - [NonSerialized] - public StreamSequenceToken? PendingStartToken; - - [NonSerialized] - public IQueueCacheCursor? RegistrationPin; - public StreamConsumerCollection(DateTime now) { queueData = new Dictionary(); @@ -39,10 +31,7 @@ public StreamConsumerCollection(DateTime now) public StreamConsumerData AddConsumer(GuidId subscriptionId, QualifiedStreamId streamId, IStreamConsumerExtension streamConsumer, string? filterData, DateTime now) { - var consumerData = new StreamConsumerData(subscriptionId, streamId, streamConsumer, filterData) - { - PendingStartToken = PendingStartToken - }; + var consumerData = new StreamConsumerData(subscriptionId, streamId, streamConsumer, filterData); queueData.Add(subscriptionId, consumerData); lastActivityTime = now; return consumerData; @@ -73,7 +62,6 @@ public IEnumerable AllConsumers() public void DisposeAll(ILogger logger) { - ReleaseRegistrationPin(); foreach (StreamConsumerData consumer in queueData.Values) { consumer.SafeDisposeCursor(logger); @@ -81,14 +69,6 @@ public void DisposeAll(ILogger logger) queueData.Clear(); } - public void ReleaseRegistrationPin() - { - RegistrationPin?.Dispose(); - RegistrationPin = null; - PendingStartToken = null; - } - - public int Count { get { return queueData.Count; } diff --git a/test/Orleans.Streaming.Tests/StreamingTests/PersistentStreamPullingAgentTests.cs b/test/Orleans.Streaming.Tests/StreamingTests/PersistentStreamPullingAgentTests.cs index ac778615f72..fcd03b9208b 100644 --- a/test/Orleans.Streaming.Tests/StreamingTests/PersistentStreamPullingAgentTests.cs +++ b/test/Orleans.Streaming.Tests/StreamingTests/PersistentStreamPullingAgentTests.cs @@ -111,46 +111,6 @@ public async Task ReadFromQueue_ClearsRegistrationTaskWhenColdStreamRegistration Assert.True(streamData.StreamRegistered); } - [TestSuite("BVT")] - [TestProvider("None")] - [TestArea("Streaming")] - [Fact, TestCategory("BVT"), TestCategory("Streaming")] - public async Task ReadFromQueue_PreservesColdStreamStartTokenForLateSubscriber() - { - var pubSub = Substitute.For(); - pubSub.RegisterProducer(default, default) - .ReturnsForAnyArgs(Task.FromResult>(new HashSet())); - - var queueId = QueueId.GetQueueId("queue", 0u, 0u); - var streamId = StreamId.Create("namespace", Guid.NewGuid()); - var qualifiedStreamId = new QualifiedStreamId("provider", streamId); - var firstToken = new EventSequenceTokenV2(1); - var receiver = Substitute.For(); - receiver.GetQueueMessagesAsync(Arg.Any(), Arg.Any()) - .Returns(Task.FromResult>( - [ - new GeneratedBatchContainer(streamId, 1, firstToken), - ])); - - var agent = CreateAgent(pubSub, queueId); - var testAccessor = (PersistentStreamPullingAgent.ITestAccessor)agent; - await InitializeAgent(agent); - - Assert.True(await testAccessor.ReadFromQueue(queueId, receiver, 1)); - - var streamData = (await testAccessor.GetPubSubCache()).Single().Value; - Assert.Equal(firstToken, streamData.PendingStartToken); - - var consumerData = streamData.AddConsumer( - GuidId.GetGuidId(Guid.NewGuid()), - qualifiedStreamId, - new RecordingConsumer(), - filterData: null, - now: DateTime.UtcNow); - - Assert.Equal(firstToken, consumerData.PendingStartToken); - } - [TestSuite("BVT")] [TestProvider("None")] [TestArea("Streaming")]