added

WriteAsync_WhenTransportExceptionOnWriteInWriterSessionThenReconnectSession_ReturnWriteResult

Author: KirillKurdyukov

src/Ydb.Sdk/src/IDriver.cs

@@ -32,7 +32,9 @@ public IBidirectionalStream<TRequest, TResponse> BidirectionalStreamCall<TReques
 public interface IBidirectionalStream<in TRequest, out TResponse> : IDisposable
 {
     public Task Write(TRequest request);
+
     public ValueTask<bool> MoveNextAsync();
+
     public TResponse Current { get; }
 }
 

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

@@ -25,7 +25,7 @@ protected async void ReconnectSession()
     {
         if (Interlocked.CompareExchange(ref _isActive, 0, 1) == 0)
         {
-            Logger.LogWarning("Skipping reconnect. A reconnect session has already been initiated");
+            Logger.LogDebug("Skipping reconnect. A reconnect session has already been initiated");
 
             return;
         }

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

@@ -23,10 +23,11 @@ internal class Writer<TValue> : IWriter<TValue>
     private readonly ILogger<Writer<TValue>> _logger;
     private readonly ISerializer<TValue> _serializer;
     private readonly ConcurrentQueue<MessageSending> _toSendBuffer = new();
+    private readonly ConcurrentQueue<MessageSending> _inFlightMessages = new();
     private readonly CancellationTokenSource _disposeTokenSource = new();
 
     private volatile TaskCompletionSource _taskWakeUpCompletionSource = new();
-    private volatile IWriteSession _session = new NotStartedWriterSession("Session not started!");
+    private volatile IWriteSession _session = null!;
 
     private int _limitBufferMaxSize;
 
@@ -60,9 +61,8 @@ public async Task<WriteResult> WriteAsync(Message<TValue> message)
 
         foreach (var metadata in message.Metadata)
         {
-            messageData.MetadataItems.Add(
-                new MetadataItem { Key = metadata.Key, Value = ByteString.CopyFrom(metadata.Value) }
-            );
+            messageData.MetadataItems.Add(new MetadataItem
+                { Key = metadata.Key, Value = ByteString.CopyFrom(metadata.Value) });
         }
 
         while (true)
@@ -77,7 +77,6 @@ public async Task<WriteResult> WriteAsync(Message<TValue> message)
                         curLimitBufferSize - data.Length, curLimitBufferSize) == curLimitBufferSize)
                 {
                     _toSendBuffer.Enqueue(new MessageSending(messageData, completeTask));
-
                     WakeUpWorker();
 
                     break;
@@ -130,6 +129,8 @@ private async Task Initialize()
         {
             _logger.LogInformation("Writer session initialization started. WriterConfig: {WriterConfig}", _config);
 
+            _session = new NotStartedWriterSession("Session not started!");
+
             var stream = _driver.BidirectionalStreamCall(
                 TopicService.StreamWriteMethod,
                 GrpcRequestSettings.DefaultInstance
@@ -194,7 +195,18 @@ private async Task Initialize()
                 return;
             }
 
-            _session = new WriterSession(_config, stream, initResponse, Initialize, _logger);
+            var newSession = new WriterSession(_config, stream, initResponse, Initialize, _logger, _inFlightMessages);
+            if (!_inFlightMessages.IsEmpty)
+            {
+                while (_inFlightMessages.TryDequeue(out var sendData))
+                {
+                    _toSendBuffer.Enqueue(sendData);
+                }
+
+                await newSession.Write(_toSendBuffer); // retry prev in flight messages
+            }
+
+            _session = newSession;
         }
         catch (Driver.TransportException e)
         {
@@ -252,7 +264,7 @@ public Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
     {
         while (toSendBuffer.TryDequeue(out var messageSending))
         {
-            messageSending.TaskCompletionSource.SetException(_reasonException);
+            messageSending.TaskCompletionSource.TrySetException(_reasonException);
         }
 
         return Task.CompletedTask;
@@ -267,7 +279,7 @@ public void Dispose()
 internal class WriterSession : TopicSession<MessageFromClient, MessageFromServer>, IWriteSession
 {
     private readonly WriterConfig _config;
-    private readonly ConcurrentQueue<MessageSending> _inFlightMessages = new();
+    private readonly ConcurrentQueue<MessageSending> _inFlightMessages;
 
     private long _seqNum;
 
@@ -276,9 +288,11 @@ public WriterSession(
         WriterStream stream,
         InitResponse initResponse,
         Func<Task> initialize,
-        ILogger logger) : base(stream, logger, initResponse.SessionId, initialize)
+        ILogger logger,
+        ConcurrentQueue<MessageSending> inFlightMessages) : base(stream, logger, initResponse.SessionId, initialize)
     {
         _config = config;
+        _inFlightMessages = inFlightMessages;
         Volatile.Write(ref _seqNum, initResponse.LastSeqNo); // happens-before for Volatile.Read
 
         RunProcessingWriteAck();
@@ -312,8 +326,6 @@ public async Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
             Logger.LogError(e, "WriterSession[{SessionId}] have error on Write, last SeqNo={SeqNo}",
                 SessionId, Volatile.Read(ref _seqNum));
 
-            ClearInFlightMessages(e);
-
             ReconnectSession();
         }
     }
@@ -366,12 +378,12 @@ Completing task on exception...
 Client SeqNo: {SeqNo}, WriteAck: {WriteAck}",
                             messageFromClient.MessageData.SeqNo, ack);
 
-                        messageFromClient.TaskCompletionSource.SetException(new WriterException(
+                        messageFromClient.TaskCompletionSource.TrySetException(new WriterException(
                             $"Client SeqNo[{messageFromClient.MessageData.SeqNo}] is less then server's WriteAck[{ack}]"));
                     }
                     else
                     {
-                        messageFromClient.TaskCompletionSource.SetResult(new WriteResult(ack));
+                        messageFromClient.TaskCompletionSource.TrySetResult(new WriteResult(ack));
                     }
 
                     _inFlightMessages.TryDequeue(out _); // Dequeue 
@@ -381,20 +393,10 @@ Completing task on exception...
         catch (Driver.TransportException e)
         {
             Logger.LogError(e, "WriterSession[{SessionId}] have error on processing writeAck", SessionId);
-
-            ClearInFlightMessages(e);
         }
         finally
         {
             ReconnectSession();
         }
     }
-
-    private void ClearInFlightMessages(Driver.TransportException e)
-    {
-        while (_inFlightMessages.TryDequeue(out var sendData))
-        {
-            sendData.TaskCompletionSource.SetException(e);
-        }
-    }
 }

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

@@ -1,4 +1,3 @@
-using System.Collections.Concurrent;
 using Grpc.Core;
 using Moq;
 using Xunit;
@@ -279,4 +278,73 @@ public async Task WriteAsyncStress_WhenBufferIsOverflow_ThrowWriterExceptionOnBu
             Assert.Equal(batchTasksSize - bufferSize / messageSize, countErrors);
         }
     }
+
+    [Fact]
+    public async Task WriteAsync_WhenTransportExceptionOnWriteInWriterSessionThenReconnectSession_ReturnWriteResult()
+    {
+        var moveFirstNextSource = new TaskCompletionSource<bool>();
+        var moveSecondNextSource = new TaskCompletionSource<bool>();
+        var moveThirdNextSource = new TaskCompletionSource<bool>();
+        var nextCompleted = new TaskCompletionSource();
+        _mockStream.SetupSequence(stream => stream.Write(It.IsAny<FromClient>()))
+            .Returns(Task.CompletedTask)
+            .Throws(() =>
+            {
+                moveFirstNextSource.SetResult(false);
+                return new Driver.TransportException(new RpcException(Grpc.Core.Status.DefaultCancelled));
+            })
+            .Returns(() =>
+            {
+                moveSecondNextSource.SetResult(true);
+                return Task.CompletedTask;
+            })
+            .Returns(() =>
+            {
+                nextCompleted.SetResult();
+                return Task.CompletedTask;
+            });
+        _mockStream.SetupSequence(stream => stream.MoveNextAsync())
+            .Returns(new ValueTask<bool>(true))
+            .Returns(new ValueTask<bool>(moveFirstNextSource.Task))
+            .Returns(new ValueTask<bool>(moveSecondNextSource.Task))
+            .Returns(new ValueTask<bool>(moveThirdNextSource.Task));
+        _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
+            })
+            .Returns(new StreamWriteMessage.Types.FromServer
+            {
+                WriteResponse = new StreamWriteMessage.Types.WriteResponse
+                {
+                    PartitionId = 1,
+                    Acks =
+                    {
+                        new StreamWriteMessage.Types.WriteResponse.Types.WriteAck
+                        {
+                            SeqNo = 1,
+                            Written = new StreamWriteMessage.Types.WriteResponse.Types.WriteAck.Types.Written
+                                { Offset = 0 }
+                        }
+                    }
+                },
+                Status = StatusIds.Types.StatusCode.Success
+            });
+        using var writer = new WriterBuilder<long>(_mockIDriver.Object, new WriterConfig("/topic")
+            { ProducerId = "producerId" }).Build();
+
+        var runTask = writer.WriteAsync(100L);
+        await nextCompleted.Task;
+        moveThirdNextSource.SetResult(true);
+
+        Assert.Equal(PersistenceStatus.Written, (await runTask).Status);
+    }
 }