fix

Author: KirillKurdyukov

src/Ydb.Sdk/src/Driver.cs

@@ -367,8 +367,7 @@ public async ValueTask<bool> MoveNextAsync()
         }
     }
 
-    internal sealed class BidirectionalStream<TRequest, TResponse> : IAsyncEnumerator<TResponse>,
-        IAsyncEnumerable<TResponse>
+    internal sealed class BidirectionalStream<TRequest, TResponse> : IDisposable
     {
         private readonly AsyncDuplexStreamingCall<TRequest, TResponse> _bidirectionalStream;
         private readonly Action _rpcErrorAction;
@@ -394,13 +393,6 @@ public async Task Write(TRequest request)
             }
         }
 
-        public ValueTask DisposeAsync()
-        {
-            _bidirectionalStream.Dispose();
-
-            return default;
-        }
-
         public async ValueTask<bool> MoveNextAsync()
         {
             try
@@ -417,9 +409,9 @@ public async ValueTask<bool> MoveNextAsync()
 
         public TResponse Current => _bidirectionalStream.ResponseStream.Current;
 
-        public IAsyncEnumerator<TResponse> GetAsyncEnumerator(CancellationToken cancellationToken = new())
+        public void Dispose()
         {
-            return this;
+            _bidirectionalStream.Dispose();
         }
     }
 

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

@@ -2,18 +2,21 @@
 
 namespace Ydb.Sdk.Services.Topic;
 
-internal abstract class TopicSession : IDisposable
+internal abstract class TopicSession<TFromClient, TFromServer> : IDisposable
 {
     private readonly Func<Task> _initialize;
 
+    protected readonly Driver.BidirectionalStream<TFromClient, TFromServer> Stream;
     protected readonly ILogger Logger;
     protected readonly string SessionId;
 
     private int _isActive = 1;
     private bool _disposed;
 
-    protected TopicSession(ILogger logger, string sessionId, Func<Task> initialize)
+    protected TopicSession(Driver.BidirectionalStream<TFromClient, TFromServer> stream, ILogger logger,
+        string sessionId, Func<Task> initialize)
     {
+        Stream = stream;
         Logger = logger;
         SessionId = sessionId;
         _initialize = initialize;
@@ -50,5 +53,7 @@ public void Dispose()
         {
             _disposed = true;
         }
+
+        Stream.Dispose();
     }
 }

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

@@ -11,6 +11,7 @@ namespace Ydb.Sdk.Services.Topic;
 using InitResponse = StreamWriteMessage.Types.InitResponse;
 using MessageData = StreamWriteMessage.Types.WriteRequest.Types.MessageData;
 using MessageFromClient = StreamWriteMessage.Types.FromClient;
+using MessageFromServer = StreamWriteMessage.Types.FromServer;
 using WriterStream = Driver.BidirectionalStream<
     StreamWriteMessage.Types.FromClient,
     StreamWriteMessage.Types.FromServer
@@ -143,10 +144,9 @@ public async Task<WriteResult> WriteAsync(Message<TValue> message)
 }
 
 // No thread safe
-internal class WriterSession : TopicSession
+internal class WriterSession : TopicSession<MessageFromClient, MessageFromServer>
 {
     private readonly WriterConfig _config;
-    private readonly WriterStream _stream;
 
     private long _seqNum;
 
@@ -155,11 +155,10 @@ public WriterSession(
         WriterStream stream,
         InitResponse initResponse,
         Func<Task> initialize,
-        ILogger logger) : base(logger, initResponse.SessionId, initialize)
+        ILogger logger) : base(stream, logger, initResponse.SessionId, initialize)
     {
         _config = config;
-        _stream = stream;
-        _seqNum = initResponse.LastSeqNo;
+        Volatile.Write(ref _seqNum, initResponse.LastSeqNo); // happens-before for Volatile.Read
     }
 
     internal async Task RunProcessingWriteAck(ConcurrentQueue<MessageSending> inFlightMessages)
@@ -168,8 +167,9 @@ internal async Task RunProcessingWriteAck(ConcurrentQueue<MessageSending> inFlig
         {
             Logger.LogInformation("WriterSession[{SessionId}] is running processing writeAck", SessionId);
 
-            await foreach (var messageFromServer in _stream)
+            while (await Stream.MoveNextAsync())
             {
+                var messageFromServer = Stream.Current;
                 var status = Status.FromProto(messageFromServer.Status, messageFromServer.Issues);
 
                 if (status.IsNotSuccess)
@@ -253,7 +253,7 @@ internal async Task Write(ConcurrentQueue<MessageSending> toSendBuffer,
             }
 
             Volatile.Write(ref _seqNum, currentSeqNum);
-            await _stream.Write(new MessageFromClient { WriteRequest = writeMessage });
+            await Stream.Write(new MessageFromClient { WriteRequest = writeMessage });
         }
         catch (TransactionException e)
         {
@@ -263,11 +263,6 @@ internal async Task Write(ConcurrentQueue<MessageSending> toSendBuffer,
             ReconnectSession();
         }
     }
-
-    public ValueTask DisposeAsync()
-    {
-        return _stream.DisposeAsync();
-    }
 }
 
 internal record MessageSending(MessageData MessageData, TaskCompletionSource<WriteResult> TaskCompletionSource);