fix: Websockets Transport now handles being disabled for scene change…

…s (#1994)

* fix: Websockets Transport now handles being disabled for scene changes

* Removed unused method

* Removed test warnings

* Changed back to await Task.Run

* adding queue to ClientJs

keeping standalone client without queue
dont need common class for now as that over complicates the fix

Co-authored-by: Chris Langsenkamp <chris@clevertech.net>

Assets/Mirror/Runtime/Transport/Websocket/Client.cs

@@ -79,6 +79,8 @@ public async void Connect(Uri uri)
             }
         }
 
+        public bool enabled;
+
         async Task ReceiveLoop(WebSocket webSocket, CancellationToken token)
         {
             byte[] buffer = new byte[MaxMessageSize];
@@ -87,6 +89,8 @@ async Task ReceiveLoop(WebSocket webSocket, CancellationToken token)
             {
                 WebSocketReceiveResult result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), token);
 
+                await Task.Run(WaitForEnabled);
+
                 if (result == null)
                     break;
                 if (result.MessageType == WebSocketMessageType.Close)
@@ -109,6 +113,17 @@ async Task ReceiveLoop(WebSocket webSocket, CancellationToken token)
             }
         }
 
+        void WaitForEnabled()
+        {
+            while (!enabled) { Task.Delay(10); }
+        }
+
+        public bool ProcessClientMessage()
+        {
+            // message in standalone client don't use queue to process
+            return false;
+        }
+
         // a message might come splitted in multiple frames
         // collect all frames
         async Task<ArraySegment<byte>> ReadFrames(WebSocketReceiveResult result, WebSocket webSocket, byte[] buffer)

Assets/Mirror/Runtime/Transport/Websocket/ClientJs.cs

@@ -1,17 +1,11 @@
 #if UNITY_WEBGL && !UNITY_EDITOR
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.Linq;
-using System.Net;
-using System.Net.Sockets;
-using System.Net.WebSockets;
 using System.Runtime.InteropServices;
 using System.Threading;
-using System.Threading.Tasks;
 using AOT;
-using Ninja.WebSockets;
-using UnityEngine;
 
 namespace Mirror.Websocket
 {
@@ -28,6 +22,9 @@ public class Client
         public event Action Disconnected;
         public event Action<Exception> ReceivedError;
 
+        readonly ConcurrentQueue<byte[]> receivedQueue = new ConcurrentQueue<byte[]>();
+
+        public bool enabled;
         public bool Connecting { get; set; }
         public bool IsConnected
         {
@@ -65,6 +62,17 @@ public void Send(ArraySegment<byte> segment)
             SocketSend(nativeRef, segment.Array, segment.Count);
         }
 
+        public bool ProcessClientMessage()
+        {
+            if (receivedQueue.TryDequeue(out byte[] data))
+            {
+                clients[id].ReceivedData(new ArraySegment<byte>(data));
+
+                return true;
+            }
+            return false;
+        }
+
 #region Javascript native functions
         [DllImport("__Internal")]
         static extern int SocketCreate(
@@ -108,7 +116,7 @@ public static void OnData(int id, IntPtr ptr, int length)
             byte[] data = new byte[length];
             Marshal.Copy(ptr, data, 0, length);
 
-            clients[id].ReceivedData(new ArraySegment<byte>(data));
+            clients[id].receivedQueue.Enqueue(data);
         }
 #endregion
     }

Assets/Mirror/Runtime/Transport/Websocket/Server.cs

@@ -178,6 +178,8 @@ bool CertVerificationCallback(object sender, X509Certificate certificate, X509Ch
             return true;
         }
 
+        public bool enabled;
+
         async Task ReceiveLoopAsync(WebSocket webSocket, CancellationToken token)
         {
             int connectionId = NextConnectionId();
@@ -194,6 +196,8 @@ async Task ReceiveLoopAsync(WebSocket webSocket, CancellationToken token)
                 {
                     WebSocketReceiveResult result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), token);
 
+                    await Task.Run(WaitForEnabled);
+
                     if (result.MessageType == WebSocketMessageType.Close)
                     {
                         Debug.Log($"Client initiated close. Status: {result.CloseStatus} Description: {result.CloseStatusDescription}");
@@ -228,6 +232,11 @@ async Task ReceiveLoopAsync(WebSocket webSocket, CancellationToken token)
             }
         }
 
+        void WaitForEnabled()
+        {
+            while (!enabled) { Task.Delay(10); }
+        }
+
         // a message might come splitted in multiple frames
         // collect all frames
         async Task<ArraySegment<byte>> ReadFrames(int connectionId, WebSocketReceiveResult result, WebSocket webSocket, byte[] buffer, CancellationToken token)

Assets/Mirror/Runtime/Transport/Websocket/WebsocketTransport.cs

@@ -60,6 +60,45 @@ public override bool Available()
             return true;
         }
 
+        void OnEnable()
+        {
+            server.enabled = true;
+            client.enabled = true;
+        }
+
+        void OnDisable()
+        {
+            server.enabled = false;
+            client.enabled = false;
+        }
+
+        void LateUpdate()
+        {
+            // note: we need to check enabled in case we set it to false
+            // when LateUpdate already started.
+            // (https://github.com/vis2k/Mirror/pull/379)
+            if (!enabled)
+                return;
+
+            // process a maximum amount of client messages per tick
+            // TODO add clientMaxReceivesPerTick same as telepathy
+            while (true)
+            {
+                // stop when there is no more message
+                if (!client.ProcessClientMessage())
+                {
+                    break;
+                }
+
+                // Some messages can disable transport
+                // If this is disabled stop processing message in queue
+                if (!enabled)
+                {
+                    break;
+                }
+            }
+        }
+
         // client
         public override bool ClientConnected() => client.IsConnected;