update writer

Author: KirillKurdyukov

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

@@ -26,9 +26,9 @@ internal class Writer<TValue> : IWriter<TValue>
     private readonly CancellationTokenSource _disposeCts = new();
 
     private volatile TaskCompletionSource _tcsWakeUp = new();
+    private volatile TaskCompletionSource _tcsBufferAvailableEvent = new();
     private volatile IWriteSession _session = null!;
-
-    private int _limitBufferMaxSize;
+    private volatile int _limitBufferMaxSize;
 
     internal Writer(IDriver driver, WriterConfig config, ISerializer<TValue> serializer)
     {
@@ -79,7 +79,7 @@ public async Task<WriteResult> WriteAsync(Message<TValue> message, CancellationT
 
         while (true)
         {
-            var curLimitBufferSize = Volatile.Read(ref _limitBufferMaxSize);
+            var curLimitBufferSize = _limitBufferMaxSize;
 
             if ( // sending one biggest message anyway
                 (curLimitBufferSize == _config.BufferMaxSize && data.Length > curLimitBufferSize)
@@ -98,15 +98,15 @@ public async Task<WriteResult> WriteAsync(Message<TValue> message, CancellationT
                 continue;
             }
 
-            _logger.LogWarning(
-                "Buffer overflow: the data size [{DataLength}] exceeds the current buffer limit ({CurLimitBufferSize}) [BufferMaxSize = {BufferMaxSize}]",
-                data.Length, curLimitBufferSize, _config.BufferMaxSize);
+            // _logger.LogWarning(
+            //     "Buffer overflow: the data size [{DataLength}] exceeds the current buffer limit ({CurLimitBufferSize}) [BufferMaxSize = {BufferMaxSize}]",
+            //     data.Length, curLimitBufferSize, _config.BufferMaxSize);
 
             try
             {
-                await Task.Delay(_config.BufferOverflowRetryTimeoutMs, cancellationToken);
+                await WaitBufferAvailable(cancellationToken);
             }
-            catch (TaskCanceledException)
+            catch (OperationCanceledException)
             {
                 throw new WriterException("Buffer overflow");
             }
@@ -121,24 +121,46 @@ public async Task<WriteResult> WriteAsync(Message<TValue> message, CancellationT
         finally
         {
             Interlocked.Add(ref _limitBufferMaxSize, data.Length);
+
+            _tcsBufferAvailableEvent.TrySetResult();
         }
     }
 
+    private async Task WaitBufferAvailable(CancellationToken cancellationToken)
+    {
+        var tcsBufferAvailableEvent = _tcsBufferAvailableEvent;
+
+        await tcsBufferAvailableEvent.Task.WaitAsync(cancellationToken);
+
+        Interlocked.CompareExchange(
+            ref _tcsBufferAvailableEvent,
+            new TaskCompletionSource(),
+            tcsBufferAvailableEvent
+        );
+    }
+
     private async void StartWriteWorker()
     {
         await Initialize();
 
-        while (!_disposeCts.Token.IsCancellationRequested)
+        try
         {
-            await _tcsWakeUp.Task;
-            _tcsWakeUp = new TaskCompletionSource();
-
-            if (_toSendBuffer.IsEmpty)
+            while (!_disposeCts.Token.IsCancellationRequested)
             {
-                continue;
-            }
+                await _tcsWakeUp.Task.WaitAsync(_disposeCts.Token);
+                _tcsWakeUp = new TaskCompletionSource();
 
-            await _session.Write(_toSendBuffer);
+                if (_toSendBuffer.IsEmpty)
+                {
+                    continue;
+                }
+
+                await _session.Write(_toSendBuffer);
+            }
+        }
+        catch (OperationCanceledException)
+        {
+            _logger.LogInformation("WriteWorker[{WriterConfig}] is disposed", _config);
         }
     }
 

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

@@ -48,27 +48,14 @@ public WriterBuilder(IDriver driver, string topicPath)
     /// </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,
-            bufferOverflowRetryTimeoutMs: BufferOverflowRetryTimeoutMs
+            partitionId: PartitionId
         );
 
         return new Writer<TValue>(

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

@@ -9,15 +9,13 @@ internal WriterConfig(
         string? producerId,
         Codec codec,
         int bufferMaxSize,
-        long? partitionId,
-        int bufferOverflowRetryTimeoutMs)
+        long? partitionId)
     {
         TopicPath = topicPath;
         ProducerId = producerId;
         Codec = codec;
         BufferMaxSize = bufferMaxSize;
         PartitionId = partitionId;
-        BufferOverflowRetryTimeoutMs = bufferOverflowRetryTimeoutMs;
     }
 
     public string TopicPath { get; }
@@ -30,17 +28,15 @@ internal WriterConfig(
 
     public long? PartitionId { get; }
 
-    public int BufferOverflowRetryTimeoutMs { get; }
-
     public override string ToString()
     {
-        var toString = new StringBuilder().Append("[TopicPath: ").Append(TopicPath);
+        var toString = new StringBuilder().Append("TopicPath: ").Append(TopicPath);
 
         if (ProducerId != null)
         {
             toString.Append(", ProducerId: ").Append(ProducerId);
         }
 
-        return toString.Append(", Codec: ").Append(Codec).Append(']').ToString();
+        return toString.Append(", Codec: ").Append(Codec).ToString();
     }
 }

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

@@ -318,7 +318,7 @@ public async Task Initialize_WhenNotSupportedCodec_ThrowWriterExceptionOnWriteAs
     }
 
     [Fact]
-    public async Task WriteAsyncStress_WhenBufferIsOverflow_ThrowWriterExceptionOnBufferOverflow()
+    public async Task WriteAsyncStress_WhenBufferIsOverflow_ThrowException()
     {
         const int countBatchSendingSize = 1000;
         const int batchTasksSize = 100;
@@ -331,77 +331,51 @@ public async Task WriteAsyncStress_WhenBufferIsOverflow_ThrowWriterExceptionOnBu
         var taskSource = new TaskCompletionSource<bool>();
         _mockStream.Setup(stream => stream.Write(It.IsAny<FromClient>()))
             .Returns(Task.CompletedTask);
-        var mockNextAsync = _mockStream.SetupSequence(stream => stream.MoveNextAsync())
+        _mockStream.SetupSequence(stream => stream.MoveNextAsync())
             .Returns(new ValueTask<bool>(true))
             .Returns(new ValueTask<bool>(taskSource.Task));
-        var sequentialResult = _mockStream.SetupSequence(stream => stream.Current)
-            .Returns(new StreamWriteMessage.Types.FromServer
-            {
-                InitResponse = new StreamWriteMessage.Types.InitResponse
-                    { LastSeqNo = 0, PartitionId = 1, SessionId = "SessionId" },
-                Status = StatusIds.Types.StatusCode.Success
-            });
         using var writer = new WriterBuilder<int>(_mockIDriver.Object, "/topic")
         {
             ProducerId = "producerId",
-            BufferMaxSize = bufferSize /* bytes */,
-            BufferOverflowRetryTimeoutMs = 1_000
+            BufferMaxSize = bufferSize /* bytes */
         }.Build();
 
         for (var attempt = 0; attempt < countBatchSendingSize; attempt++)
         {
             _testOutputHelper.WriteLine($"Processing attempt {attempt}");
             var cts = new CancellationTokenSource();
+            cts.CancelAfter(10);
 
             var tasks = new List<Task<WriteResult>>();
-            var serverAck = new StreamWriteMessage.Types.FromServer
-            {
-                WriteResponse = new StreamWriteMessage.Types.WriteResponse { PartitionId = 1 },
-                Status = StatusIds.Types.StatusCode.Success
-            };
+
             for (var i = 0; i < batchTasksSize; i++)
             {
                 tasks.Add(writer.WriteAsync(100, cts.Token));
-                serverAck.WriteResponse.Acks.Add(new StreamWriteMessage.Types.WriteResponse.Types.WriteAck
-                {
-                    SeqNo = bufferSize / messageSize * attempt + i + 1,
-                    Written = new StreamWriteMessage.Types.WriteResponse.Types.WriteAck.Types.Written
-                        { Offset = i * messageSize + bufferSize * attempt }
-                });
             }
 
-            sequentialResult.Returns(() =>
-            {
-                // ReSharper disable once AccessToModifiedClosure
-                Volatile.Write(ref taskSource, new TaskCompletionSource<bool>());
-                mockNextAsync.Returns(() =>
-                {
-                    cts.Cancel();
-                    return new ValueTask<bool>(Volatile.Read(ref taskSource).Task);
-                });
-                return serverAck;
-            });
-            taskSource.SetResult(true);
-
-            var countSuccess = 0;
-            var countErrors = 0;
+            var countErrorCancel = 0;
+            var countErrorBuffer = 0;
             foreach (var task in tasks)
             {
                 try
                 {
-                    var res = await task;
-                    countSuccess++;
-                    Assert.Equal(PersistenceStatus.Written, res.Status);
+                    await task;
                 }
                 catch (WriterException e)
                 {
-                    countErrors++;
-                    Assert.Equal("Buffer overflow", e.Message);
+                    if ("Buffer overflow" == e.Message)
+                    {
+                        countErrorBuffer++;
+                    }
+                    else
+                    {
+                        countErrorCancel++;
+                    }
                 }
             }
 
-            Assert.Equal(bufferSize / messageSize, countSuccess);
-            Assert.Equal(batchTasksSize - bufferSize / messageSize, countErrors);
+            Assert.Equal(bufferSize / messageSize, countErrorCancel);
+            Assert.Equal(batchTasksSize - bufferSize / messageSize, countErrorBuffer);
         }
     }