[ZEPPELIN-6540] Make ConnectionManager.userSocketMap thread-safe

[HwangRock]

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.

[ParkGyeongTae]

LGTM +1

[jongyoul]

ConcurrentHashMap rejects null keys, but the WebSocket close path can pass one here. NotebookServer.onOpen() only adds the socket to connectedSockets; if the connection closes before its first message, NotebookSocket.getUser() is still null and onClose() calls removeUserConnection(null, socket). The old HashMap.containsKey(null) path only logged an absent connection, while the new computeIfPresent(null, ...) throws NullPointerException during close handling. Please add a null-user guard and a regression test for close-before-user-assignment.

[HwangRock]

@jongyoul Thanks for the catch — agreed the close path should never be able to reach computeIfPresent(null, ...). Added the null guard and regression tests in 73906a1.

One note from tracing the path: NotebookSocket initializes user to StringUtils.EMPTY in its constructor, so onClose() before the first message currently passes "" rather than null, and a null principal cannot reach addUserConnection because ticket validation rejects it first. So the null case looks unreachable today, but setUser is public and the guard costs nothing, so guarding explicitly makes sense regardless. The tests cover both the close-before-user-assignment path ("") and the literal null case.

[jongyoul]

Re-reviewed at exact head. The null-user guard and regression tests address the prior concern, and the concurrent map/queue operations preserve the required invariants. The sole Selenium failure is unrelated and also recurs on master.