ReadyTask & DisconnectedTask Null Reference Fix

src/NexNet/Internals/NexNetSessionConfigurations.cs

@@ -1,4 +1,5 @@
-using NexNet.Cache;
+using System.Threading.Tasks;
+using NexNet.Cache;
 using NexNet.Invocation;
 using NexNet.Transports;
 
@@ -21,4 +22,6 @@ internal readonly struct NexNetSessionConfigurations<THub, TProxy>
     public required long Id { get; init; }
 
     public required THub Hub { get; init; }
+    public TaskCompletionSource? ReadyTaskCompletionSource { get; init; }
+    public TaskCompletionSource? DisconnectedTaskCompletionSource { get; init; }
 }

src/NexNet/NexNetClient.cs

@@ -23,6 +23,7 @@ public sealed class NexNetClient<TClientHub, TServerProxy> : INexNetClient
     private readonly SessionCacheManager<TServerProxy> _cacheManager;
     private readonly TClientHub _hub;
     private NexNetSession<TClientHub, TServerProxy>? _session;
+    private TaskCompletionSource? _disconnectedTaskCompletionSource;
 
     internal NexNetSession<TClientHub, TServerProxy>? Session => _session;
 
@@ -45,12 +46,12 @@ public sealed class NexNetClient<TClientHub, TServerProxy> : INexNetClient
     /// Task which completes upon the completed connection and optional authentication of the client.
     /// Null when the client is has not started connection or after disconnection.
     /// </summary>
-    public Task? ReadyTask => _session?.ReadyTask;
+    public Task? ReadyTask { get; private set; }
 
     /// <summary>
     /// Task which completes upon the disconnection of the client.
     /// </summary>
-    public Task DisconnectedTask => _session?.DisconnectedTask ?? Task.CompletedTask;
+    public Task DisconnectedTask => _disconnectedTaskCompletionSource?.Task ?? Task.CompletedTask;
 
     /// <summary>
     /// Creates a NexNet client for communication with a matching NexNet server.
@@ -79,6 +80,14 @@ public async Task ConnectAsync()
         if (_session != null)
             throw new InvalidOperationException("Client is already connected.");
 
+        // Set the ready task completion source now and get the task since the ConnectTransport
+        // call below can/will await and cause any usages of the ReadyTask used outside this function call
+        // to be null.
+        var readyTaskCompletionSource = new TaskCompletionSource();
+        var disconnectedTaskCompletionSource = new TaskCompletionSource();
+        ReadyTask = readyTaskCompletionSource.Task;
+        _disconnectedTaskCompletionSource = disconnectedTaskCompletionSource;
+
         var client = await _config.ConnectTransport().ConfigureAwait(false);
 
         _config.InternalOnClientConnect?.Invoke();
@@ -91,7 +100,9 @@ public async Task ConnectAsync()
             SessionManager = null,
             IsServer = false,
             Id = 0,
-            Hub = _hub
+            Hub = _hub,
+            ReadyTaskCompletionSource = readyTaskCompletionSource,
+            DisconnectedTaskCompletionSource = disconnectedTaskCompletionSource
         };
 
         _session = new NexNetSession<TClientHub, TServerProxy>(config)

src/NexNet/NexNetSession.cs

@@ -44,8 +44,8 @@ internal class NexNetSession<THub, TProxy> : INexNetSession<TProxy>
     private bool _isReconnected;
     private DisconnectReason _registeredDisconnectReason = DisconnectReason.None;
 
-    private readonly TaskCompletionSource _readyTaskCompletionSource = new TaskCompletionSource();
-    private readonly TaskCompletionSource _disconnectedTaskCompletionSource = new TaskCompletionSource();
+    private readonly TaskCompletionSource? _readyTaskCompletionSource;
+    private readonly TaskCompletionSource? _disconnectedTaskCompletionSource;
 
     public long Id { get; }
 
@@ -66,10 +66,6 @@ internal class NexNetSession<THub, TProxy> : INexNetSession<TProxy>
 
     public ConnectionState State { get; private set; }
 
-    public Task ReadyTask => _readyTaskCompletionSource.Task;
-
-    public Task DisconnectedTask => _disconnectedTaskCompletionSource.Task;
-
     public NexNetSession(in NexNetSessionConfigurations<THub, TProxy> configurations)
     {
         State = ConnectionState.Connecting;
@@ -86,6 +82,9 @@ public NexNetSession(in NexNetSessionConfigurations<THub, TProxy> configurations
             ? new ServerSessionContext<TProxy>(this, _sessionManager!)
             : new ClientSessionContext<TProxy>(this);
 
+        _readyTaskCompletionSource = configurations.ReadyTaskCompletionSource;
+        _disconnectedTaskCompletionSource = configurations.DisconnectedTaskCompletionSource;
+
         SessionInvocationStateManager = new SessionInvocationStateManager(configurations.Cache);
         SessionStore = new SessionStore();
         _invocationSemaphore = new SemaphoreSlim(configurations.Configs.MaxConcurrentConnectionInvocations,
@@ -469,7 +468,7 @@ static async void InvokeOnConnected(object? sessionObj)
 
             _ = Task.Factory.StartNew(InvokeOnConnected, this);
 
-            _readyTaskCompletionSource.TrySetResult();
+            _readyTaskCompletionSource?.TrySetResult();
         }
         else if (message is ServerGreetingMessage)
         {
@@ -479,7 +478,7 @@ static async void InvokeOnConnected(object? sessionObj)
 
             _ = Task.Factory.StartNew(InvokeOnConnected, this);
 
-            _readyTaskCompletionSource.TrySetResult();
+            _readyTaskCompletionSource?.TrySetResult();
         }
         else if (message is InvocationRequestMessage invocationRequestMessage)
         {
@@ -689,8 +688,8 @@ private async Task DisconnectCore(DisconnectReason reason, bool sendDisconnect)
         _invocationTaskArgumentsPool.Clear();
 
         // Let any waiting tasks pass.
-        _disconnectedTaskCompletionSource.TrySetResult();
-        _readyTaskCompletionSource.TrySetResult();
+        _disconnectedTaskCompletionSource?.TrySetResult();
+        _readyTaskCompletionSource?.TrySetResult();
     }
 
     private class InvocationTaskArguments