Correct closing web socket connections in case of fatal server stopping

serpion · serpion · commit ebed7cf85756 · 2022-03-25T18:07:01.000+02:00
diff --git a/src/main/java/org/java_websocket/server/WebSocketServer.java b/src/main/java/org/java_websocket/server/WebSocketServer.java
@@ -248,6 +248,10 @@ public void start() {
     new Thread(this).start();
   }
 
+  public void stop(int timeout) throws InterruptedException {
+    stop(timeout, "");
+  }
+
   /**
    * Closes all connected clients sockets, then closes the underlying ServerSocketChannel,
    * effectively killing the server socket selectorthread, freeing the port the server was bound to
@@ -257,10 +261,11 @@ public void start() {
    *
    * @param timeout Specifies how many milliseconds the overall close handshaking may take
    *                altogether before the connections are closed without proper close
-   *                handshaking.<br>
+   *                handshaking.
+   * @param closeMessage Specifies message for remote client<br>
    * @throws InterruptedException Interrupt
    */
-  public void stop(int timeout) throws InterruptedException {
+  public void stop(int timeout, String closeMessage) throws InterruptedException {
     if (!isclosed.compareAndSet(false,
         true)) { // this also makes sure that no further connections will be added to this.connections
       return;
@@ -274,7 +279,7 @@ public void stop(int timeout) throws InterruptedException {
     }
 
     for (WebSocket ws : socketsToClose) {
-      ws.close(CloseFrame.GOING_AWAY);
+      ws.close(CloseFrame.GOING_AWAY, closeMessage);
     }
 
     wsf.close();
@@ -680,6 +685,17 @@ private void handleIOException(SelectionKey key, WebSocket conn, IOException ex)
   private void handleFatal(WebSocket conn, Exception e) {
     log.error("Shutdown due to fatal error", e);
     onError(conn, e);
+
+    String causeMessage = e.getCause() != null ? " caused by " + e.getCause().getClass().getName() : "";
+    String errorMessage = "Got error on server side: " + e.getClass().getName() + causeMessage;
+    try {
+      stop(0, errorMessage);
+    } catch (InterruptedException e1) {
+      Thread.currentThread().interrupt();
+      log.error("Interrupt during stop", e);
+      onError(null, e1);
+    }
+
     //Shutting down WebSocketWorkers, see #222
     if (decoders != null) {
       for (WebSocketWorker w : decoders) {
@@ -689,13 +705,6 @@ private void handleFatal(WebSocket conn, Exception e) {
     if (selectorthread != null) {
       selectorthread.interrupt();
     }
-    try {
-      stop();
-    } catch (InterruptedException e1) {
-      Thread.currentThread().interrupt();
-      log.error("Interrupt during stop", e);
-      onError(null, e1);
-    }
   }
 
   @Override
