fix deadlock (maybe) in reversed WebSocket server

fix #67
MrXiaoM · Jul 6, 2024 · 8b00b45 · 8b00b45
1 parent 58b7fe4
commit 8b00b45
Showing 1 changed file with 12 additions and 18 deletions.
diff --git a/onebot/src/main/kotlin/client/connection/WSServer.kt b/onebot/src/main/kotlin/client/connection/WSServer.kt
@@ -3,9 +3,7 @@ package cn.evolvefield.onebot.client.connection
 import cn.evolvefield.onebot.client.config.BotConfig
 import cn.evolvefield.onebot.client.core.Bot
 import cn.evolvefield.onebot.client.handler.ActionHandler
-import kotlinx.coroutines.CompletableDeferred
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.InternalCoroutinesApi
+import kotlinx.coroutines.*
 import org.java_websocket.WebSocket
 import org.java_websocket.framing.CloseFrame
 import org.java_websocket.handshake.ClientHandshake
@@ -28,13 +26,7 @@ class WSServer(
     private val token: String
 ) : WebSocketServer(address), IAdapter {
     private var bot: Bot? = null
-    @OptIn(InternalCoroutinesApi::class)
-    val def = CompletableDeferred<Bot>(config.parentJob).apply {
-        invokeOnCompletion(
-            onCancelling = true,
-            invokeImmediately = true
-        ) { stop() }
-    }
+    val connectDef = CompletableDeferred<Bot>(config.parentJob)
 
     override fun onStart() {
         logger.info("▌ 反向 WebSocket 服务端已在 $address 启动")
@@ -67,10 +59,8 @@ class WSServer(
             }
         }
         logger.info("▌ 反向 WebSocket 客户端 ${conn.remoteSocketAddress} 已连接 ┈━═☆")
-        (bot ?: Bot(conn, config, actionHandler).also {
-            bot = it
-            def.complete(it)
-        }).conn = conn
+        bot = Bot(conn, config, actionHandler).also { it.conn = conn }
+        connectDef.complete(bot!!)
     }
 
     override fun onMessage(conn: WebSocket, message: String) = onReceiveMessage(message)
@@ -89,10 +79,14 @@ class WSServer(
     }
 
     companion object {
-        suspend fun createAndWaitConnect(scope: CoroutineScope, config: BotConfig, address: InetSocketAddress, logger: Logger, actionHandler: ActionHandler, token: String): Pair<WSServer, Bot> {
-            val ws = WSServer(scope, config, address, logger, actionHandler, token)
-            ws.start()
-            return ws to ws.def.await()
+        suspend fun createAndWaitConnect(
+            scope: CoroutineScope, config: BotConfig,
+            address: InetSocketAddress, logger: Logger,
+            actionHandler: ActionHandler, token: String
+        ): Pair<WSServer, Bot> {
+            val ws = WSServer(scope, config, address, logger, actionHandler, token).also { it.start() }
+            val bot = ws.connectDef.await()
+            return ws to bot
         }
     }
 }