HwangRock opened a new pull request, #5306: URL: https://github.com/apache/zeppelin/pull/5306
### 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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
