max / makenotwork
1 file changed,
+19 insertions,
-20 deletions
| @@ -34,28 +34,27 @@ struct SseConnectionGuard { | |||
| 34 | 34 | ||
| 35 | 35 | impl Drop for SseConnectionGuard { | |
| 36 | 36 | fn drop(&mut self) { | |
| 37 | - | // Decrement connection counter, then remove entry if it hit zero. | |
| 38 | - | // Must drop the read guard before calling remove() to avoid deadlocking | |
| 39 | - | // on the same DashMap shard. | |
| 40 | - | let should_remove_connection = self.sse_connections.get(&self.user_id).map(|counter| { | |
| 41 | - | let prev = counter.value().fetch_sub(1, Ordering::AcqRel); | |
| 42 | - | prev <= 1 | |
| 43 | - | }).unwrap_or(false); | |
| 44 | - | ||
| 45 | - | if should_remove_connection { | |
| 46 | - | self.sse_connections.remove(&self.user_id); | |
| 37 | + | // Decrement the connection counter, then drop the entry if it hit zero. | |
| 38 | + | // `remove_if` runs the predicate under the shard lock, so a connection | |
| 39 | + | // arriving concurrently cannot be counted into an entry we then remove. | |
| 40 | + | // It also avoids the same-shard deadlock a get() + remove() pair would | |
| 41 | + | // hit, which is what the earlier read/drop/remove split was working | |
| 42 | + | // around. | |
| 43 | + | if let Some(counter) = self.sse_connections.get(&self.user_id) { | |
| 44 | + | counter.value().fetch_sub(1, Ordering::AcqRel); | |
| 47 | 45 | } | |
| 48 | - | ||
| 49 | - | // Prune sync_notify channel if no receivers remain. | |
| 50 | - | // Same pattern: read first, drop guard, then remove. | |
| 46 | + | self.sse_connections | |
| 47 | + | .remove_if(&self.user_id, |_, counter| counter.load(Ordering::Acquire) == 0); | |
| 48 | + | ||
| 49 | + | // Prune the sync_notify channel once the last receiver drops. This must | |
| 50 | + | // be atomic with respect to `Sync::subscribe_channel`: reading | |
| 51 | + | // receiver_count(), dropping the guard, then removing lets a subscriber | |
| 52 | + | // land in between and hold a receiver on a sender we then discard. The | |
| 53 | + | // next push lazily creates a *different* channel, and that client | |
| 54 | + | // silently receives no further SSE events until it reconnects. | |
| 51 | 55 | let key = (self.app_id, self.user_id); | |
| 52 | - | let should_remove_notify = self.sync_notify.get(&key) | |
| 53 | - | .map(|entry| entry.value().receiver_count() == 0) | |
| 54 | - | .unwrap_or(false); | |
| 55 | - | ||
| 56 | - | if should_remove_notify { | |
| 57 | - | self.sync_notify.remove(&key); | |
| 58 | - | } | |
| 56 | + | self.sync_notify | |
| 57 | + | .remove_if(&key, |_, sender| sender.receiver_count() == 0); | |
| 59 | 58 | } | |
| 60 | 59 | } | |
| 61 | 60 |