From 30c9b7d8e8259a7516fafcbd26179a0ab6a91ec9 Mon Sep 17 00:00:00 2001
From: Chen Si Quan <406908+chsqn@users.noreply.github.com>
Date: Fri, 4 Sep 2020 02:52:42 +0800
Subject: [PATCH] fixed websocket threading error onclose (#119)

Co-authored-by: Chen Si Quan <06908+chsqn@users.noreply.github.com>
---
 Assets/Plugins/WebSocket/WebSocket.cs | 78 ++++++++++++++++++++++++++-
 1 file changed, 76 insertions(+), 2 deletions(-)

diff --git a/Assets/Plugins/WebSocket/WebSocket.cs b/Assets/Plugins/WebSocket/WebSocket.cs
index b6d632d5..a4dbae08 100644
--- a/Assets/Plugins/WebSocket/WebSocket.cs
+++ b/Assets/Plugins/WebSocket/WebSocket.cs
@@ -9,6 +9,75 @@
 
 using AOT;
 using System.Runtime.InteropServices;
+using UnityEngine;
+using System.Collections;
+
+public class MainThreadUtil : MonoBehaviour
+{
+    public static MainThreadUtil Instance { get; private set; }
+    public static SynchronizationContext synchronizationContext { get; private set; }
+
+    [RuntimeInitializeOnLoadMethod(RuntimeInitializeLoadType.BeforeSceneLoad)]
+    public static void Setup()
+    {
+        Instance = new GameObject("MainThreadUtil")
+            .AddComponent<MainThreadUtil>();
+        synchronizationContext = SynchronizationContext.Current;
+    }
+
+    public static void Run(IEnumerator waitForUpdate)
+    {
+        synchronizationContext.Post(_ => Instance.StartCoroutine(
+                    waitForUpdate), null);
+    }
+
+    void Awake()
+    {
+        gameObject.hideFlags = HideFlags.HideAndDontSave;
+        DontDestroyOnLoad(gameObject);
+    }
+}
+
+public class WaitForUpdate : CustomYieldInstruction
+{
+    public override bool keepWaiting
+    {
+        get { return false; }
+    }
+
+    public MainThreadAwaiter GetAwaiter()
+    {
+        var awaiter = new MainThreadAwaiter();
+        MainThreadUtil.Run(CoroutineWrapper(this, awaiter));
+        return awaiter;
+    }
+
+    public class MainThreadAwaiter : INotifyCompletion
+    {
+        Action continuation;
+
+        public bool IsCompleted { get; set; }
+
+        public void GetResult() { }
+
+        public void Complete()
+        {
+            IsCompleted = true;
+            continuation?.Invoke();
+        }
+
+        void INotifyCompletion.OnCompleted(Action continuation)
+        {
+            this.continuation = continuation;
+        }
+    }
+
+    public static IEnumerator CoroutineWrapper(IEnumerator theWorker, MainThreadAwaiter awaiter)
+    {
+        yield return theWorker;
+        awaiter.Complete();
+    }
+}
 
 namespace NativeWebSocket
 {
@@ -490,6 +559,7 @@ public void DispatchMessageQueue()
 
         public async Task Receive()
         {
+            WebSocketCloseCode closeCode = WebSocketCloseCode.Abnormal;
             await new WaitForBackgroundThread();
 
             ArraySegment<byte> buffer = new ArraySegment<byte>(new byte[8192]);
@@ -531,7 +601,7 @@ public async Task Receive()
                         else if (result.MessageType == WebSocketMessageType.Close)
                         {
                             await Close();
-                            OnClose?.Invoke(WebSocketHelpers.ParseCloseCodeEnum((int)result.CloseStatus));
+                            closeCode = WebSocketHelpers.ParseCloseCodeEnum((int)result.CloseStatus);
                             break;
                         }
                     }
@@ -540,7 +610,11 @@ public async Task Receive()
             catch (Exception)
             {
                 m_TokenSource.Cancel();
-                OnClose?.Invoke(WebSocketCloseCode.Abnormal);
+            }
+            finally
+            {
+                await new WaitForUpdate();
+                OnClose?.Invoke(closeCode);
             }
         }