Func<NatsUri, NatsOpts, CancellationToken, ValueTask<ISocketConnection>>? SocketConnectionFactory (#804)

Author: wolfkor

Diff for: src/NATS.Client.Core/DefaultSocketConnectionFactory.cs

@@ -61,6 +61,11 @@ public interface INatsConnection : INatsClient
     /// </summary>
     Func<(string Host, int Port), ValueTask<(string Host, int Port)>>? OnConnectingAsync { get; set; }
 
+    /// <summary>
+    /// Hook when socket is available.
+    /// </summary>
+    Func<ISocketConnection, ValueTask<ISocketConnection>>? OnSocketAvailableAsync { get; set; }
+
     /// <summary>
     /// Publishes a serializable message payload to the given subject name, optionally supplying a reply subject.
     /// </summary>

Diff for: src/NATS.Client.Core/INatsConnection.cs

@@ -3,39 +3,36 @@
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using Microsoft.Extensions.Logging;
-using NATS.Client.Core.Internal;
 
-namespace NATS.Client.Core;
+namespace NATS.Client.Core.Internal;
 
-public sealed class SocketClosedException : Exception
+internal sealed class SocketClosedException : Exception
 {
     public SocketClosedException(Exception? innerException)
         : base("Socket has been closed.", innerException)
     {
     }
 }
 
-public class TcpConnection : ISocketConnection
+internal sealed class TcpConnection : ISocketConnection
 {
+    private readonly ILogger _logger;
+    private readonly Socket _socket;
     private readonly TaskCompletionSource<Exception> _waitForClosedSource = new();
     private int _disposed;
 
     public TcpConnection(ILogger logger)
     {
-        Logger = logger;
-        Socket = new Socket(Socket.OSSupportsIPv6 ? AddressFamily.InterNetworkV6 : AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
+        _logger = logger;
+        _socket = new Socket(Socket.OSSupportsIPv6 ? AddressFamily.InterNetworkV6 : AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
         if (Socket.OSSupportsIPv6)
         {
-            Socket.DualMode = true;
+            _socket.DualMode = true;
         }
 
-        Socket.NoDelay = true;
+        _socket.NoDelay = true;
     }
 
-    public ILogger Logger { get; }
-
-    public Socket Socket { get; }
-
     public Task<Exception> WaitForClosed => _waitForClosedSource.Task;
 
     // CancellationToken is not used, operation lifetime is completely same as socket.
@@ -51,9 +48,9 @@ public TcpConnection(ILogger logger)
     public ValueTask ConnectAsync(string host, int port, CancellationToken cancellationToken)
     {
 #if NETSTANDARD
-        return new ValueTask(Socket.ConnectAsync(host, port).WaitAsync(Timeout.InfiniteTimeSpan, cancellationToken));
+        return new ValueTask(_socket.ConnectAsync(host, port).WaitAsync(Timeout.InfiniteTimeSpan, cancellationToken));
 #else
-        return Socket.ConnectAsync(host, port, cancellationToken);
+        return _socket.ConnectAsync(host, port, cancellationToken);
 #endif
     }
 
@@ -66,9 +63,9 @@ public async ValueTask ConnectAsync(string host, int port, TimeSpan timeout)
         try
         {
 #if NETSTANDARD
-            await Socket.ConnectAsync(host, port).WaitAsync(timeout, cts.Token).ConfigureAwait(false);
+            await _socket.ConnectAsync(host, port).WaitAsync(timeout, cts.Token).ConfigureAwait(false);
 #else
-                await Socket.ConnectAsync(host, port, cts.Token).ConfigureAwait(false);
+            await _socket.ConnectAsync(host, port, cts.Token).ConfigureAwait(false);
 #endif
         }
         catch (Exception ex)
@@ -94,9 +91,9 @@ public ValueTask<int> SendAsync(ReadOnlyMemory<byte> buffer)
             segment = new ArraySegment<byte>(buffer.ToArray());
         }
 
-        return new ValueTask<int>(Socket.SendAsync(segment, SocketFlags.None));
+        return new ValueTask<int>(_socket.SendAsync(segment, SocketFlags.None));
 #else
-        return Socket.SendAsync(buffer, SocketFlags.None, CancellationToken.None);
+        return _socket.SendAsync(buffer, SocketFlags.None, CancellationToken.None);
 #endif
     }
 
@@ -109,19 +106,19 @@ public ValueTask<int> ReceiveAsync(Memory<byte> buffer)
             ThrowHelper.ThrowInvalidOperationException("Can't get underlying array");
         }
 
-        return new ValueTask<int>(Socket.ReceiveAsync(segment, SocketFlags.None));
+        return new ValueTask<int>(_socket.ReceiveAsync(segment, SocketFlags.None));
 #else
-        return Socket.ReceiveAsync(buffer, SocketFlags.None, CancellationToken.None);
+        return _socket.ReceiveAsync(buffer, SocketFlags.None, CancellationToken.None);
 #endif
     }
 
     public ValueTask AbortConnectionAsync(CancellationToken cancellationToken)
     {
 #if NETSTANDARD
-        Socket.Disconnect(false);
+        _socket.Disconnect(false);
         return default;
 #else
-        return Socket.DisconnectAsync(false, cancellationToken);
+        return _socket.DisconnectAsync(false, cancellationToken);
 #endif
     }
 
@@ -135,19 +132,17 @@ public ValueTask DisposeAsync()
             }
             catch
             {
-                // ignored
             }
 
             try
             {
-                Socket.Shutdown(SocketShutdown.Both);
+                _socket.Shutdown(SocketShutdown.Both);
             }
             catch
             {
-                // ignored
             }
 
-            Socket.Dispose();
+            _socket.Dispose();
         }
 
         return default;
@@ -161,13 +156,13 @@ public void SignalDisconnected(Exception exception)
 
     // NetworkStream will own the Socket, so mark as disposed
     // in order to skip socket.Dispose() in DisposeAsync
-    internal SslStreamConnection UpgradeToSslStreamConnection(NatsTlsOpts tlsOpts)
+    public SslStreamConnection UpgradeToSslStreamConnection(NatsTlsOpts tlsOpts)
     {
         if (Interlocked.Increment(ref _disposed) == 1)
         {
             return new SslStreamConnection(
-                Logger,
-                Socket,
+                _logger,
+                _socket,
                 tlsOpts,
                 _waitForClosedSource);
         }

Diff for: src/NATS.Client.Core/ISocketConnectionFactory.cs

@@ -6,15 +6,19 @@
 using System.Runtime.InteropServices;
 #endif
 
-namespace NATS.Client.Core;
+namespace NATS.Client.Core.Internal;
 
-public class WebSocketConnection : ISocketConnection
+internal sealed class WebSocketConnection : ISocketConnection
 {
+    private readonly ClientWebSocket _socket;
     private readonly TaskCompletionSource<Exception> _waitForClosedSource = new();
     private readonly TimeSpan _socketCloseTimeout = TimeSpan.FromSeconds(5); // matches _socketComponentDisposeTimeout in NatsConnection.cs
     private int _disposed;
 
-    public ClientWebSocket Socket { get; } = new();
+    public WebSocketConnection()
+    {
+        _socket = new ClientWebSocket();
+    }
 
     public Task<Exception> WaitForClosed => _waitForClosedSource.Task;
 
@@ -30,7 +34,7 @@ public class WebSocketConnection : ISocketConnection
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public Task ConnectAsync(Uri uri, CancellationToken cancellationToken)
     {
-        return Socket.ConnectAsync(uri, cancellationToken);
+        return _socket.ConnectAsync(uri, cancellationToken);
     }
 
     /// <summary>
@@ -41,8 +45,8 @@ public async ValueTask ConnectAsync(NatsUri uri, NatsOpts opts)
         using var cts = new CancellationTokenSource(opts.ConnectTimeout);
         try
         {
-            await opts.WebSocketOpts.ApplyClientWebSocketOptionsAsync(Socket.Options, uri, opts.TlsOpts, cts.Token).ConfigureAwait(false);
-            await Socket.ConnectAsync(uri.Uri, cts.Token).ConfigureAwait(false);
+            await opts.WebSocketOpts.ApplyClientWebSocketOptionsAsync(_socket.Options, uri, opts.TlsOpts, cts.Token).ConfigureAwait(false);
+            await _socket.ConnectAsync(uri.Uri, cts.Token).ConfigureAwait(false);
         }
         catch (Exception ex)
         {
@@ -67,9 +71,9 @@ public async ValueTask<int> SendAsync(ReadOnlyMemory<byte> buffer)
             segment = new ArraySegment<byte>(buffer.ToArray());
         }
 
-        await Socket.SendAsync(segment, WebSocketMessageType.Binary, true, CancellationToken.None).ConfigureAwait(false);
+        await _socket.SendAsync(segment, WebSocketMessageType.Binary, true, CancellationToken.None).ConfigureAwait(false);
 #else
-        await Socket.SendAsync(buffer, WebSocketMessageType.Binary, WebSocketMessageFlags.EndOfMessage, CancellationToken.None).ConfigureAwait(false);
+        await _socket.SendAsync(buffer, WebSocketMessageType.Binary, WebSocketMessageFlags.EndOfMessage, CancellationToken.None).ConfigureAwait(false);
 #endif
         return buffer.Length;
     }
@@ -83,9 +87,9 @@ public async ValueTask<int> ReceiveAsync(Memory<byte> buffer)
             ThrowHelper.ThrowInvalidOperationException("Can't get underlying array");
         }
 
-        var wsRead = await Socket.ReceiveAsync(segment, CancellationToken.None).ConfigureAwait(false);
+        var wsRead = await _socket.ReceiveAsync(segment, CancellationToken.None).ConfigureAwait(false);
 #else
-        var wsRead = await Socket.ReceiveAsync(buffer, CancellationToken.None).ConfigureAwait(false);
+        var wsRead = await _socket.ReceiveAsync(buffer, CancellationToken.None).ConfigureAwait(false);
 #endif
         return wsRead.Count;
     }
@@ -94,7 +98,7 @@ public ValueTask AbortConnectionAsync(CancellationToken cancellationToken)
     {
         // ClientWebSocket.Abort() doesn't accept a cancellation token, so check at the beginning of this method
         cancellationToken.ThrowIfCancellationRequested();
-        Socket.Abort();
+        _socket.Abort();
         return default;
     }
 
@@ -113,13 +117,13 @@ public async ValueTask DisposeAsync()
             try
             {
                 var cts = new CancellationTokenSource(_socketCloseTimeout);
-                await Socket.CloseAsync(WebSocketCloseStatus.NormalClosure, default, cts.Token).ConfigureAwait(false);
+                await _socket.CloseAsync(WebSocketCloseStatus.NormalClosure, default, cts.Token).ConfigureAwait(false);
             }
             catch
             {
             }
 
-            Socket.Dispose();
+            _socket.Dispose();
         }
     }