IWriter added cancelToken

Author: KirillKurdyukov

src/Ydb.Sdk/src/Services/Topic/IWriter.cs

@@ -4,7 +4,7 @@ namespace Ydb.Sdk.Services.Topic;
 
 public interface IWriter<TValue> : IDisposable
 {
-    public Task<WriteResult> WriteAsync(TValue data);
+    public Task<WriteResult> WriteAsync(TValue data, CancellationToken cancellationToken = default);
 
-    public Task<WriteResult> WriteAsync(Message<TValue> message);
+    public Task<WriteResult> WriteAsync(Message<TValue> message, CancellationToken cancellationToken = default);
 }

src/Ydb.Sdk/src/Services/Topic/Writer/Writer.cs

@@ -42,14 +42,14 @@ internal Writer(IDriver driver, WriterConfig config, ISerializer<TValue> seriali
         StartWriteWorker();
     }
 
-    public Task<WriteResult> WriteAsync(TValue data)
+    public Task<WriteResult> WriteAsync(TValue data, CancellationToken cancellationToken)
     {
-        return WriteAsync(new Message<TValue>(data));
+        return WriteAsync(new Message<TValue>(data), cancellationToken);
     }
 
-    public async Task<WriteResult> WriteAsync(Message<TValue> message)
+    public async Task<WriteResult> WriteAsync(Message<TValue> message, CancellationToken cancellationToken)
     {
-        TaskCompletionSource<WriteResult> completeTask = new();
+        TaskCompletionSource<WriteResult> completeTask = new(cancellationToken);
 
         var data = _serializer.Serialize(message.Data);
         var messageData = new MessageData
@@ -207,7 +207,7 @@ private async Task Initialize()
                 _logger,
                 _inFlightMessages
             );
-            
+
             if (!_inFlightMessages.IsEmpty)
             {
                 var copyInFlightMessages = new ConcurrentQueue<MessageSending>();

src/Ydb.Sdk/tests/Topic/WriterMockTests.cs

@@ -318,7 +318,7 @@ public async Task WriteAsync_WhenTransportExceptionOnWriteInWriterSessionThenRec
                 return Task.CompletedTask;
             });
         _mockStream.SetupSequence(stream => stream.MoveNextAsync())
-            .Returns(new ValueTask<bool>(true))
+            .ReturnsAsync(true)
             .Returns(new ValueTask<bool>(moveFirstNextSource.Task))
             .Returns(new ValueTask<bool>(moveSecondNextSource.Task))
             .Returns(new ValueTask<bool>(moveThirdNextSource.Task))
@@ -392,7 +392,7 @@ public async Task WriteAsync_WhenTransportExceptionOnProcessingWriteAckThenRecon
         _mockStream.Setup(stream => stream.Write(It.IsAny<FromClient>()))
             .Returns(Task.CompletedTask);
         _mockStream.SetupSequence(stream => stream.MoveNextAsync())
-            .Returns(new ValueTask<bool>(true))
+            .ReturnsAsync(true)
             .ThrowsAsync(new Driver.TransportException(new RpcException(Grpc.Core.Status.DefaultCancelled)))
             .Returns(() =>
             {
@@ -425,4 +425,56 @@ public async Task WriteAsync_WhenTransportExceptionOnProcessingWriteAckThenRecon
         _mockStream.Verify(stream => stream.MoveNextAsync(), Times.Exactly(3));
         _mockStream.Verify(stream => stream.Current, Times.Once);
     }
+
+    /*
+     * Performed invocations:
+
+       Mock<IBidirectionalStream<StreamWriteMessage.Types.FromClient, StreamWriteMessage.Types.FromServer>:1> (stream):
+
+          IBidirectionalStream<StreamWriteMessage.Types.FromClient, StreamWriteMessage.Types.FromServer>.Write({ "initRequest": { "path": "/topic", "producerId": "producerId" } })
+          IBidirectionalStream<StreamWriteMessage.Types.FromClient, StreamWriteMessage.Types.FromServer>.MoveNextAsync()
+          IBidirectionalStream<StreamWriteMessage.Types.FromClient, StreamWriteMessage.Types.FromServer>.Current
+          IBidirectionalStream<StreamWriteMessage.Types.FromClient, StreamWriteMessage.Types.FromServer>.MoveNextAsync()
+          IBidirectionalStream<StreamWriteMessage.Types.FromClient, StreamWriteMessage.Types.FromServer>.Write({ "initRequest": { "path": "/topic", "producerId": "producerId" } })
+          IBidirectionalStream<StreamWriteMessage.Types.FromClient, StreamWriteMessage.Types.FromServer>.MoveNextAsync()
+     */
+    [Fact]
+    public async Task WriteAsync_WhenStreamIsClosingOnProcessingWriteAckThenReconnectSession_ReturnWriterException()
+    {
+        var nextCompleted = new TaskCompletionSource();
+        _mockStream.Setup(stream => stream.Write(It.IsAny<FromClient>()))
+            .Returns(Task.CompletedTask);
+        _mockStream.SetupSequence(stream => stream.MoveNextAsync())
+            .ReturnsAsync(true)
+            .ReturnsAsync(false)
+            .Returns(() =>
+            {
+                nextCompleted.SetResult();
+                return new ValueTask<bool>(new TaskCompletionSource<bool>().Task);
+            }); // retry init writer session
+        _mockStream.SetupSequence(stream => stream.Current)
+            .Returns(new StreamWriteMessage.Types.FromServer
+            {
+                InitResponse = new StreamWriteMessage.Types.InitResponse
+                    { LastSeqNo = 0, PartitionId = 1, SessionId = "SessionId" },
+                Status = StatusIds.Types.StatusCode.Success
+            })
+            .Returns(new StreamWriteMessage.Types.FromServer
+            {
+                InitResponse = new StreamWriteMessage.Types.InitResponse
+                    { LastSeqNo = 0, PartitionId = 1, SessionId = "SessionId" },
+                Status = StatusIds.Types.StatusCode.Success
+            });
+        using var writer = new WriterBuilder<long>(_mockIDriver.Object, new WriterConfig("/topic")
+            { ProducerId = "producerId" }).Build();
+
+        await nextCompleted.Task;
+        var writerExceptionAfterResetSession = await Assert.ThrowsAsync<WriterException>(() => writer.WriteAsync(100));
+        Assert.Equal("WriterStream is closed", writerExceptionAfterResetSession.Message);
+        Assert.Equal(StatusCode.Unspecified, writerExceptionAfterResetSession.Status.StatusCode);
+
+        _mockStream.Verify(stream => stream.Write(It.IsAny<FromClient>()), Times.Exactly(2));
+        _mockStream.Verify(stream => stream.MoveNextAsync(), Times.Exactly(3));
+        _mockStream.Verify(stream => stream.Current, Times.Once);
+    }
 }