Bump Microsoft.Orleans.Streaming from 10.2.0 to 10.3.1

Directory.Packages.props

@@ -12,13 +12,13 @@
     <PackageVersion Include="MessagePack" Version="3.1.7" />
     <PackageVersion Include="Microsoft.Extensions.Configuration.Abstractions" Version="10.0.12" />
     <PackageVersion Include="MinVer" Version="7.0.0" />
-    <PackageVersion Include="Microsoft.Orleans.Client" Version="10.2.0" />
+    <PackageVersion Include="Microsoft.Orleans.Client" Version="10.3.1" />
     <PackageVersion Include="Microsoft.NET.Test.Sdk" Version="18.10.1" />
-    <PackageVersion Include="Microsoft.Orleans.Sdk" Version="10.2.0" />
-    <PackageVersion Include="Microsoft.Orleans.Server" Version="10.2.0" />
-    <PackageVersion Include="Microsoft.Orleans.Streaming" Version="10.2.0" />
+    <PackageVersion Include="Microsoft.Orleans.Sdk" Version="10.3.1" />
+    <PackageVersion Include="Microsoft.Orleans.Server" Version="10.3.1" />
+    <PackageVersion Include="Microsoft.Orleans.Streaming" Version="10.3.1" />
     <PackageVersion Include="MSTest.TestAdapter" Version="4.3.3" />
     <PackageVersion Include="MSTest.TestFramework" Version="4.3.3" />
-    <PackageVersion Include="Testcontainers.Kafka" Version="4.12.0" />
+    <PackageVersion Include="Testcontainers.Kafka" Version="4.15.0" />
   </ItemGroup>
 </Project>

src/Orleans.Streams.Confluent/KafkaBatchContainer.cs

@@ -91,7 +91,8 @@ public static KafkaBatchContainer FromPayload(Serializer<KafkaBatchContainer> se
         ArgumentNullException.ThrowIfNull(serializer);
         ArgumentNullException.ThrowIfNull(payload);
 
-        return serializer.Deserialize(payload);
+        return serializer.Deserialize(payload)
+            ?? throw new InvalidOperationException("The Kafka batch payload deserialized to null.");
     }
 
     /// <summary>

src/Orleans.Streams.Confluent/KafkaQueueAdapter.cs

@@ -47,7 +47,7 @@ public IQueueAdapterReceiver CreateReceiver(QueueId queueId)
         return _receivers[queueId].Value;
     }
 
-    public async Task QueueMessageBatchAsync<T>(StreamId streamId, IEnumerable<T> events, StreamSequenceToken token, Dictionary<string, object> requestContext)
+    public async Task QueueMessageBatchAsync<T>(StreamId streamId, IEnumerable<T> events, StreamSequenceToken? token, Dictionary<string, object>? requestContext)
     {
         ArgumentNullException.ThrowIfNull(events);
         var effectiveRequestContext = requestContext ?? EmptyRequestContext;

src/Orleans.Streams.Confluent/KafkaQueueAdapterReceiver.cs

@@ -66,46 +66,77 @@ public Task Initialize(TimeSpan timeout)
     }
 
     public Task<IList<IBatchContainer>> GetQueueMessagesAsync(int maxCount)
+        => GetQueueMessagesAsync(maxCount, CancellationToken.None);
+
+    public Task<IList<IBatchContainer>> GetQueueMessagesAsync(int maxCount, CancellationToken cancellationToken)
     {
+        IConsumer<Ignore, byte[]>? consumer = null;
+        var consumedOffsets = new List<TopicPartitionOffset>();
+
         try
         {
+            cancellationToken.ThrowIfCancellationRequested();
+
             lock (_consumerSync)
             {
-                var consumer = _consumer;
+                consumer = _consumer;
                 if (consumer is null)
                 {
                     LogDebugReceiverNotInitialized(queueId);
                     return Task.FromResult<IList<IBatchContainer>>([]);
                 }
 
-                var batches = new List<IBatchContainer>(maxCount);
-                while (batches.Count < maxCount)
+                try
                 {
-                    var result = consumer.Consume(TimeSpan.FromMilliseconds(50));
-                    if (result is null)
+                    var batches = new List<IBatchContainer>(maxCount);
+                    while (batches.Count < maxCount)
                     {
-                        break;
+                        cancellationToken.ThrowIfCancellationRequested();
+                        var result = consumer.Consume(TimeSpan.FromMilliseconds(50));
+                        if (result is null)
+                        {
+                            cancellationToken.ThrowIfCancellationRequested();
+                            break;
+                        }
+
+                        if (!result.IsPartitionEOF)
+                        {
+                            consumedOffsets.Add(new TopicPartitionOffset(
+                                new TopicPartition(result.Topic, result.Partition),
+                                result.Offset));
+                        }
+
+                        cancellationToken.ThrowIfCancellationRequested();
+
+                        if (result.IsPartitionEOF || result.Message?.Value is null)
+                        {
+                            continue;
+                        }
+
+                        var container = KafkaBatchContainer
+                            .FromPayload(serializer, result.Message.Value)
+                            .WithKafkaMetadata(result.Topic, result.Partition.Value, result.Offset.Value);
+                        batches.Add(container);
                     }
 
-                    if (result.IsPartitionEOF || result.Message?.Value is null)
+                    if (batches.Count > 0)
                     {
-                        continue;
+                        LogDebugMessagesReceived(queueId, batches.Count);
                     }
 
-                    var container = KafkaBatchContainer
-                        .FromPayload(serializer, result.Message.Value)
-                        .WithKafkaMetadata(result.Topic, result.Partition.Value, result.Offset.Value);
-                    batches.Add(container);
+                    return Task.FromResult<IList<IBatchContainer>>(batches);
                 }
-
-                if (batches.Count > 0)
+                catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
                 {
-                    LogDebugMessagesReceived(queueId, batches.Count);
+                    RestoreOffsetsOrFail(consumer, consumedOffsets);
+                    throw;
                 }
-
-                return Task.FromResult<IList<IBatchContainer>>(batches);
             }
         }
+        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+        {
+            throw;
+        }
         catch (Exception ex)
         {
             LogErrorReceivingMessagesFailed(queueId, maxCount, ex);
@@ -114,49 +145,130 @@ public Task<IList<IBatchContainer>> GetQueueMessagesAsync(int maxCount)
     }
 
     public Task MessagesDeliveredAsync(IList<IBatchContainer> messages)
+        => MessagesDeliveredAsync(messages, CancellationToken.None);
+
+    public Task MessagesDeliveredAsync(IList<IBatchContainer> messages, CancellationToken cancellationToken)
     {
         if (messages.Count == 0)
         {
+            cancellationToken.ThrowIfCancellationRequested();
             return Task.CompletedTask;
         }
 
+        IConsumer<Ignore, byte[]>? consumer = null;
+        var deliveredOffsets = messages
+            .OfType<KafkaBatchContainer>()
+            .Select(batch => new TopicPartitionOffset(
+                new TopicPartition(batch.Topic, new Partition(batch.Partition)),
+                new Offset(batch.Offset)))
+            .ToList();
+
         try
         {
             lock (_consumerSync)
             {
-                var consumer = _consumer;
-                if (consumer is null)
+                consumer = _consumer;
+                try
                 {
-                    return Task.CompletedTask;
-                }
+                    cancellationToken.ThrowIfCancellationRequested();
+
+                    if (consumer is null)
+                    {
+                        return Task.CompletedTask;
+                    }
 
-                var commitOffsets = messages
-                    .OfType<KafkaBatchContainer>()
-                    .GroupBy(batch => new TopicPartition(batch.Topic, new Partition(batch.Partition)))
-                    .Select(group =>
+                    var commitOffsets = messages
+                        .OfType<KafkaBatchContainer>()
+                        .GroupBy(batch => new TopicPartition(batch.Topic, new Partition(batch.Partition)))
+                        .Select(group =>
+                        {
+                            var nextOffset = group.Max(batch => batch.Offset) + 1;
+                            return new TopicPartitionOffset(group.Key, new Offset(nextOffset));
+                        })
+                        .ToList();
+
+                    if (commitOffsets.Count == 0)
                     {
-                        var nextOffset = group.Max(batch => batch.Offset) + 1;
-                        return new TopicPartitionOffset(group.Key, new Offset(nextOffset));
-                    })
-                    .ToList();
+                        return Task.CompletedTask;
+                    }
 
-                if (commitOffsets.Count == 0)
-                {
+                    cancellationToken.ThrowIfCancellationRequested();
+                    consumer.Commit(commitOffsets);
+                    LogDebugMessagesCommitted(queueId, messages.Count);
                     return Task.CompletedTask;
                 }
-
-                consumer.Commit(commitOffsets);
-                LogDebugMessagesCommitted(queueId, messages.Count);
-                return Task.CompletedTask;
+                catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+                {
+                    RestoreOffsetsOrFail(consumer, deliveredOffsets);
+                    throw;
+                }
             }
         }
+        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+        {
+            throw;
+        }
         catch (Exception ex)
         {
             LogErrorCommitFailed(queueId, messages.Count, ex);
             throw;
         }
     }
 
+    private void RestoreOffsetsOrFail(IConsumer<Ignore, byte[]>? consumer, IEnumerable<TopicPartitionOffset> offsets)
+    {
+        if (consumer is null)
+        {
+            return;
+        }
+
+        Exception? restorationFailure = null;
+
+        if (!ReferenceEquals(_consumer, consumer))
+        {
+            return;
+        }
+
+        try
+        {
+            foreach (var partitionOffsets in offsets.GroupBy(offset => offset.TopicPartition))
+            {
+                var firstOffset = partitionOffsets.Min(offset => offset.Offset.Value);
+                consumer.Seek(new TopicPartitionOffset(partitionOffsets.Key, new Offset(firstOffset)));
+            }
+        }
+        catch (Exception exception)
+        {
+            _consumer = null;
+            restorationFailure = exception;
+        }
+
+        if (restorationFailure is null)
+        {
+            return;
+        }
+
+        try
+        {
+            consumer.Close();
+        }
+        catch (Exception closeException)
+        {
+            _logger.LogError(closeException, "Failed to close Kafka receiver after offset restoration failed for queue {QueueId}", queueId);
+        }
+
+        try
+        {
+            consumer.Dispose();
+        }
+        catch (Exception disposeException)
+        {
+            _logger.LogError(disposeException, "Failed to dispose Kafka receiver after offset restoration failed for queue {QueueId}", queueId);
+        }
+
+        throw new InvalidOperationException("Kafka receiver offset restoration failed after cancellation.", restorationFailure);
+    }
+
     public Task Shutdown(TimeSpan timeout)
     {
         IConsumer<Ignore, byte[]>? consumer;

test/Orleans.Streams.Confluent.IntegrationTests/KafkaStreamProviderIntegrationTests.cs

@@ -148,6 +148,22 @@ await adapter.QueueMessageBatchAsync(
             new EventSequenceTokenV2(0),
             new Dictionary<string, object>());
 
+        using (var assertionConsumer = new ConsumerBuilder<Ignore, byte[]>(new ConsumerConfig
+        {
+            BootstrapServers = bootstrapServers,
+            GroupId = $"assert-{Guid.NewGuid():N}",
+            AutoOffsetReset = AutoOffsetReset.Earliest,
+            AllowAutoCreateTopics = false
+        }).Build())
+        {
+            assertionConsumer.Assign(new TopicPartition(topicName, new Partition((int)queueId.GetNumericId())));
+            await WaitForKafkaMessageAsync(assertionConsumer);
+        }
+
+        using var readCancellationTokenSource = new CancellationTokenSource(TimeSpan.FromMilliseconds(10));
+        Func<Task> readCanceled = () => receiver.GetQueueMessagesAsync(2, readCancellationTokenSource.Token);
+        await readCanceled.Should().ThrowAsync<OperationCanceledException>();
+
         var messages = await WaitForReceiverMessagesAsync(receiver);
         messages.Should().ContainSingle();
 
@@ -156,7 +172,16 @@ await adapter.QueueMessageBatchAsync(
         batch.Partition.Should().Be((int)queueId.GetNumericId());
         batch.GetEvents<string>().Select(tuple => tuple.Item1).Should().ContainSingle().Which.Should().Be("created");
 
-        await receiver.MessagesDeliveredAsync(messages);
+        using var cancellationTokenSource = new CancellationTokenSource();
+        cancellationTokenSource.Cancel();
+
+        Func<Task> deliverCanceled = () => receiver.MessagesDeliveredAsync(messages, cancellationTokenSource.Token);
+        await deliverCanceled.Should().ThrowAsync<OperationCanceledException>();
+
+        var replayedMessages = await WaitForReceiverMessagesAsync(receiver);
+        replayedMessages.Should().ContainSingle();
+
+        await receiver.MessagesDeliveredAsync(replayedMessages, CancellationToken.None);
         await receiver.Shutdown(TimeSpan.FromSeconds(5));
     }
 
@@ -200,7 +225,7 @@ public async Task KafkaQueueAdapterReceiver_WhenShutdown_GetQueueMessagesAsync_R
         await receiver.Initialize(TimeSpan.FromSeconds(5));
         await receiver.Shutdown(TimeSpan.FromSeconds(5));
 
-        var messages = await receiver.GetQueueMessagesAsync(10);
+        var messages = await receiver.GetQueueMessagesAsync(10, CancellationToken.None);
         messages.Should().BeEmpty();
     }
 
@@ -282,7 +307,7 @@ private static async Task<IList<IBatchContainer>> WaitForReceiverMessagesAsync(I
         var deadline = DateTime.UtcNow.AddSeconds(15);
         while (DateTime.UtcNow < deadline)
         {
-            var messages = await receiver.GetQueueMessagesAsync(10);
+            var messages = await receiver.GetQueueMessagesAsync(10, CancellationToken.None);
             if (messages.Count > 0)
             {
                 return messages;

test/Orleans.Streams.Confluent.Tests/KafkaQueueAdapterReceiverTests.cs

@@ -0,0 +1,47 @@
+using FluentAssertions;
+using Microsoft.Extensions.Logging.Abstractions;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Orleans.Streams;
+
+namespace Orleans.Streams.Confluent.Tests;
+
+[TestClass]
+public sealed class KafkaQueueAdapterReceiverTests
+{
+    [TestMethod]
+    public async Task GetQueueMessagesAsync_WhenCancellationRequested_ThrowsOperationCanceledException()
+    {
+        using var cancellationTokenSource = new CancellationTokenSource();
+        cancellationTokenSource.Cancel();
+        var receiver = CreateReceiver();
+
+        Func<Task> act = async () => await receiver.GetQueueMessagesAsync(10, cancellationTokenSource.Token);
+
+        await act.Should().ThrowAsync<OperationCanceledException>();
+    }
+
+    [TestMethod]
+    public async Task MessagesDeliveredAsync_WhenCancellationRequested_ThrowsOperationCanceledException()
+    {
+        using var cancellationTokenSource = new CancellationTokenSource();
+        cancellationTokenSource.Cancel();
+        var receiver = CreateReceiver();
+        var messages = new IBatchContainer[]
+        {
+            new KafkaBatchContainer(StreamId.Create("orders", "order-123"), ["created"], [], "orders-topic", 0, 0)
+        };
+
+        Func<Task> act = async () => await receiver.MessagesDeliveredAsync(messages, cancellationTokenSource.Token);
+
+        await act.Should().ThrowAsync<OperationCanceledException>();
+    }
+
+    private static KafkaQueueAdapterReceiver CreateReceiver()
+        => new(
+            "provider",
+            null!,
+            null!,
+            NullLogger<KafkaQueueAdapterReceiver>.Instance,
+            default,
+            "consumer-group");
+}