Fix lowlevelclient connectstate being changed from the main thread only

Assets/Plugins/StreamChat/Core/LowLevelClient/StreamChatLowLevelClient.cs

@@ -28,6 +28,7 @@
 using StreamChat.Core.LowLevelClient.Requests;
 using System.Linq;
 using StreamChat.Core.Helpers;
+using Thread = System.Threading.Thread;
 
 #if STREAM_TESTS_ENABLED || STREAM_RUNTIME_TESTS_ENABLED
 using System.Runtime.CompilerServices;
@@ -212,12 +213,12 @@ private set
                 _logs.Warning($"Connection state changed from: {previous} to: {value}");
 #endif
 
-                ConnectionStateChanged?.Invoke(previous, _connectionState);
+                RaiseConnectionStateChanged(previous, _connectionState);
 
                 if (value == ConnectionState.Disconnected)
                 {
                     _disconnectionLastEventReceivedAt = _lastEventReceivedAt;
-                    Disconnected?.Invoke();
+                    RaiseDisconnected();
                 }
             }
         }
@@ -297,6 +298,8 @@ public StreamChatLowLevelClient(AuthCredentials authCredentials, IWebsocketClien
             IHttpClient httpClient, ISerializer serializer, ITimeService timeService, INetworkMonitor networkMonitor,
             IApplicationInfo applicationInfo, ILogs logs, IStreamClientConfig config)
         {
+            _mainThreadId = Thread.CurrentThread.ManagedThreadId;
+
             _authCredentials = authCredentials;
             _websocketClient = websocketClient ?? throw new ArgumentNullException(nameof(websocketClient));
             _httpClient = httpClient ?? throw new ArgumentNullException(nameof(httpClient));
@@ -410,6 +413,7 @@ public void Update(float deltaTime)
 #endif
 
             TryHandleWebsocketsConnectionFailed();
+            TryHandleWebsocketDisconnected();
             TryToReconnect();
 
             UpdateHealthCheck();
@@ -577,6 +581,12 @@ internal async Task<OwnUserInternalDTO> ConnectUserAsync(string apiKey, string u
             new Dictionary<string, Action<string>>();
 
         private readonly object _websocketConnectionFailedFlagLock = new object();
+        private readonly object _websocketDisconnectedFlagLock = new object();
+
+        /// <summary>
+        /// Every <see cref="ConnectionState"/> write must happen on this thread
+        /// </summary>
+        private readonly int _mainThreadId;
 
         private TaskCompletionSource<OwnUserInternalDTO> _connectUserTaskSource;
         private CancellationToken _connectUserCancellationToken;
@@ -592,6 +602,7 @@ internal async Task<OwnUserInternalDTO> ConnectUserAsync(string apiKey, string u
         private bool _updateCallReceived;
 
         private bool _websocketConnectionFailed;
+        private bool _websocketDisconnected;
         private ITokenProvider _tokenProvider;
 
         /// <summary>
@@ -646,14 +657,95 @@ private void TryCancelWaitingForUserConnection()
             }
         }
 
+        /// <summary>
+        /// This event can be called by a background thread and we must propagate it on the main thread
+        /// Otherwise any call to Unity API would result in Exception. Unity API can only be called from the main thread
+        /// </summary>
         private void OnWebsocketDisconnected()
         {
 #if STREAM_DEBUG_ENABLED
             _logs.Warning("Websocket Disconnected");
 #endif
+
+            if (Thread.CurrentThread.ManagedThreadId == _mainThreadId)
+            {
+                ConnectionState = ConnectionState.Disconnected;
+                return;
+            }
+
+            lock (_websocketDisconnectedFlagLock)
+            {
+                _websocketDisconnected = true;
+            }
+        }
+
+        private void TryHandleWebsocketDisconnected()
+        {
+            lock (_websocketDisconnectedFlagLock)
+            {
+                if (!_websocketDisconnected)
+                {
+                    return;
+                }
+
+                _websocketDisconnected = false;
+            }
+
+            if (ConnectionState == ConnectionState.Closing)
+            {
+                return;
+            }
+
             ConnectionState = ConnectionState.Disconnected;
         }
 
+        /// <summary>
+        /// Subscribers are invoked one by one so that a throwing handler cannot stop the remaining ones -
+        /// including the internal reconnect scheduling - from observing the transition
+        /// </summary>
+        private void RaiseConnectionStateChanged(ConnectionState previous, ConnectionState current)
+        {
+            var handler = ConnectionStateChanged;
+            if (handler == null)
+            {
+                return;
+            }
+
+            foreach (var subscriber in handler.GetInvocationList())
+            {
+                try
+                {
+                    ((ConnectionStateChangeHandler)subscriber)(previous, current);
+                }
+                catch (Exception e)
+                {
+                    _logs.Exception(e);
+                }
+            }
+        }
+
+        /// <inheritdoc cref="RaiseConnectionStateChanged"/>
+        private void RaiseDisconnected()
+        {
+            var handler = Disconnected;
+            if (handler == null)
+            {
+                return;
+            }
+
+            foreach (var subscriber in handler.GetInvocationList())
+            {
+                try
+                {
+                    ((Action)subscriber)();
+                }
+                catch (Exception e)
+                {
+                    _logs.Exception(e);
+                }
+            }
+        }
+
         /// <summary>
         /// This event can be called by a background thread and we must propagate it on the main thread
         /// Otherwise any call to Unity API would result in Exception. Unity API can only be called from the main thread

Assets/Plugins/StreamChat/Libs/Websockets/WebsocketClient.cs

@@ -124,7 +124,9 @@ public void Update()
             }
 #endif
 
-            var disconnect = false;
+            var serverClosedConnection = Interlocked.Exchange(ref _serverClosedConnectionFlag, 0) == 1;
+
+            var disconnect = serverClosedConnection;
             while (_threadWebsocketExceptionsLog.TryDequeue(out var webSocketException))
             {
                 LogExceptionIfDebugMode(webSocketException);
@@ -134,7 +136,15 @@ public void Update()
 
             if (disconnect)
             {
-                DisconnectAsync(WebSocketCloseStatus.ProtocolError, "WebSocket thrown an exception")
+                var closeStatus = serverClosedConnection
+                    ? WebSocketCloseStatus.InternalServerError
+                    : WebSocketCloseStatus.ProtocolError;
+
+                var closeMessage = serverClosedConnection
+                    ? "Server closed the connection"
+                    : "WebSocket thrown an exception";
+
+                DisconnectAsync(closeStatus, closeMessage)
                     .ContinueWith(_ => LogExceptionIfDebugMode(_.Exception), TaskContinuationOptions.OnlyOnFaulted);
                 return;
             }
@@ -200,6 +210,9 @@ public void Dispose()
 
         private WebSocketState _lastState;
 
+        // Int rather than bool so that Update can read and reset it in a single interlocked operation
+        private int _serverClosedConnectionFlag;
+
         private async void SendMessagesCallback(object state)
         {
             if (!IsConnected || _connectionCts == null || _connectionCts.IsCancellationRequested)
@@ -283,6 +296,9 @@ private async void ReceiveMessagesCallback(object state)
 
         private async Task TryCloseAndDisposeAsync(WebSocketCloseStatus closeStatus, string closeMessage)
         {
+            // A close frame that arrived while we were tearing down must not disconnect the next connection
+            Interlocked.Exchange(ref _serverClosedConnectionFlag, 0);
+
             try
             {
 #if UNITY_2021_2_OR_NEWER
@@ -419,10 +435,12 @@ await TryCloseAndDisposeAsync(WebSocketCloseStatus.ProtocolError,
             ConnectionFailed?.Invoke();
         }
 
-        // Called from a background thread
-        private void OnReceivedCloseMessage()
-            => DisconnectAsync(WebSocketCloseStatus.InternalServerError, "Server closed the connection")
-                .ContinueWith(t => LogThreadExceptionIfDebugMode(t.Exception), TaskContinuationOptions.OnlyOnFaulted);
+        /// <summary>
+        /// Called from a background thread. The disconnect is deferred to <see cref="Update"/> so that the
+        /// <see cref="Disconnected"/> event is always raised from the main thread, like every other path that
+        /// raises it. Subscribers react to it with Unity API calls, which throw off the main thread.
+        /// </summary>
+        private void OnReceivedCloseMessage() => Interlocked.Exchange(ref _serverClosedConnectionFlag, 1);
 
         private async Task<string> TryReceiveSingleMessageAsync()
         {
@@ -478,14 +496,6 @@ private void LogExceptionIfDebugMode(Exception exception)
             }
         }
 
-        private void LogThreadExceptionIfDebugMode(Exception exception)
-        {
-            if (_isDebugMode)
-            {
-                _threadExceptionsLog.Enqueue(exception);
-            }
-        }
-
         private void LogInfoIfDebugMode(string info)
         {
             if (_isDebugMode)

Assets/Plugins/StreamChat/Tests/LowLevelClient/StreamChatLowLevelClientTests.cs

@@ -3,6 +3,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Net.WebSockets;
+using System.Threading;
 using System.Threading.Tasks;
 using NSubstitute;
 using NUnit.Framework;
@@ -292,6 +293,60 @@ public void when_stream_client_health_check_timeout_detected_expect_client_disco
             Assert.IsFalse(client.ConnectionState == ConnectionState.Connected);
         }
 
+        [Test]
+        public void when_websocket_disconnected_raised_from_background_thread_expect_it_handled_on_main_thread_in_update()
+        {
+            var client = CreateConnectedClient();
+
+            var disconnectedCount = 0;
+            var disconnectedThreadId = 0;
+            client.Disconnected += () =>
+            {
+                disconnectedCount++;
+                disconnectedThreadId = Thread.CurrentThread.ManagedThreadId;
+            };
+
+            Task.Run(() => _mockWebsocketClient.Disconnected += Raise.Event<Action>()).Wait();
+
+            Assert.AreEqual(0, disconnectedCount,
+                "Disconnected must not be raised from the thread that closed the websocket");
+            Assert.IsTrue(client.ConnectionState == ConnectionState.Connected);
+
+            client.Update(deltaTime: 0.2f);
+
+            Assert.AreEqual(1, disconnectedCount);
+            Assert.AreEqual(Thread.CurrentThread.ManagedThreadId, disconnectedThreadId);
+        }
+
+        [Test]
+        public void when_websocket_disconnected_raised_from_main_thread_expect_it_handled_immediately()
+        {
+            var client = CreateConnectedClient();
+
+            var disconnectedCount = 0;
+            client.Disconnected += () => disconnectedCount++;
+
+            _mockWebsocketClient.Disconnected += Raise.Event<Action>();
+
+            // Awaiting DisconnectAsync must keep guaranteeing that the client is no longer connected
+            Assert.AreEqual(1, disconnectedCount);
+            Assert.IsFalse(client.ConnectionState == ConnectionState.Connected);
+        }
+
+        [Test]
+        public void when_connection_state_changed_subscriber_throws_expect_remaining_subscribers_notified()
+        {
+            var client = CreateConnectedClient(Substitute.For<ILogs>());
+
+            var lastStateSeenByLateSubscriber = ConnectionState.Connected;
+            client.ConnectionStateChanged += (previous, current) => throw new Exception("Subscriber failed");
+            client.ConnectionStateChanged += (previous, current) => lastStateSeenByLateSubscriber = current;
+
+            _mockWebsocketClient.Disconnected += Raise.Event<Action>();
+
+            Assert.AreNotEqual(ConnectionState.Connected, lastStateSeenByLateSubscriber);
+        }
+
         private readonly List<IDisposable> _resourcesToDispose = new List<IDisposable>();
 
         private IStreamChatLowLevelClient _lowLevelClient;
@@ -305,6 +360,29 @@ public void when_stream_client_health_check_timeout_detected_expect_client_disco
         private INetworkMonitor _mockNetworkMonitor;
         private IHttpClient _mockHttpClient;
         private IStreamClientConfig _mockStreamClientConfig;
+
+        private StreamChatLowLevelClient CreateConnectedClient(ILogs logs = null)
+        {
+            var client = new StreamChatLowLevelClient(_authCredentials, _mockWebsocketClient, _mockHttpClient,
+                new NewtonsoftJsonSerializer(), _mockTimeService, _mockNetworkMonitor, _mockApplicationInfo,
+                logs ?? _mockLogs, _mockStreamClientConfig);
+            _resourcesToDispose.Add(client);
+
+            _mockWebsocketClient.ConnectAsync(Arg.Any<Uri>()).Returns(Task.CompletedTask);
+
+            _mockWebsocketClient.TryDequeueMessage(out Arg.Any<string>()).Returns(arg =>
+            {
+                arg[0] = "{\"connection_id\":\"fakeId\", \"type\":\"health.check\"}";
+                return true;
+            }, arg => false);
+
+            client.Connect();
+            client.Update(deltaTime: 0.2f);
+
+            Assert.IsTrue(client.ConnectionState == ConnectionState.Connected);
+
+            return client;
+        }
     }
 }
 #endif