From 25bcc3c96ce3e269fc42663782583cc21c95587e Mon Sep 17 00:00:00 2001 From: HwangRock Date: Wed, 15 Jul 2026 22:36:39 +0900 Subject: [PATCH 1/2] [ZEPPELIN-6540] Make ConnectionManager.userSocketMap thread-safe userSocketMap was a plain HashMap accessed without synchronization from multiple Jetty WebSocket threads (login onMessage, disconnect onClose, note-list broadcast). Concurrent access could throw ConcurrentModificationException while iterating keySet(), spin during HashMap resize, or NPE via the containsKey/get TOCTOU in multicastToUser. The sibling noteSocketMap is guarded by synchronized everywhere; only userSocketMap was left unprotected. Switch userSocketMap to ConcurrentHashMap and make the compound operations atomic: compute() for add, computeIfPresent() for remove, and a single get()+null check for the read paths. Iteration over keySet() is safe under the weakly-consistent iterator. noteSocketMap keeps synchronized because it needs multi-entry atomic operations that a per-key guarantee cannot cover. --- .../zeppelin/socket/ConnectionManager.java | 38 +++--- .../socket/ConnectionManagerTest.java | 122 ++++++++++++++++++ 2 files changed, 142 insertions(+), 18 deletions(-) diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java b/zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java index a348d218afb..98a543c4824 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java @@ -53,6 +53,7 @@ import java.util.Map.Entry; import java.util.Queue; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; /** @@ -71,7 +72,7 @@ public class ConnectionManager { // noteId -> connection final Map> noteSocketMap = Metrics.gaugeMapSize("zeppelin_note_sockets", Tags.empty(), new HashMap<>()); // user -> connection - final Map> userSocketMap = Metrics.gaugeMapSize("zeppelin_user_sockets", Tags.empty(), new HashMap<>()); + final Map> userSocketMap = Metrics.gaugeMapSize("zeppelin_user_sockets", Tags.empty(), new ConcurrentHashMap<>()); /** * This is a special endpoint in the notebook websocket, Every connection in this Queue @@ -153,24 +154,23 @@ public void removeConnectionFromAllNote(NotebookSocket socket) { public void addUserConnection(String user, NotebookSocket conn) { LOGGER.debug("Add user connection {} for user: {}", conn, user); conn.setUser(user); - if (userSocketMap.containsKey(user)) { - userSocketMap.get(user).add(conn); - } else { - Queue socketQueue = new ConcurrentLinkedQueue<>(); - socketQueue.add(conn); - userSocketMap.put(user, socketQueue); - } + userSocketMap.compute(user, (k, connections) -> { + Queue queue = + (connections == null) ? new ConcurrentLinkedQueue<>() : connections; + queue.add(conn); + return queue; + }); } public void removeUserConnection(String user, NotebookSocket conn) { LOGGER.debug("Remove user connection {} for user: {}", conn, user); - if (userSocketMap.containsKey(user)) { - Queue connections = userSocketMap.get(user); + boolean[] wasPresent = {false}; + userSocketMap.computeIfPresent(user, (k, connections) -> { + wasPresent[0] = true; connections.remove(conn); - if (connections.isEmpty()) { - userSocketMap.remove(user); - } - } else { + return connections.isEmpty() ? null : connections; + }); + if (!wasPresent[0]) { LOGGER.warn("Closing connection that is absent in user connections"); } } @@ -330,12 +330,13 @@ public Set getConnectedUsers() { public void multicastToUser(String user, Message m) { - if (!userSocketMap.containsKey(user)) { + Queue connections = userSocketMap.get(user); + if (connections == null) { LOGGER.warn("Multicasting to user {} that is not in connections map", user); return; } - for (NotebookSocket conn : userSocketMap.get(user)) { + for (NotebookSocket conn : connections) { unicast(m, conn); } } @@ -354,12 +355,13 @@ public void unicastParagraph(Note note, Paragraph p, String user, String msgId) return; } - if (!userSocketMap.containsKey(user)) { + Queue connections = userSocketMap.get(user); + if (connections == null) { LOGGER.warn("Failed to send unicast. user {} that is not in connections map", user); return; } - for (NotebookSocket conn : userSocketMap.get(user)) { + for (NotebookSocket conn : connections) { Message m = new Message(Message.OP.PARAGRAPH).withMsgId(msgId).put("paragraph", p); unicast(m, conn); } diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/socket/ConnectionManagerTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/socket/ConnectionManagerTest.java index 92adc93c795..96f9949b4a0 100644 --- a/zeppelin-server/src/test/java/org/apache/zeppelin/socket/ConnectionManagerTest.java +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/socket/ConnectionManagerTest.java @@ -18,16 +18,22 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; +import java.util.Queue; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import org.apache.zeppelin.conf.ZeppelinConfiguration; import org.apache.zeppelin.notebook.AuthorizationService; @@ -142,6 +148,122 @@ void removeWatcherConnectionConcurrentTest() throws InterruptedException { assertEquals(0, manager.watcherSockets.size()); } + @Test + void userSocketMapConcurrentAccessTest() throws InterruptedException { + AuthorizationService authService = mock(AuthorizationService.class); + when(authService.getRoles(anyString())).thenAnswer(invocation -> new HashSet<>()); + ConnectionManager manager = new ConnectionManager(authService, ZeppelinConfiguration.load()); + + int writerCount = 8; + int readerCount = 2; + int iterations = 1000; + + List users = new ArrayList<>(); + List sockets = new ArrayList<>(); + for (int i = 0; i < writerCount; i++) { + users.add("user-" + i); + sockets.add(mock(NotebookSocket.class)); + } + + AtomicReference failure = new AtomicReference<>(); + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(writerCount + readerCount); + ExecutorService executor = Executors.newFixedThreadPool(writerCount + readerCount); + + for (int i = 0; i < writerCount; i++) { + String user = users.get(i); + NotebookSocket socket = sockets.get(i); + executor.submit(() -> { + try { + startLatch.await(); + for (int j = 0; j < iterations; j++) { + manager.addUserConnection(user, socket); + manager.removeUserConnection(user, socket); + } + } catch (Throwable t) { + failure.compareAndSet(null, t); + } finally { + doneLatch.countDown(); + } + }); + } + + for (int i = 0; i < readerCount; i++) { + executor.submit(() -> { + try { + startLatch.await(); + for (int j = 0; j < iterations; j++) { + manager.forAllUsers((user, userAndRoles) -> { }); + } + } catch (Throwable t) { + failure.compareAndSet(null, t); + } finally { + doneLatch.countDown(); + } + }); + } + + startLatch.countDown(); + assertTrue(doneLatch.await(30, TimeUnit.SECONDS)); + executor.shutdown(); + + assertNull(failure.get(), + "Concurrent access to userSocketMap should not throw, but got: " + failure.get()); + } + + @Test + void userSocketMapConcurrentAddPreservesAllConnectionsTest() throws InterruptedException { + AuthorizationService authService = mock(AuthorizationService.class); + ConnectionManager manager = new ConnectionManager(authService, ZeppelinConfiguration.load()); + + String user = "shared-user"; + int threadCount = 16; + int iterationsPerThread = 500; + + AtomicReference failure = new AtomicReference<>(); + List addedSockets = new CopyOnWriteArrayList<>(); + + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(threadCount); + for (int i = 0; i < threadCount; i++) { + executor.submit(() -> { + try { + startLatch.await(); + for (int j = 0; j < iterationsPerThread; j++) { + NotebookSocket socket = mock(NotebookSocket.class); + manager.addUserConnection(user, socket); + addedSockets.add(socket); + } + } catch (Throwable t) { + failure.compareAndSet(null, t); + } finally { + doneLatch.countDown(); + } + }); + } + + startLatch.countDown(); + assertTrue(doneLatch.await(30, TimeUnit.SECONDS)); + executor.shutdown(); + + assertNull(failure.get(), "Concurrent add should not throw, but got: " + failure.get()); + + Queue finalConnections = manager.userSocketMap.get(user); + assertEquals(threadCount * iterationsPerThread, addedSockets.size()); + List missing = new ArrayList<>(); + for (NotebookSocket socket : addedSockets) { + if (finalConnections == null || !finalConnections.contains(socket)) { + missing.add(socket); + } + } + + assertTrue(missing.isEmpty(), + missing.size() + " of " + addedSockets.size() + + " concurrently added connections were lost from userSocketMap"); + assertEquals(addedSockets.size(), finalConnections.size()); + } + @Test void switchConnectionToWatcherAndRemove() { AuthorizationService authService = mock(AuthorizationService.class); From 73906a1145b1bca2c6feaba7c1cd02907d95c3cc Mon Sep 17 00:00:00 2001 From: HwangRock Date: Sat, 18 Jul 2026 10:39:51 +0900 Subject: [PATCH 2/2] [ZEPPELIN-6540] Guard removeUserConnection against null user --- .../zeppelin/socket/ConnectionManager.java | 4 ++++ .../socket/ConnectionManagerTest.java | 20 +++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java b/zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java index 98a543c4824..6b13613ccec 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java @@ -164,6 +164,10 @@ public void addUserConnection(String user, NotebookSocket conn) { public void removeUserConnection(String user, NotebookSocket conn) { LOGGER.debug("Remove user connection {} for user: {}", conn, user); + if (user == null) { + LOGGER.warn("Closing connection for null user"); + return; + } boolean[] wasPresent = {false}; userSocketMap.computeIfPresent(user, (k, connections) -> { wasPresent[0] = true; diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/socket/ConnectionManagerTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/socket/ConnectionManagerTest.java index 96f9949b4a0..562d0658949 100644 --- a/zeppelin-server/src/test/java/org/apache/zeppelin/socket/ConnectionManagerTest.java +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/socket/ConnectionManagerTest.java @@ -16,6 +16,7 @@ */ package org.apache.zeppelin.socket; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; @@ -290,4 +291,23 @@ void switchConnectionToWatcherAndRemove() { // Verify it's completely removed assertFalse(manager.watcherSockets.contains(socket)); } + + @Test + void removeUserConnectionWithNullUserDoesNotThrow() { + AuthorizationService authService = mock(AuthorizationService.class); + ConnectionManager manager = new ConnectionManager(authService, ZeppelinConfiguration.load()); + NotebookSocket socket = mock(NotebookSocket.class); + + assertDoesNotThrow(() -> manager.removeUserConnection(null, socket)); + } + + @Test + void removeUserConnectionBeforeUserAssignment() { + AuthorizationService authService = mock(AuthorizationService.class); + ConnectionManager manager = new ConnectionManager(authService, ZeppelinConfiguration.load()); + NotebookSocket socket = mock(NotebookSocket.class); + + assertDoesNotThrow(() -> manager.removeUserConnection("", socket)); + assertTrue(manager.userSocketMap.isEmpty()); + } }