feat: impl disposing reader

Author: KirillKurdyukov

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> : IAsyncDisposable
+public interface IBidirectionalStream<in TRequest, out TResponse> : IDisposable
 {
     public Task Write(TRequest request);
 
@@ -39,6 +39,8 @@ public interface IBidirectionalStream<in TRequest, out TResponse> : IAsyncDispos
     public TResponse Current { get; }
 
     public string? AuthToken { get; }
+
+    public Task RequestStreamComplete();
 }
 
 public abstract class BaseDriver : IDriver
@@ -225,7 +227,6 @@ 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,
@@ -259,7 +260,6 @@ public async ValueTask<bool> MoveNextAsync()
         }
         catch (RpcException e)
         {
-            _closedResponseStreamTcs.SetResult();
             _rpcErrorAction(e);
 
             throw new Driver.TransportException(e);
@@ -270,7 +270,7 @@ public async ValueTask<bool> MoveNextAsync()
 
     public string? AuthToken => _credentialsProvider.GetAuthInfo();
 
-    public async ValueTask DisposeAsync()
+    public async Task RequestStreamComplete()
     {
         try
         {
@@ -280,7 +280,10 @@ public async ValueTask DisposeAsync()
         {
             _rpcErrorAction(e);
         }
+    }
 
+    public void Dispose()
+    {
         _stream.Dispose();
     }
 }

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

@@ -2,7 +2,7 @@
 
 namespace Ydb.Sdk.Services.Topic;
 
-public interface IReader<TValue> : IDisposable
+public interface IReader<TValue> : IAsyncDisposable
 {
     public ValueTask<Message<TValue>> ReadAsync(CancellationToken cancellationToken = default);
 

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

@@ -21,7 +21,9 @@ internal class Reader<TValue> : IReader<TValue>
     private readonly ReaderConfig _config;
     private readonly IDeserializer<TValue> _deserializer;
     private readonly ILogger _logger;
-    private readonly GrpcRequestSettings _readerGrpcRequestSettings;
+    private readonly GrpcRequestSettings _readerGrpcRequestSettings = new();
+
+    private ReaderSession<TValue>? _currentReaderSession;
 
     private readonly Channel<InternalBatchMessages<TValue>> _receivedMessagesChannel =
         Channel.CreateUnbounded<InternalBatchMessages<TValue>>(
@@ -41,7 +43,6 @@ internal Reader(IDriver driver, ReaderConfig config, IDeserializer<TValue> deser
         _config = config;
         _deserializer = deserializer;
         _logger = driver.LoggerFactory.CreateLogger<Reader<TValue>>();
-        _readerGrpcRequestSettings = new GrpcRequestSettings { CancellationToken = _disposeCts.Token };
 
         _ = Initialize();
     }
@@ -68,7 +69,7 @@ public async ValueTask<Message<TValue>> ReadAsync(CancellationToken cancellation
             }
         }
 
-        throw new ObjectDisposedException("Reader");
+        throw new ReaderException("Reader is disposed");
     }
 
     public async ValueTask<BatchMessages<TValue>> ReadBatchAsync(CancellationToken cancellationToken = default)
@@ -86,7 +87,7 @@ public async ValueTask<BatchMessages<TValue>> ReadBatchAsync(CancellationToken c
             }
         }
 
-        throw new ObjectDisposedException("Reader");
+        throw new ReaderException("Reader is disposed");
     }
 
     private async Task Initialize()
@@ -185,15 +186,15 @@ await stream.Write(new MessageFromClient
                 ReadRequest = new StreamReadMessage.Types.ReadRequest { BytesSize = _config.MemoryUsageMaxBytes }
             });
 
-            new ReaderSession<TValue>(
+            _currentReaderSession = new ReaderSession<TValue>(
                 _config,
                 stream,
                 initResponse.SessionId,
                 Initialize,
                 _logger,
                 _receivedMessagesChannel.Writer,
                 _deserializer
-            ).RunProcessingTopic();
+            );
         }
         catch (Driver.TransportException e)
         {
@@ -203,18 +204,12 @@ await stream.Write(new MessageFromClient
         }
     }
 
-    public void Dispose()
+    public ValueTask DisposeAsync()
     {
-        try
-        {
-            _receivedMessagesChannel.Writer.TryComplete();
+        _receivedMessagesChannel.Writer.TryComplete();
+        _disposeCts.Cancel();
 
-            _disposeCts.Cancel();
-        }
-        finally
-        {
-            _disposeCts.Dispose();
-        }
+        return _currentReaderSession?.DisposeAsync() ?? ValueTask.CompletedTask;
     }
 }
 
@@ -247,6 +242,8 @@ internal class ReaderSession<TValue> : TopicSession<MessageFromClient, MessageFr
     private readonly ChannelWriter<InternalBatchMessages<TValue>> _channelWriter;
     private readonly CancellationTokenSource _lifecycleReaderSessionCts = new();
     private readonly IDeserializer<TValue> _deserializer;
+    private readonly Task _runProcessingStreamResponse;
+    private readonly Task _runProcessingStreamRequest;
 
     private readonly Channel<MessageFromClient> _channelFromClientMessageSending =
         Channel.CreateUnbounded<MessageFromClient>(
@@ -279,29 +276,13 @@ IDeserializer<TValue> deserializer
         _readerConfig = config;
         _channelWriter = channelWriter;
         _deserializer = deserializer;
+
+        _runProcessingStreamResponse = RunProcessingStreamResponse();
+        _runProcessingStreamRequest = RunProcessingStreamRequest();
     }
 
-    public async void RunProcessingTopic()
+    private async Task RunProcessingStreamResponse()
     {
-        _ = Task.Run(async () =>
-        {
-            try
-            {
-                await foreach (var messageFromClient in _channelFromClientMessageSending.Reader.ReadAllAsync())
-                {
-                    await SendMessage(messageFromClient);
-                }
-            }
-            catch (Driver.TransportException e)
-            {
-                Logger.LogError(e, "ReaderSession[{SessionId}] have transport error on Write", SessionId);
-
-                ReconnectSession();
-
-                _lifecycleReaderSessionCts.Cancel();
-            }
-        });
-
         try
         {
             while (await Stream.MoveNextAsync())
@@ -349,6 +330,10 @@ public async void RunProcessingTopic()
             Logger.LogError(e, "ReaderSession[{SessionId}] have transport error on processing server messages",
                 SessionId);
         }
+        catch (OperationCanceledException)
+        {
+            Logger.LogInformation("");
+        }
         finally
         {
             ReconnectSession();
@@ -357,6 +342,25 @@ public async void RunProcessingTopic()
         }
     }
 
+    private async Task RunProcessingStreamRequest()
+    {
+        try
+        {
+            await foreach (var messageFromClient in _channelFromClientMessageSending.Reader.ReadAllAsync())
+            {
+                await SendMessage(messageFromClient);
+            }
+        }
+        catch (Driver.TransportException e)
+        {
+            Logger.LogError(e, "ReaderSession[{SessionId}] have transport error on Write", SessionId);
+
+            ReconnectSession();
+            
+            _lifecycleReaderSessionCts.Cancel();
+        }
+    }
+
     internal async void TryReadRequestBytes(long bytes)
     {
         var readRequestBytes = Interlocked.Add(ref _readRequestBytes, bytes);
@@ -550,4 +554,20 @@ protected override MessageFromClient GetSendUpdateTokenRequest(string token)
             }
         };
     }
+
+    public override async ValueTask DisposeAsync()
+    {
+        _channelFromClientMessageSending.Writer.Complete();
+
+        try
+        {
+            await _runProcessingStreamRequest;
+            await Stream.RequestStreamComplete();
+            await _runProcessingStreamResponse; // waiting all ack's commits
+        }
+        finally
+        {
+            Stream.Dispose();
+        }
+    }
 }

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

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

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

@@ -224,7 +224,7 @@ private async Task Initialize()
                 _logger.LogError("Stream unexpectedly closed by YDB server. Current InitRequest: {initRequest}",
                     initRequest);
 
-                _ = Task.Run(Initialize, _disposeCts.Token);
+                _ = Task.Run(Initialize);
 
                 return;
             }
@@ -239,7 +239,7 @@ private async Task Initialize()
                 {
                     _logger.LogError("Writer initialization failed to start. Reason: {Status}", status);
 
-                    _ = Task.Run(Initialize, _disposeCts.Token);
+                    _ = Task.Run(Initialize);
                 }
                 else
                 {
@@ -323,7 +323,7 @@ private async Task Initialize()
         {
             _logger.LogError(e, "Transport error on creating WriterSession");
 
-            _ = Task.Run(Initialize, _disposeCts.Token);
+            _ = Task.Run(Initialize);
         }
         catch (OperationCanceledException)
         {
@@ -593,4 +593,13 @@ protected override MessageFromClient GetSendUpdateTokenRequest(string token)
             }
         };
     }
+
+    public override async ValueTask DisposeAsync()
+    {
+        Logger.LogInformation("TopicSession[{SessionId}] is being deleted", SessionId);
+        
+        await Stream.RequestStreamComplete();
+        
+        Stream.Dispose();
+    }
 }

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

@@ -42,7 +42,7 @@ public async Task StressTest_WhenReadingThenCommiting_ReturnMessages()
             await message.CommitAsync();
         }
 
-        reader.Dispose();
+        await reader.DisposeAsync();
 
         var readerNext = new ReaderBuilder<string>(_driver)
         {
@@ -59,7 +59,7 @@ public async Task StressTest_WhenReadingThenCommiting_ReturnMessages()
             await message.CommitAsync();
         }
 
-        readerNext.Dispose();
+        await reader.DisposeAsync();
 
         await topicClient.DropTopic(new DropTopicSettings { Path = _topicName });
     }