From 06bb5180cf009fe78bca8e9ac850fb02e3abcbec Mon Sep 17 00:00:00 2001
From: Redkale <8730487+redkale@users.noreply.github.com>
Date: Wed, 8 Jan 2020 10:36:54 +0800
Subject: [PATCH] =?UTF-8?q?=E3=80=90=E4=B8=8D=E5=85=BC=E5=AE=B9=E3=80=91We?=
 =?UTF-8?q?bSocket=E4=B8=ADonConnected=E3=80=81onClose=E6=96=B9=E6=B3=95?=
 =?UTF-8?q?=E7=9A=84=E8=BF=94=E5=9B=9E=E5=80=BC=E7=94=B1void=E6=94=B9?=
 =?UTF-8?q?=E6=88=90CompletableFuture?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/org/redkale/net/http/WebSocket.java       | 10 ++++++++--
 src/org/redkale/net/http/WebSocketRunner.java |  8 +++++---
 2 files changed, 13 insertions(+), 5 deletions(-)

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<G extends Serializable, T> {
 
     /**
      * WebSokcet连接成功后的回调方法
+     *
+     * @return Future 可以为null
      */
-    public void onConnected() {
+    public CompletableFuture onConnected() {
+        return null;
     }
 
     /**
@@ -805,8 +808,11 @@ public abstract class WebSocket<G extends Serializable, T> {
      *
      * @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<Void> 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);
         }
     }