added test on DisposeAsync waiting in flight messages

Author: KirillKurdyukov

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

@@ -60,8 +60,10 @@ protected async Task SendMessage(TFromClient fromClient)
 
     protected abstract TFromClient GetSendUpdateTokenRequest(string token);
 
-    public async ValueTask DisposeAsync()
+    public ValueTask DisposeAsync()
     {
-        await Stream.DisposeAsync();
+        Logger.LogInformation("TopicSession[{SessionId}] is being deleted", SessionId);
+
+        return Stream.DisposeAsync();
     }
 }

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

@@ -268,7 +268,7 @@ private async Task Initialize()
                 return;
             }
 
-            await _sendInFlightMessagesSemaphoreSlim.WaitAsync(_disposeCts.Token);
+            await _sendInFlightMessagesSemaphoreSlim.WaitAsync();
             try
             {
                 var copyInFlightMessages = new ConcurrentQueue<MessageSending>();
@@ -333,17 +333,22 @@ private async Task Initialize()
 
     public async ValueTask DisposeAsync()
     {
+        _logger.LogInformation("Starting Writer[{WriterConfig}] disposal process", _config);
+
         await _sendInFlightMessagesSemaphoreSlim.WaitAsync();
         try
         {
+            _logger.LogDebug("Signaling cancellation token to stop writing new messages");
+
             _disposeCts.Cancel();
         }
         finally
         {
             _sendInFlightMessagesSemaphoreSlim.Release();
         }
 
-        // wait all messages
+        _logger.LogDebug("Writer[{WriterConfig}] is waiting for all in-flight messages to complete...", _config);
+
         foreach (var inFlightMessage in _inFlightMessages)
         {
             try
@@ -352,11 +357,14 @@ public async ValueTask DisposeAsync()
             }
             catch (Exception e)
             {
-                _logger.LogWarning(e, "Failed in flight message");
+                _logger.LogError(e, "Error occurred while waiting for in-flight message SeqNo: {SeqNo}",
+                    inFlightMessage.MessageData.SeqNo);
             }
         }
 
         await _session.DisposeAsync();
+
+        _logger.LogInformation("Writer[{WriterConfig}] is disposed", _config);
     }
 }
 

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

@@ -37,6 +37,11 @@ public override string ToString()
             toString.Append(", ProducerId: ").Append(ProducerId);
         }
 
+        if (PartitionId != null)
+        {
+            toString.Append(", PartitionId: ").Append(PartitionId);
+        }
+
         return toString.Append(", Codec: ").Append(Codec).ToString();
     }
 }

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

@@ -841,6 +841,82 @@ public async Task WriteAsync_WhenTokenIsUpdatedOneTime_SuccessUpdateToken()
             msg.UpdateTokenRequest.Token == "Token2")));
     }
 
+    [Fact]
+    public async Task DisposeAsync_WhenInFlightMessages_WaitingInFlightMessages()
+    {
+        var tcsDetectedWrite = new TaskCompletionSource();
+        var writeTcs1 = new TaskCompletionSource<bool>();
+
+        _mockStream.SetupSequence(stream => stream.Write(It.IsAny<FromClient>()))
+            .Returns(Task.CompletedTask)
+            .Returns(() =>
+            {
+                tcsDetectedWrite.TrySetResult();
+                return Task.CompletedTask;
+            })
+            .Returns(Task.CompletedTask)
+            .Returns(Task.CompletedTask);
+        _mockStream.SetupSequence(stream => stream.MoveNextAsync())
+            .ReturnsAsync(true)
+            .Returns(new ValueTask<bool>(writeTcs1.Task))
+            .ReturnsAsync(true)
+            .ReturnsAsync(true)
+            .Returns(new ValueTask<bool>(new TaskCompletionSource<bool>().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
+            });
+
+        var writer = new WriterBuilder<long>(_mockIDriver.Object, "/topic-16")
+            { ProducerId = "producerId" }.Build();
+
+        var writeTask1 = writer.WriteAsync(100L);
+        
+        await tcsDetectedWrite.Task;
+        var disposedTask = writer.DisposeAsync();
+
+        Assert.False(writeTask1.IsCompleted);
+        Assert.False(disposedTask.IsCompleted);
+        writeTcs1.TrySetException(new Driver.TransportException(
+            new RpcException(new Grpc.Core.Status(Grpc.Core.StatusCode.DeadlineExceeded, "Some message"))));
+        Assert.Equal("Writer[TopicPath: /topic-16, ProducerId: producerId, Codec: Raw] is disposed",
+            (await Assert.ThrowsAsync<WriterException>(() => writer.WriteAsync(12))).Message);
+        
+        Assert.Equal(PersistenceStatus.Written, (await writeTask1).Status);
+
+        Assert.Equal("Writer[TopicPath: /topic-16, ProducerId: producerId, Codec: Raw] is disposed",
+            (await Assert.ThrowsAsync<WriterException>(() => writer.WriteAsync(12))).Message);
+
+        await disposedTask;
+    }
+
     private ISetupSequentialResult<StreamWriteMessage.Types.FromServer> SetupReadOneWriteAckMessage()
     {
         return _mockStream.SetupSequence(stream => stream.Current)