Filter out dropped clients in the send worker

Author: joshua-spacetime

crates/core/src/client/client_connection.rs

@@ -1,5 +1,6 @@
 use std::collections::VecDeque;
 use std::ops::Deref;
+use std::sync::atomic::Ordering;
 use std::sync::atomic::{AtomicBool, Ordering::Relaxed};
 use std::sync::Arc;
 use std::time::Instant;
@@ -158,6 +159,10 @@ impl ClientConnectionSender {
         Self::dummy_with_channel(id, config).0
     }
 
+    pub fn is_cancelled(&self) -> bool {
+        self.cancelled.load(Ordering::Acquire)
+    }
+
     /// Send a message to the client. For data-related messages, you should probably use
     /// `BroadcastQueue::send` to ensure that the client sees data messages in a consistent order.
     pub fn send_message(&self, message: impl Into<SerializableMessage>) -> Result<(), ClientSendError> {

crates/core/src/subscription/module_subscription_manager.rs

@@ -1197,12 +1197,7 @@ impl SubscriptionManager {
                     SingleQueryUpdate { update, num_rows }
                 }
 
-                // filter out clients that've dropped
-                let clients_for_query = qstate.all_clients().filter(|id| {
-                    self.clients
-                        .get(*id)
-                        .is_some_and(|info| !info.dropped.load(Ordering::Acquire))
-                });
+                let clients_for_query = qstate.all_clients();
 
                 match eval_delta(tx, &mut acc.metrics, plan) {
                     Err(err) => {
@@ -1281,6 +1276,16 @@ struct SendWorkerClient {
     outbound_ref: Client,
 }
 
+impl SendWorkerClient {
+    fn is_dropped(&self) -> bool {
+        self.dropped.load(Ordering::Acquire)
+    }
+
+    fn is_cancelled(&self) -> bool {
+        self.outbound_ref.is_cancelled()
+    }
+}
+
 /// Asynchronous background worker which aggregates each of the clients' updates from a [`ComputedQueries`]
 /// into `DbUpdate`s and then sends them to the clients' WebSocket workers.
 ///
@@ -1319,6 +1324,14 @@ impl Drop for SendWorker {
     }
 }
 
+impl SendWorker {
+    fn is_client_dropped_or_cancelled(&self, client_id: &ClientId) -> bool {
+        self.clients
+            .get(client_id)
+            .is_some_and(|client| client.is_cancelled() || client.is_dropped())
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct BroadcastQueue(SenderWithGauge<SendWorkerMessage>);
 
@@ -1420,6 +1433,8 @@ impl SendWorker {
 
         let mut eval = updates
             .into_iter()
+            // Filter out dropped or cancelled clients
+            .filter(|upd| !self.is_client_dropped_or_cancelled(&upd.id))
             // Filter out clients whose subscriptions failed
             .filter(|upd| !clients_with_errors.contains(&upd.id))
             // For each subscriber, aggregate all the updates for the same table.