fix: memory leak CancellationTokenRegistration (#242)

1) Do not pessimize the node on Grpc.Core.StatusCode.Cancelled and Grpc.Core.StatusCode.DeadlineExceeded.
2) Dispose of WriterSession using dispose CancellationToken.
3) BidirectionalStream is internal class.
4) Move Metadata class to Ydb.Sdk.Services.Topic.
5) Fixed memory leak CancellationTokenRegistration.
6) Cancel writing tasks after disposing of Writer.

Author: KirillKurdyukov

CHANGELOG.md

@@ -1,3 +1,10 @@
+- Do not pessimize the node on Grpc.Core.StatusCode.Cancelled and Grpc.Core.StatusCode.DeadlineExceeded.
+- Dispose of WriterSession using dispose CancellationToken.
+- BidirectionalStream is internal class.
+- Move Metadata class to Ydb.Sdk.Services.Topic.
+- Fixed memory leak CancellationTokenRegistration.
+- Cancel writing tasks after disposing of Writer.
+
 ## v0.9.3
 - Fixed bug in Topic Writer: worker is stopped by disposeCts
 - Fixed bug in sql parser ADO.NET: deduplication declare param in YQL query 

src/Ydb.Sdk/src/Driver.cs

@@ -97,6 +97,12 @@ protected override (string, GrpcChannel) GetChannel(long nodeId)
     protected override void OnRpcError(string endpoint, RpcException e)
     {
         Logger.LogWarning("gRPC error [{Status}] on channel {Endpoint}", e.Status, endpoint);
+
+        if (e.StatusCode is Grpc.Core.StatusCode.Cancelled or Grpc.Core.StatusCode.DeadlineExceeded)
+        {
+            return;
+        }
+
         if (!_endpointPool.PessimizeEndpoint(endpoint))
         {
             return;

src/Ydb.Sdk/src/GrpcRequestSettings.cs

@@ -11,6 +11,7 @@ public class GrpcRequestSettings
     public string TraceId { get; set; } = string.Empty;
     public TimeSpan TransportTimeout { get; set; } = TimeSpan.Zero;
     public ImmutableArray<string> CustomClientHeaders { get; } = new();
+    public CancellationToken CancellationToken = default;
 
     internal long NodeId { get; set; }
     internal Action<Grpc.Core.Metadata> TrailersHandler { get; set; } = _ => { };

src/Ydb.Sdk/src/IDriver.cs

@@ -144,7 +144,8 @@ protected CallOptions GetCallOptions(GrpcRequestSettings settings)
         }
 
         var options = new CallOptions(
-            headers: meta
+            headers: meta,
+            cancellationToken: settings.CancellationToken
         );
 
         if (settings.TransportTimeout != TimeSpan.Zero)
@@ -213,7 +214,7 @@ public async ValueTask<bool> MoveNextAsync()
     }
 }
 
-public class BidirectionalStream<TRequest, TResponse> : IBidirectionalStream<TRequest, TResponse>
+internal class BidirectionalStream<TRequest, TResponse> : IBidirectionalStream<TRequest, TResponse>
 {
     private readonly AsyncDuplexStreamingCall<TRequest, TResponse> _stream;
     private readonly Action<RpcException> _rpcErrorAction;

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

@@ -0,0 +1,3 @@
+namespace Ydb.Sdk.Services.Topic;
+
+public record Metadata(string Key, byte[] Value);

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

@@ -13,5 +13,3 @@ public Message(TValue data)
 
     public List<Metadata> Metadata { get; } = new();
 }
-
-public record Metadata(string Key, byte[] Value);

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

@@ -21,6 +21,7 @@ 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 ConcurrentQueue<MessageSending> _toSendBuffer = new();
     private readonly ConcurrentQueue<MessageSending> _inFlightMessages = new();
     private readonly CancellationTokenSource _disposeCts = new();
@@ -37,6 +38,7 @@ 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();
     }
@@ -49,11 +51,15 @@ public Task<WriteResult> WriteAsync(TValue data, CancellationToken cancellationT
     public async Task<WriteResult> WriteAsync(Message<TValue> message, CancellationToken cancellationToken)
     {
         TaskCompletionSource<WriteResult> tcs = new();
-        cancellationToken.Register(
+        await using var registrationUserCancellationTokenRegistration = cancellationToken.Register(
             () => tcs.TrySetException(
                 new WriterException("The write operation was canceled before it could be completed")
             ), useSynchronizationContext: false
         );
+        await using var writerDisposedCancellationTokenRegistration = _disposeCts.Token.Register(
+            () => tcs.TrySetException(new WriterException($"Writer[{_config}] is disposed")),
+            useSynchronizationContext: false
+        );
 
         byte[] data;
         try
@@ -184,10 +190,7 @@ private async Task Initialize()
 
             _logger.LogInformation("Writer session initialization started. WriterConfig: {WriterConfig}", _config);
 
-            var stream = _driver.BidirectionalStreamCall(
-                TopicService.StreamWriteMethod,
-                GrpcRequestSettings.DefaultInstance
-            );
+            var stream = _driver.BidirectionalStreamCall(TopicService.StreamWriteMethod, _writerGrpcRequestSettings);
 
             var initRequest = new StreamWriteMessage.Types.InitRequest { Path = _config.TopicPath };
             if (_config.ProducerId != null)
@@ -304,22 +307,15 @@ private async Task Initialize()
 
     public void Dispose()
     {
-        try
-        {
-            _disposeCts.Cancel();
+        _disposeCts.Cancel();
 
-            _session.Dispose();
-        }
-        finally
-        {
-            _disposeCts.Dispose();
-        }
+        _session = new NotStartedWriterSession("Writer is disposed");
     }
 }
 
 internal record MessageSending(MessageData MessageData, TaskCompletionSource<WriteResult> Tcs);
 
-internal interface IWriteSession : IDisposable
+internal interface IWriteSession
 {
     Task Write(ConcurrentQueue<MessageSending> toSendBuffer);
 }
@@ -347,11 +343,6 @@ public Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
 
         return Task.CompletedTask;
     }
-
-    public void Dispose()
-    {
-        // Do nothing 
-    }
 }
 
 internal class DummyWriterSession : IWriteSession
@@ -362,10 +353,6 @@ private DummyWriterSession()
     {
     }
 
-    public void Dispose()
-    {
-    }
-
     public Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
     {
         return Task.CompletedTask;
@@ -501,18 +488,16 @@ Completing task on exception...
                     _inFlightMessages.TryDequeue(out _); // Dequeue 
                 }
             }
+
+            Logger.LogWarning("WriterSession[{SessionId}]: stream is closed", SessionId);
         }
         catch (Driver.TransportException e)
         {
             Logger.LogError(e, "WriterSession[{SessionId}] have error on processing writeAck", SessionId);
-
+        }
+        finally
+        {
             ReconnectSession();
-
-            return;
         }
-
-        Logger.LogWarning("WriterSession[{SessionId}]: stream is closed", SessionId);
-
-        ReconnectSession();
     }
 }

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

@@ -704,6 +704,23 @@ public async Task WriteAsync_WhenInFlightBufferSendInInitialize_ReturnCompletedT
         _mockStream.Verify(stream => stream.Current, Times.Exactly(3));
     }
 
+    [Fact]
+    public async Task WriteAsync_WhenWriterIsDisposed_ThrowWriterException()
+    {
+        _mockStream.SetupSequence(stream => stream.Write(It.IsAny<FromClient>()))
+            .Returns(Task.CompletedTask);
+        _mockStream.SetupSequence(stream => stream.MoveNextAsync())
+            .ReturnsAsync(true);
+        SetupReadOneWriteAckMessage();
+
+        var writer = new WriterBuilder<string>(_mockIDriver.Object, "/topic")
+            { ProducerId = "producerId" }.Build();
+        writer.Dispose();
+
+        Assert.Equal("Writer[TopicPath: /topic, ProducerId: producerId, Codec: Raw] is disposed",
+            (await Assert.ThrowsAsync<WriterException>(() => writer.WriteAsync("abacaba"))).Message);
+    }
+
     private void SetupReadOneWriteAckMessage()
     {
         _mockStream.SetupSequence(stream => stream.Current)