feat: writer retry always

Author: KirillKurdyukov

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

@@ -4,20 +4,15 @@ public class WriterException : Exception
 {
     public WriterException(string message) : base(message)
     {
-        Status = new Status(StatusCode.Unspecified);
     }
 
     public WriterException(string message, Status status) : base(message + ": " + status)
     {
-        Status = status;
     }
 
-    public WriterException(string message, Driver.TransportException e) : base(message, e)
+    public WriterException(string message, Exception inner) : base(message, inner)
     {
-        Status = e.Status;
     }
-
-    public Status Status { get; }
 }
 
 public class ReaderException : Exception

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

@@ -5,7 +5,6 @@ namespace Ydb.Sdk.Services.Topic;
 internal abstract class TopicSession<TFromClient, TFromServer> : IDisposable
 {
     private readonly Func<Task> _initialize;
-    private readonly Action<WriterException> _resetSessionOnTransportError;
 
     protected readonly IBidirectionalStream<TFromClient, TFromServer> Stream;
     protected readonly ILogger Logger;
@@ -17,28 +16,23 @@ protected TopicSession(
         IBidirectionalStream<TFromClient, TFromServer> stream,
         ILogger logger,
         string sessionId,
-        Func<Task> initialize,
-        Action<WriterException> resetSessionOnTransportError)
+        Func<Task> initialize)
     {
         Stream = stream;
         Logger = logger;
         SessionId = sessionId;
         _initialize = initialize;
-        _resetSessionOnTransportError = resetSessionOnTransportError;
     }
 
-    protected async void ReconnectSession(WriterException exception)
-    {
+    protected async void ReconnectSession() {
         if (Interlocked.CompareExchange(ref _isActive, 0, 1) == 0)
         {
             Logger.LogDebug("Skipping reconnect. A reconnect session has already been initiated");
 
             return;
         }
 
-        _resetSessionOnTransportError(exception);
-
-        Logger.LogInformation("WriterSession[{SessionId}] has been deactivated, starting to reconnect", SessionId);
+        Logger.LogInformation("TopicSession[{SessionId}] has been deactivated, starting to reconnect", SessionId);
 
         await _initialize();
     }

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

@@ -24,10 +24,10 @@ internal class Writer<TValue> : IWriter<TValue>
     private readonly ISerializer<TValue> _serializer;
     private readonly ConcurrentQueue<MessageSending> _toSendBuffer = new();
     private readonly ConcurrentQueue<MessageSending> _inFlightMessages = new();
-    private readonly CancellationTokenSource _disposeTokenSource = new();
+    private readonly CancellationTokenSource _disposeCts = new();
 
     private volatile TaskCompletionSource _tcsWakeUp = new();
-    private volatile IWriteSession _session = new NotStartedWriterSession("Session not started!");
+    private volatile IWriteSession _session = null!;
 
     private int _limitBufferMaxSize;
 
@@ -51,18 +51,27 @@ public async Task<WriteResult> WriteAsync(Message<TValue> message, CancellationT
     {
         TaskCompletionSource<WriteResult> tcs = new();
         cancellationToken.Register(
-            () => tcs.TrySetCanceled(cancellationToken),
-            useSynchronizationContext: false
+            () => tcs.TrySetException(
+                new WriterException("The write operation was canceled before it could be completed")
+            ), useSynchronizationContext: false
         );
 
-        var data = _serializer.Serialize(message.Data);
+        byte[] data;
+        try
+        {
+            data = _serializer.Serialize(message.Data);
+        }
+        catch (Exception e)
+        {
+            throw new WriterException("Error when serializing message data", e);
+        }
+
         var messageData = new MessageData
         {
             Data = ByteString.CopyFrom(data),
             CreatedAt = Timestamp.FromDateTime(message.Timestamp.ToUniversalTime()),
             UncompressedSize = data.Length
         };
-
         foreach (var metadata in message.Metadata)
         {
             messageData.MetadataItems.Add(new MetadataItem
@@ -94,7 +103,14 @@ public async Task<WriteResult> WriteAsync(Message<TValue> message, CancellationT
                 "Buffer overflow: the data size [{DataLength}] exceeds the current buffer limit ({CurLimitBufferSize}) [BufferMaxSize = {BufferMaxSize}]",
                 data.Length, curLimitBufferSize, _config.BufferMaxSize);
 
-            throw new WriterException("Buffer overflow");
+            try
+            {
+                await Task.Delay(_config.BufferOverflowRetryTimeoutMs, cancellationToken);
+            }
+            catch (TaskCanceledException)
+            {
+                throw new WriterException("Buffer overflow");
+            }
         }
 
         try
@@ -113,7 +129,7 @@ private async void StartWriteWorker()
     {
         await Initialize();
 
-        while (!_disposeTokenSource.Token.IsCancellationRequested)
+        while (!_disposeCts.Token.IsCancellationRequested)
         {
             await _tcsWakeUp.Task;
             _tcsWakeUp = new TaskCompletionSource();
@@ -129,10 +145,14 @@ private void WakeUpWorker()
 
     private async Task Initialize()
     {
+        _session = DummyWriteSession.Instance;
+
         try
         {
-            if (_disposeTokenSource.IsCancellationRequested)
+            if (_disposeCts.IsCancellationRequested)
             {
+                _logger.LogWarning("Initialize writer is canceled because it has been disposed");
+
                 return;
             }
 
@@ -162,7 +182,7 @@ private async Task Initialize()
                 _session = new NotStartedWriterSession(
                     $"Stream unexpectedly closed by YDB server. Current InitRequest: {initRequest}");
 
-                _ = Task.Run(Initialize, _disposeTokenSource.Token);
+                _ = Task.Run(Initialize, _disposeCts.Token);
 
                 return;
             }
@@ -177,7 +197,7 @@ private async Task Initialize()
 
                 if (status.StatusCode != StatusCode.SchemeError)
                 {
-                    _ = Task.Run(Initialize, _disposeTokenSource.Token);
+                    _ = Task.Run(Initialize, _disposeCts.Token);
                 }
 
                 _logger.LogCritical("Writer initialization failed to start. Reason: {Status}", status);
@@ -207,17 +227,16 @@ private async Task Initialize()
                 stream,
                 initResponse,
                 Initialize,
-                e => { _session = new NotStartedWriterSession(e); },
                 _logger,
-                _inFlightMessages
+                _toSendBuffer
             );
 
             if (!_inFlightMessages.IsEmpty)
             {
                 var copyInFlightMessages = new ConcurrentQueue<MessageSending>();
                 while (_inFlightMessages.TryDequeue(out var sendData))
                 {
-                    if (sendData.Tcs.Task.IsCanceled)
+                    if (sendData.Tcs.Task.IsFaulted)
                     {
                         _logger.LogWarning("Message[SeqNo={SeqNo}] is cancelled", sendData.MessageData.SeqNo);
 
@@ -232,29 +251,31 @@ private async Task Initialize()
 
             _session = newSession;
             newSession.RunProcessingWriteAck();
+            if (!_toSendBuffer.IsEmpty)
+            {
+                WakeUpWorker(); // send buffer    
+            }
         }
         catch (Driver.TransportException e)
         {
-            _logger.LogError(e, "Unable to connect the session");
-
-            _session = new NotStartedWriterSession(
-                new WriterException("Transport error on creating WriterSession", e));
+            _logger.LogError(e, "Transport error on creating WriterSession");
 
-            _ = Task.Run(Initialize, _disposeTokenSource.Token);
+            _session = DummyWriteSession.Instance;
+            _ = Task.Run(Initialize, _disposeCts.Token);
         }
     }
 
     public void Dispose()
     {
         try
         {
-            _disposeTokenSource.Cancel();
+            _disposeCts.Cancel();
 
             _session.Dispose();
         }
         finally
         {
-            _disposeTokenSource.Dispose();
+            _disposeCts.Dispose();
         }
     }
 }
@@ -280,11 +301,6 @@ public NotStartedWriterSession(string reasonExceptionMessage, Status status)
         _reasonException = new WriterException(reasonExceptionMessage, status);
     }
 
-    public NotStartedWriterSession(WriterException reasonException)
-    {
-        _reasonException = reasonException;
-    }
-
     public Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
     {
         while (toSendBuffer.TryDequeue(out var messageSending))
@@ -297,6 +313,21 @@ public Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
 
     public void Dispose()
     {
+        // Do nothing 
+    }
+}
+
+internal class DummyWriteSession : IWriteSession
+{
+    internal static readonly DummyWriteSession Instance = new();
+
+    public void Dispose()
+    {
+    }
+
+    public Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
+    {
+        return Task.CompletedTask;
     }
 }
 
@@ -312,15 +343,13 @@ public WriterSession(
         WriterStream stream,
         InitResponse initResponse,
         Func<Task> initialize,
-        Action<WriterException> resetSessionOnTransportError,
         ILogger logger,
         ConcurrentQueue<MessageSending> inFlightMessages
     ) : base(
         stream,
         logger,
         initResponse.SessionId,
-        initialize,
-        resetSessionOnTransportError
+        initialize
     )
     {
         _config = config;
@@ -357,10 +386,10 @@ public async Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
         }
         catch (Driver.TransportException e)
         {
-            Logger.LogError(e, "WriterSession[{SessionId}] have error on Write, last SeqNo={SeqNo}",
+            Logger.LogError(e, "WriterSession[{SessionId}] have transport error on Write, last SeqNo={SeqNo}",
                 SessionId, Volatile.Read(ref _seqNum));
 
-            ReconnectSession(new WriterException("Transport error in the WriterSession on write messages", e));
+            ReconnectSession();
         }
     }
 
@@ -428,13 +457,13 @@ Completing task on exception...
         {
             Logger.LogError(e, "WriterSession[{SessionId}] have error on processing writeAck", SessionId);
 
-            ReconnectSession(new WriterException("Transport error in the WriterSession on processing writeAck", e));
+            ReconnectSession();
 
             return;
         }
 
         Logger.LogWarning("WriterSession[{SessionId}]: stream is closed", SessionId);
 
-        ReconnectSession(new WriterException("WriterStream is closed"));
+        ReconnectSession();
     }
 }

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

@@ -39,16 +39,36 @@ public WriterBuilder(IDriver driver, string topicPath)
     /// </summary>    
     public long? PartitionId { get; set; }
 
+    /// <summary>
+    /// The serializer to use to serialize values.
+    /// </summary>
+    /// <remarks>
+    ///     If your value serializer throws an exception, this will be
+    ///     wrapped in a WriterException with unspecified status.
+    /// </remarks>
     public ISerializer<TValue>? Serializer { get; set; }
 
+
+    /// <summary>
+    /// Represents the timeout duration, in milliseconds, used when a buffer overflow is detected.
+    /// This timeout specifies how long the system should wait before attempting to retry the operation.
+    /// </summary>
+    /// <remarks>
+    /// This timeout is important for managing system performance and stability. 
+    /// Too short a timeout could lead to rapid retry attempts, potentially causing further resource contention
+    /// and degrading system performance. Conversely, too long a timeout might delay processing significantly.
+    /// </remarks>
+    public int BufferOverflowRetryTimeoutMs { get; set; } = 10;
+
     public IWriter<TValue> Build()
     {
         var config = new WriterConfig(
             topicPath: TopicPath,
             producerId: ProducerId,
             codec: Codec,
             bufferMaxSize: BufferMaxSize,
-            partitionId: PartitionId
+            partitionId: PartitionId,
+            bufferOverflowRetryTimeoutMs: BufferOverflowRetryTimeoutMs
         );
 
         return new Writer<TValue>(

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

@@ -9,14 +9,15 @@ internal WriterConfig(
         string? producerId,
         Codec codec,
         int bufferMaxSize,
-        long? partitionId
-    )
+        long? partitionId, 
+        int bufferOverflowRetryTimeoutMs)
     {
         TopicPath = topicPath;
         ProducerId = producerId;
         Codec = codec;
         BufferMaxSize = bufferMaxSize;
         PartitionId = partitionId;
+        BufferOverflowRetryTimeoutMs = bufferOverflowRetryTimeoutMs;
     }
 
     public string TopicPath { get; }
@@ -28,6 +29,8 @@ internal WriterConfig(
     public int BufferMaxSize { get; }
 
     public long? PartitionId { get; }
+    
+    public int BufferOverflowRetryTimeoutMs { get; }
 
     public override string ToString()
     {

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

@@ -45,8 +45,8 @@ public async Task WriteAsync_WhenTopicNotFound_ReturnNotFoundException()
         using var writer = new WriterBuilder<string>(_driver, _topicName + "_not_found")
             { ProducerId = "producerId" }.Build();
 
-        Assert.Equal(StatusCode.SchemeError, (await Assert.ThrowsAsync<WriterException>(
-            () => writer.WriteAsync("hello world"))).Status.StatusCode);
+        Assert.Equal("StatusCode.SchemeError", (await Assert.ThrowsAsync<WriterException>(
+            () => writer.WriteAsync("hello world"))).Message);
     }
 
     [Fact]