feat: impl disposing Writer.cs

Author: KirillKurdyukov

slo/src/TopicService/SloTopicContext.cs

@@ -73,7 +73,7 @@ public async Task Run(RunConfig config)
                 {
                     try
                     {
-                        using var writer = new WriterBuilder<string>(driver, PathTopic)
+                        await using var writer = new WriterBuilder<string>(driver, PathTopic)
                         {
                             BufferMaxSize = 8 * 1024 * 1024,
                             ProducerId = "producer-" + partitionId,

src/Ydb.Sdk/src/IDriver.cs

@@ -30,7 +30,7 @@ public IBidirectionalStream<TRequest, TResponse> BidirectionalStreamCall<TReques
     ILoggerFactory LoggerFactory { get; }
 }
 
-public interface IBidirectionalStream<in TRequest, out TResponse> : IDisposable
+public interface IBidirectionalStream<in TRequest, out TResponse> : IAsyncDisposable
 {
     public Task Write(TRequest request);
 
@@ -225,6 +225,7 @@ internal class BidirectionalStream<TRequest, TResponse> : IBidirectionalStream<T
     private readonly AsyncDuplexStreamingCall<TRequest, TResponse> _stream;
     private readonly Action<RpcException> _rpcErrorAction;
     private readonly ICredentialsProvider _credentialsProvider;
+    private readonly TaskCompletionSource _closedResponseStreamTcs = new();
 
     internal BidirectionalStream(
         AsyncDuplexStreamingCall<TRequest, TResponse> stream,
@@ -258,6 +259,7 @@ public async ValueTask<bool> MoveNextAsync()
         }
         catch (RpcException e)
         {
+            _closedResponseStreamTcs.SetResult();
             _rpcErrorAction(e);
 
             throw new Driver.TransportException(e);
@@ -268,8 +270,17 @@ public async ValueTask<bool> MoveNextAsync()
 
     public string? AuthToken => _credentialsProvider.GetAuthInfo();
 
-    public void Dispose()
+    public async ValueTask DisposeAsync()
     {
+        try
+        {
+            await _stream.RequestStream.CompleteAsync();
+        }
+        catch (RpcException e)
+        {
+            _rpcErrorAction(e);
+        }
+
         _stream.Dispose();
     }
 }

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

@@ -2,7 +2,7 @@
 
 namespace Ydb.Sdk.Services.Topic;
 
-public interface IWriter<TValue> : IDisposable
+public interface IWriter<TValue> : IAsyncDisposable
 {
     /// <summary>
     /// Asynchronously send a data to a YDB Topic.

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

@@ -2,7 +2,7 @@
 
 namespace Ydb.Sdk.Services.Topic;
 
-internal abstract class TopicSession<TFromClient, TFromServer> : IDisposable
+internal abstract class TopicSession<TFromClient, TFromServer> : IAsyncDisposable
 {
     private readonly Func<Task> _initialize;
 
@@ -58,10 +58,10 @@ protected async Task SendMessage(TFromClient fromClient)
         await Stream.Write(fromClient);
     }
 
-    public void Dispose()
+    protected abstract TFromClient GetSendUpdateTokenRequest(string token);
+
+    public async ValueTask DisposeAsync()
     {
-        Stream.Dispose();
+        await Stream.DisposeAsync();
     }
-
-    protected abstract TFromClient GetSendUpdateTokenRequest(string token);
 }

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

@@ -21,11 +21,11 @@ internal class Writer<TValue> : IWriter<TValue>
     private readonly WriterConfig _config;
     private readonly ILogger<Writer<TValue>> _logger;
     private readonly ISerializer<TValue> _serializer;
-    private readonly GrpcRequestSettings _writerGrpcRequestSettings;
+    private readonly GrpcRequestSettings _writerGrpcRequestSettings = new();
     private readonly ConcurrentQueue<MessageSending> _toSendBuffer = new();
     private readonly ConcurrentQueue<MessageSending> _inFlightMessages = new();
     private readonly CancellationTokenSource _disposeCts = new();
-    private readonly SemaphoreSlim _clearInFlightMessagesSemaphoreSlim = new(1);
+    private readonly SemaphoreSlim _sendInFlightMessagesSemaphoreSlim = new(1);
 
     private volatile TaskCompletionSource _tcsWakeUp = new();
     private volatile TaskCompletionSource _tcsBufferAvailableEvent = new();
@@ -39,7 +39,6 @@ internal Writer(IDriver driver, WriterConfig config, ISerializer<TValue> seriali
         _logger = driver.LoggerFactory.CreateLogger<Writer<TValue>>();
         _serializer = serializer;
         _limitBufferMaxSize = config.BufferMaxSize;
-        _writerGrpcRequestSettings = new GrpcRequestSettings { CancellationToken = _disposeCts.Token };
 
         StartWriteWorker();
     }
@@ -95,7 +94,9 @@ public async Task<WriteResult> WriteAsync(Message<TValue> message, CancellationT
                 if (Interlocked.CompareExchange(ref _limitBufferMaxSize,
                         curLimitBufferSize - data.Length, curLimitBufferSize) == curLimitBufferSize)
                 {
-                    _toSendBuffer.Enqueue(new MessageSending(messageData, tcs));
+                    _toSendBuffer.Enqueue(
+                        new MessageSending(messageData, tcs, writerDisposedCancellationTokenRegistration)
+                    );
                     WakeUpWorker();
 
                     break;
@@ -162,7 +163,7 @@ private async void StartWriteWorker()
                     continue;
                 }
 
-                await _clearInFlightMessagesSemaphoreSlim.WaitAsync(_disposeCts.Token);
+                await _sendInFlightMessagesSemaphoreSlim.WaitAsync(_disposeCts.Token);
                 try
                 {
                     if (_session.IsActive)
@@ -172,7 +173,7 @@ private async void StartWriteWorker()
                 }
                 finally
                 {
-                    _clearInFlightMessagesSemaphoreSlim.Release();
+                    _sendInFlightMessagesSemaphoreSlim.Release();
                 }
             }
         }
@@ -193,7 +194,7 @@ private async Task Initialize()
 
         try
         {
-            if (_disposeCts.IsCancellationRequested)
+            if (_disposeCts.IsCancellationRequested && _inFlightMessages.IsEmpty)
             {
                 _logger.LogWarning("Initialize writer is canceled because it has been disposed");
 
@@ -267,7 +268,7 @@ private async Task Initialize()
                 return;
             }
 
-            await _clearInFlightMessagesSemaphoreSlim.WaitAsync(_disposeCts.Token);
+            await _sendInFlightMessagesSemaphoreSlim.WaitAsync(_disposeCts.Token);
             try
             {
                 var copyInFlightMessages = new ConcurrentQueue<MessageSending>();
@@ -315,7 +316,7 @@ private async Task Initialize()
             }
             finally
             {
-                _clearInFlightMessagesSemaphoreSlim.Release();
+                _sendInFlightMessagesSemaphoreSlim.Release();
             }
         }
         catch (Driver.TransportException e)
@@ -330,17 +331,42 @@ private async Task Initialize()
         }
     }
 
-    public void Dispose()
+    public async ValueTask DisposeAsync()
     {
-        _disposeCts.Cancel();
+        await _sendInFlightMessagesSemaphoreSlim.WaitAsync();
+        try
+        {
+            _disposeCts.Dispose();
+        }
+        finally
+        {
+            _sendInFlightMessagesSemaphoreSlim.Release();
+        }
 
-        _session = new NotStartedWriterSession("Writer is disposed");
+        // wait all messages
+        foreach (var inFlightMessage in _inFlightMessages)
+        {
+            try
+            {
+                await inFlightMessage.Tcs.Task;
+            }
+            catch (Exception e)
+            {
+                _logger.LogWarning(e, "Failed in flight message");
+            }
+        }
+
+        await _session.DisposeAsync();
     }
 }
 
-internal record MessageSending(MessageData MessageData, TaskCompletionSource<WriteResult> Tcs);
+internal record MessageSending(
+    MessageData MessageData,
+    TaskCompletionSource<WriteResult> Tcs,
+    CancellationTokenRegistration DisposedCtr
+);
 
-internal interface IWriteSession
+internal interface IWriteSession : IAsyncDisposable
 {
     Task Write(ConcurrentQueue<MessageSending> toSendBuffer);
 
@@ -372,6 +398,11 @@ public Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
     }
 
     public bool IsActive => true;
+
+    public ValueTask DisposeAsync()
+    {
+        return ValueTask.CompletedTask;
+    }
 }
 
 internal class DummyWriterSession : IWriteSession
@@ -388,6 +419,11 @@ public Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
     }
 
     public bool IsActive => false;
+
+    public ValueTask DisposeAsync()
+    {
+        return ValueTask.CompletedTask;
+    }
 }
 
 internal class WriterSession : TopicSession<MessageFromClient, MessageFromServer>, IWriteSession
@@ -437,6 +473,8 @@ public async Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
                     continue;
                 }
 
+                sendData.DisposedCtr.Unregister();
+
                 var messageData = sendData.MessageData;
 
                 if (messageData.SeqNo == default)

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

@@ -25,7 +25,7 @@ public async Task StressTest_WhenReadingThenCommiting_ReturnMessages()
         topicSettings.Consumers.Add(new Consumer("Consumer"));
         await topicClient.CreateTopic(topicSettings);
 
-        using var writer = new WriterBuilder<string>(_driver, _topicName)
+        await using var writer = new WriterBuilder<string>(_driver, _topicName)
             { ProducerId = "producerId" }.Build();
         var reader = new ReaderBuilder<string>(_driver)
         {

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

@@ -30,7 +30,7 @@ public async Task WriteAsync_WhenOneMessage_ReturnWritten()
         };
         await topicClient.CreateTopic(topicSettings);
 
-        using var writer = new WriterBuilder<string>(_driver, _topicName) { ProducerId = "producerId" }.Build();
+        await using var writer = new WriterBuilder<string>(_driver, _topicName) { ProducerId = "producerId" }.Build();
 
         var result = await writer.WriteAsync("abacaba");
 
@@ -42,7 +42,7 @@ public async Task WriteAsync_WhenOneMessage_ReturnWritten()
     [Fact]
     public async Task WriteAsync_WhenTopicNotFound_ReturnNotFoundException()
     {
-        using var writer = new WriterBuilder<string>(_driver, _topicName + "_not_found")
+        await using var writer = new WriterBuilder<string>(_driver, _topicName + "_not_found")
             { ProducerId = "producerId" }.Build();
 
         Assert.Contains(
@@ -61,7 +61,7 @@ public async Task WriteAsync_When1000Messages_ReturnWriteResultIsPersisted()
         topicSettings.Consumers.Add(new Consumer("Consumer"));
         await topicClient.CreateTopic(topicSettings);
 
-        using var writer = new WriterBuilder<int>(_driver, topicName)
+        await using var writer = new WriterBuilder<int>(_driver, topicName)
             { ProducerId = "producerId" }.Build();
 
         var tasks = new List<Task>();