From 6d1c4e6c8594d95f790fb556a9c683f3da47a4b2 Mon Sep 17 00:00:00 2001
From: Redkale <22250530@qq.com>
Date: Mon, 28 May 2018 10:10:59 +0800
Subject: [PATCH]

---
 src/org/redkale/net/http/WebSocket.java       |  4 +++-
 .../redkale/net/http/WebSocketServlet.java    | 23 ++++++++++++-------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/src/org/redkale/net/http/WebSocket.java b/src/org/redkale/net/http/WebSocket.java
index 6fa758218..05f28bba3 100644
--- a/src/org/redkale/net/http/WebSocket.java
+++ b/src/org/redkale/net/http/WebSocket.java
@@ -756,8 +756,10 @@ public abstract class WebSocket<G extends Serializable, T> {
     /**
      * 当Single模式下用户重复登陆时回调函数，调用完之后关闭之前的WebSocket连接
      *
+     * @return Future 可以为null
      */
-    public void onSingleRepeatConnect() {
+    public CompletableFuture<Integer> onSingleRepeatConnect() {
+        return null;
     }
 
     /**
diff --git a/src/org/redkale/net/http/WebSocketServlet.java b/src/org/redkale/net/http/WebSocketServlet.java
index e235068d0..65e72abd0 100644
--- a/src/org/redkale/net/http/WebSocketServlet.java
+++ b/src/org/redkale/net/http/WebSocketServlet.java
@@ -242,14 +242,21 @@ public abstract class WebSocketServlet extends HttpServlet implements Resourcabl
                                 if (single && !anyuser) {
                                     WebSocketServlet.this.node.existsWebSocket(userid).whenComplete((rs, nex) -> {
                                         if (rs) {
-                                            webSocket.onSingleRepeatConnect();
-                                            node.forceCloseWebSocket(userid).whenComplete((fr, fex) -> {
-                                                WebSocketServlet.this.node.localEngine.add(webSocket);
-                                                WebSocketRunner runner = new WebSocketRunner(context, webSocket, restMessageConsumer, response.removeChannel());
-                                                webSocket._runner = runner;
-                                                context.runAsync(runner);
-                                                response.finish(true);
-                                            });
+                                            CompletableFuture rcFuture = webSocket.onSingleRepeatConnect();
+                                            Runnable task = () -> {
+                                                node.forceCloseWebSocket(userid).whenComplete((fr, fex) -> {
+                                                    WebSocketServlet.this.node.localEngine.add(webSocket);
+                                                    WebSocketRunner runner = new WebSocketRunner(context, webSocket, restMessageConsumer, response.removeChannel());
+                                                    webSocket._runner = runner;
+                                                    context.runAsync(runner);
+                                                    response.finish(true);
+                                                });
+                                            };
+                                            if (rcFuture == null) {
+                                                task.run();
+                                            } else {
+                                                rcFuture.whenComplete((r, e) -> task.run());
+                                            }
                                         } else {
                                             WebSocketServlet.this.node.localEngine.add(webSocket);
                                             WebSocketRunner runner = new WebSocketRunner(context, webSocket, restMessageConsumer, response.removeChannel());