This is an automated email from the ASF dual-hosted git repository.
ParkGyeongTae pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/zeppelin.git
The following commit(s) were added to refs/heads/master by this push:
new fb7f9dc86f [ZEPPELIN-6540] Make ConnectionManager.userSocketMap
thread-safe
fb7f9dc86f is described below
commit fb7f9dc86f2966395610cf8fa7942f6a56abf274
Author: HwangRock <[email protected]>
AuthorDate: Tue Jul 21 22:48:54 2026 +0900
[ZEPPELIN-6540] Make ConnectionManager.userSocketMap thread-safe
### What is this PR for?
`ConnectionManager.userSocketMap` (`user -> Queue<NotebookSocket>`) is a
plain `HashMap`, but every access to it is unsynchronized:
`addUserConnection`/`removeUserConnection` (writes) and
`multicastToUser`/`unicastParagraph`/`forAllUsers`/`broadcastNoteListExcept`
(reads and `keySet()` iteration). The sibling field `noteSocketMap` is guarded
by `synchronized (noteSocketMap)` at every access — only `userSocketMap` was
left unprotected.
These paths run on different Jetty WebSocket threads (login `onMessage`,
disconnect `onClose`, note-list broadcast `broadcastNoteListUpdate`). Under
concurrent connect/disconnect — e.g. a burst of client reconnects after a
server restart — this leads to:
- `ConcurrentModificationException` while iterating `keySet()`, aborting
the note-list broadcast so some users stop receiving updates
- possible CPU spin / lost entries during `HashMap` resize
- `NullPointerException` from the `containsKey` + `get` TOCTOU in
`multicastToUser`/`unicastParagraph`
The map value is already a `ConcurrentLinkedQueue`, so only the map itself
was unprotected.
Fix: switch `userSocketMap` to `ConcurrentHashMap` and make the compound
operations atomic:
- `addUserConnection`: `compute(...)` so the queue create-or-reuse and the
`add` happen in one atomic map operation (closing the add-after-remove window
that a bare `computeIfAbsent(...).add(...)` would leave open)
- `removeUserConnection`: `computeIfPresent(...)`, removing the key when
the queue becomes empty
- `multicastToUser` / `unicastParagraph`: a single `get()` + null check,
removing the TOCTOU/NPE
- `forAllUsers` / `broadcastNoteListExcept`: unchanged — `keySet()`
iteration is safe under `ConcurrentHashMap`'s weakly-consistent iterator
`noteSocketMap` intentionally keeps `synchronized`: it needs multi-entry
atomic operations (`removeConnectionFromAllNote`, `checkCollaborativeStatus`)
that a per-key `ConcurrentHashMap` guarantee does not cover, so a
single-strategy migration would not be correct there.
### What type of PR is it?
Bug Fix
### Todos
* [ ] - none
### What is the Jira issue?
* https://issues.apache.org/jira/browse/ZEPPELIN-6540
### How should this be tested?
Added two unit tests in `ConnectionManagerTest`:
- `userSocketMapConcurrentAccessTest`: 8 writer threads add/remove
connections while 2 reader threads iterate via `forAllUsers`. On the old
`HashMap` this reproduces `ConcurrentModificationException` /
`NullPointerException`; with the fix it passes with no thrown exception.
- `userSocketMapConcurrentAddPreservesAllConnectionsTest`: 16 threads
concurrently add unique sockets for the same user, then assert every socket is
present and the final queue size matches — validates the atomic publish
`compute()` guarantees.
Ran the full `ConnectionManagerTest` 5× consecutively — stable, no
intermittent failures.
Note on the add-after-remove window: it is extremely narrow and did not
reproduce as a deterministic failing test even under heavy contention (16 churn
threads, 16×5000 iterations). The `compute()` fix closes it by construction
(atomic per-key create-and-add under `ConcurrentHashMap`); the correctness
rests on that plus the concurrent-add invariant test rather than a RED→GREEN of
the exact interleaving.
### Screenshots (if appropriate)
### Questions:
* Does the license files need to update? No.
* Is there breaking changes for older versions? No — public method
signatures are unchanged.
* Does this needs documentation? No.
Closes #5306 from HwangRock/ZEPPELIN-6540.
Signed-off-by: ParkGyeongTae <[email protected]>
---
.../apache/zeppelin/socket/ConnectionManager.java | 42 +++---
.../zeppelin/socket/ConnectionManagerTest.java | 142 +++++++++++++++++++++
2 files changed, 166 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 a348d218af..6b13613cce 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;
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<String, Set<NotebookSocket>> noteSocketMap =
Metrics.gaugeMapSize("zeppelin_note_sockets", Tags.empty(), new HashMap<>());
// user -> connection
- final Map<String, Queue<NotebookSocket>> userSocketMap =
Metrics.gaugeMapSize("zeppelin_user_sockets", Tags.empty(), new HashMap<>());
+ final Map<String, Queue<NotebookSocket>> 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,27 @@ public class ConnectionManager {
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<NotebookSocket> socketQueue = new ConcurrentLinkedQueue<>();
- socketQueue.add(conn);
- userSocketMap.put(user, socketQueue);
- }
+ userSocketMap.compute(user, (k, connections) -> {
+ Queue<NotebookSocket> 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<NotebookSocket> connections = userSocketMap.get(user);
+ if (user == null) {
+ LOGGER.warn("Closing connection for null user");
+ return;
+ }
+ 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 +334,13 @@ public class ConnectionManager {
public void multicastToUser(String user, Message m) {
- if (!userSocketMap.containsKey(user)) {
+ Queue<NotebookSocket> 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 +359,13 @@ public class ConnectionManager {
return;
}
- if (!userSocketMap.containsKey(user)) {
+ Queue<NotebookSocket> 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 92adc93c79..562d065894 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,18 +16,25 @@
*/
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;
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 +149,122 @@ class ConnectionManagerTest {
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<String> users = new ArrayList<>();
+ List<NotebookSocket> sockets = new ArrayList<>();
+ for (int i = 0; i < writerCount; i++) {
+ users.add("user-" + i);
+ sockets.add(mock(NotebookSocket.class));
+ }
+
+ AtomicReference<Throwable> 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<Throwable> failure = new AtomicReference<>();
+ List<NotebookSocket> 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<NotebookSocket> finalConnections = manager.userSocketMap.get(user);
+ assertEquals(threadCount * iterationsPerThread, addedSockets.size());
+ List<NotebookSocket> 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);
@@ -168,4 +291,23 @@ class ConnectionManagerTest {
// 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());
+ }
}
