diff --git a/src/org/redkale/net/http/WebSocket.java b/src/org/redkale/net/http/WebSocket.java index 33183643b..6f8ce6fa0 100644 --- a/src/org/redkale/net/http/WebSocket.java +++ b/src/org/redkale/net/http/WebSocket.java @@ -740,8 +740,11 @@ public abstract class WebSocket { /** * WebSokcet连接成功后的回调方法 + * + * @return Future 可以为null */ - public void onConnected() { + public CompletableFuture onConnected() { + return null; } /** @@ -805,8 +808,11 @@ public abstract class WebSocket { * * @param code 结果码,非0表示非正常关闭 * @param reason 关闭原因 + * + * @return Future 可以为null */ - public void onClose(int code, String reason) { + public CompletableFuture onClose(int code, String reason) { + return null; } /** diff --git a/src/org/redkale/net/http/WebSocketRunner.java b/src/org/redkale/net/http/WebSocketRunner.java index b1455879d..2b6718570 100644 --- a/src/org/redkale/net/http/WebSocketRunner.java +++ b/src/org/redkale/net/http/WebSocketRunner.java @@ -59,7 +59,8 @@ class WebSocketRunner implements Runnable { final boolean debug = context.getLogger().isLoggable(Level.FINEST); final WebSocketRunner self = this; try { - webSocket.onConnected(); + CompletableFuture connectfFuture = webSocket.onConnected(); + if (connectfFuture != null) connectfFuture.join(); webSocket._channel.setReadTimeoutSeconds(300); //读取超时5分钟 if (webSocket._channel.isOpen()) { final int wsmaxbody = webSocket._engine.wsmaxbody; @@ -307,8 +308,9 @@ class WebSocketRunner implements Runnable { closed = true; CompletableFuture future = engine.removeLocalThenClose(webSocket); webSocket._channel.dispose(); - webSocket.onClose(code, reason); - return future; + CompletableFuture closeFuture = webSocket.onClose(code, reason); + if (closeFuture == null) return future; + return CompletableFuture.allOf(future, closeFuture); } }