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]

Reply via email to