From e90d588d9c2732ffd64729d177bdfe93646e9505 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Mon, 22 Jan 2024 17:09:43 +0000 Subject: [PATCH 01/25] Send buffer refinements --- tests/NATS.Client.Core.Tests/ProtocolTest.cs | 48 ++++++++++++++++ .../InMemoryTestLoggerFactory.cs | 55 +++++++++++++++++++ 2 files changed, 103 insertions(+) create mode 100644 tests/NATS.Client.TestUtilities/InMemoryTestLoggerFactory.cs diff --git a/tests/NATS.Client.Core.Tests/ProtocolTest.cs b/tests/NATS.Client.Core.Tests/ProtocolTest.cs index 030ece422..e97fef5b3 100644 --- a/tests/NATS.Client.Core.Tests/ProtocolTest.cs +++ b/tests/NATS.Client.Core.Tests/ProtocolTest.cs @@ -1,5 +1,7 @@ using System.Buffers; using System.Text; +using Microsoft.Extensions.Logging; +using NATS.Client.TestUtilities; namespace NATS.Client.Core.Tests; @@ -338,6 +340,52 @@ await Retry.Until( await nats.DisposeAsync(); } + [Fact] + public async Task Protocol_parser_under_load() + { + await using var server = NatsServer.Start(); + var logger = new InMemoryTestLoggerFactory(LogLevel.Error); + var opts = server.ClientOpts(NatsOpts.Default) with { LoggerFactory = logger }; + var nats = new NatsConnection(opts); + + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + + var signal = new WaitSignal(); + + _ = Task.Run( + async () => + { + var count = 0; + await foreach (var unused in nats.SubscribeAsync("x", cancellationToken: cts.Token)) + { + if (++count > 10_000) + signal.Pulse(); + } + }, + cts.Token); + + _ = Task.Run( + async () => + { + while (!cts.Token.IsCancellationRequested) + await nats.PublishAsync("x", "x", cancellationToken: cts.Token); + }, + cts.Token); + + await signal; + + for (var i = 0; i < 3; i++) + { + await Task.Delay(1_000, cts.Token); + await server.RestartAsync(); + } + + foreach (var log in logger.Logs.Where(x => x.EventId == NatsLogEvents.Protocol && x.LogLevel == LogLevel.Error)) + { + Assert.DoesNotContain("Unknown Protocol Operation", log.Message); + } + } + private sealed class NatsSubReconnectTest : NatsSubBase { private readonly Action _callback; diff --git a/tests/NATS.Client.TestUtilities/InMemoryTestLoggerFactory.cs b/tests/NATS.Client.TestUtilities/InMemoryTestLoggerFactory.cs new file mode 100644 index 000000000..1f47e0f3f --- /dev/null +++ b/tests/NATS.Client.TestUtilities/InMemoryTestLoggerFactory.cs @@ -0,0 +1,55 @@ +using Microsoft.Extensions.Logging; + +namespace NATS.Client.TestUtilities; + +public class InMemoryTestLoggerFactory(LogLevel level) : ILoggerFactory +{ + private readonly List _messages = new(); + + public IReadOnlyList Logs + { + get + { + lock (_messages) + return _messages.ToList(); + } + } + + public ILogger CreateLogger(string categoryName) => new TestLogger(categoryName, level, this); + + public void AddProvider(ILoggerProvider provider) + { + } + + public void Dispose() + { + } + + private void Log(string categoryName, LogLevel logLevel, EventId eventId, Exception? exception, string message) + { + lock (_messages) + _messages.Add(new LogMessage(categoryName, logLevel, eventId, exception, message)); + } + + public record LogMessage(string Category, LogLevel LogLevel, EventId EventId, Exception? Exception, string Message); + + private class TestLogger(string categoryName, LogLevel level, InMemoryTestLoggerFactory logger) : ILogger + { + public void Log(LogLevel logLevel, EventId eventId, TState state, Exception? exception, Func formatter) + { + if (logLevel >= level) + logger.Log(categoryName, logLevel, eventId, exception, formatter(state, exception)); + } + + public bool IsEnabled(LogLevel logLevel) => logLevel >= level; + + public IDisposable BeginScope(TState state) => new NullDisposable(); + + private class NullDisposable : IDisposable + { + public void Dispose() + { + } + } + } +} From 2604370ab41e1145ef93ced0a1f33a674616b92c Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Mon, 22 Jan 2024 22:10:14 +0000 Subject: [PATCH 02/25] Command writer pipeline encapsulation --- .../Commands/CommandWriter.cs | 794 +++++++++--------- src/NATS.Client.Core/Commands/PingCommand.cs | 52 +- .../Commands/ProtocolWriter.cs | 280 +++--- src/NATS.Client.Core/Internal/HeaderWriter.cs | 37 +- .../NatsPipeliningWriteProtocolProcessor.cs | 618 +++++++------- .../Internal/NatsReadProtocolProcessor.cs | 4 +- src/NATS.Client.Core/NatsConnection.Ping.cs | 14 +- src/NATS.Client.Core/NatsConnection.cs | 21 +- src/NATS.Client.Core/NatsLogEvents.cs | 1 + .../CancellationTest.cs | 2 +- .../NATS.Client.Core.Tests/NatsHeaderTest.cs | 21 +- 11 files changed, 939 insertions(+), 905 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 818b7b766..4e78bca61 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -1,40 +1,45 @@ +using System.Buffers; using System.IO.Pipelines; +using System.Net.Sockets; +using System.Numerics; using System.Runtime.CompilerServices; -using System.Threading.Channels; +using Microsoft.Extensions.Logging; using NATS.Client.Core.Internal; namespace NATS.Client.Core.Commands; /// -/// Used to track commands that have been enqueued to the PipeReader -/// -internal readonly record struct QueuedCommand(int Size, int Trim = 0); - -/// -/// Sets up a Pipe, and provides methods to write to the PipeWriter +/// Sets up a buffer (Pipe), and provides methods to write protocol messages to the buffer /// When methods complete, they have been queued for sending /// and further cancellation is not possible /// /// /// These methods are in the hot path, and have all been -/// optimized to eliminate allocations and make an initial attempt -/// to run synchronously without the async state machine +/// optimized to eliminate allocations and minimize copying /// internal sealed class CommandWriter : IAsyncDisposable { + private readonly ILogger _logger; + private readonly ObjectPool _pool; + private readonly object _lock = new(); + private readonly CancellationTokenSource _cts; private readonly ConnectionStatsCounter _counter; private readonly TimeSpan _defaultCommandTimeout; private readonly Action _enqueuePing; private readonly NatsOpts _opts; - private readonly PipeWriter _pipeWriter; private readonly ProtocolWriter _protocolWriter; - private readonly ChannelWriter _queuedCommandsWriter; private readonly SemaphoreSlim _semLock; - private Task? _flushTask; + private readonly Task _readerLoopTask; + private readonly HeaderWriter _headerWriter; + private ISocketConnection? _socketConnection; + private PipeReader? _pipeReader; + private PipeWriter? _pipeWriter; private bool _disposed; - public CommandWriter(NatsOpts opts, ConnectionStatsCounter counter, Action enqueuePing, TimeSpan? overrideCommandTimeout = default) + public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter counter, Action enqueuePing, TimeSpan? overrideCommandTimeout = default) { + _logger = opts.LoggerFactory.CreateLogger(); + _pool = pool; _counter = counter; _defaultCommandTimeout = overrideCommandTimeout ?? opts.CommandTimeout; _enqueuePing = enqueuePing; @@ -44,23 +49,39 @@ public CommandWriter(NatsOpts opts, ConnectionStatsCounter counter, Action(new UnboundedChannelOptions { SingleWriter = true, SingleReader = true }); + _protocolWriter = new ProtocolWriter(opts.SubjectEncoding); _semLock = new SemaphoreSlim(1); - QueuedCommandsReader = channel.Reader; - _queuedCommandsWriter = channel.Writer; + _headerWriter = new HeaderWriter(_opts.HeaderEncoding); + _cts = new CancellationTokenSource(); + _readerLoopTask = Task.Run(ReaderLoopAsync); } - public PipeReader PipeReader { get; } - - public ChannelReader QueuedCommandsReader { get; } + public void Reset(ISocketConnection socketConnection) + { + var pipe = new Pipe(new PipeOptions( + pauseWriterThreshold: _opts.WriterBufferSize, // flush will block after hitting + resumeWriterThreshold: _opts.WriterBufferSize / 2, + useSynchronizationContext: false)); - public Queue InFlightCommands { get; } = new(); + lock (_lock) + { + _pipeWriter?.Complete(); + _pipeReader?.Complete(); + _socketConnection = socketConnection; + _pipeReader = pipe.Reader; + _pipeWriter = pipe.Writer; + } + } public async ValueTask DisposeAsync() { +#if NET6_0 + _cts.Cancel(); +#else + await _cts.CancelAsync().ConfigureAwait(false); +#endif + await _semLock.WaitAsync().ConfigureAwait(false); try { @@ -70,8 +91,13 @@ public async ValueTask DisposeAsync() } _disposed = true; - _queuedCommandsWriter.Complete(); - await _pipeWriter.CompleteAsync().ConfigureAwait(false); + + if (_pipeWriter != null) + await _pipeWriter.CompleteAsync().ConfigureAwait(false); + if (_pipeReader != null) + await _pipeReader.CompleteAsync().ConfigureAwait(false); + + await _readerLoopTask.ConfigureAwait(false); } finally { @@ -79,24 +105,10 @@ public async ValueTask DisposeAsync() } } - public NatsPipeliningWriteProtocolProcessor CreateNatsPipeliningWriteProtocolProcessor(ISocketConnection socketConnection) => new(socketConnection, this, _opts, _counter); - - public ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken cancellationToken) + public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken cancellationToken) { -#pragma warning disable CA2016 -#pragma warning disable VSTHRD103 - if (!_semLock.Wait(0)) -#pragma warning restore VSTHRD103 -#pragma warning restore CA2016 - { - return ConnectStateMachineAsync(false, connectOpts, cancellationToken); - } - - if (_flushTask is { IsCompletedSuccessfully: false }) - { - return ConnectStateMachineAsync(true, connectOpts, cancellationToken); - } - + Interlocked.Add(ref _counter.PendingMessages, 1); + await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { if (_disposed) @@ -104,41 +116,28 @@ public ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken cancella throw new ObjectDisposedException(nameof(CommandWriter)); } - var success = false; - try - { - _protocolWriter.WriteConnect(connectOpts); - success = true; - } - finally + PipeWriter bw; + lock (_lock) { - EnqueueCommand(success); + if (_pipeWriter == null) + ThrowOnDisconnected(); + bw = _pipeWriter!; } + + _protocolWriter.WriteConnect(bw, connectOpts!); + await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { _semLock.Release(); + Interlocked.Add(ref _counter.PendingMessages, -1); } - - return ValueTask.CompletedTask; } - public ValueTask PingAsync(PingCommand pingCommand, CancellationToken cancellationToken) + public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken cancellationToken) { -#pragma warning disable CA2016 -#pragma warning disable VSTHRD103 - if (!_semLock.Wait(0)) -#pragma warning restore VSTHRD103 -#pragma warning restore CA2016 - { - return PingStateMachineAsync(false, pingCommand, cancellationToken); - } - - if (_flushTask is { IsCompletedSuccessfully: false }) - { - return PingStateMachineAsync(true, pingCommand, cancellationToken); - } - + Interlocked.Add(ref _counter.PendingMessages, 1); + await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { if (_disposed) @@ -146,42 +145,30 @@ public ValueTask PingAsync(PingCommand pingCommand, CancellationToken cancellati throw new ObjectDisposedException(nameof(CommandWriter)); } - var success = false; - try - { - _protocolWriter.WritePing(); - _enqueuePing(pingCommand); - success = true; - } - finally + PipeWriter bw; + lock (_lock) { - EnqueueCommand(success); + if (_pipeWriter == null) + ThrowOnDisconnected(); + bw = _pipeWriter!; } + + _enqueuePing(pingCommand); + + _protocolWriter.WritePing(bw); + await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { _semLock.Release(); + Interlocked.Add(ref _counter.PendingMessages, -1); } - - return ValueTask.CompletedTask; } - public ValueTask PongAsync(CancellationToken cancellationToken = default) + public async ValueTask PongAsync(CancellationToken cancellationToken = default) { -#pragma warning disable CA2016 -#pragma warning disable VSTHRD103 - if (!_semLock.Wait(0)) -#pragma warning restore VSTHRD103 -#pragma warning restore CA2016 - { - return PongStateMachineAsync(false, cancellationToken); - } - - if (_flushTask is { IsCompletedSuccessfully: false }) - { - return PongStateMachineAsync(true, cancellationToken); - } - + Interlocked.Add(ref _counter.PendingMessages, 1); + await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { if (_disposed) @@ -189,41 +176,46 @@ public ValueTask PongAsync(CancellationToken cancellationToken = default) throw new ObjectDisposedException(nameof(CommandWriter)); } - var success = false; - try + PipeWriter bw; + lock (_lock) { - _protocolWriter.WritePong(); - success = true; - } - finally - { - EnqueueCommand(success); + if (_pipeWriter == null) + ThrowOnDisconnected(); + bw = _pipeWriter!; } + + _protocolWriter.WritePong(bw); + await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { _semLock.Release(); } - - return ValueTask.CompletedTask; } public ValueTask PublishAsync(string subject, T? value, NatsHeaders? headers, string? replyTo, INatsSerialize serializer, CancellationToken cancellationToken) { -#pragma warning disable CA2016 -#pragma warning disable VSTHRD103 - if (!_semLock.Wait(0)) -#pragma warning restore VSTHRD103 -#pragma warning restore CA2016 + NatsPooledBufferWriter? headersBuffer = null; + if (headers != null) { - return PublishStateMachineAsync(false, subject, value, headers, replyTo, serializer, cancellationToken); + if (!_pool.TryRent(out headersBuffer)) + headersBuffer = new NatsPooledBufferWriter(); + _headerWriter.Write(headersBuffer, headers); } - if (_flushTask is { IsCompletedSuccessfully: false }) - { - return PublishStateMachineAsync(true, subject, value, headers, replyTo, serializer, cancellationToken); - } + NatsPooledBufferWriter payloadBuffer; + if (!_pool.TryRent(out payloadBuffer!)) + payloadBuffer = new NatsPooledBufferWriter(); + if (value != null) + serializer.Serialize(payloadBuffer, value); + + return PublishLockedAsync(subject, replyTo, payloadBuffer, headersBuffer, cancellationToken); + } + public async ValueTask SubscribeAsync(int sid, string subject, string? queueGroup, int? maxMsgs, CancellationToken cancellationToken) + { + Interlocked.Add(ref _counter.PendingMessages, 1); + await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { if (_disposed) @@ -231,42 +223,28 @@ public ValueTask PublishAsync(string subject, T? value, NatsHeaders? headers, throw new ObjectDisposedException(nameof(CommandWriter)); } - var trim = 0; - var success = false; - try - { - trim = _protocolWriter.WritePublish(subject, value, headers, replyTo, serializer); - success = true; - } - finally + PipeWriter bw; + lock (_lock) { - EnqueueCommand(success, trim: trim); + if (_pipeWriter == null) + ThrowOnDisconnected(); + bw = _pipeWriter!; } + + _protocolWriter.WriteSubscribe(bw, sid, subject, queueGroup, maxMsgs); + await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { _semLock.Release(); + Interlocked.Add(ref _counter.PendingMessages, -1); } - - return ValueTask.CompletedTask; } - public ValueTask SubscribeAsync(int sid, string subject, string? queueGroup, int? maxMsgs, CancellationToken cancellationToken) + public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken cancellationToken) { -#pragma warning disable CA2016 -#pragma warning disable VSTHRD103 - if (!_semLock.Wait(0)) -#pragma warning restore VSTHRD103 -#pragma warning restore CA2016 - { - return SubscribeStateMachineAsync(false, sid, subject, queueGroup, maxMsgs, cancellationToken); - } - - if (_flushTask is { IsCompletedSuccessfully: false }) - { - return SubscribeStateMachineAsync(true, sid, subject, queueGroup, maxMsgs, cancellationToken); - } - + Interlocked.Add(ref _counter.PendingMessages, 1); + await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { if (_disposed) @@ -274,385 +252,361 @@ public ValueTask SubscribeAsync(int sid, string subject, string? queueGroup, int throw new ObjectDisposedException(nameof(CommandWriter)); } - var success = false; - try - { - _protocolWriter.WriteSubscribe(sid, subject, queueGroup, maxMsgs); - success = true; - } - finally + PipeWriter bw; + lock (_lock) { - EnqueueCommand(success); + if (_pipeWriter == null) + ThrowOnDisconnected(); + bw = _pipeWriter!; } + + _protocolWriter.WriteUnsubscribe(bw, sid, maxMsgs); + await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { _semLock.Release(); + Interlocked.Add(ref _counter.PendingMessages, -1); } - - return ValueTask.CompletedTask; } - public ValueTask UnsubscribeAsync(int sid, CancellationToken cancellationToken) - { -#pragma warning disable CA2016 -#pragma warning disable VSTHRD103 - if (!_semLock.Wait(0)) -#pragma warning restore VSTHRD103 -#pragma warning restore CA2016 - { - return UnsubscribeStateMachineAsync(false, sid, cancellationToken); - } - - if (_flushTask is { IsCompletedSuccessfully: false }) - { - return UnsubscribeStateMachineAsync(true, sid, cancellationToken); - } + [MethodImpl(MethodImplOptions.NoInlining)] + private static void ThrowOnDisconnected() => throw new NatsException("Connection hasn't been established yet."); + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] + private async ValueTask PublishLockedAsync(string subject, string? replyTo, NatsPooledBufferWriter payloadBuffer, NatsPooledBufferWriter? headersBuffer, CancellationToken cancellationToken) + { + // Interlocked.Add(ref _counter.PendingMessages, 1); + await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { - if (_disposed) - { - throw new ObjectDisposedException(nameof(CommandWriter)); - } + var payload = payloadBuffer.WrittenMemory; + var headers = headersBuffer?.WrittenMemory; - var success = false; - try - { - _protocolWriter.WriteUnsubscribe(sid, null); - success = true; - } - finally + // if (_disposed) + // { + // throw new ObjectDisposedException(nameof(CommandWriter)); + // } + + PipeWriter bw; + lock (_lock) { - EnqueueCommand(success); + if (_pipeWriter == null) + ThrowOnDisconnected(); + bw = _pipeWriter!; } - } - finally - { - _semLock.Release(); - } - return ValueTask.CompletedTask; - } + _protocolWriter.WritePublish(bw, subject, replyTo, headers, payload); - // only used for internal testing - internal async Task TestStallFlushAsync(TimeSpan timeSpan) - { - await _semLock.WaitAsync().ConfigureAwait(false); + payloadBuffer.Reset(); + _pool.Return(payloadBuffer); - try - { - if (_flushTask is { IsCompletedSuccessfully: false }) + if (headersBuffer != null) { - await _flushTask.ConfigureAwait(false); + headersBuffer.Reset(); + _pool.Return(headersBuffer); } - _flushTask = Task.Delay(timeSpan); + await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { _semLock.Release(); + // Interlocked.Add(ref _counter.PendingMessages, -1); } } - private async ValueTask ConnectStateMachineAsync(bool lockHeld, ClientOpts connectOpts, CancellationToken cancellationToken) + private async Task ReaderLoopAsync() { - if (!lockHeld) - { - if (!await _semLock.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false)) - { - throw new TimeoutException(); - } - } + // 8520 should fit into 6 packets on 1500 MTU TLS connection or 1 packet on 9000 MTU TLS connection + // assuming 40 bytes TCP overhead + 40 bytes TLS overhead per packet + const int maxSendMemoryLength = 8520; + var sendMemory = new Memory(new byte[maxSendMemoryLength]); - try + var cancellationToken = _cts.Token; + while (cancellationToken.IsCancellationRequested == false) { - if (_disposed) - { - throw new ObjectDisposedException(nameof(CommandWriter)); - } - - if (_flushTask is { IsCompletedSuccessfully: false }) - { - await _flushTask.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false); - } - - var success = false; try { - _protocolWriter.WriteConnect(connectOpts); - success = true; - } - finally - { - EnqueueCommand(success); + ISocketConnection? connection; + PipeReader? pipeReader; + lock (_lock) + { + connection = _socketConnection; + pipeReader = _pipeReader; + } + + if (connection == null || pipeReader == null) + { + await Task.Delay(10, cancellationToken).ConfigureAwait(false); + continue; + } + + var result = await pipeReader.ReadAsync(cancellationToken).ConfigureAwait(false); + var buffer = result.Buffer; + try + { + if (!buffer.IsEmpty) + { + Memory memory; + var length = (int)buffer.Length; + + if (length > maxSendMemoryLength) + { + length = maxSendMemoryLength; + buffer = buffer.Slice(0, buffer.GetPosition(length)); + memory = sendMemory; + } + else + { + memory = sendMemory.Slice(0, length); + } + + buffer.CopyTo(memory.Span); + + try + { + await connection.SendAsync(memory).ConfigureAwait(false); + } + catch (SocketException e) + { + connection.SignalDisconnected(e); + _logger.LogWarning(NatsLogEvents.TcpSocket, e, "Error while sending data"); + } + catch (Exception e) + { + connection.SignalDisconnected(e); + _logger.LogError(NatsLogEvents.TcpSocket, e, "Unexpected error while sending data"); + } + } + } + finally + { + pipeReader.AdvanceTo(buffer.End); + } + } + catch (InvalidOperationException) + { + // We might still be using the previous pipe reader which might be completed already + } + catch (OperationCanceledException) + { + // Expected during shutdown + } + catch (Exception e) + { + _logger.LogError(NatsLogEvents.Buffer, e, "Unexpected error in send buffer reader loop"); } } - finally - { - _semLock.Release(); - } } +} + +internal sealed class PriorityCommandWriter : IAsyncDisposable +{ + private int _disposed; + + public PriorityCommandWriter(ObjectPool pool, ISocketConnection socketConnection, NatsOpts opts, ConnectionStatsCounter counter, Action enqueuePing) + { + CommandWriter = new CommandWriter(pool, opts, counter, enqueuePing, overrideCommandTimeout: TimeSpan.MaxValue); + CommandWriter.Reset(socketConnection); + } + + public CommandWriter CommandWriter { get; } - private async ValueTask PingStateMachineAsync(bool lockHeld, PingCommand pingCommand, CancellationToken cancellationToken) + public async ValueTask DisposeAsync() { - if (!lockHeld) + if (Interlocked.Increment(ref _disposed) == 1) { - if (!await _semLock.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false)) - { - throw new TimeoutException(); - } + // disposing command writer marks pipe writer as complete + await CommandWriter.DisposeAsync().ConfigureAwait(false); } + } +} - try - { - if (_disposed) - { - throw new ObjectDisposedException(nameof(CommandWriter)); - } +internal sealed class NatsPooledBufferWriter : IBufferWriter, IObjectPoolNode> +{ + private const int DefaultInitialBufferSize = 256; - if (_flushTask is { IsCompletedSuccessfully: false }) - { - await _flushTask.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false); - } + private readonly ArrayPool _pool; + private T[]? _array; + private int _index; + private NatsPooledBufferWriter? _next; - var success = false; - try - { - _protocolWriter.WritePing(); - _enqueuePing(pingCommand); - success = true; - } - finally - { - EnqueueCommand(success); - } - } - finally - { - _semLock.Release(); - } + public NatsPooledBufferWriter() + { + _pool = ArrayPool.Shared; + _array = _pool.Rent(DefaultInitialBufferSize); + _index = 0; } - private async ValueTask PongStateMachineAsync(bool lockHeld, CancellationToken cancellationToken) + public ref NatsPooledBufferWriter? NextNode => ref _next; + + /// + /// Gets the data written to the underlying buffer so far, as a . + /// + public ReadOnlyMemory WrittenMemory { - if (!lockHeld) + get { - if (!await _semLock.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false)) + var array = _array; + + if (array is null) { - throw new TimeoutException(); + ThrowObjectDisposedException(); } + + return array!.AsMemory(0, _index); } + } - try + /// + /// Gets the data written to the underlying buffer so far, as a . + /// + public ReadOnlySpan WrittenSpan + { + get { - if (_disposed) - { - throw new ObjectDisposedException(nameof(CommandWriter)); - } + var array = _array; - if (_flushTask is { IsCompletedSuccessfully: false }) + if (array is null) { - await _flushTask.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false); + ThrowObjectDisposedException(); } - var success = false; - try - { - _protocolWriter.WritePong(); - success = true; - } - finally - { - EnqueueCommand(success); - } - } - finally - { - _semLock.Release(); + return array!.AsSpan(0, _index); } } - private async ValueTask PublishStateMachineAsync(bool lockHeld, string subject, T? value, NatsHeaders? headers, string? replyTo, INatsSerialize serializer, CancellationToken cancellationToken) + /// + /// Gets the amount of data written to the underlying buffer so far. + /// + public int WrittenCount + { + get => _index; + } + + /// + public void Advance(int count) { - if (!lockHeld) + var array = _array; + + if (array is null) { - if (!await _semLock.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false)) - { - throw new TimeoutException(); - } + ThrowObjectDisposedException(); } - try + if (count < 0) { - if (_disposed) - { - throw new ObjectDisposedException(nameof(CommandWriter)); - } - - if (_flushTask is { IsCompletedSuccessfully: false }) - { - await _flushTask.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false); - } - - var trim = 0; - var success = false; - try - { - trim = _protocolWriter.WritePublish(subject, value, headers, replyTo, serializer); - success = true; - } - finally - { - EnqueueCommand(success, trim: trim); - } + ThrowArgumentOutOfRangeExceptionForNegativeCount(); } - finally + + if (_index > array!.Length - count) { - _semLock.Release(); + ThrowArgumentExceptionForAdvancedTooFar(); } + + _index += count; } - private async ValueTask SubscribeStateMachineAsync(bool lockHeld, int sid, string subject, string? queueGroup, int? maxMsgs, CancellationToken cancellationToken) + /// + public Memory GetMemory(int sizeHint = 0) { - if (!lockHeld) - { - if (!await _semLock.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false)) - { - throw new TimeoutException(); - } - } + CheckBufferAndEnsureCapacity(sizeHint); - try - { - if (_disposed) - { - throw new ObjectDisposedException(nameof(CommandWriter)); - } + return _array.AsMemory(_index); + } - if (_flushTask is { IsCompletedSuccessfully: false }) - { - await _flushTask.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false); - } + /// + public Span GetSpan(int sizeHint = 0) + { + CheckBufferAndEnsureCapacity(sizeHint); - var success = false; - try - { - _protocolWriter.WriteSubscribe(sid, subject, queueGroup, maxMsgs); - success = true; - } - finally - { - EnqueueCommand(success); - } - } - finally - { - _semLock.Release(); - } + return _array.AsSpan(_index); + } + + public void Reset() + { + if (_array != null) + _pool.Return(_array); + _array = _pool.Rent(DefaultInitialBufferSize); + _index = 0; } - private async ValueTask UnsubscribeStateMachineAsync(bool lockHeld, int sid, CancellationToken cancellationToken) + /// + public override string ToString() { - if (!lockHeld) + // See comments in MemoryOwner about this + if (typeof(T) == typeof(char) && + _array is char[] chars) { - if (!await _semLock.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false)) - { - throw new TimeoutException(); - } + return new(chars, 0, _index); } - try - { - if (_disposed) - { - throw new ObjectDisposedException(nameof(CommandWriter)); - } + // Same representation used in Span + return $"NatsPooledBufferWriter<{typeof(T)}>[{_index}]"; + } - if (_flushTask is { IsCompletedSuccessfully: false }) - { - await _flushTask.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false); - } + [MethodImpl(MethodImplOptions.NoInlining)] + private static void ThrowArgumentOutOfRangeExceptionForNegativeCount() => throw new ArgumentOutOfRangeException("count", "The count can't be a negative value."); - var success = false; - try - { - _protocolWriter.WriteUnsubscribe(sid, null); - success = true; - } - finally - { - EnqueueCommand(success); - } - } - finally - { - _semLock.Release(); - } - } + [MethodImpl(MethodImplOptions.NoInlining)] + private static void ThrowArgumentOutOfRangeExceptionForNegativeSizeHint() => throw new ArgumentOutOfRangeException("sizeHint", "The size hint can't be a negative value."); + + [MethodImpl(MethodImplOptions.NoInlining)] + private static void ThrowArgumentExceptionForAdvancedTooFar() => throw new ArgumentException("The buffer writer has advanced too far."); + + [MethodImpl(MethodImplOptions.NoInlining)] + private static void ThrowObjectDisposedException() => throw new ObjectDisposedException("The current buffer has already been disposed."); /// - /// Enqueues a command, and kicks off a flush + /// Ensures that has enough free space to contain a given number of new items. /// - /// - /// Whether the command was successful - /// If true, it will be sent on the wire - /// If false, it will be thrown out - /// - /// - /// Number of bytes to skip from beginning of message - /// when sending on the wire - /// + /// The minimum number of items to ensure space for in . [MethodImpl(MethodImplOptions.AggressiveInlining)] - private void EnqueueCommand(bool success, int trim = 0) + private void CheckBufferAndEnsureCapacity(int sizeHint) { - var size = (int)_pipeWriter.UnflushedBytes; - if (size == 0) + var array = _array; + + if (array is null) { - // no unflushed bytes means no command was produced - _flushTask = null; - return; + ThrowObjectDisposedException(); } - if (success) + if (sizeHint < 0) { - Interlocked.Add(ref _counter.PendingMessages, 1); + ThrowArgumentOutOfRangeExceptionForNegativeSizeHint(); } - _queuedCommandsWriter.TryWrite(new QueuedCommand(Size: size, Trim: success ? trim : size)); - var flush = _pipeWriter.FlushAsync(); - _flushTask = flush.IsCompletedSuccessfully ? null : flush.AsTask(); - } -} - -internal sealed class PriorityCommandWriter : IAsyncDisposable -{ - private readonly NatsPipeliningWriteProtocolProcessor _natsPipeliningWriteProtocolProcessor; - private int _disposed; + if (sizeHint == 0) + { + sizeHint = 1; + } - public PriorityCommandWriter(ISocketConnection socketConnection, NatsOpts opts, ConnectionStatsCounter counter, Action enqueuePing) - { - CommandWriter = new CommandWriter(opts, counter, enqueuePing, overrideCommandTimeout: TimeSpan.MaxValue); - _natsPipeliningWriteProtocolProcessor = CommandWriter.CreateNatsPipeliningWriteProtocolProcessor(socketConnection); + if (sizeHint > array!.Length - _index) + { + ResizeBuffer(sizeHint); + } } - public CommandWriter CommandWriter { get; } - - public async ValueTask DisposeAsync() + /// + /// Resizes to ensure it can fit the specified number of new items. + /// + /// The minimum number of items to ensure space for in . + [MethodImpl(MethodImplOptions.NoInlining)] + private void ResizeBuffer(int sizeHint) { - if (Interlocked.Increment(ref _disposed) == 1) + var minimumSize = (uint)_index + (uint)sizeHint; + + // The ArrayPool class has a maximum threshold of 1024 * 1024 for the maximum length of + // pooled arrays, and once this is exceeded it will just allocate a new array every time + // of exactly the requested size. In that case, we manually round up the requested size to + // the nearest power of two, to ensure that repeated consecutive writes when the array in + // use is bigger than that threshold don't end up causing a resize every single time. + if (minimumSize > 1024 * 1024) { - // disposing command writer marks pipe writer as complete - await CommandWriter.DisposeAsync().ConfigureAwait(false); - try - { - // write loop will complete once pipe reader completes - await _natsPipeliningWriteProtocolProcessor.WriteLoop.ConfigureAwait(false); - } - finally - { - await _natsPipeliningWriteProtocolProcessor.DisposeAsync().ConfigureAwait(false); - } + minimumSize = BitOperations.RoundUpToPowerOf2(minimumSize); } + + _pool.Resize(ref _array, (int)minimumSize); } } diff --git a/src/NATS.Client.Core/Commands/PingCommand.cs b/src/NATS.Client.Core/Commands/PingCommand.cs index 1861adc0b..a8c6236d4 100644 --- a/src/NATS.Client.Core/Commands/PingCommand.cs +++ b/src/NATS.Client.Core/Commands/PingCommand.cs @@ -1,12 +1,56 @@ +using System.Threading.Tasks.Sources; +using NATS.Client.Core.Internal; + namespace NATS.Client.Core.Commands; -internal struct PingCommand +internal class PingCommand : IValueTaskSource, IObjectPoolNode { - public PingCommand() + private readonly ObjectPool? _pool; + private DateTimeOffset _start; + private ManualResetValueTaskSourceCore _core; + private PingCommand? _next; + + public PingCommand(ObjectPool? pool) { + _pool = pool; + _core = new ManualResetValueTaskSourceCore + { + RunContinuationsAsynchronously = true, + }; + _start = DateTimeOffset.MinValue; + } + + public ref PingCommand? NextNode => ref _next; + + public void Start() => _start = DateTimeOffset.UtcNow; + + public void SetResult() => _core.SetResult(DateTimeOffset.UtcNow - _start); + + public void SetCanceled() => _core.SetException(new OperationCanceledException()); + + public void Reset() + { + _start = DateTimeOffset.MinValue; + _core.Reset(); + } + + public ValueTask RunAsync() => new(this, _core.Version); + + public TimeSpan GetResult(short token) + { + var result = _core.GetResult(token); + + if (_pool is not null) + { + Reset(); + _pool.Return(this); + } + + return result; } - public DateTimeOffset WriteTime { get; } = DateTimeOffset.UtcNow; + public ValueTaskSourceStatus GetStatus(short token) => _core.GetStatus(token); - public TaskCompletionSource TaskCompletionSource { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + public void OnCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) + => _core.OnCompleted(continuation, state, token, flags); } diff --git a/src/NATS.Client.Core/Commands/ProtocolWriter.cs b/src/NATS.Client.Core/Commands/ProtocolWriter.cs index 4ddc439d3..264b035ef 100644 --- a/src/NATS.Client.Core/Commands/ProtocolWriter.cs +++ b/src/NATS.Client.Core/Commands/ProtocolWriter.cs @@ -1,6 +1,6 @@ +using System.Buffers; using System.Buffers.Binary; using System.Buffers.Text; -using System.IO.Pipelines; using System.Runtime.CompilerServices; using System.Text; using System.Text.Json; @@ -36,47 +36,40 @@ internal sealed class ProtocolWriter private static readonly ulong PongNewLine = BinaryPrimitives.ReadUInt64LittleEndian("PONG\r\n "u8); private static readonly ulong UnsubSpace = BinaryPrimitives.ReadUInt64LittleEndian("UNSUB "u8); - private readonly PipeWriter _writer; - private readonly HeaderWriter _headerWriter; private readonly Encoding _subjectEncoding; - public ProtocolWriter(PipeWriter writer, Encoding subjectEncoding, Encoding headerEncoding) - { - _writer = writer; - _subjectEncoding = subjectEncoding; - _headerWriter = new HeaderWriter(writer, headerEncoding); - } + public ProtocolWriter(Encoding subjectEncoding) => _subjectEncoding = subjectEncoding; // https://docs.nats.io/reference/reference-protocols/nats-protocol#connect // CONNECT {["option_name":option_value],...} - public void WriteConnect(ClientOpts opts) + public void WriteConnect(IBufferWriter writer, ClientOpts opts) { - var span = _writer.GetSpan(UInt64Length); + var span = writer.GetSpan(UInt64Length); BinaryPrimitives.WriteUInt64LittleEndian(span, ConnectSpace); - _writer.Advance(ConnectSpaceLength); + writer.Advance(ConnectSpaceLength); - var jsonWriter = new Utf8JsonWriter(_writer); + var jsonWriter = new Utf8JsonWriter(writer); JsonSerializer.Serialize(jsonWriter, opts, JsonContext.Default.ClientOpts); - span = _writer.GetSpan(UInt16Length); + span = writer.GetSpan(UInt16Length); BinaryPrimitives.WriteUInt16LittleEndian(span, NewLine); - _writer.Advance(NewLineLength); + writer.Advance(NewLineLength); } // https://docs.nats.io/reference/reference-protocols/nats-protocol#ping-pong - public void WritePing() + public void WritePing(IBufferWriter writer) { - var span = _writer.GetSpan(UInt64Length); + var span = writer.GetSpan(UInt64Length); BinaryPrimitives.WriteUInt64LittleEndian(span, PingNewLine); - _writer.Advance(PingNewLineLength); + writer.Advance(PingNewLineLength); } // https://docs.nats.io/reference/reference-protocols/nats-protocol#ping-pong - public void WritePong() + public void WritePong(IBufferWriter writer) { - var span = _writer.GetSpan(UInt64Length); + var span = writer.GetSpan(UInt64Length); BinaryPrimitives.WriteUInt64LittleEndian(span, PongNewLine); - _writer.Advance(PongNewLineLength); + writer.Advance(PongNewLineLength); } // https://docs.nats.io/reference/reference-protocols/nats-protocol#pub @@ -84,124 +77,21 @@ public void WritePong() // or // https://docs.nats.io/reference/reference-protocols/nats-protocol#hpub // HPUB [reply-to] <#header bytes> <#total bytes>\r\n[headers]\r\n\r\n[payload]\r\n - // - // returns the number of bytes that should be skipped when writing to the wire - public int WritePublish(string subject, T? value, NatsHeaders? headers, string? replyTo, INatsSerialize serializer) + public void WritePublish(IBufferWriter writer, string subject, string? replyTo, ReadOnlyMemory? headers, ReadOnlyMemory payload) { - int ctrlLen; - if (headers == null) - { - // 'PUB ' + subject +' '+ payload len +'\r\n' - ctrlLen = PubSpaceLength + _subjectEncoding.GetByteCount(subject) + 1 + MaxIntStringLength + NewLineLength; - } - else - { - // 'HPUB ' + subject +' '+ header len +' '+ payload len +'\r\n' - ctrlLen = HpubSpaceLength + _subjectEncoding.GetByteCount(subject) + 1 + MaxIntStringLength + 1 + MaxIntStringLength + NewLineLength; - } - - if (replyTo != null) - { - // len += replyTo +' ' - ctrlLen += _subjectEncoding.GetByteCount(replyTo) + 1; - } - - var ctrlSpan = _writer.GetSpan(ctrlLen); - var span = ctrlSpan; - if (headers == null) - { - BinaryPrimitives.WriteUInt32LittleEndian(span, PubSpace); - span = span[PubSpaceLength..]; - } - else - { - BinaryPrimitives.WriteUInt64LittleEndian(span, HpubSpace); - span = span[HpubSpaceLength..]; - } - - var written = _subjectEncoding.GetBytes(subject, span); - span[written] = (byte)' '; - span = span[(written + 1)..]; - - if (replyTo != null) - { - written = _subjectEncoding.GetBytes(replyTo, span); - span[written] = (byte)' '; - span = span[(written + 1)..]; - } - - Span lenSpan; if (headers == null) { - // len = payload len - lenSpan = span[..MaxIntStringLength]; - span = span[lenSpan.Length..]; + WritePub(writer, subject, replyTo, payload); } else { - // len = header len +' '+ payload len - lenSpan = span[..(MaxIntStringLength + 1 + MaxIntStringLength)]; - span = span[lenSpan.Length..]; + WriteHpub(writer, subject, replyTo, headers.Value, payload); } - - BinaryPrimitives.WriteUInt16LittleEndian(span, NewLine); - _writer.Advance(ctrlLen); - - var headersLength = 0L; - var totalLength = 0L; - if (headers != null) - { - headersLength = _headerWriter.Write(headers); - totalLength += headersLength; - } - - // Consider null as empty payload. This way we are able to transmit null values as sentinels. - // Another point is serializer behaviour. For instance JSON serializer seems to serialize null - // as a string "null", others might throw exception. - if (value != null) - { - var initialCount = _writer.UnflushedBytes; - serializer.Serialize(_writer, value); - totalLength += _writer.UnflushedBytes - initialCount; - } - - span = _writer.GetSpan(UInt16Length); - BinaryPrimitives.WriteUInt16LittleEndian(span, NewLine); - _writer.Advance(NewLineLength); - - // write the length - var lenWritten = 0; - if (headers != null) - { - if (!Utf8Formatter.TryFormat(headersLength, lenSpan, out lenWritten)) - { - ThrowOnUtf8FormatFail(); - } - - lenSpan[lenWritten] = (byte)' '; - lenWritten += 1; - } - - if (!Utf8Formatter.TryFormat(totalLength, lenSpan[lenWritten..], out var tLen)) - { - ThrowOnUtf8FormatFail(); - } - - lenWritten += tLen; - var trim = lenSpan.Length - lenWritten; - if (trim > 0) - { - // shift right - ctrlSpan[..(ctrlLen - trim - NewLineLength)].CopyTo(ctrlSpan[trim..]); - ctrlSpan[..trim].Clear(); - } - - return trim; } // https://docs.nats.io/reference/reference-protocols/nats-protocol#sub // SUB [queue group] - public void WriteSubscribe(int sid, string subject, string? queueGroup, int? maxMsgs) + public void WriteSubscribe(IBufferWriter writer, int sid, string subject, string? queueGroup, int? maxMsgs) { // 'SUB ' + subject +' '+ sid +'\r\n' var ctrlLen = SubSpaceLength + _subjectEncoding.GetByteCount(subject) + 1 + MaxIntStringLength + NewLineLength; @@ -212,7 +102,7 @@ public void WriteSubscribe(int sid, string subject, string? queueGroup, int? max ctrlLen += _subjectEncoding.GetByteCount(queueGroup) + 1; } - var span = _writer.GetSpan(ctrlLen); + var span = writer.GetSpan(ctrlLen); BinaryPrimitives.WriteUInt32LittleEndian(span, SubSpace); var size = SubSpaceLength; span = span[SubSpaceLength..]; @@ -241,20 +131,20 @@ public void WriteSubscribe(int sid, string subject, string? queueGroup, int? max BinaryPrimitives.WriteUInt16LittleEndian(span, NewLine); size += NewLineLength; - _writer.Advance(size); + writer.Advance(size); // Immediately send UNSUB to minimize the risk of // receiving more messages than in case they are published // between our SUB and UNSUB calls. if (maxMsgs != null) { - WriteUnsubscribe(sid, maxMsgs); + WriteUnsubscribe(writer, sid, maxMsgs); } } // https://docs.nats.io/reference/reference-protocols/nats-protocol#unsub // UNSUB [max_msgs] - public void WriteUnsubscribe(int sid, int? maxMessages) + public void WriteUnsubscribe(IBufferWriter writer, int sid, int? maxMessages) { // 'UNSUB ' + sid +'\r\n' var ctrlLen = UnsubSpaceLength + MaxIntStringLength + NewLineLength; @@ -264,7 +154,7 @@ public void WriteUnsubscribe(int sid, int? maxMessages) ctrlLen += 1 + MaxIntStringLength; } - var span = _writer.GetSpan(ctrlLen); + var span = writer.GetSpan(ctrlLen); BinaryPrimitives.WriteUInt64LittleEndian(span, UnsubSpace); var size = UnsubSpaceLength; span = span[UnsubSpaceLength..]; @@ -291,10 +181,132 @@ public void WriteUnsubscribe(int sid, int? maxMessages) BinaryPrimitives.WriteUInt16LittleEndian(span, NewLine); size += NewLineLength; - _writer.Advance(size); + writer.Advance(size); } // optimization detailed here: https://github.com/nats-io/nats.net.v2/issues/320#issuecomment-1886165748 [MethodImpl(MethodImplOptions.NoInlining)] private static void ThrowOnUtf8FormatFail() => throw new NatsException("Can not format integer."); + + // PUB [reply-to] <#bytes>\r\n[payload]\r\n + private void WritePub(IBufferWriter writer, string subject, string? replyTo, ReadOnlyMemory payload) + { + Span spanPayloadLength = stackalloc byte[MaxIntStringLength]; + if (!Utf8Formatter.TryFormat(payload.Length, spanPayloadLength, out var payloadLengthWritten)) + { + ThrowOnUtf8FormatFail(); + } + + spanPayloadLength = spanPayloadLength.Slice(0, payloadLengthWritten); + + var total = PubSpaceLength; + + var subjectSpaceLength = _subjectEncoding.GetByteCount(subject) + 1; + total += subjectSpaceLength; + + var replyToLengthSpace = 0; + if (replyTo != null) + { + replyToLengthSpace = _subjectEncoding.GetByteCount(replyTo) + 1; + total += replyToLengthSpace; + } + + total += spanPayloadLength.Length + NewLineLength + payload.Length + NewLineLength; + + var span = writer.GetSpan(total); + + BinaryPrimitives.WriteUInt32LittleEndian(span, PubSpace); + span = span.Slice(PubSpaceLength); + + _subjectEncoding.GetBytes(subject, span); + span[subjectSpaceLength - 1] = (byte)' '; + span = span.Slice(subjectSpaceLength); + + if (replyTo != null) + { + _subjectEncoding.GetBytes(replyTo, span); + span[replyToLengthSpace - 1] = (byte)' '; + span = span.Slice(replyToLengthSpace); + } + + spanPayloadLength.CopyTo(span); + span = span.Slice(spanPayloadLength.Length); + + BinaryPrimitives.WriteUInt16LittleEndian(span, NewLine); + span = span.Slice(NewLineLength); + + payload.Span.CopyTo(span); + span = span.Slice(payload.Length); + + BinaryPrimitives.WriteUInt16LittleEndian(span, NewLine); + + writer.Advance(total); + } + + // HPUB [reply-to] <#header bytes> <#total bytes>\r\n[headers]\r\n\r\n[payload]\r\n + private void WriteHpub(IBufferWriter writer, string subject, string? replyTo, ReadOnlyMemory headers, ReadOnlyMemory payload) + { + Span spanPayloadLength = stackalloc byte[MaxIntStringLength]; + if (!Utf8Formatter.TryFormat(payload.Length, spanPayloadLength, out var payloadLengthWritten)) + { + ThrowOnUtf8FormatFail(); + } + + spanPayloadLength = spanPayloadLength.Slice(0, payloadLengthWritten); + + Span spanHeadersLength = stackalloc byte[MaxIntStringLength + 1]; + if (!Utf8Formatter.TryFormat(headers.Length, spanHeadersLength, out var headersLengthWritten)) + { + ThrowOnUtf8FormatFail(); + } + + spanHeadersLength = spanHeadersLength.Slice(0, headersLengthWritten + 1); + spanHeadersLength[headersLengthWritten] = (byte)' '; + + var total = HpubSpaceLength; + + var subjectSpaceLength = _subjectEncoding.GetByteCount(subject) + 1; + total += subjectSpaceLength; + + var replyToLengthSpace = 0; + if (replyTo != null) + { + replyToLengthSpace = _subjectEncoding.GetByteCount(replyTo) + 1; + total += replyToLengthSpace; + } + + total += spanHeadersLength.Length + spanPayloadLength.Length + NewLineLength + payload.Length + NewLineLength; + + var span = writer.GetSpan(total); + + BinaryPrimitives.WriteUInt64LittleEndian(span, HpubSpace); + span = span.Slice(HpubSpaceLength); + + _subjectEncoding.GetBytes(subject, span); + span[subjectSpaceLength - 1] = (byte)' '; + span = span.Slice(subjectSpaceLength); + + if (replyTo != null) + { + _subjectEncoding.GetBytes(replyTo, span); + span[replyToLengthSpace - 1] = (byte)' '; + span = span.Slice(replyToLengthSpace); + } + + spanHeadersLength.CopyTo(span); + span = span.Slice(spanHeadersLength.Length); + + spanPayloadLength.CopyTo(span); + span = span.Slice(spanPayloadLength.Length); + + BinaryPrimitives.WriteUInt16LittleEndian(span, NewLine); + span = span.Slice(NewLineLength); + + payload.Span.CopyTo(span); + span = span.Slice(payload.Length); + + BinaryPrimitives.WriteUInt16LittleEndian(span, NewLine); + + writer.Advance(total); + } } diff --git a/src/NATS.Client.Core/Internal/HeaderWriter.cs b/src/NATS.Client.Core/Internal/HeaderWriter.cs index 47b0d8c2e..8bb62694d 100644 --- a/src/NATS.Client.Core/Internal/HeaderWriter.cs +++ b/src/NATS.Client.Core/Internal/HeaderWriter.cs @@ -12,23 +12,19 @@ internal class HeaderWriter private const byte ByteColon = (byte)':'; private const byte ByteSpace = (byte)' '; private const byte ByteDel = 127; - private readonly PipeWriter _pipeWriter; + private readonly Encoding _encoding; - public HeaderWriter(PipeWriter pipeWriter, Encoding encoding) - { - _pipeWriter = pipeWriter; - _encoding = encoding; - } + public HeaderWriter(Encoding encoding) => _encoding = encoding; private static ReadOnlySpan CrLf => new[] { ByteCr, ByteLf }; private static ReadOnlySpan ColonSpace => new[] { ByteColon, ByteSpace }; - internal long Write(NatsHeaders headers) + internal long Write(IBufferWriter bufferWriter, NatsHeaders headers) { - var initialCount = _pipeWriter.UnflushedBytes; - _pipeWriter.WriteSpan(CommandConstants.NatsHeaders10NewLine); + bufferWriter.WriteSpan(CommandConstants.NatsHeaders10NewLine); + var len = CommandConstants.NatsHeaders10NewLine.Length; foreach (var kv in headers) { @@ -38,7 +34,7 @@ internal long Write(NatsHeaders headers) { // write key var keyLength = _encoding.GetByteCount(kv.Key); - var keySpan = _pipeWriter.GetSpan(keyLength); + var keySpan = bufferWriter.GetSpan(keyLength); _encoding.GetBytes(kv.Key, keySpan); if (!ValidateKey(keySpan.Slice(0, keyLength))) { @@ -46,20 +42,26 @@ internal long Write(NatsHeaders headers) $"Invalid header key '{kv.Key}': contains colon, space, or other non-printable ASCII characters"); } - _pipeWriter.Advance(keyLength); - _pipeWriter.Write(ColonSpace); + bufferWriter.Advance(keyLength); + len += keyLength; + + bufferWriter.Write(ColonSpace); + len += ColonSpace.Length; // write values var valueLength = _encoding.GetByteCount(value); - var valueSpan = _pipeWriter.GetSpan(valueLength); + var valueSpan = bufferWriter.GetSpan(valueLength); _encoding.GetBytes(value, valueSpan); if (!ValidateValue(valueSpan.Slice(0, valueLength))) { throw new NatsException($"Invalid header value for key '{kv.Key}': contains CRLF"); } - _pipeWriter.Advance(valueLength); - _pipeWriter.Write(CrLf); + bufferWriter.Advance(valueLength); + len += valueLength; + + bufferWriter.Write(CrLf); + len += CrLf.Length; } } } @@ -67,9 +69,10 @@ internal long Write(NatsHeaders headers) // Even empty header needs to terminate. // We will send NATS/1.0 version line // even if there are no headers. - _pipeWriter.Write(CrLf); + bufferWriter.Write(CrLf); + len += CrLf.Length; - return _pipeWriter.UnflushedBytes - initialCount; + return len; } // cannot contain ASCII Bytes <=32, 58, or 127 diff --git a/src/NATS.Client.Core/Internal/NatsPipeliningWriteProtocolProcessor.cs b/src/NATS.Client.Core/Internal/NatsPipeliningWriteProtocolProcessor.cs index b26e144ec..922cb5b95 100644 --- a/src/NATS.Client.Core/Internal/NatsPipeliningWriteProtocolProcessor.cs +++ b/src/NATS.Client.Core/Internal/NatsPipeliningWriteProtocolProcessor.cs @@ -1,309 +1,309 @@ -using System.Buffers; -using System.Diagnostics; -using System.IO.Pipelines; -using System.Threading.Channels; -using Microsoft.Extensions.Logging; -using NATS.Client.Core.Commands; - -namespace NATS.Client.Core.Internal; - -internal sealed class NatsPipeliningWriteProtocolProcessor : IAsyncDisposable -{ - private readonly CancellationTokenSource _cts; - private readonly ConnectionStatsCounter _counter; - private readonly NatsOpts _opts; - private readonly PipeReader _pipeReader; - private readonly Queue _inFlightCommands; - private readonly ChannelReader _queuedCommandReader; - private readonly ISocketConnection _socketConnection; - private readonly Stopwatch _stopwatch = new Stopwatch(); - private int _disposed; - - public NatsPipeliningWriteProtocolProcessor(ISocketConnection socketConnection, CommandWriter commandWriter, NatsOpts opts, ConnectionStatsCounter counter) - { - _cts = new CancellationTokenSource(); - _counter = counter; - _inFlightCommands = commandWriter.InFlightCommands; - _opts = opts; - _pipeReader = commandWriter.PipeReader; - _queuedCommandReader = commandWriter.QueuedCommandsReader; - _socketConnection = socketConnection; - WriteLoop = Task.Run(WriteLoopAsync); - } - - public Task WriteLoop { get; } - - public async ValueTask DisposeAsync() - { - if (Interlocked.Increment(ref _disposed) == 1) - { -#if NET6_0 - _cts.Cancel(); -#else - await _cts.CancelAsync().ConfigureAwait(false); -#endif - try - { - await WriteLoop.ConfigureAwait(false); // wait to drain writer - } - catch - { - // ignore - } - } - } - - private async Task WriteLoopAsync() - { - var logger = _opts.LoggerFactory.CreateLogger(); - var isEnabledTraceLogging = logger.IsEnabled(LogLevel.Trace); - var cancellationToken = _cts.Token; - var pending = 0; - var trimming = 0; - var examinedOffset = 0; - - // memory segment used to consolidate multiple small memory chunks - // should <= (minimumSegmentSize * 0.5) in CommandWriter - // 8520 should fit into 6 packets on 1500 MTU TLS connection or 1 packet on 9000 MTU TLS connection - // assuming 40 bytes TCP overhead + 40 bytes TLS overhead per packet - var consolidateMem = new Memory(new byte[8520]); - - // add up in flight command sum - var inFlightSum = 0; - foreach (var command in _inFlightCommands) - { - inFlightSum += command.Size; - } - - try - { - while (true) - { - var result = await _pipeReader.ReadAsync(cancellationToken).ConfigureAwait(false); - if (result.IsCanceled) - { - break; - } - - var consumedPos = result.Buffer.Start; - var examinedPos = result.Buffer.Start; - try - { - var buffer = result.Buffer.Slice(examinedOffset); - while (inFlightSum < result.Buffer.Length) - { - QueuedCommand queuedCommand; - while (!_queuedCommandReader.TryRead(out queuedCommand)) - { - await _queuedCommandReader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); - } - - _inFlightCommands.Enqueue(queuedCommand); - inFlightSum += queuedCommand.Size; - } - - while (buffer.Length > 0) - { - if (pending == 0) - { - var peek = _inFlightCommands.Peek(); - pending = peek.Size; - trimming = peek.Trim; - } - - if (trimming > 0) - { - var trimmed = Math.Min(trimming, (int)buffer.Length); - consumedPos = buffer.GetPosition(trimmed); - examinedPos = buffer.GetPosition(trimmed); - examinedOffset = 0; - buffer = buffer.Slice(trimmed); - pending -= trimmed; - trimming -= trimmed; - if (pending == 0) - { - // the entire command was trimmed (canceled) - inFlightSum -= _inFlightCommands.Dequeue().Size; - } - - continue; - } - - var sendMem = buffer.First; - var maxSize = 0; - var maxSizeCap = Math.Max(sendMem.Length, consolidateMem.Length); - var doTrim = false; - foreach (var command in _inFlightCommands) - { - if (maxSize == 0) - { - // first command; set to pending - maxSize = pending; - continue; - } - - if (maxSize > maxSizeCap) - { - // over cap - break; - } - - if (command.Trim > 0) - { - // will have to trim - doTrim = true; - break; - } - - maxSize += command.Size; - } - - if (sendMem.Length > maxSize) - { - sendMem = sendMem[..maxSize]; - } - - var bufferIter = buffer; - if (doTrim || (bufferIter.Length > sendMem.Length && sendMem.Length < consolidateMem.Length)) - { - var memIter = consolidateMem; - var trimmedSize = 0; - foreach (var command in _inFlightCommands) - { - if (bufferIter.Length == 0 || memIter.Length == 0) - { - break; - } - - int write; - if (trimmedSize == 0) - { - // first command, only write pending data - write = pending; - } - else if (command.Trim == 0) - { - write = command.Size; - } - else - { - if (bufferIter.Length < command.Trim + 1) - { - // not enough bytes to start writing the next command - break; - } - - bufferIter = bufferIter.Slice(command.Trim); - write = command.Size - command.Trim; - if (write == 0) - { - // the entire command was trimmed (canceled) - continue; - } - } - - write = Math.Min(memIter.Length, write); - write = Math.Min((int)bufferIter.Length, write); - bufferIter.Slice(0, write).CopyTo(memIter.Span); - memIter = memIter[write..]; - bufferIter = bufferIter.Slice(write); - trimmedSize += write; - } - - sendMem = consolidateMem[..trimmedSize]; - } - - // perform send - _stopwatch.Restart(); - var sent = await _socketConnection.SendAsync(sendMem).ConfigureAwait(false); - _stopwatch.Stop(); - Interlocked.Add(ref _counter.SentBytes, sent); - if (isEnabledTraceLogging) - { - logger.LogTrace("Socket.SendAsync. Size: {0} Elapsed: {1}ms", sent, _stopwatch.Elapsed.TotalMilliseconds); - } - - var consumed = 0; - var sentAndTrimmed = sent; - while (consumed < sentAndTrimmed) - { - if (pending == 0) - { - var peek = _inFlightCommands.Peek(); - pending = peek.Size - peek.Trim; - consumed += peek.Trim; - sentAndTrimmed += peek.Trim; - - if (pending == 0) - { - // the entire command was trimmed (canceled) - inFlightSum -= _inFlightCommands.Dequeue().Size; - continue; - } - } - - if (pending <= sentAndTrimmed - consumed) - { - // the entire command was sent - inFlightSum -= _inFlightCommands.Dequeue().Size; - Interlocked.Add(ref _counter.PendingMessages, -1); - Interlocked.Add(ref _counter.SentMessages, 1); - - // mark the bytes as consumed, and reset pending - consumed += pending; - pending = 0; - } - else - { - // the entire command was not sent; decrement pending by - // the number of bytes from the command that was sent - pending += consumed - sentAndTrimmed; - break; - } - } - - if (consumed > 0) - { - // mark fully sent commands as consumed - consumedPos = buffer.GetPosition(consumed); - examinedOffset = sentAndTrimmed - consumed; - } - else - { - // no commands were consumed - examinedOffset += sentAndTrimmed; - } - - // lop off sent bytes for next iteration - examinedPos = buffer.GetPosition(sentAndTrimmed); - buffer = buffer.Slice(sentAndTrimmed); - } - } - finally - { - _pipeReader.AdvanceTo(consumedPos, examinedPos); - } - - if (result.IsCompleted) - { - break; - } - } - } - catch (OperationCanceledException) - { - // ignore, intentionally disposed - } - catch (SocketClosedException) - { - // ignore, will be handled in read loop - } - catch (Exception ex) - { - logger.LogError(ex, "Unexpected error occured in write loop"); - throw; - } - - logger.LogDebug("WriteLoop finished."); - } -} +// using System.Buffers; +// using System.Diagnostics; +// using System.IO.Pipelines; +// using System.Threading.Channels; +// using Microsoft.Extensions.Logging; +// using NATS.Client.Core.Commands; +// +// namespace NATS.Client.Core.Internal; +// +// internal sealed class NatsPipeliningWriteProtocolProcessor : IAsyncDisposable +// { +// private readonly CancellationTokenSource _cts; +// private readonly ConnectionStatsCounter _counter; +// private readonly NatsOpts _opts; +// private readonly PipeReader _pipeReader; +// private readonly Queue _inFlightCommands; +// private readonly ChannelReader _queuedCommandReader; +// private readonly ISocketConnection _socketConnection; +// private readonly Stopwatch _stopwatch = new Stopwatch(); +// private int _disposed; +// +// public NatsPipeliningWriteProtocolProcessor(ISocketConnection socketConnection, CommandWriter commandWriter, NatsOpts opts, ConnectionStatsCounter counter) +// { +// _cts = new CancellationTokenSource(); +// _counter = counter; +// _inFlightCommands = commandWriter.InFlightCommands; +// _opts = opts; +// _pipeReader = commandWriter.PipeReader; +// _queuedCommandReader = commandWriter.QueuedCommandsReader; +// _socketConnection = socketConnection; +// WriteLoop = Task.Run(WriteLoopAsync); +// } +// +// public Task WriteLoop { get; } +// +// public async ValueTask DisposeAsync() +// { +// if (Interlocked.Increment(ref _disposed) == 1) +// { +// #if NET6_0 +// _cts.Cancel(); +// #else +// await _cts.CancelAsync().ConfigureAwait(false); +// #endif +// try +// { +// await WriteLoop.ConfigureAwait(false); // wait to drain writer +// } +// catch +// { +// // ignore +// } +// } +// } +// +// private async Task WriteLoopAsync() +// { +// var logger = _opts.LoggerFactory.CreateLogger(); +// var isEnabledTraceLogging = logger.IsEnabled(LogLevel.Trace); +// var cancellationToken = _cts.Token; +// var pending = 0; +// var trimming = 0; +// var examinedOffset = 0; +// +// // memory segment used to consolidate multiple small memory chunks +// // should <= (minimumSegmentSize * 0.5) in CommandWriter +// // 8520 should fit into 6 packets on 1500 MTU TLS connection or 1 packet on 9000 MTU TLS connection +// // assuming 40 bytes TCP overhead + 40 bytes TLS overhead per packet +// var consolidateMem = new Memory(new byte[8520]); +// +// // add up in flight command sum +// var inFlightSum = 0; +// foreach (var command in _inFlightCommands) +// { +// inFlightSum += command.Size; +// } +// +// try +// { +// while (true) +// { +// var result = await _pipeReader.ReadAsync(cancellationToken).ConfigureAwait(false); +// if (result.IsCanceled) +// { +// break; +// } +// +// var consumedPos = result.Buffer.Start; +// var examinedPos = result.Buffer.Start; +// try +// { +// var buffer = result.Buffer.Slice(examinedOffset); +// while (inFlightSum < result.Buffer.Length) +// { +// QueuedCommand queuedCommand; +// while (!_queuedCommandReader.TryRead(out queuedCommand)) +// { +// await _queuedCommandReader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); +// } +// +// _inFlightCommands.Enqueue(queuedCommand); +// inFlightSum += queuedCommand.Size; +// } +// +// while (buffer.Length > 0) +// { +// if (pending == 0) +// { +// var peek = _inFlightCommands.Peek(); +// pending = peek.Size; +// trimming = peek.Trim; +// } +// +// if (trimming > 0) +// { +// var trimmed = Math.Min(trimming, (int)buffer.Length); +// consumedPos = buffer.GetPosition(trimmed); +// examinedPos = buffer.GetPosition(trimmed); +// examinedOffset = 0; +// buffer = buffer.Slice(trimmed); +// pending -= trimmed; +// trimming -= trimmed; +// if (pending == 0) +// { +// // the entire command was trimmed (canceled) +// inFlightSum -= _inFlightCommands.Dequeue().Size; +// } +// +// continue; +// } +// +// var sendMem = buffer.First; +// var maxSize = 0; +// var maxSizeCap = Math.Max(sendMem.Length, consolidateMem.Length); +// var doTrim = false; +// foreach (var command in _inFlightCommands) +// { +// if (maxSize == 0) +// { +// // first command; set to pending +// maxSize = pending; +// continue; +// } +// +// if (maxSize > maxSizeCap) +// { +// // over cap +// break; +// } +// +// if (command.Trim > 0) +// { +// // will have to trim +// doTrim = true; +// break; +// } +// +// maxSize += command.Size; +// } +// +// if (sendMem.Length > maxSize) +// { +// sendMem = sendMem[..maxSize]; +// } +// +// var bufferIter = buffer; +// if (doTrim || (bufferIter.Length > sendMem.Length && sendMem.Length < consolidateMem.Length)) +// { +// var memIter = consolidateMem; +// var trimmedSize = 0; +// foreach (var command in _inFlightCommands) +// { +// if (bufferIter.Length == 0 || memIter.Length == 0) +// { +// break; +// } +// +// int write; +// if (trimmedSize == 0) +// { +// // first command, only write pending data +// write = pending; +// } +// else if (command.Trim == 0) +// { +// write = command.Size; +// } +// else +// { +// if (bufferIter.Length < command.Trim + 1) +// { +// // not enough bytes to start writing the next command +// break; +// } +// +// bufferIter = bufferIter.Slice(command.Trim); +// write = command.Size - command.Trim; +// if (write == 0) +// { +// // the entire command was trimmed (canceled) +// continue; +// } +// } +// +// write = Math.Min(memIter.Length, write); +// write = Math.Min((int)bufferIter.Length, write); +// bufferIter.Slice(0, write).CopyTo(memIter.Span); +// memIter = memIter[write..]; +// bufferIter = bufferIter.Slice(write); +// trimmedSize += write; +// } +// +// sendMem = consolidateMem[..trimmedSize]; +// } +// +// // perform send +// _stopwatch.Restart(); +// var sent = await _socketConnection.SendAsync(sendMem).ConfigureAwait(false); +// _stopwatch.Stop(); +// Interlocked.Add(ref _counter.SentBytes, sent); +// if (isEnabledTraceLogging) +// { +// logger.LogTrace("Socket.SendAsync. Size: {0} Elapsed: {1}ms", sent, _stopwatch.Elapsed.TotalMilliseconds); +// } +// +// var consumed = 0; +// var sentAndTrimmed = sent; +// while (consumed < sentAndTrimmed) +// { +// if (pending == 0) +// { +// var peek = _inFlightCommands.Peek(); +// pending = peek.Size - peek.Trim; +// consumed += peek.Trim; +// sentAndTrimmed += peek.Trim; +// +// if (pending == 0) +// { +// // the entire command was trimmed (canceled) +// inFlightSum -= _inFlightCommands.Dequeue().Size; +// continue; +// } +// } +// +// if (pending <= sentAndTrimmed - consumed) +// { +// // the entire command was sent +// inFlightSum -= _inFlightCommands.Dequeue().Size; +// Interlocked.Add(ref _counter.PendingMessages, -1); +// Interlocked.Add(ref _counter.SentMessages, 1); +// +// // mark the bytes as consumed, and reset pending +// consumed += pending; +// pending = 0; +// } +// else +// { +// // the entire command was not sent; decrement pending by +// // the number of bytes from the command that was sent +// pending += consumed - sentAndTrimmed; +// break; +// } +// } +// +// if (consumed > 0) +// { +// // mark fully sent commands as consumed +// consumedPos = buffer.GetPosition(consumed); +// examinedOffset = sentAndTrimmed - consumed; +// } +// else +// { +// // no commands were consumed +// examinedOffset += sentAndTrimmed; +// } +// +// // lop off sent bytes for next iteration +// examinedPos = buffer.GetPosition(sentAndTrimmed); +// buffer = buffer.Slice(sentAndTrimmed); +// } +// } +// finally +// { +// _pipeReader.AdvanceTo(consumedPos, examinedPos); +// } +// +// if (result.IsCompleted) +// { +// break; +// } +// } +// } +// catch (OperationCanceledException) +// { +// // ignore, intentionally disposed +// } +// catch (SocketClosedException) +// { +// // ignore, will be handled in read loop +// } +// catch (Exception ex) +// { +// logger.LogError(ex, "Unexpected error occured in write loop"); +// throw; +// } +// +// logger.LogDebug("WriteLoop finished."); +// } +// } diff --git a/src/NATS.Client.Core/Internal/NatsReadProtocolProcessor.cs b/src/NATS.Client.Core/Internal/NatsReadProtocolProcessor.cs index 05c9dda23..fd77311d6 100644 --- a/src/NATS.Client.Core/Internal/NatsReadProtocolProcessor.cs +++ b/src/NATS.Client.Core/Internal/NatsReadProtocolProcessor.cs @@ -52,7 +52,7 @@ public async ValueTask DisposeAsync() await _readLoop.ConfigureAwait(false); // wait for drain buffer. foreach (var item in _pingCommands) { - item.TaskCompletionSource.SetCanceled(); + item.SetCanceled(); } _waitForInfoSignal.TrySetCanceled(); @@ -329,7 +329,7 @@ private async ValueTask> DispatchCommandAsync(int code, R if (_pingCommands.TryDequeue(out var pingCommand)) { - pingCommand.TaskCompletionSource.SetResult(DateTimeOffset.UtcNow - pingCommand.WriteTime); + pingCommand.SetResult(); } if (length < PongSize) diff --git a/src/NATS.Client.Core/NatsConnection.Ping.cs b/src/NATS.Client.Core/NatsConnection.Ping.cs index 9dac84b51..891954564 100644 --- a/src/NATS.Client.Core/NatsConnection.Ping.cs +++ b/src/NATS.Client.Core/NatsConnection.Ping.cs @@ -12,9 +12,17 @@ public async ValueTask PingAsync(CancellationToken cancellationToken = await ConnectAsync().AsTask().WaitAsync(cancellationToken).ConfigureAwait(false); } - var pingCommand = new PingCommand(); + PingCommand pingCommand; + if (!_pool.TryRent(out pingCommand!)) + { + pingCommand = new PingCommand(_pool); + } + + pingCommand.Start(); + await CommandWriter.PingAsync(pingCommand, cancellationToken).ConfigureAwait(false); - return await pingCommand.TaskCompletionSource.Task.ConfigureAwait(false); + + return await pingCommand.RunAsync().ConfigureAwait(false); } /// @@ -27,6 +35,6 @@ public async ValueTask PingAsync(CancellationToken cancellationToken = /// representing the asynchronous operation private ValueTask PingOnlyAsync(CancellationToken cancellationToken = default) => ConnectionState == NatsConnectionState.Open - ? CommandWriter.PingAsync(new PingCommand(), cancellationToken) + ? CommandWriter.PingAsync(new PingCommand(_pool), cancellationToken) : ValueTask.CompletedTask; } diff --git a/src/NATS.Client.Core/NatsConnection.cs b/src/NATS.Client.Core/NatsConnection.cs index 8dbfe3a81..8e4dd7753 100644 --- a/src/NATS.Client.Core/NatsConnection.cs +++ b/src/NATS.Client.Core/NatsConnection.cs @@ -1,6 +1,5 @@ using System.Buffers; using System.Diagnostics; -using System.IO.Pipelines; using System.Threading.Channels; using Microsoft.Extensions.Logging; using NATS.Client.Core.Commands; @@ -47,7 +46,6 @@ public partial class NatsConnection : INatsConnection private volatile NatsUri? _currentConnectUri; private volatile NatsUri? _lastSeedConnectUri; private NatsReadProtocolProcessor? _socketReader; - private NatsPipeliningWriteProtocolProcessor? _socketWriter; private TaskCompletionSource _waitForOpenConnection; private TlsCerts? _tlsCerts; private UserCredentials? _userCredentials; @@ -71,7 +69,7 @@ public NatsConnection(NatsOpts opts) _cancellationTimerPool = new CancellationTimerPool(_pool, _disposedCancellationTokenSource.Token); _name = opts.Name; Counter = new ConnectionStatsCounter(); - CommandWriter = new CommandWriter(Opts, Counter, EnqueuePing); + CommandWriter = new CommandWriter(_pool, Opts, Counter, EnqueuePing); InboxPrefix = NewInbox(opts.InboxPrefix); SubscriptionManager = new SubscriptionManager(this, InboxPrefix); _logger = opts.LoggerFactory.CreateLogger(); @@ -197,7 +195,8 @@ internal ValueTask UnsubscribeAsync(int sid) { try { - return CommandWriter.UnsubscribeAsync(sid, CancellationToken.None); + // TODO: use maxMsgs in INatsSub to unsubscribe. + return CommandWriter.UnsubscribeAsync(sid, null, CancellationToken.None); } catch (Exception ex) { @@ -410,11 +409,11 @@ private async ValueTask SetupReaderWriterAsync(bool reconnect) // Authentication _userCredentials?.Authenticate(_clientOpts, WritableServerInfo); - await using (var priorityCommandWriter = new PriorityCommandWriter(_socket!, Opts, Counter, EnqueuePing)) + await using (var priorityCommandWriter = new PriorityCommandWriter(_pool, _socket!, Opts, Counter, EnqueuePing)) { // add CONNECT and PING command to priority lane await priorityCommandWriter.CommandWriter.ConnectAsync(_clientOpts, CancellationToken.None).ConfigureAwait(false); - await priorityCommandWriter.CommandWriter.PingAsync(new PingCommand(), CancellationToken.None).ConfigureAwait(false); + await priorityCommandWriter.CommandWriter.PingAsync(new PingCommand(_pool), CancellationToken.None).ConfigureAwait(false); Task? reconnectTask = null; if (reconnect) @@ -434,7 +433,7 @@ private async ValueTask SetupReaderWriterAsync(bool reconnect) } // create the socket writer - _socketWriter = CommandWriter.CreateNatsPipeliningWriteProtocolProcessor(_socket!); + CommandWriter.Reset(_socket!); lock (_gate) { @@ -710,7 +709,7 @@ private void EnqueuePing(PingCommand pingCommand) } // Can not add PING, set fail. - pingCommand.TaskCompletionSource.SetCanceled(); + pingCommand.SetCanceled(); } [System.Runtime.CompilerServices.MethodImpl(System.Runtime.CompilerServices.MethodImplOptions.AggressiveInlining)] @@ -740,11 +739,7 @@ private async ValueTask DisposeSocketComponentAsync(IAsyncDisposable component, private async ValueTask DisposeSocketAsync(bool asyncReaderDispose) { // writer's internal buffer/channel is not thread-safe, must wait until complete. - if (_socketWriter != null) - { - await DisposeSocketComponentAsync(_socketWriter, "socket writer").ConfigureAwait(false); - _socketWriter = null; - } + await CommandWriter.DisposeAsync().ConfigureAwait(false); if (_socket != null) { diff --git a/src/NATS.Client.Core/NatsLogEvents.cs b/src/NATS.Client.Core/NatsLogEvents.cs index fc1cd0fe2..15456e79a 100644 --- a/src/NATS.Client.Core/NatsLogEvents.cs +++ b/src/NATS.Client.Core/NatsLogEvents.cs @@ -11,4 +11,5 @@ public static class NatsLogEvents public static readonly EventId Protocol = new(1005, nameof(Protocol)); public static readonly EventId TcpSocket = new(1006, nameof(TcpSocket)); public static readonly EventId Internal = new(1006, nameof(Internal)); + public static readonly EventId Buffer = new(1007, nameof(Buffer)); } diff --git a/tests/NATS.Client.Core.Tests/CancellationTest.cs b/tests/NATS.Client.Core.Tests/CancellationTest.cs index 1148cd569..5c2d5c4fb 100644 --- a/tests/NATS.Client.Core.Tests/CancellationTest.cs +++ b/tests/NATS.Client.Core.Tests/CancellationTest.cs @@ -16,7 +16,7 @@ public async Task CommandTimeoutTest() await conn.ConnectAsync(); // stall the flush task - await conn.CommandWriter.TestStallFlushAsync(TimeSpan.FromSeconds(5)); + // TODO: await conn.CommandWriter.TestStallFlushAsync(TimeSpan.FromSeconds(5)); // commands that call ConnectAsync throw OperationCanceledException await Assert.ThrowsAsync(() => conn.PingAsync().AsTask()); diff --git a/tests/NATS.Client.Core.Tests/NatsHeaderTest.cs b/tests/NATS.Client.Core.Tests/NatsHeaderTest.cs index 6b74776af..e2081e740 100644 --- a/tests/NATS.Client.Core.Tests/NatsHeaderTest.cs +++ b/tests/NATS.Client.Core.Tests/NatsHeaderTest.cs @@ -21,8 +21,8 @@ public async Task WriterTests() ["key"] = "a-long-header-value", }; var pipe = new Pipe(new PipeOptions(pauseWriterThreshold: 0)); - var writer = new HeaderWriter(pipe.Writer, Encoding.UTF8); - var written = writer.Write(headers); + var writer = new HeaderWriter(Encoding.UTF8); + var written = writer.Write(pipe.Writer, headers); var text = "NATS/1.0\r\nk1: v1\r\nk2: v2-0\r\nk2: v2-1\r\na-long-header-key: value\r\nkey: a-long-header-value\r\n\r\n"; var expected = new ReadOnlySequence(Encoding.UTF8.GetBytes(text)); @@ -30,7 +30,24 @@ public async Task WriterTests() Assert.Equal(expected.Length, written); await pipe.Writer.FlushAsync(); var result = await pipe.Reader.ReadAtLeastAsync((int)written); + Assert.True(expected.ToSpan().SequenceEqual(result.Buffer.ToSpan())); + _output.WriteLine($"Buffer:\n{result.Buffer.FirstSpan.Dump()}"); + } + [Fact] + public async Task WriterEmptyTests() + { + var headers = new NatsHeaders(); + var pipe = new Pipe(new PipeOptions(pauseWriterThreshold: 0)); + var writer = new HeaderWriter(Encoding.UTF8); + var written = writer.Write(pipe.Writer, headers); + + var text = "NATS/1.0\r\n\r\n"; + var expected = new ReadOnlySequence(Encoding.UTF8.GetBytes(text)); + + Assert.Equal(expected.Length, written); + await pipe.Writer.FlushAsync(); + var result = await pipe.Reader.ReadAtLeastAsync((int)written); Assert.True(expected.ToSpan().SequenceEqual(result.Buffer.ToSpan())); _output.WriteLine($"Buffer:\n{result.Buffer.FirstSpan.Dump()}"); } From 85972f5078b8cce6f7b8f625ca658a79c7423eeb Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Mon, 22 Jan 2024 23:34:00 +0000 Subject: [PATCH 03/25] Default write buffer size --- src/NATS.Client.Core/NatsOpts.cs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/NATS.Client.Core/NatsOpts.cs b/src/NATS.Client.Core/NatsOpts.cs index c44ffda56..3fb4861e2 100644 --- a/src/NATS.Client.Core/NatsOpts.cs +++ b/src/NATS.Client.Core/NatsOpts.cs @@ -31,7 +31,9 @@ public sealed record NatsOpts public ILoggerFactory LoggerFactory { get; init; } = NullLoggerFactory.Instance; - public int WriterBufferSize { get; init; } = 1048576; + // Same as default pipelines pause writer size + // Performing better compared to nats bench on localhost + public int WriterBufferSize { get; init; } = 65536; public int ReaderBufferSize { get; init; } = 1048576; From 78f0392afd4286ef31cab61eb4a43b36e3f9616d Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Tue, 23 Jan 2024 00:09:43 +0000 Subject: [PATCH 04/25] Revert tmp comment outs --- src/NATS.Client.Core/Commands/CommandWriter.cs | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 4e78bca61..c557ec0aa 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -44,12 +44,6 @@ public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter coun _defaultCommandTimeout = overrideCommandTimeout ?? opts.CommandTimeout; _enqueuePing = enqueuePing; _opts = opts; - var pipe = new Pipe(new PipeOptions( - pauseWriterThreshold: opts.WriterBufferSize, // flush will block after hitting - resumeWriterThreshold: opts.WriterBufferSize / 2, // will start flushing again after catching up - minimumSegmentSize: 16384, // segment that is part of an uninterrupted payload can be sent using socket.send - useSynchronizationContext: false)); - _pipeWriter = pipe.Writer; _protocolWriter = new ProtocolWriter(opts.SubjectEncoding); _semLock = new SemaphoreSlim(1); _headerWriter = new HeaderWriter(_opts.HeaderEncoding); @@ -276,17 +270,17 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] private async ValueTask PublishLockedAsync(string subject, string? replyTo, NatsPooledBufferWriter payloadBuffer, NatsPooledBufferWriter? headersBuffer, CancellationToken cancellationToken) { - // Interlocked.Add(ref _counter.PendingMessages, 1); + Interlocked.Add(ref _counter.PendingMessages, 1); await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { var payload = payloadBuffer.WrittenMemory; var headers = headersBuffer?.WrittenMemory; - // if (_disposed) - // { - // throw new ObjectDisposedException(nameof(CommandWriter)); - // } + if (_disposed) + { + throw new ObjectDisposedException(nameof(CommandWriter)); + } PipeWriter bw; lock (_lock) @@ -312,7 +306,7 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat finally { _semLock.Release(); - // Interlocked.Add(ref _counter.PendingMessages, -1); + Interlocked.Add(ref _counter.PendingMessages, -1); } } From abf863cbec6eaff77c18ffbc5398c07771a6d3ff Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Tue, 23 Jan 2024 00:54:08 +0000 Subject: [PATCH 05/25] HPUB writer fix --- src/NATS.Client.Core/Commands/CommandWriter.cs | 2 +- src/NATS.Client.Core/Commands/ProtocolWriter.cs | 7 +++++-- src/NATS.Client.Core/NatsConnection.cs | 2 +- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index c557ec0aa..290a7edc1 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -51,7 +51,7 @@ public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter coun _readerLoopTask = Task.Run(ReaderLoopAsync); } - public void Reset(ISocketConnection socketConnection) + public void Reset(ISocketConnection? socketConnection) { var pipe = new Pipe(new PipeOptions( pauseWriterThreshold: _opts.WriterBufferSize, // flush will block after hitting diff --git a/src/NATS.Client.Core/Commands/ProtocolWriter.cs b/src/NATS.Client.Core/Commands/ProtocolWriter.cs index 264b035ef..6d477c0d9 100644 --- a/src/NATS.Client.Core/Commands/ProtocolWriter.cs +++ b/src/NATS.Client.Core/Commands/ProtocolWriter.cs @@ -247,7 +247,7 @@ private void WritePub(IBufferWriter writer, string subject, string? replyT private void WriteHpub(IBufferWriter writer, string subject, string? replyTo, ReadOnlyMemory headers, ReadOnlyMemory payload) { Span spanPayloadLength = stackalloc byte[MaxIntStringLength]; - if (!Utf8Formatter.TryFormat(payload.Length, spanPayloadLength, out var payloadLengthWritten)) + if (!Utf8Formatter.TryFormat(payload.Length + headers.Length, spanPayloadLength, out var payloadLengthWritten)) { ThrowOnUtf8FormatFail(); } @@ -275,7 +275,7 @@ private void WriteHpub(IBufferWriter writer, string subject, string? reply total += replyToLengthSpace; } - total += spanHeadersLength.Length + spanPayloadLength.Length + NewLineLength + payload.Length + NewLineLength; + total += spanHeadersLength.Length + spanPayloadLength.Length + NewLineLength + headers.Length + payload.Length + NewLineLength; var span = writer.GetSpan(total); @@ -302,6 +302,9 @@ private void WriteHpub(IBufferWriter writer, string subject, string? reply BinaryPrimitives.WriteUInt16LittleEndian(span, NewLine); span = span.Slice(NewLineLength); + headers.Span.CopyTo(span); + span = span.Slice(headers.Length); + payload.Span.CopyTo(span); span = span.Slice(payload.Length); diff --git a/src/NATS.Client.Core/NatsConnection.cs b/src/NATS.Client.Core/NatsConnection.cs index 8e4dd7753..abeda8889 100644 --- a/src/NATS.Client.Core/NatsConnection.cs +++ b/src/NATS.Client.Core/NatsConnection.cs @@ -739,7 +739,7 @@ private async ValueTask DisposeSocketComponentAsync(IAsyncDisposable component, private async ValueTask DisposeSocketAsync(bool asyncReaderDispose) { // writer's internal buffer/channel is not thread-safe, must wait until complete. - await CommandWriter.DisposeAsync().ConfigureAwait(false); + CommandWriter.Reset(_socket); if (_socket != null) { From 7b87f225f26d890fa95b596bc5212dce77d6a7cd Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Tue, 23 Jan 2024 15:34:37 +0000 Subject: [PATCH 06/25] Use channel instead of semaphore for locks --- .../Commands/CommandWriter.cs | 119 +++++++++++++----- 1 file changed, 91 insertions(+), 28 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 290a7edc1..d3538ac12 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -3,6 +3,7 @@ using System.Net.Sockets; using System.Numerics; using System.Runtime.CompilerServices; +using System.Threading.Channels; using Microsoft.Extensions.Logging; using NATS.Client.Core.Internal; @@ -28,9 +29,9 @@ internal sealed class CommandWriter : IAsyncDisposable private readonly Action _enqueuePing; private readonly NatsOpts _opts; private readonly ProtocolWriter _protocolWriter; - private readonly SemaphoreSlim _semLock; private readonly Task _readerLoopTask; private readonly HeaderWriter _headerWriter; + private readonly Channel _channelLock; private ISocketConnection? _socketConnection; private PipeReader? _pipeReader; private PipeWriter? _pipeWriter; @@ -45,7 +46,7 @@ public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter coun _enqueuePing = enqueuePing; _opts = opts; _protocolWriter = new ProtocolWriter(opts.SubjectEncoding); - _semLock = new SemaphoreSlim(1); + _channelLock = Channel.CreateBounded(1); _headerWriter = new HeaderWriter(_opts.HeaderEncoding); _cts = new CancellationTokenSource(); _readerLoopTask = Task.Run(ReaderLoopAsync); @@ -76,7 +77,11 @@ public async ValueTask DisposeAsync() await _cts.CancelAsync().ConfigureAwait(false); #endif - await _semLock.WaitAsync().ConfigureAwait(false); + while (!_channelLock.Writer.TryWrite(1)) + { + await _channelLock.Writer.WaitToWriteAsync().ConfigureAwait(false); + } + try { if (_disposed) @@ -95,14 +100,22 @@ public async ValueTask DisposeAsync() } finally { - _semLock.Release(); + while (!_channelLock.Reader.TryRead(out _)) + { + await _channelLock.Reader.WaitToReadAsync().ConfigureAwait(false); + } } } public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken cancellationToken) { - Interlocked.Add(ref _counter.PendingMessages, 1); - await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); + Interlocked.Increment(ref _counter.PendingMessages); + + while (!_channelLock.Writer.TryWrite(1)) + { + await _channelLock.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false); + } + try { if (_disposed) @@ -118,20 +131,29 @@ public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken ca bw = _pipeWriter!; } - _protocolWriter.WriteConnect(bw, connectOpts!); + _protocolWriter.WriteConnect(bw, connectOpts); await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { - _semLock.Release(); - Interlocked.Add(ref _counter.PendingMessages, -1); + while (!_channelLock.Reader.TryRead(out _)) + { + await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); + } + + Interlocked.Decrement(ref _counter.PendingMessages); } } public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken cancellationToken) { - Interlocked.Add(ref _counter.PendingMessages, 1); - await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); + Interlocked.Increment(ref _counter.PendingMessages); + + while (!_channelLock.Writer.TryWrite(1)) + { + await _channelLock.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false); + } + try { if (_disposed) @@ -154,15 +176,24 @@ public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken canc } finally { - _semLock.Release(); - Interlocked.Add(ref _counter.PendingMessages, -1); + while (!_channelLock.Reader.TryRead(out _)) + { + await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); + } + + Interlocked.Decrement(ref _counter.PendingMessages); } } public async ValueTask PongAsync(CancellationToken cancellationToken = default) { - Interlocked.Add(ref _counter.PendingMessages, 1); - await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); + Interlocked.Increment(ref _counter.PendingMessages); + + while (!_channelLock.Writer.TryWrite(1)) + { + await _channelLock.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false); + } + try { if (_disposed) @@ -183,7 +214,12 @@ public async ValueTask PongAsync(CancellationToken cancellationToken = default) } finally { - _semLock.Release(); + while (!_channelLock.Reader.TryRead(out _)) + { + await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); + } + + Interlocked.Decrement(ref _counter.PendingMessages); } } @@ -208,8 +244,13 @@ public ValueTask PublishAsync(string subject, T? value, NatsHeaders? headers, public async ValueTask SubscribeAsync(int sid, string subject, string? queueGroup, int? maxMsgs, CancellationToken cancellationToken) { - Interlocked.Add(ref _counter.PendingMessages, 1); - await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); + Interlocked.Increment(ref _counter.PendingMessages); + + while (!_channelLock.Writer.TryWrite(1)) + { + await _channelLock.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false); + } + try { if (_disposed) @@ -230,15 +271,24 @@ public async ValueTask SubscribeAsync(int sid, string subject, string? queueGrou } finally { - _semLock.Release(); - Interlocked.Add(ref _counter.PendingMessages, -1); + while (!_channelLock.Reader.TryRead(out _)) + { + await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); + } + + Interlocked.Decrement(ref _counter.PendingMessages); } } public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken cancellationToken) { - Interlocked.Add(ref _counter.PendingMessages, 1); - await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); + Interlocked.Increment(ref _counter.PendingMessages); + + while (!_channelLock.Writer.TryWrite(1)) + { + await _channelLock.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false); + } + try { if (_disposed) @@ -259,8 +309,12 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken } finally { - _semLock.Release(); - Interlocked.Add(ref _counter.PendingMessages, -1); + while (!_channelLock.Reader.TryRead(out _)) + { + await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); + } + + Interlocked.Decrement(ref _counter.PendingMessages); } } @@ -270,8 +324,13 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] private async ValueTask PublishLockedAsync(string subject, string? replyTo, NatsPooledBufferWriter payloadBuffer, NatsPooledBufferWriter? headersBuffer, CancellationToken cancellationToken) { - Interlocked.Add(ref _counter.PendingMessages, 1); - await _semLock.WaitAsync(cancellationToken).ConfigureAwait(false); + Interlocked.Increment(ref _counter.PendingMessages); + + while (!_channelLock.Writer.TryWrite(1)) + { + await _channelLock.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false); + } + try { var payload = payloadBuffer.WrittenMemory; @@ -305,8 +364,12 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat } finally { - _semLock.Release(); - Interlocked.Add(ref _counter.PendingMessages, -1); + while (!_channelLock.Reader.TryRead(out _)) + { + await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); + } + + Interlocked.Decrement(ref _counter.PendingMessages); } } From 2e15356c68783fd5e0b4b79b1861794d27957e58 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Fri, 26 Jan 2024 09:58:38 +0000 Subject: [PATCH 07/25] Reuse same pipeline avoiding data loss --- .../Commands/CommandWriter.cs | 225 +++++++++--------- 1 file changed, 109 insertions(+), 116 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index d3538ac12..0b88cb307 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -32,9 +32,9 @@ internal sealed class CommandWriter : IAsyncDisposable private readonly Task _readerLoopTask; private readonly HeaderWriter _headerWriter; private readonly Channel _channelLock; + private readonly PipeReader _pipeReader; + private readonly PipeWriter _pipeWriter; private ISocketConnection? _socketConnection; - private PipeReader? _pipeReader; - private PipeWriter? _pipeWriter; private bool _disposed; public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter counter, Action enqueuePing, TimeSpan? overrideCommandTimeout = default) @@ -49,23 +49,20 @@ public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter coun _channelLock = Channel.CreateBounded(1); _headerWriter = new HeaderWriter(_opts.HeaderEncoding); _cts = new CancellationTokenSource(); - _readerLoopTask = Task.Run(ReaderLoopAsync); - } - - public void Reset(ISocketConnection? socketConnection) - { var pipe = new Pipe(new PipeOptions( pauseWriterThreshold: _opts.WriterBufferSize, // flush will block after hitting resumeWriterThreshold: _opts.WriterBufferSize / 2, useSynchronizationContext: false)); + _pipeReader = pipe.Reader; + _pipeWriter = pipe.Writer; + _readerLoopTask = Task.Run(ReaderLoopAsync); + } + public void Reset(ISocketConnection? socketConnection) + { lock (_lock) { - _pipeWriter?.Complete(); - _pipeReader?.Complete(); _socketConnection = socketConnection; - _pipeReader = pipe.Reader; - _pipeWriter = pipe.Writer; } } @@ -91,10 +88,8 @@ public async ValueTask DisposeAsync() _disposed = true; - if (_pipeWriter != null) - await _pipeWriter.CompleteAsync().ConfigureAwait(false); - if (_pipeReader != null) - await _pipeReader.CompleteAsync().ConfigureAwait(false); + await _pipeWriter.CompleteAsync().ConfigureAwait(false); + await _pipeReader.CompleteAsync().ConfigureAwait(false); await _readerLoopTask.ConfigureAwait(false); } @@ -111,9 +106,17 @@ public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken ca { Interlocked.Increment(ref _counter.PendingMessages); - while (!_channelLock.Writer.TryWrite(1)) + try { - await _channelLock.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false); + await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); + } + catch (OperationCanceledException) + { + return; + } + catch (ChannelClosedException) + { + return; } try @@ -123,16 +126,8 @@ public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken ca throw new ObjectDisposedException(nameof(CommandWriter)); } - PipeWriter bw; - lock (_lock) - { - if (_pipeWriter == null) - ThrowOnDisconnected(); - bw = _pipeWriter!; - } - - _protocolWriter.WriteConnect(bw, connectOpts); - await bw.FlushAsync(cancellationToken).ConfigureAwait(false); + _protocolWriter.WriteConnect(_pipeWriter, connectOpts); + await _pipeWriter.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { @@ -149,10 +144,19 @@ public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken canc { Interlocked.Increment(ref _counter.PendingMessages); - while (!_channelLock.Writer.TryWrite(1)) + try + { + await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); + } + catch (OperationCanceledException) { - await _channelLock.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false); + return; } + catch (ChannelClosedException) + { + return; + } + try { @@ -161,18 +165,10 @@ public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken canc throw new ObjectDisposedException(nameof(CommandWriter)); } - PipeWriter bw; - lock (_lock) - { - if (_pipeWriter == null) - ThrowOnDisconnected(); - bw = _pipeWriter!; - } - _enqueuePing(pingCommand); - _protocolWriter.WritePing(bw); - await bw.FlushAsync(cancellationToken).ConfigureAwait(false); + _protocolWriter.WritePing(_pipeWriter); + await _pipeWriter.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { @@ -189,10 +185,19 @@ public async ValueTask PongAsync(CancellationToken cancellationToken = default) { Interlocked.Increment(ref _counter.PendingMessages); - while (!_channelLock.Writer.TryWrite(1)) + try { - await _channelLock.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false); + await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); } + catch (OperationCanceledException) + { + return; + } + catch (ChannelClosedException) + { + return; + } + try { @@ -201,16 +206,8 @@ public async ValueTask PongAsync(CancellationToken cancellationToken = default) throw new ObjectDisposedException(nameof(CommandWriter)); } - PipeWriter bw; - lock (_lock) - { - if (_pipeWriter == null) - ThrowOnDisconnected(); - bw = _pipeWriter!; - } - - _protocolWriter.WritePong(bw); - await bw.FlushAsync(cancellationToken).ConfigureAwait(false); + _protocolWriter.WritePong(_pipeWriter); + await _pipeWriter.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { @@ -246,10 +243,19 @@ public async ValueTask SubscribeAsync(int sid, string subject, string? queueGrou { Interlocked.Increment(ref _counter.PendingMessages); - while (!_channelLock.Writer.TryWrite(1)) + try { - await _channelLock.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false); + await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); } + catch (OperationCanceledException) + { + return; + } + catch (ChannelClosedException) + { + return; + } + try { @@ -258,16 +264,8 @@ public async ValueTask SubscribeAsync(int sid, string subject, string? queueGrou throw new ObjectDisposedException(nameof(CommandWriter)); } - PipeWriter bw; - lock (_lock) - { - if (_pipeWriter == null) - ThrowOnDisconnected(); - bw = _pipeWriter!; - } - - _protocolWriter.WriteSubscribe(bw, sid, subject, queueGroup, maxMsgs); - await bw.FlushAsync(cancellationToken).ConfigureAwait(false); + _protocolWriter.WriteSubscribe(_pipeWriter, sid, subject, queueGroup, maxMsgs); + await _pipeWriter.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { @@ -284,9 +282,17 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken { Interlocked.Increment(ref _counter.PendingMessages); - while (!_channelLock.Writer.TryWrite(1)) + try + { + await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); + } + catch (OperationCanceledException) + { + return; + } + catch (ChannelClosedException) { - await _channelLock.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false); + return; } try @@ -296,16 +302,8 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken throw new ObjectDisposedException(nameof(CommandWriter)); } - PipeWriter bw; - lock (_lock) - { - if (_pipeWriter == null) - ThrowOnDisconnected(); - bw = _pipeWriter!; - } - - _protocolWriter.WriteUnsubscribe(bw, sid, maxMsgs); - await bw.FlushAsync(cancellationToken).ConfigureAwait(false); + _protocolWriter.WriteUnsubscribe(_pipeWriter, sid, maxMsgs); + await _pipeWriter.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { @@ -326,10 +324,19 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat { Interlocked.Increment(ref _counter.PendingMessages); - while (!_channelLock.Writer.TryWrite(1)) + try { - await _channelLock.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false); + await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); } + catch (OperationCanceledException) + { + return; + } + catch (ChannelClosedException) + { + return; + } + try { @@ -341,15 +348,7 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat throw new ObjectDisposedException(nameof(CommandWriter)); } - PipeWriter bw; - lock (_lock) - { - if (_pipeWriter == null) - ThrowOnDisconnected(); - bw = _pipeWriter!; - } - - _protocolWriter.WritePublish(bw, subject, replyTo, headers, payload); + _protocolWriter.WritePublish(_pipeWriter, subject, replyTo, headers, payload); payloadBuffer.Reset(); _pool.Return(payloadBuffer); @@ -360,7 +359,7 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat _pool.Return(headersBuffer); } - await bw.FlushAsync(cancellationToken).ConfigureAwait(false); + await _pipeWriter.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { @@ -375,55 +374,44 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat private async Task ReaderLoopAsync() { - // 8520 should fit into 6 packets on 1500 MTU TLS connection or 1 packet on 9000 MTU TLS connection - // assuming 40 bytes TCP overhead + 40 bytes TLS overhead per packet - const int maxSendMemoryLength = 8520; - var sendMemory = new Memory(new byte[maxSendMemoryLength]); - var cancellationToken = _cts.Token; while (cancellationToken.IsCancellationRequested == false) { try { ISocketConnection? connection; - PipeReader? pipeReader; lock (_lock) { connection = _socketConnection; - pipeReader = _pipeReader; } - if (connection == null || pipeReader == null) + if (connection == null) { await Task.Delay(10, cancellationToken).ConfigureAwait(false); continue; } - var result = await pipeReader.ReadAsync(cancellationToken).ConfigureAwait(false); + var result = await _pipeReader.ReadAsync(cancellationToken).ConfigureAwait(false); + + if (result.IsCanceled) + { + break; + } + var buffer = result.Buffer; + var completed = buffer.Start; try { if (!buffer.IsEmpty) { - Memory memory; - var length = (int)buffer.Length; - - if (length > maxSendMemoryLength) - { - length = maxSendMemoryLength; - buffer = buffer.Slice(0, buffer.GetPosition(length)); - memory = sendMemory; - } - else - { - memory = sendMemory.Slice(0, length); - } - - buffer.CopyTo(memory.Span); + var bytes = ArrayPool.Shared.Rent((int)buffer.Length); + buffer.CopyTo(bytes); + var memory = bytes.AsMemory(0, (int)buffer.Length); try { await connection.SendAsync(memory).ConfigureAwait(false); + completed = buffer.End; } catch (SocketException e) { @@ -435,16 +423,21 @@ private async Task ReaderLoopAsync() connection.SignalDisconnected(e); _logger.LogError(NatsLogEvents.TcpSocket, e, "Unexpected error while sending data"); } + finally + { + ArrayPool.Shared.Return(bytes); + } } } finally { - pipeReader.AdvanceTo(buffer.End); + _pipeReader.AdvanceTo(completed); + } + + if (result.IsCompleted) + { + break; } - } - catch (InvalidOperationException) - { - // We might still be using the previous pipe reader which might be completed already } catch (OperationCanceledException) { From 1d11ccf60ef845fd2f673cd3649e460357e913c4 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Fri, 26 Jan 2024 10:38:06 +0000 Subject: [PATCH 08/25] Tidy up --- .../Commands/CommandWriter.cs | 380 +++--------------- .../Commands/NatsPooledBufferWriter.cs | 195 +++++++++ .../Commands/PriorityCommandWriter.cs | 25 ++ 3 files changed, 276 insertions(+), 324 deletions(-) create mode 100644 src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs create mode 100644 src/NATS.Client.Core/Commands/PriorityCommandWriter.cs diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 0b88cb307..6909934b8 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -1,7 +1,6 @@ using System.Buffers; using System.IO.Pipelines; using System.Net.Sockets; -using System.Numerics; using System.Runtime.CompilerServices; using System.Threading.Channels; using Microsoft.Extensions.Logging; @@ -35,7 +34,7 @@ internal sealed class CommandWriter : IAsyncDisposable private readonly PipeReader _pipeReader; private readonly PipeWriter _pipeWriter; private ISocketConnection? _socketConnection; - private bool _disposed; + private volatile bool _disposed; public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter counter, Action enqueuePing, TimeSpan? overrideCommandTimeout = default) { @@ -68,38 +67,26 @@ public void Reset(ISocketConnection? socketConnection) public async ValueTask DisposeAsync() { + if (_disposed) + { + return; + } + + _disposed = true; + #if NET6_0 _cts.Cancel(); #else await _cts.CancelAsync().ConfigureAwait(false); #endif - while (!_channelLock.Writer.TryWrite(1)) - { - await _channelLock.Writer.WaitToWriteAsync().ConfigureAwait(false); - } + await _channelLock.Writer.WriteAsync(1).ConfigureAwait(false); + _channelLock.Writer.TryComplete(); - try - { - if (_disposed) - { - return; - } + await _pipeWriter.CompleteAsync().ConfigureAwait(false); + await _pipeReader.CompleteAsync().ConfigureAwait(false); - _disposed = true; - - await _pipeWriter.CompleteAsync().ConfigureAwait(false); - await _pipeReader.CompleteAsync().ConfigureAwait(false); - - await _readerLoopTask.ConfigureAwait(false); - } - finally - { - while (!_channelLock.Reader.TryRead(out _)) - { - await _channelLock.Reader.WaitToReadAsync().ConfigureAwait(false); - } - } + await _readerLoopTask.ConfigureAwait(false); } public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken cancellationToken) @@ -142,22 +129,11 @@ public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken ca public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken cancellationToken) { - Interlocked.Increment(ref _counter.PendingMessages); - - try - { - await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); - } - catch (OperationCanceledException) - { - return; - } - catch (ChannelClosedException) + if (!await LockAsync(cancellationToken).ConfigureAwait(false)) { return; } - try { if (_disposed) @@ -172,32 +148,16 @@ public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken canc } finally { - while (!_channelLock.Reader.TryRead(out _)) - { - await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); - } - - Interlocked.Decrement(ref _counter.PendingMessages); + await UnLockAsync(cancellationToken).ConfigureAwait(true); } } public async ValueTask PongAsync(CancellationToken cancellationToken = default) { - Interlocked.Increment(ref _counter.PendingMessages); - - try - { - await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); - } - catch (OperationCanceledException) + if (!await LockAsync(cancellationToken).ConfigureAwait(false)) { return; } - catch (ChannelClosedException) - { - return; - } - try { @@ -211,12 +171,7 @@ public async ValueTask PongAsync(CancellationToken cancellationToken = default) } finally { - while (!_channelLock.Reader.TryRead(out _)) - { - await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); - } - - Interlocked.Decrement(ref _counter.PendingMessages); + await UnLockAsync(cancellationToken).ConfigureAwait(true); } } @@ -241,22 +196,11 @@ public ValueTask PublishAsync(string subject, T? value, NatsHeaders? headers, public async ValueTask SubscribeAsync(int sid, string subject, string? queueGroup, int? maxMsgs, CancellationToken cancellationToken) { - Interlocked.Increment(ref _counter.PendingMessages); - - try - { - await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); - } - catch (OperationCanceledException) - { - return; - } - catch (ChannelClosedException) + if (!await LockAsync(cancellationToken).ConfigureAwait(false)) { return; } - try { if (_disposed) @@ -269,28 +213,13 @@ public async ValueTask SubscribeAsync(int sid, string subject, string? queueGrou } finally { - while (!_channelLock.Reader.TryRead(out _)) - { - await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); - } - - Interlocked.Decrement(ref _counter.PendingMessages); + await UnLockAsync(cancellationToken).ConfigureAwait(true); } } public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken cancellationToken) { - Interlocked.Increment(ref _counter.PendingMessages); - - try - { - await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); - } - catch (OperationCanceledException) - { - return; - } - catch (ChannelClosedException) + if (!await LockAsync(cancellationToken).ConfigureAwait(false)) { return; } @@ -307,37 +236,18 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken } finally { - while (!_channelLock.Reader.TryRead(out _)) - { - await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); - } - - Interlocked.Decrement(ref _counter.PendingMessages); + await UnLockAsync(cancellationToken).ConfigureAwait(true); } } - [MethodImpl(MethodImplOptions.NoInlining)] - private static void ThrowOnDisconnected() => throw new NatsException("Connection hasn't been established yet."); - [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] private async ValueTask PublishLockedAsync(string subject, string? replyTo, NatsPooledBufferWriter payloadBuffer, NatsPooledBufferWriter? headersBuffer, CancellationToken cancellationToken) { - Interlocked.Increment(ref _counter.PendingMessages); - - try - { - await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); - } - catch (OperationCanceledException) - { - return; - } - catch (ChannelClosedException) + if (!await LockAsync(cancellationToken).ConfigureAwait(false)) { return; } - try { var payload = payloadBuffer.WrittenMemory; @@ -363,13 +273,45 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat } finally { - while (!_channelLock.Reader.TryRead(out _)) + await UnLockAsync(cancellationToken).ConfigureAwait(true); + } + } + + private async Task UnLockAsync(CancellationToken cancellationToken) + { + while (!_channelLock.Reader.TryRead(out _)) + { + try { await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); } + catch (OperationCanceledException) + { + break; + } + } - Interlocked.Decrement(ref _counter.PendingMessages); + Interlocked.Decrement(ref _counter.PendingMessages); + } + + private async Task LockAsync(CancellationToken cancellationToken) + { + Interlocked.Increment(ref _counter.PendingMessages); + + try + { + await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); } + catch (OperationCanceledException) + { + return false; + } + catch (ChannelClosedException) + { + return false; + } + + return true; } private async Task ReaderLoopAsync() @@ -450,213 +392,3 @@ private async Task ReaderLoopAsync() } } } - -internal sealed class PriorityCommandWriter : IAsyncDisposable -{ - private int _disposed; - - public PriorityCommandWriter(ObjectPool pool, ISocketConnection socketConnection, NatsOpts opts, ConnectionStatsCounter counter, Action enqueuePing) - { - CommandWriter = new CommandWriter(pool, opts, counter, enqueuePing, overrideCommandTimeout: TimeSpan.MaxValue); - CommandWriter.Reset(socketConnection); - } - - public CommandWriter CommandWriter { get; } - - public async ValueTask DisposeAsync() - { - if (Interlocked.Increment(ref _disposed) == 1) - { - // disposing command writer marks pipe writer as complete - await CommandWriter.DisposeAsync().ConfigureAwait(false); - } - } -} - -internal sealed class NatsPooledBufferWriter : IBufferWriter, IObjectPoolNode> -{ - private const int DefaultInitialBufferSize = 256; - - private readonly ArrayPool _pool; - private T[]? _array; - private int _index; - private NatsPooledBufferWriter? _next; - - public NatsPooledBufferWriter() - { - _pool = ArrayPool.Shared; - _array = _pool.Rent(DefaultInitialBufferSize); - _index = 0; - } - - public ref NatsPooledBufferWriter? NextNode => ref _next; - - /// - /// Gets the data written to the underlying buffer so far, as a . - /// - public ReadOnlyMemory WrittenMemory - { - get - { - var array = _array; - - if (array is null) - { - ThrowObjectDisposedException(); - } - - return array!.AsMemory(0, _index); - } - } - - /// - /// Gets the data written to the underlying buffer so far, as a . - /// - public ReadOnlySpan WrittenSpan - { - get - { - var array = _array; - - if (array is null) - { - ThrowObjectDisposedException(); - } - - return array!.AsSpan(0, _index); - } - } - - /// - /// Gets the amount of data written to the underlying buffer so far. - /// - public int WrittenCount - { - get => _index; - } - - /// - public void Advance(int count) - { - var array = _array; - - if (array is null) - { - ThrowObjectDisposedException(); - } - - if (count < 0) - { - ThrowArgumentOutOfRangeExceptionForNegativeCount(); - } - - if (_index > array!.Length - count) - { - ThrowArgumentExceptionForAdvancedTooFar(); - } - - _index += count; - } - - /// - public Memory GetMemory(int sizeHint = 0) - { - CheckBufferAndEnsureCapacity(sizeHint); - - return _array.AsMemory(_index); - } - - /// - public Span GetSpan(int sizeHint = 0) - { - CheckBufferAndEnsureCapacity(sizeHint); - - return _array.AsSpan(_index); - } - - public void Reset() - { - if (_array != null) - _pool.Return(_array); - _array = _pool.Rent(DefaultInitialBufferSize); - _index = 0; - } - - /// - public override string ToString() - { - // See comments in MemoryOwner about this - if (typeof(T) == typeof(char) && - _array is char[] chars) - { - return new(chars, 0, _index); - } - - // Same representation used in Span - return $"NatsPooledBufferWriter<{typeof(T)}>[{_index}]"; - } - - [MethodImpl(MethodImplOptions.NoInlining)] - private static void ThrowArgumentOutOfRangeExceptionForNegativeCount() => throw new ArgumentOutOfRangeException("count", "The count can't be a negative value."); - - [MethodImpl(MethodImplOptions.NoInlining)] - private static void ThrowArgumentOutOfRangeExceptionForNegativeSizeHint() => throw new ArgumentOutOfRangeException("sizeHint", "The size hint can't be a negative value."); - - [MethodImpl(MethodImplOptions.NoInlining)] - private static void ThrowArgumentExceptionForAdvancedTooFar() => throw new ArgumentException("The buffer writer has advanced too far."); - - [MethodImpl(MethodImplOptions.NoInlining)] - private static void ThrowObjectDisposedException() => throw new ObjectDisposedException("The current buffer has already been disposed."); - - /// - /// Ensures that has enough free space to contain a given number of new items. - /// - /// The minimum number of items to ensure space for in . - [MethodImpl(MethodImplOptions.AggressiveInlining)] - private void CheckBufferAndEnsureCapacity(int sizeHint) - { - var array = _array; - - if (array is null) - { - ThrowObjectDisposedException(); - } - - if (sizeHint < 0) - { - ThrowArgumentOutOfRangeExceptionForNegativeSizeHint(); - } - - if (sizeHint == 0) - { - sizeHint = 1; - } - - if (sizeHint > array!.Length - _index) - { - ResizeBuffer(sizeHint); - } - } - - /// - /// Resizes to ensure it can fit the specified number of new items. - /// - /// The minimum number of items to ensure space for in . - [MethodImpl(MethodImplOptions.NoInlining)] - private void ResizeBuffer(int sizeHint) - { - var minimumSize = (uint)_index + (uint)sizeHint; - - // The ArrayPool class has a maximum threshold of 1024 * 1024 for the maximum length of - // pooled arrays, and once this is exceeded it will just allocate a new array every time - // of exactly the requested size. In that case, we manually round up the requested size to - // the nearest power of two, to ensure that repeated consecutive writes when the array in - // use is bigger than that threshold don't end up causing a resize every single time. - if (minimumSize > 1024 * 1024) - { - minimumSize = BitOperations.RoundUpToPowerOf2(minimumSize); - } - - _pool.Resize(ref _array, (int)minimumSize); - } -} diff --git a/src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs b/src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs new file mode 100644 index 000000000..d84698f10 --- /dev/null +++ b/src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs @@ -0,0 +1,195 @@ +using System.Buffers; +using System.Numerics; +using System.Runtime.CompilerServices; +using NATS.Client.Core.Internal; + +namespace NATS.Client.Core.Commands; + +// adapted from https://github.com/CommunityToolkit/dotnet/blob/v8.2.2/src/CommunityToolkit.HighPerformance/Buffers/ArrayPoolBufferWriter%7BT%7D.cs +internal sealed class NatsPooledBufferWriter : IBufferWriter, IObjectPoolNode> +{ + private const int DefaultInitialBufferSize = 256; + + private readonly ArrayPool _pool; + private T[]? _array; + private int _index; + private NatsPooledBufferWriter? _next; + + public NatsPooledBufferWriter() + { + _pool = ArrayPool.Shared; + _array = _pool.Rent(DefaultInitialBufferSize); + _index = 0; + } + + public ref NatsPooledBufferWriter? NextNode => ref _next; + + /// + /// Gets the data written to the underlying buffer so far, as a . + /// + public ReadOnlyMemory WrittenMemory + { + get + { + var array = _array; + + if (array is null) + { + ThrowObjectDisposedException(); + } + + return array!.AsMemory(0, _index); + } + } + + /// + /// Gets the data written to the underlying buffer so far, as a . + /// + public ReadOnlySpan WrittenSpan + { + get + { + var array = _array; + + if (array is null) + { + ThrowObjectDisposedException(); + } + + return array!.AsSpan(0, _index); + } + } + + /// + /// Gets the amount of data written to the underlying buffer so far. + /// + public int WrittenCount + { + get => _index; + } + + /// + public void Advance(int count) + { + var array = _array; + + if (array is null) + { + ThrowObjectDisposedException(); + } + + if (count < 0) + { + ThrowArgumentOutOfRangeExceptionForNegativeCount(); + } + + if (_index > array!.Length - count) + { + ThrowArgumentExceptionForAdvancedTooFar(); + } + + _index += count; + } + + /// + public Memory GetMemory(int sizeHint = 0) + { + CheckBufferAndEnsureCapacity(sizeHint); + + return _array.AsMemory(_index); + } + + /// + public Span GetSpan(int sizeHint = 0) + { + CheckBufferAndEnsureCapacity(sizeHint); + + return _array.AsSpan(_index); + } + + public void Reset() + { + if (_array != null) + _pool.Return(_array); + _array = _pool.Rent(DefaultInitialBufferSize); + _index = 0; + } + + /// + public override string ToString() + { + // See comments in MemoryOwner about this + if (typeof(T) == typeof(char) && + _array is char[] chars) + { + return new(chars, 0, _index); + } + + // Same representation used in Span + return $"NatsPooledBufferWriter<{typeof(T)}>[{_index}]"; + } + + [MethodImpl(MethodImplOptions.NoInlining)] + private static void ThrowArgumentOutOfRangeExceptionForNegativeCount() => throw new ArgumentOutOfRangeException("count", "The count can't be a negative value."); + + [MethodImpl(MethodImplOptions.NoInlining)] + private static void ThrowArgumentOutOfRangeExceptionForNegativeSizeHint() => throw new ArgumentOutOfRangeException("sizeHint", "The size hint can't be a negative value."); + + [MethodImpl(MethodImplOptions.NoInlining)] + private static void ThrowArgumentExceptionForAdvancedTooFar() => throw new ArgumentException("The buffer writer has advanced too far."); + + [MethodImpl(MethodImplOptions.NoInlining)] + private static void ThrowObjectDisposedException() => throw new ObjectDisposedException("The current buffer has already been disposed."); + + /// + /// Ensures that has enough free space to contain a given number of new items. + /// + /// The minimum number of items to ensure space for in . + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private void CheckBufferAndEnsureCapacity(int sizeHint) + { + var array = _array; + + if (array is null) + { + ThrowObjectDisposedException(); + } + + if (sizeHint < 0) + { + ThrowArgumentOutOfRangeExceptionForNegativeSizeHint(); + } + + if (sizeHint == 0) + { + sizeHint = 1; + } + + if (sizeHint > array!.Length - _index) + { + ResizeBuffer(sizeHint); + } + } + + /// + /// Resizes to ensure it can fit the specified number of new items. + /// + /// The minimum number of items to ensure space for in . + [MethodImpl(MethodImplOptions.NoInlining)] + private void ResizeBuffer(int sizeHint) + { + var minimumSize = (uint)_index + (uint)sizeHint; + + // The ArrayPool class has a maximum threshold of 1024 * 1024 for the maximum length of + // pooled arrays, and once this is exceeded it will just allocate a new array every time + // of exactly the requested size. In that case, we manually round up the requested size to + // the nearest power of two, to ensure that repeated consecutive writes when the array in + // use is bigger than that threshold don't end up causing a resize every single time. + if (minimumSize > 1024 * 1024) + { + minimumSize = BitOperations.RoundUpToPowerOf2(minimumSize); + } + + _pool.Resize(ref _array, (int)minimumSize); + } +} diff --git a/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs b/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs new file mode 100644 index 000000000..efec61c50 --- /dev/null +++ b/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs @@ -0,0 +1,25 @@ +using NATS.Client.Core.Internal; + +namespace NATS.Client.Core.Commands; + +internal sealed class PriorityCommandWriter : IAsyncDisposable +{ + private int _disposed; + + public PriorityCommandWriter(ObjectPool pool, ISocketConnection socketConnection, NatsOpts opts, ConnectionStatsCounter counter, Action enqueuePing) + { + CommandWriter = new CommandWriter(pool, opts, counter, enqueuePing, overrideCommandTimeout: TimeSpan.MaxValue); + CommandWriter.Reset(socketConnection); + } + + public CommandWriter CommandWriter { get; } + + public async ValueTask DisposeAsync() + { + if (Interlocked.Increment(ref _disposed) == 1) + { + // disposing command writer marks pipe writer as complete + await CommandWriter.DisposeAsync().ConfigureAwait(false); + } + } +} From 693266d9e4e03a1f6c1b5e981d1187ded8e599f9 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Fri, 26 Jan 2024 10:45:31 +0000 Subject: [PATCH 09/25] Reverted ping command --- src/NATS.Client.Core/Commands/PingCommand.cs | 52 ++----------------- .../Internal/NatsReadProtocolProcessor.cs | 4 +- src/NATS.Client.Core/NatsConnection.Ping.cs | 14 ++--- src/NATS.Client.Core/NatsConnection.cs | 4 +- 4 files changed, 11 insertions(+), 63 deletions(-) diff --git a/src/NATS.Client.Core/Commands/PingCommand.cs b/src/NATS.Client.Core/Commands/PingCommand.cs index a8c6236d4..1861adc0b 100644 --- a/src/NATS.Client.Core/Commands/PingCommand.cs +++ b/src/NATS.Client.Core/Commands/PingCommand.cs @@ -1,56 +1,12 @@ -using System.Threading.Tasks.Sources; -using NATS.Client.Core.Internal; - namespace NATS.Client.Core.Commands; -internal class PingCommand : IValueTaskSource, IObjectPoolNode +internal struct PingCommand { - private readonly ObjectPool? _pool; - private DateTimeOffset _start; - private ManualResetValueTaskSourceCore _core; - private PingCommand? _next; - - public PingCommand(ObjectPool? pool) + public PingCommand() { - _pool = pool; - _core = new ManualResetValueTaskSourceCore - { - RunContinuationsAsynchronously = true, - }; - _start = DateTimeOffset.MinValue; - } - - public ref PingCommand? NextNode => ref _next; - - public void Start() => _start = DateTimeOffset.UtcNow; - - public void SetResult() => _core.SetResult(DateTimeOffset.UtcNow - _start); - - public void SetCanceled() => _core.SetException(new OperationCanceledException()); - - public void Reset() - { - _start = DateTimeOffset.MinValue; - _core.Reset(); - } - - public ValueTask RunAsync() => new(this, _core.Version); - - public TimeSpan GetResult(short token) - { - var result = _core.GetResult(token); - - if (_pool is not null) - { - Reset(); - _pool.Return(this); - } - - return result; } - public ValueTaskSourceStatus GetStatus(short token) => _core.GetStatus(token); + public DateTimeOffset WriteTime { get; } = DateTimeOffset.UtcNow; - public void OnCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) - => _core.OnCompleted(continuation, state, token, flags); + public TaskCompletionSource TaskCompletionSource { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); } diff --git a/src/NATS.Client.Core/Internal/NatsReadProtocolProcessor.cs b/src/NATS.Client.Core/Internal/NatsReadProtocolProcessor.cs index fd77311d6..05c9dda23 100644 --- a/src/NATS.Client.Core/Internal/NatsReadProtocolProcessor.cs +++ b/src/NATS.Client.Core/Internal/NatsReadProtocolProcessor.cs @@ -52,7 +52,7 @@ public async ValueTask DisposeAsync() await _readLoop.ConfigureAwait(false); // wait for drain buffer. foreach (var item in _pingCommands) { - item.SetCanceled(); + item.TaskCompletionSource.SetCanceled(); } _waitForInfoSignal.TrySetCanceled(); @@ -329,7 +329,7 @@ private async ValueTask> DispatchCommandAsync(int code, R if (_pingCommands.TryDequeue(out var pingCommand)) { - pingCommand.SetResult(); + pingCommand.TaskCompletionSource.SetResult(DateTimeOffset.UtcNow - pingCommand.WriteTime); } if (length < PongSize) diff --git a/src/NATS.Client.Core/NatsConnection.Ping.cs b/src/NATS.Client.Core/NatsConnection.Ping.cs index 891954564..9dac84b51 100644 --- a/src/NATS.Client.Core/NatsConnection.Ping.cs +++ b/src/NATS.Client.Core/NatsConnection.Ping.cs @@ -12,17 +12,9 @@ public async ValueTask PingAsync(CancellationToken cancellationToken = await ConnectAsync().AsTask().WaitAsync(cancellationToken).ConfigureAwait(false); } - PingCommand pingCommand; - if (!_pool.TryRent(out pingCommand!)) - { - pingCommand = new PingCommand(_pool); - } - - pingCommand.Start(); - + var pingCommand = new PingCommand(); await CommandWriter.PingAsync(pingCommand, cancellationToken).ConfigureAwait(false); - - return await pingCommand.RunAsync().ConfigureAwait(false); + return await pingCommand.TaskCompletionSource.Task.ConfigureAwait(false); } /// @@ -35,6 +27,6 @@ public async ValueTask PingAsync(CancellationToken cancellationToken = /// representing the asynchronous operation private ValueTask PingOnlyAsync(CancellationToken cancellationToken = default) => ConnectionState == NatsConnectionState.Open - ? CommandWriter.PingAsync(new PingCommand(_pool), cancellationToken) + ? CommandWriter.PingAsync(new PingCommand(), cancellationToken) : ValueTask.CompletedTask; } diff --git a/src/NATS.Client.Core/NatsConnection.cs b/src/NATS.Client.Core/NatsConnection.cs index 5018fc6b9..2c3c585ff 100644 --- a/src/NATS.Client.Core/NatsConnection.cs +++ b/src/NATS.Client.Core/NatsConnection.cs @@ -432,7 +432,7 @@ private async ValueTask SetupReaderWriterAsync(bool reconnect) { // add CONNECT and PING command to priority lane await priorityCommandWriter.CommandWriter.ConnectAsync(_clientOpts, CancellationToken.None).ConfigureAwait(false); - await priorityCommandWriter.CommandWriter.PingAsync(new PingCommand(_pool), CancellationToken.None).ConfigureAwait(false); + await priorityCommandWriter.CommandWriter.PingAsync(new PingCommand(), CancellationToken.None).ConfigureAwait(false); Task? reconnectTask = null; if (reconnect) @@ -768,7 +768,7 @@ private void EnqueuePing(PingCommand pingCommand) } // Can not add PING, set fail. - pingCommand.SetCanceled(); + pingCommand.TaskCompletionSource.SetCanceled(); } [System.Runtime.CompilerServices.MethodImpl(System.Runtime.CompilerServices.MethodImplOptions.AggressiveInlining)] From 2ee24c059d1dabae8df0d6bdda82410c8092cfa5 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Fri, 26 Jan 2024 13:22:33 +0000 Subject: [PATCH 10/25] Revert to clearing buffer Other NATS clients do the same --- .../Commands/CommandWriter.cs | 90 ++++++++++++------- tests/NATS.Client.Core.Tests/ProtocolTest.cs | 28 ++++-- 2 files changed, 80 insertions(+), 38 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 6909934b8..d5191244d 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -31,8 +31,8 @@ internal sealed class CommandWriter : IAsyncDisposable private readonly Task _readerLoopTask; private readonly HeaderWriter _headerWriter; private readonly Channel _channelLock; - private readonly PipeReader _pipeReader; - private readonly PipeWriter _pipeWriter; + private PipeReader? _pipeReader; + private PipeWriter? _pipeWriter; private ISocketConnection? _socketConnection; private volatile bool _disposed; @@ -48,19 +48,22 @@ public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter coun _channelLock = Channel.CreateBounded(1); _headerWriter = new HeaderWriter(_opts.HeaderEncoding); _cts = new CancellationTokenSource(); - var pipe = new Pipe(new PipeOptions( - pauseWriterThreshold: _opts.WriterBufferSize, // flush will block after hitting - resumeWriterThreshold: _opts.WriterBufferSize / 2, - useSynchronizationContext: false)); - _pipeReader = pipe.Reader; - _pipeWriter = pipe.Writer; _readerLoopTask = Task.Run(ReaderLoopAsync); } public void Reset(ISocketConnection? socketConnection) { + var pipe = new Pipe(new PipeOptions( + pauseWriterThreshold: _opts.WriterBufferSize, // flush will block after hitting + resumeWriterThreshold: _opts.WriterBufferSize / 2, + useSynchronizationContext: false)); + lock (_lock) { + _pipeReader?.Complete(); + _pipeWriter?.Complete(); + _pipeReader = pipe.Reader; + _pipeWriter = pipe.Writer; _socketConnection = socketConnection; } } @@ -83,8 +86,11 @@ public async ValueTask DisposeAsync() await _channelLock.Writer.WriteAsync(1).ConfigureAwait(false); _channelLock.Writer.TryComplete(); - await _pipeWriter.CompleteAsync().ConfigureAwait(false); - await _pipeReader.CompleteAsync().ConfigureAwait(false); + if (_pipeWriter != null) + await _pipeWriter.CompleteAsync().ConfigureAwait(false); + + if (_pipeReader != null) + await _pipeReader.CompleteAsync().ConfigureAwait(false); await _readerLoopTask.ConfigureAwait(false); } @@ -113,8 +119,9 @@ public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken ca throw new ObjectDisposedException(nameof(CommandWriter)); } - _protocolWriter.WriteConnect(_pipeWriter, connectOpts); - await _pipeWriter.FlushAsync(cancellationToken).ConfigureAwait(false); + var bw = GetWriter(); + _protocolWriter.WriteConnect(bw, connectOpts); + await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { @@ -143,8 +150,9 @@ public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken canc _enqueuePing(pingCommand); - _protocolWriter.WritePing(_pipeWriter); - await _pipeWriter.FlushAsync(cancellationToken).ConfigureAwait(false); + var bw = GetWriter(); + _protocolWriter.WritePing(bw); + await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { @@ -166,8 +174,9 @@ public async ValueTask PongAsync(CancellationToken cancellationToken = default) throw new ObjectDisposedException(nameof(CommandWriter)); } - _protocolWriter.WritePong(_pipeWriter); - await _pipeWriter.FlushAsync(cancellationToken).ConfigureAwait(false); + var bw = GetWriter(); + _protocolWriter.WritePong(bw); + await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { @@ -208,8 +217,9 @@ public async ValueTask SubscribeAsync(int sid, string subject, string? queueGrou throw new ObjectDisposedException(nameof(CommandWriter)); } - _protocolWriter.WriteSubscribe(_pipeWriter, sid, subject, queueGroup, maxMsgs); - await _pipeWriter.FlushAsync(cancellationToken).ConfigureAwait(false); + var bw = GetWriter(); + _protocolWriter.WriteSubscribe(bw, sid, subject, queueGroup, maxMsgs); + await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { @@ -231,8 +241,9 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken throw new ObjectDisposedException(nameof(CommandWriter)); } - _protocolWriter.WriteUnsubscribe(_pipeWriter, sid, maxMsgs); - await _pipeWriter.FlushAsync(cancellationToken).ConfigureAwait(false); + var bw = GetWriter(); + _protocolWriter.WriteUnsubscribe(bw, sid, maxMsgs); + await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { @@ -240,6 +251,19 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken } } + [MethodImpl(MethodImplOptions.NoInlining)] + private static void ThrowOnDisconnected() => throw new NatsException("Connection hasn't been established yet."); + + private PipeWriter GetWriter() + { + lock (_lock) + { + if (_pipeWriter == null) + ThrowOnDisconnected(); + return _pipeWriter!; + } + } + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] private async ValueTask PublishLockedAsync(string subject, string? replyTo, NatsPooledBufferWriter payloadBuffer, NatsPooledBufferWriter? headersBuffer, CancellationToken cancellationToken) { @@ -258,7 +282,8 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat throw new ObjectDisposedException(nameof(CommandWriter)); } - _protocolWriter.WritePublish(_pipeWriter, subject, replyTo, headers, payload); + var bw = GetWriter(); + _protocolWriter.WritePublish(bw, subject, replyTo, headers, payload); payloadBuffer.Reset(); _pool.Return(payloadBuffer); @@ -269,7 +294,7 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat _pool.Return(headersBuffer); } - await _pipeWriter.FlushAsync(cancellationToken).ConfigureAwait(false); + await bw.FlushAsync(cancellationToken).ConfigureAwait(false); } finally { @@ -322,18 +347,20 @@ private async Task ReaderLoopAsync() try { ISocketConnection? connection; + PipeReader? pipeReader; lock (_lock) { connection = _socketConnection; + pipeReader = _pipeReader; } - if (connection == null) + if (connection == null || pipeReader == null) { await Task.Delay(10, cancellationToken).ConfigureAwait(false); continue; } - var result = await _pipeReader.ReadAsync(cancellationToken).ConfigureAwait(false); + var result = await pipeReader.ReadAsync(cancellationToken).ConfigureAwait(false); if (result.IsCanceled) { @@ -352,8 +379,8 @@ private async Task ReaderLoopAsync() try { - await connection.SendAsync(memory).ConfigureAwait(false); - completed = buffer.End; + var sent = await connection.SendAsync(memory).ConfigureAwait(false); + completed = buffer.GetPosition(sent); } catch (SocketException e) { @@ -373,18 +400,17 @@ private async Task ReaderLoopAsync() } finally { - _pipeReader.AdvanceTo(completed); - } - - if (result.IsCompleted) - { - break; + pipeReader.AdvanceTo(completed); } } catch (OperationCanceledException) { // Expected during shutdown } + catch (InvalidOperationException) + { + // We might still be using the previous pipe reader which might be completed already + } catch (Exception e) { _logger.LogError(NatsLogEvents.Buffer, e, "Unexpected error in send buffer reader loop"); diff --git a/tests/NATS.Client.Core.Tests/ProtocolTest.cs b/tests/NATS.Client.Core.Tests/ProtocolTest.cs index 35a64e266..f7b7feab5 100644 --- a/tests/NATS.Client.Core.Tests/ProtocolTest.cs +++ b/tests/NATS.Client.Core.Tests/ProtocolTest.cs @@ -1,4 +1,5 @@ using System.Buffers; +using System.Collections.Concurrent; using System.Text; using Microsoft.Extensions.Logging; using NATS.Client.TestUtilities; @@ -344,8 +345,11 @@ await Retry.Until( await nats.DisposeAsync(); } - [Fact] - public async Task Protocol_parser_under_load() + [Theory] + [InlineData(1)] + [InlineData(1024)] + [InlineData(1024 * 1024)] + public async Task Protocol_parser_under_load(int size) { await using var server = NatsServer.Start(); var logger = new InMemoryTestLoggerFactory(LogLevel.Error); @@ -355,24 +359,27 @@ public async Task Protocol_parser_under_load() using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); var signal = new WaitSignal(); - + var counts = new ConcurrentDictionary(); _ = Task.Run( async () => { var count = 0; - await foreach (var unused in nats.SubscribeAsync("x", cancellationToken: cts.Token)) + await foreach (var msg in nats.SubscribeAsync("x.*", cancellationToken: cts.Token)) { - if (++count > 10_000) + if (++count > 100) signal.Pulse(); + counts.AddOrUpdate(msg.Subject, 1, (_, c) => c + 1); } }, cts.Token); + var payload = new byte[size]; + var r = 0; _ = Task.Run( async () => { while (!cts.Token.IsCancellationRequested) - await nats.PublishAsync("x", "x", cancellationToken: cts.Token); + await nats.PublishAsync($"x.{Volatile.Read(ref r)}", payload, cancellationToken: cts.Token); }, cts.Token); @@ -382,12 +389,21 @@ public async Task Protocol_parser_under_load() { await Task.Delay(1_000, cts.Token); await server.RestartAsync(); + Interlocked.Increment(ref r); + await Task.Delay(1_000, cts.Token); } foreach (var log in logger.Logs.Where(x => x.EventId == NatsLogEvents.Protocol && x.LogLevel == LogLevel.Error)) { Assert.DoesNotContain("Unknown Protocol Operation", log.Message); } + + foreach (var (key, value) in counts) + { + _output.WriteLine($"{key} {value}"); + } + + counts.Count.Should().BeGreaterOrEqualTo(3); } private sealed class NatsSubReconnectTest : NatsSubBase From 3efe37c9047b0464ab44c02cbf394d84c1dc9bd7 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Sun, 28 Jan 2024 18:43:32 +0000 Subject: [PATCH 11/25] Lock function --- .../Commands/CommandWriter.cs | 58 +++---------------- 1 file changed, 9 insertions(+), 49 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index d5191244d..4b076922a 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -136,10 +136,7 @@ public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken ca public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken cancellationToken) { - if (!await LockAsync(cancellationToken).ConfigureAwait(false)) - { - return; - } + await LockAsync(cancellationToken).ConfigureAwait(false); try { @@ -162,10 +159,7 @@ public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken canc public async ValueTask PongAsync(CancellationToken cancellationToken = default) { - if (!await LockAsync(cancellationToken).ConfigureAwait(false)) - { - return; - } + await LockAsync(cancellationToken).ConfigureAwait(false); try { @@ -205,10 +199,7 @@ public ValueTask PublishAsync(string subject, T? value, NatsHeaders? headers, public async ValueTask SubscribeAsync(int sid, string subject, string? queueGroup, int? maxMsgs, CancellationToken cancellationToken) { - if (!await LockAsync(cancellationToken).ConfigureAwait(false)) - { - return; - } + await LockAsync(cancellationToken).ConfigureAwait(false); try { @@ -229,10 +220,7 @@ public async ValueTask SubscribeAsync(int sid, string subject, string? queueGrou public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken cancellationToken) { - if (!await LockAsync(cancellationToken).ConfigureAwait(false)) - { - return; - } + await LockAsync(cancellationToken).ConfigureAwait(false); try { @@ -267,10 +255,7 @@ private PipeWriter GetWriter() [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] private async ValueTask PublishLockedAsync(string subject, string? replyTo, NatsPooledBufferWriter payloadBuffer, NatsPooledBufferWriter? headersBuffer, CancellationToken cancellationToken) { - if (!await LockAsync(cancellationToken).ConfigureAwait(false)) - { - return; - } + await LockAsync(cancellationToken).ConfigureAwait(false); try { @@ -302,41 +287,16 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat } } - private async Task UnLockAsync(CancellationToken cancellationToken) + private ValueTask UnLockAsync(CancellationToken cancellationToken) { - while (!_channelLock.Reader.TryRead(out _)) - { - try - { - await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); - } - catch (OperationCanceledException) - { - break; - } - } - Interlocked.Decrement(ref _counter.PendingMessages); + return _channelLock.Reader.ReadAsync(cancellationToken); } - private async Task LockAsync(CancellationToken cancellationToken) + private ValueTask LockAsync(CancellationToken cancellationToken) { Interlocked.Increment(ref _counter.PendingMessages); - - try - { - await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); - } - catch (OperationCanceledException) - { - return false; - } - catch (ChannelClosedException) - { - return false; - } - - return true; + return _channelLock.Writer.WriteAsync(1, cancellationToken); } private async Task ReaderLoopAsync() From 2404ff3dfcbe01c017a30ad3c7e220069f42c57a Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Mon, 29 Jan 2024 10:50:32 +0000 Subject: [PATCH 12/25] Added inlining --- src/NATS.Client.Core/Commands/CommandWriter.cs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 4b076922a..a44892ddc 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -242,6 +242,7 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken [MethodImpl(MethodImplOptions.NoInlining)] private static void ThrowOnDisconnected() => throw new NatsException("Connection hasn't been established yet."); + [MethodImpl(MethodImplOptions.AggressiveInlining)] private PipeWriter GetWriter() { lock (_lock) @@ -287,12 +288,14 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat } } + [MethodImpl(MethodImplOptions.AggressiveInlining)] private ValueTask UnLockAsync(CancellationToken cancellationToken) { Interlocked.Decrement(ref _counter.PendingMessages); return _channelLock.Reader.ReadAsync(cancellationToken); } + [MethodImpl(MethodImplOptions.AggressiveInlining)] private ValueTask LockAsync(CancellationToken cancellationToken) { Interlocked.Increment(ref _counter.PendingMessages); From 722150fbbb8248f484be8a74487d670c2d49f586 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Tue, 30 Jan 2024 14:58:57 +0000 Subject: [PATCH 13/25] Command timeouts --- .../Commands/CommandWriter.cs | 26 ++++++++++++------- .../Internal/CancellationTimer.cs | 6 ++++- 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index a44892ddc..1e1407646 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -31,6 +31,7 @@ internal sealed class CommandWriter : IAsyncDisposable private readonly Task _readerLoopTask; private readonly HeaderWriter _headerWriter; private readonly Channel _channelLock; + private readonly CancellationTimerPool _ctPool; private PipeReader? _pipeReader; private PipeWriter? _pipeWriter; private ISocketConnection? _socketConnection; @@ -49,6 +50,9 @@ public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter coun _headerWriter = new HeaderWriter(_opts.HeaderEncoding); _cts = new CancellationTokenSource(); _readerLoopTask = Task.Run(ReaderLoopAsync); + + // _ctPool = new CancellationTimerPool(_pool, _cts.Token); + _ctPool = new CancellationTimerPool(_pool, CancellationToken.None); } public void Reset(ISocketConnection? socketConnection) @@ -97,11 +101,12 @@ public async ValueTask DisposeAsync() public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken cancellationToken) { + var cancellationTimer = _ctPool.Start(_defaultCommandTimeout, cancellationToken); Interlocked.Increment(ref _counter.PendingMessages); try { - await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); + await _channelLock.Writer.WriteAsync(1, cancellationTimer.Token).ConfigureAwait(false); } catch (OperationCanceledException) { @@ -121,7 +126,7 @@ public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken ca var bw = GetWriter(); _protocolWriter.WriteConnect(bw, connectOpts); - await bw.FlushAsync(cancellationToken).ConfigureAwait(false); + await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); } finally { @@ -131,6 +136,7 @@ public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken ca } Interlocked.Decrement(ref _counter.PendingMessages); + cancellationTimer.TryReturn(); } } @@ -153,7 +159,7 @@ public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken canc } finally { - await UnLockAsync(cancellationToken).ConfigureAwait(true); + await UnLockAsync(cancellationToken).ConfigureAwait(false); } } @@ -174,7 +180,7 @@ public async ValueTask PongAsync(CancellationToken cancellationToken = default) } finally { - await UnLockAsync(cancellationToken).ConfigureAwait(true); + await UnLockAsync(cancellationToken).ConfigureAwait(false); } } @@ -214,7 +220,7 @@ public async ValueTask SubscribeAsync(int sid, string subject, string? queueGrou } finally { - await UnLockAsync(cancellationToken).ConfigureAwait(true); + await UnLockAsync(cancellationToken).ConfigureAwait(false); } } @@ -235,7 +241,7 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken } finally { - await UnLockAsync(cancellationToken).ConfigureAwait(true); + await UnLockAsync(cancellationToken).ConfigureAwait(false); } } @@ -256,7 +262,8 @@ private PipeWriter GetWriter() [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] private async ValueTask PublishLockedAsync(string subject, string? replyTo, NatsPooledBufferWriter payloadBuffer, NatsPooledBufferWriter? headersBuffer, CancellationToken cancellationToken) { - await LockAsync(cancellationToken).ConfigureAwait(false); + var cancellationTimer = _ctPool.Start(_defaultCommandTimeout, cancellationToken); + await LockAsync(cancellationTimer.Token).ConfigureAwait(false); try { @@ -280,11 +287,12 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat _pool.Return(headersBuffer); } - await bw.FlushAsync(cancellationToken).ConfigureAwait(false); + await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); } finally { - await UnLockAsync(cancellationToken).ConfigureAwait(true); + await UnLockAsync(cancellationTimer.Token).ConfigureAwait(false); + cancellationTimer.TryReturn(); } } diff --git a/src/NATS.Client.Core/Internal/CancellationTimer.cs b/src/NATS.Client.Core/Internal/CancellationTimer.cs index 537184ec6..5966ba345 100644 --- a/src/NATS.Client.Core/Internal/CancellationTimer.cs +++ b/src/NATS.Client.Core/Internal/CancellationTimer.cs @@ -66,7 +66,11 @@ public static CancellationTimer Start(ObjectPool pool, CancellationToken rootTok } self._timeout = timeout; - self._cancellationTokenSource.CancelAfter(timeout); + if (timeout != TimeSpan.MaxValue) + { + self._cancellationTokenSource.CancelAfter(timeout); + } + return self; } From a8a5d5d8ce7a795fb7f6dc82e196943df99c621c Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Wed, 31 Jan 2024 11:49:28 +0000 Subject: [PATCH 14/25] Cancellation fixes --- .../Commands/CommandWriter.cs | 141 ++++++++++-------- .../Commands/PriorityCommandWriter.cs | 2 +- .../Internal/CancellationTimer.cs | 2 +- .../CancellationTest.cs | 18 +++ .../NATS.Client.Core.Tests/SendBufferTest.cs | 78 ++++++++++ tests/NATS.Client.TestUtilities/MockServer.cs | 112 ++++++++++++++ .../TmpFileLogger.cs | 18 +++ 7 files changed, 310 insertions(+), 61 deletions(-) create mode 100644 tests/NATS.Client.Core.Tests/SendBufferTest.cs create mode 100644 tests/NATS.Client.TestUtilities/MockServer.cs create mode 100644 tests/NATS.Client.TestUtilities/TmpFileLogger.cs diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 1e1407646..af24d300f 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -51,8 +51,12 @@ public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter coun _cts = new CancellationTokenSource(); _readerLoopTask = Task.Run(ReaderLoopAsync); - // _ctPool = new CancellationTimerPool(_pool, _cts.Token); - _ctPool = new CancellationTimerPool(_pool, CancellationToken.None); + // We need a new ObjectPool here because of the root token (_cts.Token). + // When the root token is cancelled as this object is disposed, cancellation + // objects in the pooled CancellationTimer should not be reused since the + // root token would already be cancelled which means CancellationTimer tokens + // would always be in a cancelled state. + _ctPool = new CancellationTimerPool(new ObjectPool(opts.ObjectPoolSize), _cts.Token); } public void Reset(ISocketConnection? socketConnection) @@ -102,21 +106,7 @@ public async ValueTask DisposeAsync() public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken cancellationToken) { var cancellationTimer = _ctPool.Start(_defaultCommandTimeout, cancellationToken); - Interlocked.Increment(ref _counter.PendingMessages); - - try - { - await _channelLock.Writer.WriteAsync(1, cancellationTimer.Token).ConfigureAwait(false); - } - catch (OperationCanceledException) - { - return; - } - catch (ChannelClosedException) - { - return; - } - + await LockAsync(cancellationTimer.Token).ConfigureAwait(false); try { if (_disposed) @@ -126,24 +116,23 @@ public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken ca var bw = GetWriter(); _protocolWriter.WriteConnect(bw, connectOpts); - await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + var result = await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + if (result.IsCanceled) + { + throw new OperationCanceledException(); + } } finally { - while (!_channelLock.Reader.TryRead(out _)) - { - await _channelLock.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); - } - - Interlocked.Decrement(ref _counter.PendingMessages); + await UnLockAsync().ConfigureAwait(false); cancellationTimer.TryReturn(); } } public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken cancellationToken) { - await LockAsync(cancellationToken).ConfigureAwait(false); - + var cancellationTimer = _ctPool.Start(_defaultCommandTimeout, cancellationToken); + await LockAsync(cancellationTimer.Token).ConfigureAwait(false); try { if (_disposed) @@ -155,18 +144,23 @@ public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken canc var bw = GetWriter(); _protocolWriter.WritePing(bw); - await bw.FlushAsync(cancellationToken).ConfigureAwait(false); + var result = await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + if (result.IsCanceled) + { + throw new OperationCanceledException(); + } } finally { - await UnLockAsync(cancellationToken).ConfigureAwait(false); + await UnLockAsync().ConfigureAwait(false); + cancellationTimer.TryReturn(); } } public async ValueTask PongAsync(CancellationToken cancellationToken = default) { - await LockAsync(cancellationToken).ConfigureAwait(false); - + var cancellationTimer = _ctPool.Start(_defaultCommandTimeout, cancellationToken); + await LockAsync(cancellationTimer.Token).ConfigureAwait(false); try { if (_disposed) @@ -176,11 +170,16 @@ public async ValueTask PongAsync(CancellationToken cancellationToken = default) var bw = GetWriter(); _protocolWriter.WritePong(bw); - await bw.FlushAsync(cancellationToken).ConfigureAwait(false); + var result = await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + if (result.IsCanceled) + { + throw new OperationCanceledException(); + } } finally { - await UnLockAsync(cancellationToken).ConfigureAwait(false); + await UnLockAsync().ConfigureAwait(false); + cancellationTimer.TryReturn(); } } @@ -205,8 +204,8 @@ public ValueTask PublishAsync(string subject, T? value, NatsHeaders? headers, public async ValueTask SubscribeAsync(int sid, string subject, string? queueGroup, int? maxMsgs, CancellationToken cancellationToken) { - await LockAsync(cancellationToken).ConfigureAwait(false); - + var cancellationTimer = _ctPool.Start(_defaultCommandTimeout, cancellationToken); + await LockAsync(cancellationTimer.Token).ConfigureAwait(false); try { if (_disposed) @@ -216,18 +215,23 @@ public async ValueTask SubscribeAsync(int sid, string subject, string? queueGrou var bw = GetWriter(); _protocolWriter.WriteSubscribe(bw, sid, subject, queueGroup, maxMsgs); - await bw.FlushAsync(cancellationToken).ConfigureAwait(false); + var result = await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + if (result.IsCanceled) + { + throw new OperationCanceledException(); + } } finally { - await UnLockAsync(cancellationToken).ConfigureAwait(false); + await UnLockAsync().ConfigureAwait(false); + cancellationTimer.TryReturn(); } } public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken cancellationToken) { - await LockAsync(cancellationToken).ConfigureAwait(false); - + var cancellationTimer = _ctPool.Start(_defaultCommandTimeout, cancellationToken); + await LockAsync(cancellationTimer.Token).ConfigureAwait(false); try { if (_disposed) @@ -237,34 +241,27 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken var bw = GetWriter(); _protocolWriter.WriteUnsubscribe(bw, sid, maxMsgs); - await bw.FlushAsync(cancellationToken).ConfigureAwait(false); + var result = await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + if (result.IsCanceled) + { + throw new OperationCanceledException(); + } } finally { - await UnLockAsync(cancellationToken).ConfigureAwait(false); + await UnLockAsync().ConfigureAwait(false); + cancellationTimer.TryReturn(); } } [MethodImpl(MethodImplOptions.NoInlining)] private static void ThrowOnDisconnected() => throw new NatsException("Connection hasn't been established yet."); - [MethodImpl(MethodImplOptions.AggressiveInlining)] - private PipeWriter GetWriter() - { - lock (_lock) - { - if (_pipeWriter == null) - ThrowOnDisconnected(); - return _pipeWriter!; - } - } - [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] private async ValueTask PublishLockedAsync(string subject, string? replyTo, NatsPooledBufferWriter payloadBuffer, NatsPooledBufferWriter? headersBuffer, CancellationToken cancellationToken) { var cancellationTimer = _ctPool.Start(_defaultCommandTimeout, cancellationToken); await LockAsync(cancellationTimer.Token).ConfigureAwait(false); - try { var payload = payloadBuffer.WrittenMemory; @@ -287,27 +284,53 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nat _pool.Return(headersBuffer); } - await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + var result = await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + if (result.IsCanceled) + { + throw new OperationCanceledException(); + } } finally { - await UnLockAsync(cancellationTimer.Token).ConfigureAwait(false); + await UnLockAsync().ConfigureAwait(false); cancellationTimer.TryReturn(); } } [MethodImpl(MethodImplOptions.AggressiveInlining)] - private ValueTask UnLockAsync(CancellationToken cancellationToken) + private async ValueTask LockAsync(CancellationToken cancellationToken) + { + Interlocked.Increment(ref _counter.PendingMessages); + try + { + await _channelLock.Writer.WriteAsync(1, cancellationToken).ConfigureAwait(false); + } + catch (TaskCanceledException) + { + throw new OperationCanceledException(); + } + catch (ChannelClosedException) + { + throw new OperationCanceledException(); + } + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private ValueTask UnLockAsync() { Interlocked.Decrement(ref _counter.PendingMessages); - return _channelLock.Reader.ReadAsync(cancellationToken); + return _channelLock.Reader.ReadAsync(_cts.Token); } [MethodImpl(MethodImplOptions.AggressiveInlining)] - private ValueTask LockAsync(CancellationToken cancellationToken) + private PipeWriter GetWriter() { - Interlocked.Increment(ref _counter.PendingMessages); - return _channelLock.Writer.WriteAsync(1, cancellationToken); + lock (_lock) + { + if (_pipeWriter == null) + ThrowOnDisconnected(); + return _pipeWriter!; + } } private async Task ReaderLoopAsync() diff --git a/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs b/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs index efec61c50..8880c393f 100644 --- a/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs +++ b/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs @@ -8,7 +8,7 @@ internal sealed class PriorityCommandWriter : IAsyncDisposable public PriorityCommandWriter(ObjectPool pool, ISocketConnection socketConnection, NatsOpts opts, ConnectionStatsCounter counter, Action enqueuePing) { - CommandWriter = new CommandWriter(pool, opts, counter, enqueuePing, overrideCommandTimeout: TimeSpan.MaxValue); + CommandWriter = new CommandWriter(pool, opts, counter, enqueuePing, overrideCommandTimeout: Timeout.InfiniteTimeSpan); CommandWriter.Reset(socketConnection); } diff --git a/src/NATS.Client.Core/Internal/CancellationTimer.cs b/src/NATS.Client.Core/Internal/CancellationTimer.cs index 5966ba345..6fe7b280d 100644 --- a/src/NATS.Client.Core/Internal/CancellationTimer.cs +++ b/src/NATS.Client.Core/Internal/CancellationTimer.cs @@ -66,7 +66,7 @@ public static CancellationTimer Start(ObjectPool pool, CancellationToken rootTok } self._timeout = timeout; - if (timeout != TimeSpan.MaxValue) + if (timeout != Timeout.InfiniteTimeSpan) { self._cancellationTokenSource.CancelAfter(timeout); } diff --git a/tests/NATS.Client.Core.Tests/CancellationTest.cs b/tests/NATS.Client.Core.Tests/CancellationTest.cs index 5c2d5c4fb..0e451064e 100644 --- a/tests/NATS.Client.Core.Tests/CancellationTest.cs +++ b/tests/NATS.Client.Core.Tests/CancellationTest.cs @@ -75,4 +75,22 @@ await Assert.ThrowsAsync(async () => } }); } + + [Fact] + public async Task Cancellation_timer() + { + var objectPool = new ObjectPool(10); + var cancellationTimerPool = new CancellationTimerPool(objectPool, CancellationToken.None); + var cancellationTimer = cancellationTimerPool.Start(TimeSpan.FromSeconds(2), CancellationToken.None); + + try + { + await Task.Delay(TimeSpan.FromSeconds(4), cancellationTimer.Token); + _output.WriteLine($"delayed 4 seconds"); + } + catch (Exception e) + { + _output.WriteLine($"Exception: {e.GetType().Name}"); + } + } } diff --git a/tests/NATS.Client.Core.Tests/SendBufferTest.cs b/tests/NATS.Client.Core.Tests/SendBufferTest.cs new file mode 100644 index 000000000..fdf1f75b2 --- /dev/null +++ b/tests/NATS.Client.Core.Tests/SendBufferTest.cs @@ -0,0 +1,78 @@ +using System.Diagnostics; +using NATS.Client.TestUtilities; + +namespace NATS.Client.Core.Tests; + +public class SendBufferTest +{ + private readonly ITestOutputHelper _output; + + public SendBufferTest(ITestOutputHelper output) => _output = output; + + [Fact] + public async Task Send_cancel() + { + void Log(string m) => TmpFileLogger.Log(m); + + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + + await using var server = new MockServer( + async (s, cmd) => + { + if (cmd.Name == "PUB" && cmd.Subject == "pause") + { + s.Log("[S] pause"); + await Task.Delay(10_000, cts.Token); + } + }, + Log, + cts.Token); + + Log("__________________________________"); + + await using var nats = new NatsConnection(new NatsOpts { Url = server.Url }); + + Log($"[C] connect {server.Url}"); + await nats.ConnectAsync(); + + Log($"[C] ping"); + var rtt = await nats.PingAsync(cts.Token); + Log($"[C] ping rtt={rtt}"); + + server.Log($"[C] publishing pause..."); + await nats.PublishAsync("pause", "x", cancellationToken: cts.Token); + + server.Log($"[C] publishing 1M..."); + var payload = new byte[1024 * 1024]; + var tasks = new List(); + for (var i = 0; i < 10; i++) + { + var i1 = i; + tasks.Add(Task.Run(async () => + { + var stopwatch = Stopwatch.StartNew(); + + try + { + Log($"[C] ({i1}) publish..."); + await nats.PublishAsync("x", payload, cancellationToken: cts.Token); + } + catch (Exception e) + { + stopwatch.Stop(); + Log($"[C] ({i1}) publish cancelled after {stopwatch.Elapsed.TotalSeconds:n0} s (exception: {e.GetType()})"); + return; + } + + stopwatch.Stop(); + Log($"[C] ({i1}) publish took {stopwatch.Elapsed.TotalSeconds:n3} s"); + })); + } + + for (var i = 0; i < 10; i++) + { + Log($"[C] await tasks {i}..."); + await tasks[i]; + } + } +} diff --git a/tests/NATS.Client.TestUtilities/MockServer.cs b/tests/NATS.Client.TestUtilities/MockServer.cs new file mode 100644 index 000000000..ea7e5927e --- /dev/null +++ b/tests/NATS.Client.TestUtilities/MockServer.cs @@ -0,0 +1,112 @@ +using System.Net; +using System.Net.Sockets; +using System.Text; +using System.Text.RegularExpressions; + +namespace NATS.Client.TestUtilities; + +public class MockServer : IAsyncDisposable +{ + private readonly Action _logger; + private readonly TcpListener _server; + private readonly Task _accept; + + public MockServer( + Func handler, + Action logger, + CancellationToken cancellationToken) + { + _logger = logger; + _server = new TcpListener(IPAddress.Parse("127.0.0.1"), 0); + _server.Start(); + Port = ((IPEndPoint)_server.LocalEndpoint).Port; + + _accept = Task.Run( + async () => + { + var client = await _server.AcceptTcpClientAsync(); + + var stream = client.GetStream(); + + var sw = new StreamWriter(stream, Encoding.ASCII); + await sw.WriteAsync("INFO {}\r\n"); + await sw.FlushAsync(); + + var sr = new StreamReader(stream, Encoding.ASCII); + + while (!cancellationToken.IsCancellationRequested) + { + Log($"[S] >>> READ LINE"); + var line = (await sr.ReadLineAsync())!; + + if (line.StartsWith("CONNECT")) + { + Log($"[S] RCV CONNECT"); + } + else if (line.StartsWith("PING")) + { + Log($"[S] RCV PING"); + await sw.WriteAsync("PONG\r\n"); + await sw.FlushAsync(); + Log($"[S] SND PONG"); + } + else if (line.StartsWith("SUB")) + { + var m = Regex.Match(line, @"^SUB\s+(?\S+)"); + var subject = m.Groups["subject"].Value; + Log($"[S] RCV SUB {subject}"); + await handler(this, new Cmd("SUB", subject, 0)); + } + else if (line.StartsWith("PUB") || line.StartsWith("HPUB")) + { + var m = Regex.Match(line, @"^(H?PUB)\s+(?\S+).*?(?\d+)$"); + var size = int.Parse(m.Groups["size"].Value); + var subject = m.Groups["subject"].Value; + Log($"[S] RCV PUB {subject} {size}"); + var read = 0; + var buffer = new byte[size]; + while (read < size) + { + var received = await stream.ReadAsync(buffer, read, size - read); + read += received; + Log($"[S] RCV {received} bytes (size={size} read={read})"); + } + + await handler(this, new Cmd("PUB", subject, size)); + await sr.ReadLineAsync(); + } + else + { + Log($"[S] RCV LINE: {line}"); + } + } + }, + cancellationToken); + } + + public int Port { get; } + + public string Url => $"127.0.0.1:{Port}"; + + public async ValueTask DisposeAsync() + { + _server.Stop(); + try + { + await _accept; + } + catch (OperationCanceledException) + { + } + catch (SocketException) + { + } + catch (IOException) + { + } + } + + public void Log(string m) => _logger(m); + + public record Cmd(string Name, string Subject, int Size); +} diff --git a/tests/NATS.Client.TestUtilities/TmpFileLogger.cs b/tests/NATS.Client.TestUtilities/TmpFileLogger.cs new file mode 100644 index 000000000..72d8e8ee0 --- /dev/null +++ b/tests/NATS.Client.TestUtilities/TmpFileLogger.cs @@ -0,0 +1,18 @@ +using System.Text; + +namespace NATS.Client.TestUtilities; + +public static class TmpFileLogger +{ + private static readonly object Gate = new(); + + public static void Log(string m) + { + lock (Gate) + { + using var fs = new FileStream("/tmp/test.log", FileMode.Append, FileAccess.Write, FileShare.ReadWrite); + using var sw = new StreamWriter(fs, Encoding.UTF8); + sw.WriteLine($"{DateTime.Now:HH:mm:ss.fff} {m}"); + } + } +} From 37871332abf365270515efa756c1ed881faf1bdf Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Wed, 31 Jan 2024 16:14:33 +0000 Subject: [PATCH 15/25] Derive pool rent size from buffer size option --- .../Commands/CommandWriter.cs | 10 +++++++-- .../Commands/NatsPooledBufferWriter.cs | 21 +++++++++++++++---- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index af24d300f..7c46e9e13 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -21,6 +21,7 @@ internal sealed class CommandWriter : IAsyncDisposable { private readonly ILogger _logger; private readonly ObjectPool _pool; + private readonly int _arrayPoolInitialSize; private readonly object _lock = new(); private readonly CancellationTokenSource _cts; private readonly ConnectionStatsCounter _counter; @@ -41,6 +42,11 @@ public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter coun { _logger = opts.LoggerFactory.CreateLogger(); _pool = pool; + + // Derive ArrayPool rent size from buffer size to + // avoid defining another option. + _arrayPoolInitialSize = opts.WriterBufferSize / 256; + _counter = counter; _defaultCommandTimeout = overrideCommandTimeout ?? opts.CommandTimeout; _enqueuePing = enqueuePing; @@ -189,13 +195,13 @@ public ValueTask PublishAsync(string subject, T? value, NatsHeaders? headers, if (headers != null) { if (!_pool.TryRent(out headersBuffer)) - headersBuffer = new NatsPooledBufferWriter(); + headersBuffer = new NatsPooledBufferWriter(_arrayPoolInitialSize); _headerWriter.Write(headersBuffer, headers); } NatsPooledBufferWriter payloadBuffer; if (!_pool.TryRent(out payloadBuffer!)) - payloadBuffer = new NatsPooledBufferWriter(); + payloadBuffer = new NatsPooledBufferWriter(_arrayPoolInitialSize); if (value != null) serializer.Serialize(payloadBuffer, value); diff --git a/src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs b/src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs index d84698f10..5b3760aa8 100644 --- a/src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs +++ b/src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs @@ -8,17 +8,30 @@ namespace NATS.Client.Core.Commands; // adapted from https://github.com/CommunityToolkit/dotnet/blob/v8.2.2/src/CommunityToolkit.HighPerformance/Buffers/ArrayPoolBufferWriter%7BT%7D.cs internal sealed class NatsPooledBufferWriter : IBufferWriter, IObjectPoolNode> { - private const int DefaultInitialBufferSize = 256; + private const int DefaultInitialMinBufferSize = 256; + private const int DefaultInitialMaxBufferSize = 65536; private readonly ArrayPool _pool; + private readonly int _size; private T[]? _array; private int _index; private NatsPooledBufferWriter? _next; - public NatsPooledBufferWriter() + public NatsPooledBufferWriter(int size) { + if (size < DefaultInitialMinBufferSize) + { + size = DefaultInitialMinBufferSize; + } + + if (size > DefaultInitialMaxBufferSize) + { + size = DefaultInitialMaxBufferSize; + } + + _size = size; _pool = ArrayPool.Shared; - _array = _pool.Rent(DefaultInitialBufferSize); + _array = _pool.Rent(size); _index = 0; } @@ -111,7 +124,7 @@ public void Reset() { if (_array != null) _pool.Return(_array); - _array = _pool.Rent(DefaultInitialBufferSize); + _array = _pool.Rent(_size); _index = 0; } From 45411a62415fb87914d87b2e86263aaaffd44881 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Wed, 31 Jan 2024 16:17:05 +0000 Subject: [PATCH 16/25] Fixed format --- src/NATS.Client.Core/Commands/CommandWriter.cs | 2 +- src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs | 2 +- src/NATS.Client.Core/Commands/PriorityCommandWriter.cs | 2 +- tests/NATS.Client.Core.Tests/SendBufferTest.cs | 2 +- tests/NATS.Client.TestUtilities/MockServer.cs | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 7c46e9e13..682df2939 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -264,7 +264,7 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken private static void ThrowOnDisconnected() => throw new NatsException("Connection hasn't been established yet."); [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] - private async ValueTask PublishLockedAsync(string subject, string? replyTo, NatsPooledBufferWriter payloadBuffer, NatsPooledBufferWriter? headersBuffer, CancellationToken cancellationToken) + private async ValueTask PublishLockedAsync(string subject, string? replyTo, NatsPooledBufferWriter payloadBuffer, NatsPooledBufferWriter? headersBuffer, CancellationToken cancellationToken) { var cancellationTimer = _ctPool.Start(_defaultCommandTimeout, cancellationToken); await LockAsync(cancellationTimer.Token).ConfigureAwait(false); diff --git a/src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs b/src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs index 5b3760aa8..b62d1c617 100644 --- a/src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs +++ b/src/NATS.Client.Core/Commands/NatsPooledBufferWriter.cs @@ -1,4 +1,4 @@ -using System.Buffers; +using System.Buffers; using System.Numerics; using System.Runtime.CompilerServices; using NATS.Client.Core.Internal; diff --git a/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs b/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs index 8880c393f..3020e84f4 100644 --- a/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs +++ b/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs @@ -1,4 +1,4 @@ -using NATS.Client.Core.Internal; +using NATS.Client.Core.Internal; namespace NATS.Client.Core.Commands; diff --git a/tests/NATS.Client.Core.Tests/SendBufferTest.cs b/tests/NATS.Client.Core.Tests/SendBufferTest.cs index fdf1f75b2..d6d098804 100644 --- a/tests/NATS.Client.Core.Tests/SendBufferTest.cs +++ b/tests/NATS.Client.Core.Tests/SendBufferTest.cs @@ -1,4 +1,4 @@ -using System.Diagnostics; +using System.Diagnostics; using NATS.Client.TestUtilities; namespace NATS.Client.Core.Tests; diff --git a/tests/NATS.Client.TestUtilities/MockServer.cs b/tests/NATS.Client.TestUtilities/MockServer.cs index ea7e5927e..d9a0f273d 100644 --- a/tests/NATS.Client.TestUtilities/MockServer.cs +++ b/tests/NATS.Client.TestUtilities/MockServer.cs @@ -1,4 +1,4 @@ -using System.Net; +using System.Net; using System.Net.Sockets; using System.Text; using System.Text.RegularExpressions; From 01b742bdd2a9fd2d6c0ac6756ec53b1cb6d61a2b Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Wed, 31 Jan 2024 16:45:51 +0000 Subject: [PATCH 17/25] Command timeout test --- src/NATS.Client.Core/Commands/CommandWriter.cs | 4 +++- tests/NATS.Client.Core.Tests/CancellationTest.cs | 13 ++++++++----- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 682df2939..1387d7c0f 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -97,7 +97,6 @@ public async ValueTask DisposeAsync() await _cts.CancelAsync().ConfigureAwait(false); #endif - await _channelLock.Writer.WriteAsync(1).ConfigureAwait(false); _channelLock.Writer.TryComplete(); if (_pipeWriter != null) @@ -260,6 +259,9 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken } } + // only used for internal testing + internal bool TestStallFlush() => _channelLock.Writer.TryWrite(1); + [MethodImpl(MethodImplOptions.NoInlining)] private static void ThrowOnDisconnected() => throw new NatsException("Connection hasn't been established yet."); diff --git a/tests/NATS.Client.Core.Tests/CancellationTest.cs b/tests/NATS.Client.Core.Tests/CancellationTest.cs index 0e451064e..732c07f84 100644 --- a/tests/NATS.Client.Core.Tests/CancellationTest.cs +++ b/tests/NATS.Client.Core.Tests/CancellationTest.cs @@ -15,15 +15,18 @@ public async Task CommandTimeoutTest() await using var conn = server.CreateClientConnection(NatsOpts.Default with { CommandTimeout = TimeSpan.FromMilliseconds(1) }); await conn.ConnectAsync(); + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + var cancellationToken = cts.Token; + // stall the flush task - // TODO: await conn.CommandWriter.TestStallFlushAsync(TimeSpan.FromSeconds(5)); + Assert.True(conn.CommandWriter.TestStallFlush()); // commands that call ConnectAsync throw OperationCanceledException - await Assert.ThrowsAsync(() => conn.PingAsync().AsTask()); - await Assert.ThrowsAsync(() => conn.PublishAsync("test").AsTask()); - await Assert.ThrowsAsync(async () => + await Assert.ThrowsAsync(() => conn.PingAsync(cancellationToken).AsTask()); + await Assert.ThrowsAsync(() => conn.PublishAsync("test", cancellationToken: cancellationToken).AsTask()); + await Assert.ThrowsAsync(async () => { - await foreach (var unused in conn.SubscribeAsync("test")) + await foreach (var unused in conn.SubscribeAsync("test", cancellationToken: cancellationToken)) { } }); From 535d697ec594cd7bf896f56b28f62c5b2effc4c4 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Wed, 31 Jan 2024 17:14:19 +0000 Subject: [PATCH 18/25] Test debug --- tests/NATS.Client.Core.Tests/ProtocolTest.cs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/NATS.Client.Core.Tests/ProtocolTest.cs b/tests/NATS.Client.Core.Tests/ProtocolTest.cs index f7b7feab5..704f2242c 100644 --- a/tests/NATS.Client.Core.Tests/ProtocolTest.cs +++ b/tests/NATS.Client.Core.Tests/ProtocolTest.cs @@ -336,6 +336,12 @@ await Retry.Until( () => proxy.ClientFrames.Any(f => f.Message.StartsWith("PUB foo"))); var frames = proxy.ClientFrames.Select(f => f.Message).ToList(); + + foreach (var frame in frames) + { + _output.WriteLine($"frame: {frame}"); + } + Assert.StartsWith("SUB foo", frames[0]); Assert.StartsWith("PUB bar1", frames[1]); Assert.StartsWith("PUB bar2", frames[2]); From 51fa05d5d55db0e58f51530a17c3c7e282d47c21 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Fri, 2 Feb 2024 13:57:44 +0000 Subject: [PATCH 19/25] Keep send buffer at message boundaries --- .../Commands/CommandWriter.cs | 303 ++++++++++-------- .../Commands/PriorityCommandWriter.cs | 1 + src/NATS.Client.Core/NatsConnection.cs | 4 +- 3 files changed, 171 insertions(+), 137 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 1387d7c0f..6edb8c2d6 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -1,6 +1,5 @@ using System.Buffers; using System.IO.Pipelines; -using System.Net.Sockets; using System.Runtime.CompilerServices; using System.Threading.Channels; using Microsoft.Extensions.Logging; @@ -27,15 +26,16 @@ internal sealed class CommandWriter : IAsyncDisposable private readonly ConnectionStatsCounter _counter; private readonly TimeSpan _defaultCommandTimeout; private readonly Action _enqueuePing; - private readonly NatsOpts _opts; private readonly ProtocolWriter _protocolWriter; - private readonly Task _readerLoopTask; private readonly HeaderWriter _headerWriter; private readonly Channel _channelLock; + private readonly Channel _channelSize; private readonly CancellationTimerPool _ctPool; - private PipeReader? _pipeReader; - private PipeWriter? _pipeWriter; + private readonly PipeReader _pipeReader; + private readonly PipeWriter _pipeWriter; private ISocketConnection? _socketConnection; + private Task? _readerLoopTask; + private CancellationTokenSource? _ctsReader; private volatile bool _disposed; public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter counter, Action enqueuePing, TimeSpan? overrideCommandTimeout = default) @@ -50,12 +50,18 @@ public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter coun _counter = counter; _defaultCommandTimeout = overrideCommandTimeout ?? opts.CommandTimeout; _enqueuePing = enqueuePing; - _opts = opts; _protocolWriter = new ProtocolWriter(opts.SubjectEncoding); _channelLock = Channel.CreateBounded(1); - _headerWriter = new HeaderWriter(_opts.HeaderEncoding); + _channelSize = Channel.CreateUnbounded(new UnboundedChannelOptions { SingleWriter = true, SingleReader = true }); + _headerWriter = new HeaderWriter(opts.HeaderEncoding); _cts = new CancellationTokenSource(); - _readerLoopTask = Task.Run(ReaderLoopAsync); + + var pipe = new Pipe(new PipeOptions( + pauseWriterThreshold: opts.WriterBufferSize, // flush will block after hitting + resumeWriterThreshold: opts.WriterBufferSize / 2, + useSynchronizationContext: false)); + _pipeReader = pipe.Reader; + _pipeWriter = pipe.Writer; // We need a new ObjectPool here because of the root token (_cts.Token). // When the root token is cancelled as this object is disposed, cancellation @@ -65,23 +71,43 @@ public CommandWriter(ObjectPool pool, NatsOpts opts, ConnectionStatsCounter coun _ctPool = new CancellationTimerPool(new ObjectPool(opts.ObjectPoolSize), _cts.Token); } - public void Reset(ISocketConnection? socketConnection) + public void Reset(ISocketConnection socketConnection) { - var pipe = new Pipe(new PipeOptions( - pauseWriterThreshold: _opts.WriterBufferSize, // flush will block after hitting - resumeWriterThreshold: _opts.WriterBufferSize / 2, - useSynchronizationContext: false)); - lock (_lock) { - _pipeReader?.Complete(); - _pipeWriter?.Complete(); - _pipeReader = pipe.Reader; - _pipeWriter = pipe.Writer; _socketConnection = socketConnection; + _ctsReader = CancellationTokenSource.CreateLinkedTokenSource(_cts.Token); + + _readerLoopTask = Task.Run(async () => + { + await ReaderLoopAsync(_logger, _socketConnection, _pipeReader, _channelSize, _ctsReader.Token).ConfigureAwait(false); + }); } } + public async Task FlushAsync() + { + CancellationTokenSource? cts; + Task? readerTask; + lock (_lock) + { + cts = _ctsReader; + readerTask = _readerLoopTask; + } + + if (cts != null) + { +#if NET6_0 + cts.Cancel(); +#else + await cts.CancelAsync().ConfigureAwait(false); +#endif + } + + if (readerTask != null) + await readerTask.WaitAsync(TimeSpan.FromSeconds(3), _cts.Token).ConfigureAwait(false); + } + public async ValueTask DisposeAsync() { if (_disposed) @@ -98,14 +124,8 @@ public async ValueTask DisposeAsync() #endif _channelLock.Writer.TryComplete(); - - if (_pipeWriter != null) - await _pipeWriter.CompleteAsync().ConfigureAwait(false); - - if (_pipeReader != null) - await _pipeReader.CompleteAsync().ConfigureAwait(false); - - await _readerLoopTask.ConfigureAwait(false); + _channelSize.Writer.TryComplete(); + await _pipeWriter.CompleteAsync().ConfigureAwait(false); } public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken cancellationToken) @@ -119,9 +139,12 @@ public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken ca throw new ObjectDisposedException(nameof(CommandWriter)); } - var bw = GetWriter(); - _protocolWriter.WriteConnect(bw, connectOpts); - var result = await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + _protocolWriter.WriteConnect(_pipeWriter, connectOpts); + + var size = (int)_pipeWriter.UnflushedBytes; + _channelSize.Writer.TryWrite(size); + + var result = await _pipeWriter.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); if (result.IsCanceled) { throw new OperationCanceledException(); @@ -147,9 +170,12 @@ public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken canc _enqueuePing(pingCommand); - var bw = GetWriter(); - _protocolWriter.WritePing(bw); - var result = await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + _protocolWriter.WritePing(_pipeWriter); + + var size = (int)_pipeWriter.UnflushedBytes; + _channelSize.Writer.TryWrite(size); + + var result = await _pipeWriter.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); if (result.IsCanceled) { throw new OperationCanceledException(); @@ -173,9 +199,12 @@ public async ValueTask PongAsync(CancellationToken cancellationToken = default) throw new ObjectDisposedException(nameof(CommandWriter)); } - var bw = GetWriter(); - _protocolWriter.WritePong(bw); - var result = await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + _protocolWriter.WritePong(_pipeWriter); + + var size = (int)_pipeWriter.UnflushedBytes; + _channelSize.Writer.TryWrite(size); + + var result = await _pipeWriter.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); if (result.IsCanceled) { throw new OperationCanceledException(); @@ -218,9 +247,12 @@ public async ValueTask SubscribeAsync(int sid, string subject, string? queueGrou throw new ObjectDisposedException(nameof(CommandWriter)); } - var bw = GetWriter(); - _protocolWriter.WriteSubscribe(bw, sid, subject, queueGroup, maxMsgs); - var result = await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + _protocolWriter.WriteSubscribe(_pipeWriter, sid, subject, queueGroup, maxMsgs); + + var size = (int)_pipeWriter.UnflushedBytes; + _channelSize.Writer.TryWrite(size); + + var result = await _pipeWriter.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); if (result.IsCanceled) { throw new OperationCanceledException(); @@ -244,9 +276,12 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken throw new ObjectDisposedException(nameof(CommandWriter)); } - var bw = GetWriter(); - _protocolWriter.WriteUnsubscribe(bw, sid, maxMsgs); - var result = await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + _protocolWriter.WriteUnsubscribe(_pipeWriter, sid, maxMsgs); + + var size = (int)_pipeWriter.UnflushedBytes; + _channelSize.Writer.TryWrite(size); + + var result = await _pipeWriter.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); if (result.IsCanceled) { throw new OperationCanceledException(); @@ -262,8 +297,92 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken // only used for internal testing internal bool TestStallFlush() => _channelLock.Writer.TryWrite(1); - [MethodImpl(MethodImplOptions.NoInlining)] - private static void ThrowOnDisconnected() => throw new NatsException("Connection hasn't been established yet."); + private static async Task ReaderLoopAsync(ILogger logger, ISocketConnection connection, PipeReader pipeReader, Channel channelSize, CancellationToken cancellationToken) + { + try + { + while (true) + { + var result = await pipeReader.ReadAsync(cancellationToken).ConfigureAwait(false); + + if (result.IsCanceled) + { + break; + } + + var buffer = result.Buffer; + var completed = buffer.Start; + try + { + if (!buffer.IsEmpty) + { + var bufferLength = (int)buffer.Length; + + var bytes = ArrayPool.Shared.Rent(bufferLength); + buffer.CopyTo(bytes); + var memory = bytes.AsMemory(0, (int)buffer.Length); + + try + { + completed = buffer.Start; + var totalSent = 0; + var totalSize = 0; + while (totalSent < bufferLength) + { + var sent = await connection.SendAsync(memory).ConfigureAwait(false); + + totalSent += sent; + memory = memory[sent..]; + + while (totalSize < totalSent) + { + int size; + while (!channelSize.Reader.TryRead(out size)) + { + await channelSize.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); + } + + totalSize += size; + } + + // make sure to mark the buffer only at message boundaries. + // if there was a message sent only partially, we consider it sent even if it might not be + // due to a socket error in following socket send iteration. this is to avoid re-sending + // the same message again, ensuring at-most-once delivery. + completed = buffer.GetPosition(totalSize); + } + } + finally + { + ArrayPool.Shared.Return(bytes); + } + } + } + finally + { + // Always examine to the end to potentially unblock writer + pipeReader.AdvanceTo(completed, buffer.End); + } + + if (result.IsCompleted) + { + break; + } + } + } + catch (OperationCanceledException) + { + // Expected during shutdown + } + catch (InvalidOperationException) + { + // We might still be using the previous pipe reader which might be completed already + } + catch (Exception e) + { + logger.LogError(NatsLogEvents.Buffer, e, "Unexpected error in send buffer reader loop"); + } + } [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] private async ValueTask PublishLockedAsync(string subject, string? replyTo, NatsPooledBufferWriter payloadBuffer, NatsPooledBufferWriter? headersBuffer, CancellationToken cancellationToken) @@ -280,8 +399,7 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nats throw new ObjectDisposedException(nameof(CommandWriter)); } - var bw = GetWriter(); - _protocolWriter.WritePublish(bw, subject, replyTo, headers, payload); + _protocolWriter.WritePublish(_pipeWriter, subject, replyTo, headers, payload); payloadBuffer.Reset(); _pool.Return(payloadBuffer); @@ -292,7 +410,10 @@ private async ValueTask PublishLockedAsync(string subject, string? replyTo, Nats _pool.Return(headersBuffer); } - var result = await bw.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); + var size = (int)_pipeWriter.UnflushedBytes; + _channelSize.Writer.TryWrite(size); + + var result = await _pipeWriter.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); if (result.IsCanceled) { throw new OperationCanceledException(); @@ -329,94 +450,4 @@ private ValueTask UnLockAsync() Interlocked.Decrement(ref _counter.PendingMessages); return _channelLock.Reader.ReadAsync(_cts.Token); } - - [MethodImpl(MethodImplOptions.AggressiveInlining)] - private PipeWriter GetWriter() - { - lock (_lock) - { - if (_pipeWriter == null) - ThrowOnDisconnected(); - return _pipeWriter!; - } - } - - private async Task ReaderLoopAsync() - { - var cancellationToken = _cts.Token; - while (cancellationToken.IsCancellationRequested == false) - { - try - { - ISocketConnection? connection; - PipeReader? pipeReader; - lock (_lock) - { - connection = _socketConnection; - pipeReader = _pipeReader; - } - - if (connection == null || pipeReader == null) - { - await Task.Delay(10, cancellationToken).ConfigureAwait(false); - continue; - } - - var result = await pipeReader.ReadAsync(cancellationToken).ConfigureAwait(false); - - if (result.IsCanceled) - { - break; - } - - var buffer = result.Buffer; - var completed = buffer.Start; - try - { - if (!buffer.IsEmpty) - { - var bytes = ArrayPool.Shared.Rent((int)buffer.Length); - buffer.CopyTo(bytes); - var memory = bytes.AsMemory(0, (int)buffer.Length); - - try - { - var sent = await connection.SendAsync(memory).ConfigureAwait(false); - completed = buffer.GetPosition(sent); - } - catch (SocketException e) - { - connection.SignalDisconnected(e); - _logger.LogWarning(NatsLogEvents.TcpSocket, e, "Error while sending data"); - } - catch (Exception e) - { - connection.SignalDisconnected(e); - _logger.LogError(NatsLogEvents.TcpSocket, e, "Unexpected error while sending data"); - } - finally - { - ArrayPool.Shared.Return(bytes); - } - } - } - finally - { - pipeReader.AdvanceTo(completed); - } - } - catch (OperationCanceledException) - { - // Expected during shutdown - } - catch (InvalidOperationException) - { - // We might still be using the previous pipe reader which might be completed already - } - catch (Exception e) - { - _logger.LogError(NatsLogEvents.Buffer, e, "Unexpected error in send buffer reader loop"); - } - } - } } diff --git a/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs b/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs index 3020e84f4..4bd0cce06 100644 --- a/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs +++ b/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs @@ -18,6 +18,7 @@ public async ValueTask DisposeAsync() { if (Interlocked.Increment(ref _disposed) == 1) { + await CommandWriter.FlushAsync().ConfigureAwait(false); // disposing command writer marks pipe writer as complete await CommandWriter.DisposeAsync().ConfigureAwait(false); } diff --git a/src/NATS.Client.Core/NatsConnection.cs b/src/NATS.Client.Core/NatsConnection.cs index 954f6e426..1afe66f92 100644 --- a/src/NATS.Client.Core/NatsConnection.cs +++ b/src/NATS.Client.Core/NatsConnection.cs @@ -500,6 +500,8 @@ private async void ReconnectLoop() // If dispose this client, WaitForClosed throws OperationCanceledException so stop reconnect-loop correctly. await _socket!.WaitForClosed.ConfigureAwait(false); + await CommandWriter.FlushAsync().ConfigureAwait(false); + _logger.LogTrace(NatsLogEvents.Connection, "Connection {Name} is closed. Will cleanup and reconnect", _name); lock (_gate) { @@ -801,7 +803,7 @@ private async ValueTask DisposeSocketComponentAsync(IAsyncDisposable component, private async ValueTask DisposeSocketAsync(bool asyncReaderDispose) { // writer's internal buffer/channel is not thread-safe, must wait until complete. - CommandWriter.Reset(_socket); + await CommandWriter.FlushAsync().ConfigureAwait(false); if (_socket != null) { From 54c007fd246eb22a6d2a1a1e2f77acd1cb2a8e99 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Fri, 2 Feb 2024 14:25:28 +0000 Subject: [PATCH 20/25] Flush buffers cleanly on dispose --- src/NATS.Client.Core/Commands/CommandWriter.cs | 14 ++++++++++++-- .../Commands/PriorityCommandWriter.cs | 1 - src/NATS.Client.Core/NatsConnection.cs | 5 +---- tests/NATS.Client.Core.Tests/ProtocolTest.cs | 18 +++++++++++------- 4 files changed, 24 insertions(+), 14 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 6edb8c2d6..98f887615 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -76,7 +76,7 @@ public void Reset(ISocketConnection socketConnection) lock (_lock) { _socketConnection = socketConnection; - _ctsReader = CancellationTokenSource.CreateLinkedTokenSource(_cts.Token); + _ctsReader = new CancellationTokenSource(); _readerLoopTask = Task.Run(async () => { @@ -85,7 +85,7 @@ public void Reset(ISocketConnection socketConnection) } } - public async Task FlushAsync() + public async Task CancelReaderLoopAsync() { CancellationTokenSource? cts; Task? readerTask; @@ -126,6 +126,16 @@ public async ValueTask DisposeAsync() _channelLock.Writer.TryComplete(); _channelSize.Writer.TryComplete(); await _pipeWriter.CompleteAsync().ConfigureAwait(false); + + Task? readerTask; + lock (_lock) + { + readerTask = _readerLoopTask; + } + + if (readerTask != null) + await readerTask.ConfigureAwait(false); + } public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken cancellationToken) diff --git a/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs b/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs index 4bd0cce06..3020e84f4 100644 --- a/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs +++ b/src/NATS.Client.Core/Commands/PriorityCommandWriter.cs @@ -18,7 +18,6 @@ public async ValueTask DisposeAsync() { if (Interlocked.Increment(ref _disposed) == 1) { - await CommandWriter.FlushAsync().ConfigureAwait(false); // disposing command writer marks pipe writer as complete await CommandWriter.DisposeAsync().ConfigureAwait(false); } diff --git a/src/NATS.Client.Core/NatsConnection.cs b/src/NATS.Client.Core/NatsConnection.cs index 1afe66f92..0a285d285 100644 --- a/src/NATS.Client.Core/NatsConnection.cs +++ b/src/NATS.Client.Core/NatsConnection.cs @@ -500,7 +500,7 @@ private async void ReconnectLoop() // If dispose this client, WaitForClosed throws OperationCanceledException so stop reconnect-loop correctly. await _socket!.WaitForClosed.ConfigureAwait(false); - await CommandWriter.FlushAsync().ConfigureAwait(false); + await CommandWriter.CancelReaderLoopAsync().ConfigureAwait(false); _logger.LogTrace(NatsLogEvents.Connection, "Connection {Name} is closed. Will cleanup and reconnect", _name); lock (_gate) @@ -802,9 +802,6 @@ private async ValueTask DisposeSocketComponentAsync(IAsyncDisposable component, // Dispose Reader(Drain read buffers but no reads more) private async ValueTask DisposeSocketAsync(bool asyncReaderDispose) { - // writer's internal buffer/channel is not thread-safe, must wait until complete. - await CommandWriter.FlushAsync().ConfigureAwait(false); - if (_socket != null) { await DisposeSocketComponentAsync(_socket, "socket").ConfigureAwait(false); diff --git a/tests/NATS.Client.Core.Tests/ProtocolTest.cs b/tests/NATS.Client.Core.Tests/ProtocolTest.cs index 704f2242c..6c7228c35 100644 --- a/tests/NATS.Client.Core.Tests/ProtocolTest.cs +++ b/tests/NATS.Client.Core.Tests/ProtocolTest.cs @@ -343,10 +343,13 @@ await Retry.Until( } Assert.StartsWith("SUB foo", frames[0]); - Assert.StartsWith("PUB bar1", frames[1]); - Assert.StartsWith("PUB bar2", frames[2]); - Assert.StartsWith("PUB bar3", frames[3]); - Assert.StartsWith("PUB foo", frames[4]); + + for (var i = 0; i < 100; i++) + { + Assert.StartsWith($"PUB bar{i}", frames[i + 1]); + } + + Assert.StartsWith("PUB foo", frames[101]); await nats.DisposeAsync(); } @@ -425,9 +428,10 @@ internal override async ValueTask WriteReconnectCommandsAsync(CommandWriter comm await base.WriteReconnectCommandsAsync(commandWriter, sid); // Any additional commands to send on reconnect - await commandWriter.PublishAsync("bar1", default, default, default, NatsRawSerializer.Default, default); - await commandWriter.PublishAsync("bar2", default, default, default, NatsRawSerializer.Default, default); - await commandWriter.PublishAsync("bar3", default, default, default, NatsRawSerializer.Default, default); + for (var i = 0; i < 100; i++) + { + await commandWriter.PublishAsync($"bar{i}", default, default, default, NatsRawSerializer.Default, default); + } } protected override ValueTask ReceiveInternalAsync(string subject, string? replyTo, ReadOnlySequence? headersBuffer, ReadOnlySequence payloadBuffer) From cce3450d9e7cbcb8037daefb09b6f41c4df8d825 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Fri, 2 Feb 2024 16:33:02 +0000 Subject: [PATCH 21/25] Fixing buffer msg boundry issue --- .../Commands/CommandWriter.cs | 30 +++++++++++++------ .../ConnectionRetryTest.cs | 2 ++ 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 98f887615..4352e974f 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -311,6 +311,7 @@ private static async Task ReaderLoopAsync(ILogger logger, ISocket { try { + var examinedOffset = 0; while (true) { var result = await pipeReader.ReadAsync(cancellationToken).ConfigureAwait(false); @@ -321,7 +322,11 @@ private static async Task ReaderLoopAsync(ILogger logger, ISocket } var buffer = result.Buffer; - var completed = buffer.Start; + var consumed = buffer.Start; + + buffer = buffer.Slice(examinedOffset); + var examined = buffer.Start; + try { if (!buffer.IsEmpty) @@ -334,7 +339,6 @@ private static async Task ReaderLoopAsync(ILogger logger, ISocket try { - completed = buffer.Start; var totalSent = 0; var totalSize = 0; while (totalSent < bufferLength) @@ -346,20 +350,28 @@ private static async Task ReaderLoopAsync(ILogger logger, ISocket while (totalSize < totalSent) { - int size; - while (!channelSize.Reader.TryRead(out size)) + int peek; + while (!channelSize.Reader.TryPeek(out peek)) { await channelSize.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); } + // Don't just mark the message as complete if we have more data to send + if (totalSize + peek > totalSent) + { + break; + } + + var size = await channelSize.Reader.ReadAsync(cancellationToken).ConfigureAwait(false); + totalSize += size; + examinedOffset = 0; } // make sure to mark the buffer only at message boundaries. - // if there was a message sent only partially, we consider it sent even if it might not be - // due to a socket error in following socket send iteration. this is to avoid re-sending - // the same message again, ensuring at-most-once delivery. - completed = buffer.GetPosition(totalSize); + consumed = buffer.GetPosition(totalSize); + examined = buffer.GetPosition(totalSent); + examinedOffset += totalSent - totalSize; } } finally @@ -371,7 +383,7 @@ private static async Task ReaderLoopAsync(ILogger logger, ISocket finally { // Always examine to the end to potentially unblock writer - pipeReader.AdvanceTo(completed, buffer.End); + pipeReader.AdvanceTo(consumed, examined); } if (result.IsCompleted) diff --git a/tests/NATS.Client.Core.Tests/ConnectionRetryTest.cs b/tests/NATS.Client.Core.Tests/ConnectionRetryTest.cs index 261337c60..dadd708a1 100644 --- a/tests/NATS.Client.Core.Tests/ConnectionRetryTest.cs +++ b/tests/NATS.Client.Core.Tests/ConnectionRetryTest.cs @@ -131,6 +131,8 @@ public async Task Reconnect_doesnt_drop_partially_sent_msgs() Assert.True(reconnects > 0, "connection did not reconnect"); Assert.True(received <= sent, $"duplicate messages sent on wire- {sent} sent, {received} received"); + _output.WriteLine($"reconnects: {reconnects}, sent: {sent}, received: {received}"); + // some messages may still be lost, as socket could have been disconnected // after socket.WriteAsync returned, but before OS sent // check to ensure that the loss was < 1% From d3c0288dcd6307be1c0fb9ab8aea3ae70f47665f Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Fri, 2 Feb 2024 17:22:05 +0000 Subject: [PATCH 22/25] Format fixed --- src/NATS.Client.Core/Commands/CommandWriter.cs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 4352e974f..803ee5a42 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -135,7 +135,6 @@ public async ValueTask DisposeAsync() if (readerTask != null) await readerTask.ConfigureAwait(false); - } public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken cancellationToken) From 6334b5747ca6968c139e28b87065a4ba1f17cebb Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Fri, 2 Feb 2024 17:32:35 +0000 Subject: [PATCH 23/25] Buffer position --- src/NATS.Client.Core/Commands/CommandWriter.cs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 803ee5a42..79e64429e 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -322,19 +322,18 @@ private static async Task ReaderLoopAsync(ILogger logger, ISocket var buffer = result.Buffer; var consumed = buffer.Start; - - buffer = buffer.Slice(examinedOffset); - var examined = buffer.Start; + var examined = buffer.GetPosition(examinedOffset); try { if (!buffer.IsEmpty) { - var bufferLength = (int)buffer.Length; + var readBuffer = buffer.Slice(examinedOffset); + var bufferLength = (int)readBuffer.Length; var bytes = ArrayPool.Shared.Rent(bufferLength); - buffer.CopyTo(bytes); - var memory = bytes.AsMemory(0, (int)buffer.Length); + readBuffer.CopyTo(bytes); + var memory = bytes.AsMemory(0, bufferLength); try { From ad5c0b2bfd2ebc42e8a7e8dbf98e0cb9bb571057 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Fri, 2 Feb 2024 18:02:50 +0000 Subject: [PATCH 24/25] Handle socket exception --- src/NATS.Client.Core/Commands/CommandWriter.cs | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index 79e64429e..e211435e9 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -1,5 +1,7 @@ using System.Buffers; using System.IO.Pipelines; +using System.Linq.Expressions; +using System.Net.Sockets; using System.Runtime.CompilerServices; using System.Threading.Channels; using Microsoft.Extensions.Logging; @@ -323,12 +325,12 @@ private static async Task ReaderLoopAsync(ILogger logger, ISocket var buffer = result.Buffer; var consumed = buffer.Start; var examined = buffer.GetPosition(examinedOffset); + var readBuffer = buffer.Slice(examinedOffset); try { - if (!buffer.IsEmpty) + if (!buffer.IsEmpty && !readBuffer.IsEmpty) { - var readBuffer = buffer.Slice(examinedOffset); var bufferLength = (int)readBuffer.Length; var bytes = ArrayPool.Shared.Rent(bufferLength); @@ -341,7 +343,15 @@ private static async Task ReaderLoopAsync(ILogger logger, ISocket var totalSize = 0; while (totalSent < bufferLength) { - var sent = await connection.SendAsync(memory).ConfigureAwait(false); + int sent; + try + { + sent = await connection.SendAsync(memory).ConfigureAwait(false); + } + catch (SocketException) + { + break; + } totalSent += sent; memory = memory[sent..]; From 72f12137bc27f5b9977ab4b3155f0cb3f35ea364 Mon Sep 17 00:00:00 2001 From: Caleb Lloyd <2414837+caleblloyd@users.noreply.github.com> Date: Fri, 2 Feb 2024 14:28:01 -0500 Subject: [PATCH 25/25] throw away bytes in send buffer after a failed send (#368) Signed-off-by: Caleb Lloyd --- .../Commands/CommandWriter.cs | 119 +++++++++++------- .../ConnectionRetryTest.cs | 4 +- 2 files changed, 75 insertions(+), 48 deletions(-) diff --git a/src/NATS.Client.Core/Commands/CommandWriter.cs b/src/NATS.Client.Core/Commands/CommandWriter.cs index e211435e9..31cc150d1 100644 --- a/src/NATS.Client.Core/Commands/CommandWriter.cs +++ b/src/NATS.Client.Core/Commands/CommandWriter.cs @@ -1,7 +1,5 @@ using System.Buffers; using System.IO.Pipelines; -using System.Linq.Expressions; -using System.Net.Sockets; using System.Runtime.CompilerServices; using System.Threading.Channels; using Microsoft.Extensions.Logging; @@ -20,6 +18,9 @@ namespace NATS.Client.Core.Commands; /// internal sealed class CommandWriter : IAsyncDisposable { + // set to a reasonable socket write mem size + private const int MaxSendSize = 16384; + private readonly ILogger _logger; private readonly ObjectPool _pool; private readonly int _arrayPoolInitialSize; @@ -36,6 +37,7 @@ internal sealed class CommandWriter : IAsyncDisposable private readonly PipeReader _pipeReader; private readonly PipeWriter _pipeWriter; private ISocketConnection? _socketConnection; + private Task? _flushTask; private Task? _readerLoopTask; private CancellationTokenSource? _ctsReader; private volatile bool _disposed; @@ -150,16 +152,16 @@ public async ValueTask ConnectAsync(ClientOpts connectOpts, CancellationToken ca throw new ObjectDisposedException(nameof(CommandWriter)); } - _protocolWriter.WriteConnect(_pipeWriter, connectOpts); - - var size = (int)_pipeWriter.UnflushedBytes; - _channelSize.Writer.TryWrite(size); - - var result = await _pipeWriter.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); - if (result.IsCanceled) + if (_flushTask is { IsCompletedSuccessfully: false }) { - throw new OperationCanceledException(); + await _flushTask.WaitAsync(cancellationTimer.Token).ConfigureAwait(false); } + + _protocolWriter.WriteConnect(_pipeWriter, connectOpts); + + _channelSize.Writer.TryWrite((int)_pipeWriter.UnflushedBytes); + var flush = _pipeWriter.FlushAsync(CancellationToken.None); + _flushTask = flush.IsCompletedSuccessfully ? null : flush.AsTask(); } finally { @@ -179,18 +181,17 @@ public async ValueTask PingAsync(PingCommand pingCommand, CancellationToken canc throw new ObjectDisposedException(nameof(CommandWriter)); } - _enqueuePing(pingCommand); + if (_flushTask is { IsCompletedSuccessfully: false }) + { + await _flushTask.WaitAsync(cancellationTimer.Token).ConfigureAwait(false); + } + _enqueuePing(pingCommand); _protocolWriter.WritePing(_pipeWriter); - var size = (int)_pipeWriter.UnflushedBytes; - _channelSize.Writer.TryWrite(size); - - var result = await _pipeWriter.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); - if (result.IsCanceled) - { - throw new OperationCanceledException(); - } + _channelSize.Writer.TryWrite((int)_pipeWriter.UnflushedBytes); + var flush = _pipeWriter.FlushAsync(CancellationToken.None); + _flushTask = flush.IsCompletedSuccessfully ? null : flush.AsTask(); } finally { @@ -210,16 +211,16 @@ public async ValueTask PongAsync(CancellationToken cancellationToken = default) throw new ObjectDisposedException(nameof(CommandWriter)); } - _protocolWriter.WritePong(_pipeWriter); - - var size = (int)_pipeWriter.UnflushedBytes; - _channelSize.Writer.TryWrite(size); - - var result = await _pipeWriter.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); - if (result.IsCanceled) + if (_flushTask is { IsCompletedSuccessfully: false }) { - throw new OperationCanceledException(); + await _flushTask.WaitAsync(cancellationTimer.Token).ConfigureAwait(false); } + + _protocolWriter.WritePong(_pipeWriter); + + _channelSize.Writer.TryWrite((int)_pipeWriter.UnflushedBytes); + var flush = _pipeWriter.FlushAsync(CancellationToken.None); + _flushTask = flush.IsCompletedSuccessfully ? null : flush.AsTask(); } finally { @@ -258,16 +259,16 @@ public async ValueTask SubscribeAsync(int sid, string subject, string? queueGrou throw new ObjectDisposedException(nameof(CommandWriter)); } - _protocolWriter.WriteSubscribe(_pipeWriter, sid, subject, queueGroup, maxMsgs); - - var size = (int)_pipeWriter.UnflushedBytes; - _channelSize.Writer.TryWrite(size); - - var result = await _pipeWriter.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); - if (result.IsCanceled) + if (_flushTask is { IsCompletedSuccessfully: false }) { - throw new OperationCanceledException(); + await _flushTask.WaitAsync(cancellationTimer.Token).ConfigureAwait(false); } + + _protocolWriter.WriteSubscribe(_pipeWriter, sid, subject, queueGroup, maxMsgs); + + _channelSize.Writer.TryWrite((int)_pipeWriter.UnflushedBytes); + var flush = _pipeWriter.FlushAsync(CancellationToken.None); + _flushTask = flush.IsCompletedSuccessfully ? null : flush.AsTask(); } finally { @@ -287,16 +288,16 @@ public async ValueTask UnsubscribeAsync(int sid, int? maxMsgs, CancellationToken throw new ObjectDisposedException(nameof(CommandWriter)); } - _protocolWriter.WriteUnsubscribe(_pipeWriter, sid, maxMsgs); - - var size = (int)_pipeWriter.UnflushedBytes; - _channelSize.Writer.TryWrite(size); - - var result = await _pipeWriter.FlushAsync(cancellationTimer.Token).ConfigureAwait(false); - if (result.IsCanceled) + if (_flushTask is { IsCompletedSuccessfully: false }) { - throw new OperationCanceledException(); + await _flushTask.WaitAsync(cancellationTimer.Token).ConfigureAwait(false); } + + _protocolWriter.WriteUnsubscribe(_pipeWriter, sid, maxMsgs); + + _channelSize.Writer.TryWrite((int)_pipeWriter.UnflushedBytes); + var flush = _pipeWriter.FlushAsync(CancellationToken.None); + _flushTask = flush.IsCompletedSuccessfully ? null : flush.AsTask(); } finally { @@ -343,14 +344,26 @@ private static async Task ReaderLoopAsync(ILogger logger, ISocket var totalSize = 0; while (totalSent < bufferLength) { + var sendMemory = memory; + if (sendMemory.Length > MaxSendSize) + { + // cap the send size, the OS can only handle so much in a send buffer at a time + // also if the send fails, we have to throw this many bytes away + sendMemory = memory[..MaxSendSize]; + } + int sent; + Exception? sendEx = null; try { - sent = await connection.SendAsync(memory).ConfigureAwait(false); + sent = await connection.SendAsync(sendMemory).ConfigureAwait(false); } - catch (SocketException) + catch (Exception ex) { - break; + // we have no idea how many bytes were actually sent, so we have to assume they all were + // this could result in message loss, but is consistent with at-most once delivery + sendEx = ex; + sent = sendMemory.Length; } totalSent += sent; @@ -361,6 +374,7 @@ private static async Task ReaderLoopAsync(ILogger logger, ISocket int peek; while (!channelSize.Reader.TryPeek(out peek)) { + // should never happen; channel sizes are written before flush is called await channelSize.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); } @@ -370,7 +384,12 @@ private static async Task ReaderLoopAsync(ILogger logger, ISocket break; } - var size = await channelSize.Reader.ReadAsync(cancellationToken).ConfigureAwait(false); + int size; + while (!channelSize.Reader.TryRead(out size)) + { + // should never happen; channel sizes are written before flush is called (plus we just peeked) + await channelSize.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); + } totalSize += size; examinedOffset = 0; @@ -380,6 +399,12 @@ private static async Task ReaderLoopAsync(ILogger logger, ISocket consumed = buffer.GetPosition(totalSize); examined = buffer.GetPosition(totalSent); examinedOffset += totalSent - totalSize; + + // throw if there was a send failure + if (sendEx != null) + { + throw sendEx; + } } } finally diff --git a/tests/NATS.Client.Core.Tests/ConnectionRetryTest.cs b/tests/NATS.Client.Core.Tests/ConnectionRetryTest.cs index dadd708a1..132b727c3 100644 --- a/tests/NATS.Client.Core.Tests/ConnectionRetryTest.cs +++ b/tests/NATS.Client.Core.Tests/ConnectionRetryTest.cs @@ -58,6 +58,7 @@ public async Task Retry_and_connect_after_disconnected() [Fact] public async Task Reconnect_doesnt_drop_partially_sent_msgs() { + const int msgSize = 1048576; // 1MiB await using var server = NatsServer.Start(); await using var pubConn = server.CreateClientConnection(); @@ -86,6 +87,7 @@ public async Task Reconnect_doesnt_drop_partially_sent_msgs() } else { + Assert.Equal(msgSize, msg.Data.Length); Interlocked.Increment(ref received); } } @@ -99,7 +101,7 @@ public async Task Reconnect_doesnt_drop_partially_sent_msgs() } var sent = 0; - var data = new byte[1048576]; // 1MiB + var data = new byte[msgSize]; var sendTask = Task.Run(async () => { while (!stopCts.IsCancellationRequested)