ws.broadcast broadcasts to all clients instead of broadcasting to all clients on a specific path/route fix jooby-project#889

jknack · jknack · commit 49631739b213 · 2017-10-17T12:00:00.000-03:00
diff --git a/jooby/src/main/java/org/jooby/WebSocket.java b/jooby/src/main/java/org/jooby/WebSocket.java
@@ -962,7 +962,6 @@ default void broadcast(final Object data, final SuccessCallback success) throws
    */
   default void broadcast(final Object data, final OnError err) throws Exception {
     broadcast(data, SUCCESS, err);
-    ;
   }
 
   /**
diff --git a/jooby/src/main/java/org/jooby/internal/WebSocketImpl.java b/jooby/src/main/java/org/jooby/internal/WebSocketImpl.java
@@ -222,12 +222,15 @@
 
 import java.nio.channels.ClosedChannelException;
 import java.nio.charset.StandardCharsets;
+import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.NoSuchElementException;
-import java.util.Queue;
-import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.CopyOnWriteArrayList;
 
 @SuppressWarnings("unchecked")
 public class WebSocketImpl implements WebSocket {
@@ -243,7 +246,7 @@ public class WebSocketImpl implements WebSocket {
   private final Logger log = LoggerFactory.getLogger(WebSocket.class);
 
   /** All connected websocket. */
-  private static final Queue<WebSocket> sessions = new ConcurrentLinkedQueue<>();
+  private static final ConcurrentMap<String, List<WebSocket>> sessions = new ConcurrentHashMap<>();
 
   private Locale locale;
 
@@ -290,7 +293,7 @@ public WebSocketImpl(final OnOpen handler, final String path,
 
   @Override
   public void close(final CloseStatus status) {
-    sessions.remove(this);
+    removeSession(this);
     synchronized (this) {
       open = false;
       ws.close(status.code(), status.reason());
@@ -299,7 +302,7 @@ public void close(final CloseStatus status) {
 
   @Override
   public void resume() {
-    sessions.add(this);
+    addSession(this);
     synchronized (this) {
       if (suspended) {
         ws.resume();
@@ -310,7 +313,7 @@ public void resume() {
 
   @Override
   public void pause() {
-    sessions.remove(this);
+    removeSession(this);
     synchronized (this) {
       if (!suspended) {
         ws.pause();
@@ -321,7 +324,7 @@ public void pause() {
 
   @Override
   public void terminate() throws Exception {
-    sessions.remove(this);
+    removeSession(this);
     synchronized (this) {
       open = false;
       ws.terminate();
@@ -336,7 +339,7 @@ public boolean isOpen() {
   @Override
   public void broadcast(final Object data, final SuccessCallback success, final OnError err)
       throws Exception {
-    for (WebSocket ws : sessions) {
+    for (WebSocket ws : sessions.getOrDefault(this.pattern, Collections.emptyList())) {
       try {
         ws.send(data, success, err);
       } catch (Exception ex) {
@@ -394,7 +397,7 @@ public void connect(final Injector injector, final Request req, final NativeWebS
         .onFailure(this::handleErr));
 
     ws.onCloseMessage((code, reason) -> {
-      sessions.remove(this);
+      removeSession(this);
 
       Try.run(sync(() -> {
         this.open = false;
@@ -410,7 +413,7 @@ public void connect(final Injector injector, final Request req, final NativeWebS
 
     // connect now
     try {
-      sessions.add(this);
+      addSession(this);
       handler.onOpen(req, this);
     } catch (Throwable ex) {
       handleErr(ex);
@@ -515,5 +518,11 @@ private Throwing.Runnable sync(final Throwing.Runnable task) {
     };
   }
 
-  ;
+  private static void addSession(WebSocketImpl ws) {
+    sessions.computeIfAbsent(ws.pattern, k -> new CopyOnWriteArrayList<>()).add(ws);
+  }
+
+  private static  void removeSession(WebSocketImpl ws) {
+    Optional.ofNullable(sessions.get(ws.pattern)).ifPresent(list-> list.remove(ws));
+  }
 }
diff --git a/jooby/src/test/java/org/jooby/internal/WebSocketImplTest.java b/jooby/src/test/java/org/jooby/internal/WebSocketImplTest.java
@@ -124,7 +124,7 @@ public void sendString() throws Exception {
   public void resetSessions() throws Exception {
     Field field = WebSocketImpl.class.getDeclaredField("sessions");
     field.setAccessible(true);
-    Queue<WebSocket> sessions = (Queue<WebSocket>) field.get(null);
+    Map<String, List<WebSocket>> sessions = (Map<String, List<WebSocket>>) field.get(null);
     sessions.clear();
   }
 

Original file line number	Diff line number	Diff line change
`@@ -962,7 +962,6 @@ default void broadcast(final Object data, final SuccessCallback success) throws`
`962`	`962`	`*/`
`963`	`963`	`default void broadcast(final Object data, final OnError err) throws Exception {`
`964`	`964`	`broadcast(data, SUCCESS, err);`
`965`		`- ;`
`966`	`965`	`}`
`967`	`966`
`968`	`967`	`/**`
Original file line number	Diff line number	Diff line change
`@@ -124,7 +124,7 @@ public void sendString() throws Exception {`
`124`	`124`	`public void resetSessions() throws Exception {`
`125`	`125`	`Field field = WebSocketImpl.class.getDeclaredField("sessions");`
`126`	`126`	`field.setAccessible(true);`
`127`		`- Queue<WebSocket> sessions = (Queue<WebSocket>) field.get(null);`
	`127`	`+ Map<String, List<WebSocket>> sessions = (Map<String, List<WebSocket>>) field.get(null);`
`128`	`128`	`sessions.clear();`
`129`	`129`	`}`
`130`	`130`