feat: initial commit

Author: KirillKurdyukov

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

@@ -2,7 +2,7 @@
 
 namespace Ydb.Sdk.Services.Topic;
 
-public interface IWriter<TValue>
+public interface IWriter<TValue> : IDisposable
 {
     public Task<WriteResult> WriteAsync(TValue data);
 

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

@@ -23,22 +23,109 @@ internal class Writer<TValue> : IWriter<TValue>
     private readonly WriterConfig _config;
     private readonly ILogger<Writer<TValue>> _logger;
     private readonly ISerializer<TValue> _serializer;
-
-    private readonly ConcurrentQueue<MessageSending> _inFlightMessages = new();
     private readonly ConcurrentQueue<MessageSending> _toSendBuffer = new();
-    private readonly SemaphoreSlim _writeSemaphoreSlim = new(1);
 
-    private volatile WriterSession _session = null!;
+    private volatile TaskCompletionSource _taskWakeUpCompletionSource = new();
+    private volatile IWriteSession _session = null!;
+    private volatile bool _disposed;
+
+    private int _limitBufferMaxSize;
 
     internal Writer(IDriver driver, WriterConfig config, ISerializer<TValue> serializer)
     {
         _driver = driver;
         _config = config;
-        _serializer = serializer;
         _logger = driver.LoggerFactory.CreateLogger<Writer<TValue>>();
+        _serializer = serializer;
+        _limitBufferMaxSize = config.BufferMaxSize;
+
+        StartWriteWorker();
+    }
+
+    public Task<WriteResult> WriteAsync(TValue data)
+    {
+        return WriteAsync(new Message<TValue>(data));
+    }
+
+    public async Task<WriteResult> WriteAsync(Message<TValue> message)
+    {
+        TaskCompletionSource<WriteResult> completeTask = new();
+
+        var data = _serializer.Serialize(message.Data);
+        var messageData = new MessageData
+        {
+            Data = ByteString.CopyFrom(data),
+            CreatedAt = Timestamp.FromDateTime(message.Timestamp),
+            UncompressedSize = data.Length
+        };
+
+        foreach (var metadata in message.Metadata)
+        {
+            messageData.MetadataItems.Add(
+                new MetadataItem { Key = metadata.Key, Value = ByteString.CopyFrom(metadata.Value) }
+            );
+        }
+
+        while (true)
+        {
+            var curLimitBufferSize = Volatile.Read(ref _limitBufferMaxSize);
+
+            if ( // sending one biggest message anyway
+                curLimitBufferSize == _config.BufferMaxSize && data.Length > curLimitBufferSize
+                || curLimitBufferSize >= data.Length)
+            {
+                if (Interlocked.CompareExchange(ref _limitBufferMaxSize, curLimitBufferSize,
+                        curLimitBufferSize - data.Length) == curLimitBufferSize)
+                {
+                    _toSendBuffer.Enqueue(new MessageSending(messageData, completeTask));
+
+                    WakeUpWorker();
+
+                    break;
+                }
+
+                // Next try on race condition
+                continue;
+            }
+
+            _logger.LogWarning(
+                "Buffer overflow: the data size [{DataLength}] exceeds the current buffer limit ({CurLimitBufferSize}) [BufferMaxSize = {BufferMaxSize}]",
+                data.Length, curLimitBufferSize, _config.BufferMaxSize);
+
+            throw new YdbWriterException("Buffer overflow");
+        }
+
+        try
+        {
+            var writeResult = await completeTask.Task;
+
+            return writeResult;
+        }
+        finally
+        {
+            Interlocked.Add(ref _limitBufferMaxSize, data.Length);
+        }
     }
 
-    internal async Task Initialize()
+    private async void StartWriteWorker()
+    {
+        await Initialize();
+
+        while (!_disposed)
+        {
+            await _taskWakeUpCompletionSource.Task;
+            _taskWakeUpCompletionSource = new TaskCompletionSource();
+
+            await _session.Write(_toSendBuffer);
+        }
+    }
+
+    private void WakeUpWorker()
+    {
+        _taskWakeUpCompletionSource.TrySetResult();
+    }
+
+    private async Task Initialize()
     {
         _logger.LogInformation("Writer session initialization started. WriterConfig: {WriterConfig}", _config);
 
@@ -63,13 +150,26 @@ internal async Task Initialize()
         await stream.Write(new MessageFromClient { InitRequest = initRequest });
         if (!await stream.MoveNextAsync())
         {
-            throw new YdbWriterException(
+            _session = new NotStartedWriterSession(
                 $"Stream unexpectedly closed by YDB server. Current InitRequest: {initRequest}");
+
+            _ = Task.Run(Initialize);
+
+            return;
         }
 
         var receivedInitMessage = stream.Current;
 
-        Status.FromProto(receivedInitMessage.Status, receivedInitMessage.Issues).EnsureSuccess();
+        var status = Status.FromProto(receivedInitMessage.Status, receivedInitMessage.Issues);
+
+        if (status.IsNotSuccess)
+        {
+            _session = new NotStartedWriterSession(status.ToString());
+
+            _ = Task.Run(Initialize);
+
+            return;
+        }
 
         var initResponse = receivedInitMessage.InitResponse;
 
@@ -78,75 +178,54 @@ internal async Task Initialize()
 
         if (!initResponse.SupportedCodecs.Codecs.Contains((int)_config.Codec))
         {
-            throw new YdbWriterException($"Topic[{_config.TopicPath}] is not supported codec: {_config.Codec}");
-        }
-
-        _session = new WriterSession(_config, stream, initResponse, Initialize, _logger);
+            _logger.LogCritical("Topic[{TopicPath}] is not supported codec: {Codec}", _config.TopicPath, _config.Codec);
 
-        await _writeSemaphoreSlim.WaitAsync();
-        try
-        {
-            _logger.LogDebug("Retrying to send pending in-flight messages after stream restart");
-
-            await _session.Write(_inFlightMessages, _inFlightMessages);
-        }
-        finally
-        {
-            _writeSemaphoreSlim.Release();
+            _session = new NotStartedWriterSession(
+                $"Topic[{_config.TopicPath}] is not supported codec: {_config.Codec}");
+            return;
         }
 
-        _ = _session.RunProcessingWriteAck(_inFlightMessages);
+        _session = new WriterSession(_config, stream, initResponse, Initialize, _logger);
     }
 
-    public Task<WriteResult> WriteAsync(TValue data)
+    public void Dispose()
     {
-        return WriteAsync(new Message<TValue>(data));
+        _disposed = true;
     }
+}
 
-    public async Task<WriteResult> WriteAsync(Message<TValue> message)
-    {
-        TaskCompletionSource<WriteResult> completeTask = new();
-
-        var data = _serializer.Serialize(message.Data);
-        var messageData = new MessageData
-        {
-            Data = ByteString.CopyFrom(data),
-            CreatedAt = Timestamp.FromDateTime(message.Timestamp),
-            UncompressedSize = data.Length
-        };
+internal record MessageSending(MessageData MessageData, TaskCompletionSource<WriteResult> TaskCompletionSource);
 
-        foreach (var metadata in message.Metadata)
-        {
-            messageData.MetadataItems.Add(
-                new MetadataItem { Key = metadata.Key, Value = ByteString.CopyFrom(metadata.Value) }
-            );
-        }
+internal interface IWriteSession
+{
+    Task Write(ConcurrentQueue<MessageSending> toSendBuffer);
+}
 
-        _toSendBuffer.Enqueue(new MessageSending(messageData, completeTask));
+internal class NotStartedWriterSession : IWriteSession
+{
+    private readonly YdbWriterException _reasonException;
 
-        if (_toSendBuffer.IsEmpty) // concurrent sending
-        {
-            return await completeTask.Task;
-        }
+    public NotStartedWriterSession(string reasonExceptionMessage)
+    {
+        _reasonException = new YdbWriterException(reasonExceptionMessage);
+    }
 
-        await _writeSemaphoreSlim.WaitAsync();
-        try
-        {
-            await _session.Write(_toSendBuffer, _inFlightMessages);
-        }
-        finally
+    public Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
+    {
+        foreach (var messageSending in toSendBuffer)
         {
-            _writeSemaphoreSlim.Release();
+            messageSending.TaskCompletionSource.SetException(_reasonException);
         }
 
-        return await completeTask.Task;
+        return Task.CompletedTask;
     }
 }
 
 // No thread safe
-internal class WriterSession : TopicSession<MessageFromClient, MessageFromServer>
+internal class WriterSession : TopicSession<MessageFromClient, MessageFromServer>, IWriteSession
 {
     private readonly WriterConfig _config;
+    private readonly ConcurrentQueue<MessageSending> _inFlightMessages = new();
 
     private long _seqNum;
 
@@ -159,9 +238,48 @@ public WriterSession(
     {
         _config = config;
         Volatile.Write(ref _seqNum, initResponse.LastSeqNo); // happens-before for Volatile.Read
+
+        RunProcessingWriteAck();
     }
 
-    internal async Task RunProcessingWriteAck(ConcurrentQueue<MessageSending> inFlightMessages)
+    public async Task Write(ConcurrentQueue<MessageSending> toSendBuffer)
+    {
+        try
+        {
+            var writeMessage = new StreamWriteMessage.Types.WriteRequest
+            {
+                Codec = (int)_config.Codec
+            };
+
+            var currentSeqNum = Volatile.Read(ref _seqNum);
+
+            while (toSendBuffer.TryDequeue(out var sendData))
+            {
+                var messageData = sendData.MessageData;
+
+                messageData.SeqNo = ++currentSeqNum;
+                writeMessage.Messages.Add(messageData);
+                _inFlightMessages.Enqueue(sendData);
+            }
+
+            Volatile.Write(ref _seqNum, currentSeqNum);
+            await Stream.Write(new MessageFromClient { WriteRequest = writeMessage });
+        }
+        catch (TransactionException e)
+        {
+            Logger.LogError(e, "WriterSession[{SessionId}] have error on Write, last SeqNo={SeqNo}",
+                SessionId, Volatile.Read(ref _seqNum));
+
+            ReconnectSession();
+
+            while (_inFlightMessages.TryDequeue(out var sendData))
+            {
+                sendData.TaskCompletionSource.SetException(e);
+            }
+        }
+    }
+
+    private async void RunProcessingWriteAck()
     {
         try
         {
@@ -182,7 +300,7 @@ internal async Task RunProcessingWriteAck(ConcurrentQueue<MessageSending> inFlig
 
                 foreach (var ack in messageFromServer.WriteResponse.Acks)
                 {
-                    if (!inFlightMessages.TryPeek(out var messageFromClient))
+                    if (!_inFlightMessages.TryPeek(out var messageFromClient))
                     {
                         Logger.LogCritical("No client message was found upon receipt of an acknowledgement: {WriteAck}",
                             ack);
@@ -217,7 +335,7 @@ Completing task on exception...
                         messageFromClient.TaskCompletionSource.SetResult(new WriteResult(ack));
                     }
 
-                    inFlightMessages.TryDequeue(out _); // Dequeue 
+                    _inFlightMessages.TryDequeue(out _); // Dequeue 
                 }
             }
         }
@@ -230,39 +348,4 @@ Completing task on exception...
             ReconnectSession();
         }
     }
-
-    internal async Task Write(ConcurrentQueue<MessageSending> toSendBuffer,
-        ConcurrentQueue<MessageSending> inFlightMessages)
-    {
-        try
-        {
-            var writeMessage = new StreamWriteMessage.Types.WriteRequest
-            {
-                Codec = (int)_config.Codec
-            };
-
-            var currentSeqNum = Volatile.Read(ref _seqNum);
-
-            while (toSendBuffer.TryDequeue(out var sendData))
-            {
-                var messageData = sendData.MessageData;
-
-                messageData.SeqNo = ++currentSeqNum;
-                writeMessage.Messages.Add(messageData);
-                inFlightMessages.Enqueue(sendData);
-            }
-
-            Volatile.Write(ref _seqNum, currentSeqNum);
-            await Stream.Write(new MessageFromClient { WriteRequest = writeMessage });
-        }
-        catch (TransactionException e)
-        {
-            Logger.LogError(e, "WriterSession[{SessionId}] have error on Write, last SeqNo={SeqNo}",
-                SessionId, Volatile.Read(ref _seqNum));
-
-            ReconnectSession();
-        }
-    }
 }
-
-internal record MessageSending(MessageData MessageData, TaskCompletionSource<WriteResult> TaskCompletionSource);

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

@@ -13,9 +13,9 @@ public WriterBuilder(Driver driver, WriterConfig config)
 
     public ISerializer<TValue>? Serializer { get; set; }
 
-    public async Task<IWriter<TValue>> Build()
+    public IWriter<TValue> Build()
     {
-        var writer = new Writer<TValue>(
+        return new Writer<TValue>(
             _driver,
             _config,
             Serializer ?? (ISerializer<TValue>)(
@@ -24,9 +24,5 @@ public async Task<IWriter<TValue>> Build()
                     : throw new YdbWriterException("The serializer is not set")
             )
         );
-
-        await writer.Initialize();
-
-        return writer;
     }
 }