fix(WS): fix closed socket exceptions

Author: MukjepScarlet

src/main/kotlin/net/ccbluex/netty/http/HttpServer.kt

@@ -25,14 +25,16 @@ import io.netty.channel.ChannelOption
 import io.netty.channel.EventLoopGroup
 import io.netty.handler.logging.LogLevel
 import io.netty.handler.logging.LoggingHandler
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
+import net.ccbluex.netty.http.coroutines.awaitSuspend
+import net.ccbluex.netty.http.coroutines.syncSuspend
 import net.ccbluex.netty.http.middleware.Middleware
 import net.ccbluex.netty.http.rest.RouteController
 import net.ccbluex.netty.http.util.TransportType
 import net.ccbluex.netty.http.websocket.WebSocketController
 import org.apache.logging.log4j.LogManager
 import java.net.InetSocketAddress
-import java.util.concurrent.locks.ReentrantLock
-import kotlin.concurrent.withLock
 
 
 /**
@@ -43,15 +45,16 @@ import kotlin.concurrent.withLock
 class HttpServer {
 
     val routeController = RouteController()
-    val webSocketController = WebSocketController()
 
-    private val lock = ReentrantLock()
+    private val lock = Mutex()
 
     internal val middlewares = mutableListOf<Middleware>()
 
     private var bossGroup: EventLoopGroup? = null
     private var workerGroup: EventLoopGroup? = null
     private var serverChannel: Channel? = null
+    var webSocketController: WebSocketController? = null
+        private set
 
     companion object {
         internal val logger = LogManager.getLogger("HttpServer")
@@ -69,7 +72,7 @@ class HttpServer {
      *
      * @return actual port of server.
      */
-    fun start(port: Int, useNativeTransport: Boolean = true): Int = lock.withLock {
+    suspend fun start(port: Int, useNativeTransport: Boolean = true): Int = lock.withLock {
         val b = ServerBootstrap()
 
         val groups = TransportType.apply(b, useNativeTransport)
@@ -81,8 +84,9 @@ class HttpServer {
             b.option(ChannelOption.SO_BACKLOG, 1024)
                 .handler(LoggingHandler(LogLevel.INFO))
                 .childHandler(HttpChannelInitializer(this))
-            val ch = b.bind(port).sync().channel()
+            val ch = b.bind(port).syncSuspend().channel()
             serverChannel = ch
+            webSocketController = WebSocketController(ch)
 
             logger.info("Netty server started on port $port.")
 
@@ -98,12 +102,13 @@ class HttpServer {
     /**
      * Stops the Netty server gracefully.
      */
-    fun stop() = lock.withLock {
+    suspend fun stop() = lock.withLock {
         logger.info("Shutting down Netty server...")
         try {
-            serverChannel?.close()?.sync()
-            bossGroup?.shutdownGracefully()?.sync()
-            workerGroup?.shutdownGracefully()?.sync()
+            webSocketController?.disconnect()
+            serverChannel?.close()?.awaitSuspend()
+            bossGroup?.shutdownGracefully()?.awaitSuspend()
+            workerGroup?.shutdownGracefully()?.awaitSuspend()
         } catch (e: Exception) {
             logger.warn("Error during shutdown", e)
         } finally {

src/main/kotlin/net/ccbluex/netty/http/HttpServerHandler.kt

@@ -75,7 +75,7 @@ internal class HttpServerHandler(private val server: HttpServer) : ChannelInboun
                     + CoroutineName("${ctx.name()}#${ctx.channel().id().asShortText()}")
                     + exceptionHandler
         )
-        ctx.channel().closeFuture().addListener { channelScope.cancel() }
+        ctx.channel().closeFuture().addListener { channelScope.cancel("Channel closed") }
     }
 
     /**
@@ -123,7 +123,7 @@ internal class HttpServerHandler(private val server: HttpServer) : ChannelInboun
                         handshaker.handshake(ctx.channel(), msg)
                     }
 
-                    server.webSocketController.addContext(ctx)
+                    server.webSocketController!!.addContext(ctx)
                 } else {
                     val requestContext = RequestContext(
                         msg.method(),

src/main/kotlin/net/ccbluex/netty/http/coroutines/NettyCoroutine.kt

@@ -3,37 +3,67 @@ package net.ccbluex.netty.http.coroutines
 import io.netty.util.concurrent.Future
 import io.netty.util.concurrent.GenericFutureListener
 import kotlinx.coroutines.CancellableContinuation
-import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.suspendCancellableCoroutine
+import java.util.concurrent.CancellationException
 
 /**
- * Suspend until this Netty Future completes.
+ * Suspend until this Netty Future completes,
+ * and rethrows the cause of the failure if this future failed.
  *
- * Returns the Future result. Throws on failure or cancellation.
+ * Returns the Future itself.
  */
-suspend fun <V, F : Future<V>> F.suspend(): V {
+suspend fun <V, F : Future<V>> F.syncSuspend(): F {
     if (isDone) return unwrapDone().getOrThrow()
 
     return suspendCancellableCoroutine { cont ->
-        addListener(futureContinuationListener(cont))
+        addListener(FutureResultContListener(cont))
 
         cont.invokeOnCancellation {
             this.cancel(false)
         }
     }
 }
 
-private fun <V, F : Future<V>> futureContinuationListener(
-    cont: CancellableContinuation<V>
-): GenericFutureListener<F> = GenericFutureListener { future ->
-    if (cont.isActive) {
-        cont.resumeWith(future.unwrapDone())
+/**
+ * Suspend until this Netty Future completes.
+ *
+ * Returns the Future itself.
+ */
+suspend fun <F : Future<*>> F.awaitSuspend(): F {
+    if (isDone) return this
+
+    return suspendCancellableCoroutine { cont ->
+        addListener(FutureContListener(cont))
+
+        cont.invokeOnCancellation {
+            this.cancel(false)
+        }
+    }
+}
+
+private class FutureContListener<V, F : Future<V>>(
+    private val cont: CancellableContinuation<F>
+): GenericFutureListener<F> {
+    override fun operationComplete(future: F) {
+        if (cont.isActive) {
+            cont.resumeWith(Result.success(future))
+        }
+    }
+}
+
+private class FutureResultContListener<V, F : Future<V>>(
+    private val cont: CancellableContinuation<F>
+): GenericFutureListener<F> {
+    override fun operationComplete(future: F) {
+        if (cont.isActive) {
+            cont.resumeWith(future.unwrapDone())
+        }
     }
 }
 
-private fun <V, F : Future<V>> F.unwrapDone(): Result<V> =
+private fun <V, F : Future<V>> F.unwrapDone(): Result<F> =
     when {
-        isSuccess -> Result.success(this.now)
+        isSuccess -> Result.success(this)
         isCancelled -> Result.failure(CancellationException("Netty Future was cancelled"))
         else -> Result.failure(
             this.cause() ?: IllegalStateException("Future failed without cause")

src/main/kotlin/net/ccbluex/netty/http/websocket/WebSocketController.kt

@@ -19,42 +19,66 @@
  */
 package net.ccbluex.netty.http.websocket
 
+import io.netty.channel.Channel
 import io.netty.channel.ChannelHandlerContext
 import io.netty.handler.codec.http.websocketx.TextWebSocketFrame
-import java.util.concurrent.CopyOnWriteArrayList
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.asCoroutineDispatcher
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.joinAll
+import kotlinx.coroutines.launch
+import net.ccbluex.netty.http.coroutines.awaitSuspend
+import net.ccbluex.netty.http.coroutines.syncSuspend
+import java.util.concurrent.CopyOnWriteArraySet
 import java.util.function.BiConsumer
 
 /**
  * Controller for handling websocket connections.
  */
-class WebSocketController {
+class WebSocketController(
+    private val serverChannel: Channel,
+) {
+
+    private val scope = CoroutineScope(
+        serverChannel.eventLoop().asCoroutineDispatcher() + SupervisorJob()
+    )
+
+    init {
+        serverChannel.closeFuture().addListener { scope.cancel("Channel closed") }
+    }
 
     /**
      * Keeps track of all connected websocket connections to the server.
      * This is used to broadcast messages to all connected clients.
      */
-    private val activeContexts = CopyOnWriteArrayList<ChannelHandlerContext>()
+    private val activeContexts = CopyOnWriteArraySet<ChannelHandlerContext>()
 
     /**
      * Broadcasts a message to all connected clients.
      *
      * @param text The message to broadcast.
      * @param onFailure The action to take if a failure occurs.
      */
-    fun broadcast(text: String, onFailure: BiConsumer<ChannelHandlerContext, Throwable>? = null) {
-        val frame = TextWebSocketFrame(text)
-        for (handlerContext in activeContexts) {
-            val channelFuture = handlerContext.channel().writeAndFlush(frame.retainedDuplicate())
-            if (onFailure != null) {
-                channelFuture.addListener {
-                    if (!it.isSuccess) {
-                        onFailure.accept(handlerContext, it.cause())
+    fun broadcast(text: String, onFailure: BiConsumer<ChannelHandlerContext, Throwable>? = null): Job =
+        scope.launch {
+            val frame = TextWebSocketFrame(text)
+
+            activeContexts.map { handlerContext ->
+                launch {
+                    try {
+                        handlerContext.channel()
+                            .writeAndFlush(frame.retainedDuplicate())
+                            .syncSuspend()
+                    } catch (e: Exception) {
+                        onFailure?.accept(handlerContext, e)
                     }
                 }
-            }
+            }.joinAll()
+
+            frame.release()
         }
-        frame.release()
-    }
 
     /**
      * Closes all active contexts.
@@ -74,6 +98,9 @@ class WebSocketController {
      */
     fun addContext(context: ChannelHandlerContext) {
         activeContexts.add(context)
+        context.channel().closeFuture().addListener {
+            removeContext(context)
+        }
     }
 
     /**

src/main/kotlin/net/ccbluex/netty/http/websocket/WebSocketHandler.kt

@@ -53,7 +53,7 @@ internal class WebSocketHandler(private val server: HttpServer) : ChannelInbound
                         .writeAndFlush(msg.retainedDuplicate())
                         .addListener(ChannelFutureListener.CLOSE)
 
-                    server.webSocketController.removeContext(ctx)
+                    server.webSocketController!!.removeContext(ctx)
                     logger.debug("WebSocket closed due to ${msg.reasonText()} (${msg.statusCode()})")
                 }
                 else -> logger.error("Unknown WebSocketFrame type: ${msg.javaClass.name}")