Socket Connection Factory

src/NATS.Client.Core/Commands/CommandWriter.cs

@@ -49,7 +49,7 @@ internal sealed class CommandWriter : IAsyncDisposable
     private readonly PipeWriter _pipeWriter;
     private readonly SemaphoreSlim _semLock = new(1);
     private readonly PartialSendFailureCounter _partialSendFailureCounter = new();
-    private ISocketConnection? _socketConnection;
+    private SocketConnectionWrapper? _socketConnection;
     private Task? _flushTask;
     private Task? _readerLoopTask;
     private CancellationTokenSource? _ctsReader;
@@ -87,7 +87,7 @@ public CommandWriter(string name, NatsConnection connection, ObjectPool pool, Na
         _logger.LogDebug(NatsLogEvents.Buffer, "Created {Name}", _name);
     }
 
-    public void Reset(ISocketConnection socketConnection)
+    public void Reset(SocketConnectionWrapper socketConnection)
     {
         _logger.LogDebug(NatsLogEvents.Buffer, "Resetting {Name}", _name);
 
@@ -481,7 +481,7 @@ internal async Task TestStallFlushAsync(TimeSpan timeSpan, CancellationToken can
 
     private static async Task ReaderLoopAsync(
         ILogger<CommandWriter> logger,
-        ISocketConnection connection,
+        SocketConnectionWrapper socketConnection,
         PipeReader pipeReader,
         Channel<int> channelSize,
         Memory<byte> consolidateMem,
@@ -535,7 +535,7 @@ private static async Task ReaderLoopAsync(
                                 stopwatch.Restart();
                             }
 
-                            sent = await connection.SendAsync(sendMem).ConfigureAwait(false);
+                            sent = await socketConnection.SendAsync(sendMem).ConfigureAwait(false);
 
                             if (trace)
                             {
@@ -664,7 +664,7 @@ private static async Task ReaderLoopAsync(
                 // We signal the connection to disconnect, which will trigger a reconnect
                 // in the connection loop.  This is necessary because the connection may
                 // be half-open, and we can't rely on the reader loop to detect that.
-                connection.SignalDisconnected(e);
+                socketConnection.SignalDisconnected(e);
             }
             catch (Exception e1)
             {

src/NATS.Client.Core/Commands/PriorityCommandWriter.cs

@@ -6,7 +6,7 @@ internal sealed class PriorityCommandWriter : IAsyncDisposable
 {
     private int _disposed;
 
-    public PriorityCommandWriter(NatsConnection connection, ObjectPool pool, ISocketConnection socketConnection, NatsOpts opts, ConnectionStatsCounter counter, Action<PingCommand> enqueuePing)
+    public PriorityCommandWriter(NatsConnection connection, ObjectPool pool, SocketConnectionWrapper socketConnection, NatsOpts opts, ConnectionStatsCounter counter, Action<PingCommand> enqueuePing)
     {
         CommandWriter = new CommandWriter("init", connection, pool, opts, counter, enqueuePing);
         CommandWriter.Reset(socketConnection);

src/NATS.Client.Core/INatsConnection.cs

@@ -64,7 +64,7 @@ public interface INatsConnection : INatsClient
     /// <summary>
     /// Hook when socket is available.
     /// </summary>
-    Func<ISocketConnection, ValueTask<ISocketConnection>>? OnSocketAvailableAsync { get; set; }
+    Func<INatsSocketConnection, ValueTask<INatsSocketConnection>>? OnSocketAvailableAsync { get; set; }
 
     /// <summary>
     /// Publishes a serializable message payload to the given subject name, optionally supplying a reply subject.

src/NATS.Client.Core/INatsSocketConnection.cs

@@ -0,0 +1,56 @@
+using System.Net.Sockets;
+
+namespace NATS.Client.Core;
+
+/// <summary>
+/// Represents a socket connection to a NATS server.
+/// </summary>
+/// <remarks>
+/// This interface defines the contract for low-level socket connections to NATS servers,
+/// providing methods for sending and receiving data and disposing the socket.
+/// Disposing should attempt to perform graceful shutdown of the socket.
+/// </remarks>
+public interface INatsSocketConnection : IAsyncDisposable
+{
+    /// <summary>
+    /// Sends data asynchronously over the connection.
+    /// </summary>
+    /// <param name="buffer">The buffer containing the data to send.</param>
+    /// <returns>A task representing the asynchronous send operation with the number of bytes sent.</returns>
+    ValueTask<int> SendAsync(ReadOnlyMemory<byte> buffer);
+
+    /// <summary>
+    /// Receives data asynchronously from the connection.
+    /// </summary>
+    /// <param name="buffer">The buffer to store the received data.</param>
+    /// <returns>A task representing the asynchronous receive operation with the number of bytes received.</returns>
+    ValueTask<int> ReceiveAsync(Memory<byte> buffer);
+}
+
+/// <summary>
+/// Obsolete, use <see cref="INatsSocketConnection"/> instead
+/// </summary>
+[Obsolete("This interface is obsolete, use INatsSocketConnection instead.", error: false)]
+public interface ISocketConnection : INatsSocketConnection
+{
+    Task<Exception> WaitForClosed { get; }
+
+    void SignalDisconnected(Exception exception);
+
+    ValueTask AbortConnectionAsync(CancellationToken cancellationToken);
+}
+
+/// <summary>
+/// Represents a NATS socket connection that can be upgraded to use TLS.
+/// </summary>
+/// <remarks>
+/// This interface extends <see cref="INatsSocketConnection"/> to provide access to the underlying
+/// socket that is used to create a SslStream.
+/// </remarks>
+public interface INatsTlsUpgradeableSocketConnection : INatsSocketConnection
+{
+    /// <summary>
+    /// Gets the underlying Socket instance for the connection.
+    /// </summary>
+    Socket Socket { get; }
+}

src/NATS.Client.Core/INatsSocketConnectionFactory.cs

@@ -0,0 +1,23 @@
+namespace NATS.Client.Core;
+
+/// <summary>
+/// Factory interface for creating socket connections to NATS servers.
+/// </summary>
+/// <remarks>
+/// This interface abstracts the creation of socket connections to NATS servers,
+/// allowing for custom implementations of connection logic.
+/// </remarks>
+public interface INatsSocketConnectionFactory
+{
+    /// <summary>
+    /// Establishes a connection to a NATS server at the specified URI.
+    /// </summary>
+    /// <param name="uri">The URI of the NATS server to connect to.</param>
+    /// <param name="opts">The Options associated with the NATS Client.</param>
+    /// <param name="cancellationToken">A token to cancel the asynchronous operation.</param>
+    /// <returns>
+    /// A task that represents the asynchronous connect operation. The task result contains
+    /// the socket connection.
+    /// </returns>
+    ValueTask<INatsSocketConnection> ConnectAsync(Uri uri, NatsOpts opts, CancellationToken cancellationToken);
+}

src/NATS.Client.Core/Internal/NatsReadProtocolProcessor.cs

@@ -24,7 +24,7 @@ internal sealed class NatsReadProtocolProcessor : IAsyncDisposable
     private readonly bool _trace;
     private int _disposed;
 
-    public NatsReadProtocolProcessor(ISocketConnection socketConnection, NatsConnection connection, TaskCompletionSource waitForInfoSignal, TaskCompletionSource waitForPongOrErrorSignal, Task infoParsed)
+    public NatsReadProtocolProcessor(SocketConnectionWrapper socketConnection, NatsConnection connection, TaskCompletionSource waitForInfoSignal, TaskCompletionSource waitForPongOrErrorSignal, Task infoParsed)
     {
         _connection = connection;
         _logger = connection.Opts.LoggerFactory.CreateLogger<NatsReadProtocolProcessor>();

src/NATS.Client.Core/Internal/NatsUri.cs

@@ -1,6 +1,6 @@
 namespace NATS.Client.Core.Internal;
 
-internal sealed class NatsUri : IEquatable<NatsUri>
+internal sealed record NatsUri
 {
     public const string DefaultScheme = "nats";
 
@@ -57,7 +57,7 @@ public NatsUri(string urlString, bool isSeed, string defaultScheme = DefaultSche
         _redacted = IsWebSocket && Uri.AbsolutePath != "/" ? uriBuilder.Uri.ToString() : uriBuilder.Uri.ToString().Trim('/');
     }
 
-    public Uri Uri { get; }
+    public Uri Uri { get; init; }
 
     public bool IsSeed { get; }
 
@@ -69,28 +69,5 @@ public NatsUri(string urlString, bool isSeed, string defaultScheme = DefaultSche
 
     public int Port => Uri.Port;
 
-    public NatsUri CloneWith(string host, int? port = default)
-    {
-        var newUri = new UriBuilder(Uri)
-        {
-            Host = host,
-            Port = port ?? Port,
-        }.Uri.ToString();
-
-        return new NatsUri(newUri, IsSeed);
-    }
-
     public override string ToString() => _redacted;
-
-    public override int GetHashCode() => Uri.GetHashCode();
-
-    public bool Equals(NatsUri? other)
-    {
-        if (other == null)
-            return false;
-        if (ReferenceEquals(this, other))
-            return true;
-
-        return Uri.Equals(other.Uri);
-    }
 }

src/NATS.Client.Core/Internal/ServerInfo.cs

@@ -7,71 +7,71 @@ namespace NATS.Client.Core.Internal;
 internal sealed record ServerInfo : INatsServerInfo
 {
     [JsonPropertyName("server_id")]
-    public string Id { get; set; } = string.Empty;
+    public string Id { get; init; } = string.Empty;
 
     [JsonPropertyName("server_name")]
-    public string Name { get; set; } = string.Empty;
+    public string Name { get; init; } = string.Empty;
 
     [JsonPropertyName("version")]
-    public string Version { get; set; } = string.Empty;
+    public string Version { get; init; } = string.Empty;
 
     [JsonPropertyName("proto")]
-    public long ProtocolVersion { get; set; }
+    public long ProtocolVersion { get; init; }
 
     [JsonPropertyName("git_commit")]
-    public string GitCommit { get; set; } = string.Empty;
+    public string GitCommit { get; init; } = string.Empty;
 
     [JsonPropertyName("go")]
-    public string GoVersion { get; set; } = string.Empty;
+    public string GoVersion { get; init; } = string.Empty;
 
     [JsonPropertyName("host")]
-    public string Host { get; set; } = string.Empty;
+    public string Host { get; init; } = string.Empty;
 
     [JsonPropertyName("port")]
-    public int Port { get; set; }
+    public int Port { get; init; }
 
     [JsonPropertyName("headers")]
-    public bool HeadersSupported { get; set; }
+    public bool HeadersSupported { get; init; }
 
     [JsonPropertyName("auth_required")]
-    public bool AuthRequired { get; set; }
+    public bool AuthRequired { get; init; }
 
     [JsonPropertyName("tls_required")]
-    public bool TlsRequired { get; set; }
+    public bool TlsRequired { get; init; }
 
     [JsonPropertyName("tls_verify")]
-    public bool TlsVerify { get; set; }
+    public bool TlsVerify { get; init; }
 
     [JsonPropertyName("tls_available")]
-    public bool TlsAvailable { get; set; }
+    public bool TlsAvailable { get; init; }
 
     [JsonPropertyName("max_payload")]
-    public int MaxPayload { get; set; }
+    public int MaxPayload { get; init; }
 
     [JsonPropertyName("jetstream")]
-    public bool JetStreamAvailable { get; set; }
+    public bool JetStreamAvailable { get; init; }
 
     [JsonPropertyName("client_id")]
-    public ulong ClientId { get; set; }
+    public ulong ClientId { get; init; }
 
     [JsonPropertyName("client_ip")]
-    public string ClientIp { get; set; } = string.Empty;
+    public string ClientIp { get; init; } = string.Empty;
 
     [JsonPropertyName("nonce")]
-    public string? Nonce { get; set; }
+    public string? Nonce { get; init; }
 
     [JsonPropertyName("cluster")]
-    public string? Cluster { get; set; }
+    public string? Cluster { get; init; }
 
     [JsonPropertyName("cluster_dynamic")]
-    public bool ClusterDynamic { get; set; }
+    public bool ClusterDynamic { get; init; }
 
     [JsonPropertyName("connect_urls")]
-    public string[]? ClientConnectUrls { get; set; }
+    public string[]? ClientConnectUrls { get; init; }
 
     [JsonPropertyName("ws_connect_urls")]
-    public string[]? WebSocketConnectUrls { get; set; }
+    public string[]? WebSocketConnectUrls { get; init; }
 
     [JsonPropertyName("ldm")]
-    public bool LameDuckMode { get; set; }
+    public bool LameDuckMode { get; init; }
 }

src/NATS.Client.Core/Internal/SocketConnectionWrapper.cs

@@ -0,0 +1,56 @@
+namespace NATS.Client.Core.Internal;
+
+// wraps INatsSocketConnection and signals on disconnect
+internal record SocketConnectionWrapper(INatsSocketConnection InnerSocket) : INatsSocketConnection
+{
+    private readonly TaskCompletionSource _waitForClosedSource =
+#if NETSTANDARD
+        new(TaskCreationOptions.None);
+#else
+        new();
+#endif
+
+    private readonly SemaphoreSlim _sem = new(1);
+
+    public Task WaitForClosed => _waitForClosedSource.Task;
+
+    public void SignalDisconnected(Exception exception)
+    {
+        // guard with semaphore so this doesn't race DisposeAsync
+        _sem.Wait();
+        try
+        {
+            _waitForClosedSource.TrySetException(exception);
+        }
+        finally
+        {
+            _sem.Release(1);
+        }
+    }
+
+    public ValueTask<int> SendAsync(ReadOnlyMemory<byte> buffer) => InnerSocket.SendAsync(buffer);
+
+    public ValueTask<int> ReceiveAsync(Memory<byte> buffer) => InnerSocket.ReceiveAsync(buffer);
+
+    public async ValueTask DisposeAsync()
+    {
+        // guard with semaphore in case SignalDisconnected races
+        await _sem.WaitAsync().ConfigureAwait(false);
+        try
+        {
+            // dispose first, then signal
+            try
+            {
+                await InnerSocket.DisposeAsync().ConfigureAwait(false);
+            }
+            finally
+            {
+                _waitForClosedSource.TrySetResult();
+            }
+        }
+        finally
+        {
+            _sem.Release(1);
+        }
+    }
+}

src/NATS.Client.Core/Internal/SocketReader.cs

@@ -14,11 +14,11 @@ internal sealed class SocketReader
     private readonly Stopwatch _stopwatch = new Stopwatch();
     private readonly ILogger<SocketReader> _logger;
     private readonly bool _isTraceLogging;
-    private ISocketConnection _socketConnection;
+    private readonly SocketConnectionWrapper _socketConnection;
 
     private Memory<byte> _availableMemory;
 
-    public SocketReader(ISocketConnection socketConnection, int minimumBufferSize, ConnectionStatsCounter counter, ILoggerFactory loggerFactory)
+    public SocketReader(SocketConnectionWrapper socketConnection, int minimumBufferSize, ConnectionStatsCounter counter, ILoggerFactory loggerFactory)
     {
         _socketConnection = socketConnection;
         _minimumBufferSize = minimumBufferSize;