wse-server 2.0.2__tar.gz → 2.0.7__tar.gz

{wse_server-2.0.2 → wse_server-2.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wse-server
-Version: 2.0.2
+Version: 2.0.7
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License

{wse_server-2.0.2 → wse_server-2.0.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "wse-server"
-version = "2.0.2"
+version = "2.0.7"
 description = "Rust-powered WebSocket engine for Python. Up to 5M fan-out del/s. Zero-GIL JWT."
 readme = "README.md"
 license = { text = "MIT" }

{wse_server-2.0.2 → wse_server-2.0.7}/rust/Cargo.lock

@@ -1912,7 +1912,7 @@ dependencies = [
 
 [[package]]
 name = "wse-accel"
-version = "2.0.2"
+version = "2.0.7"
 dependencies = [
  "aes-gcm",
  "ahash",

{wse_server-2.0.2 → wse_server-2.0.7}/rust/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "wse-accel"
-version = "2.0.2"
+version = "2.0.7"
 edition = "2024"
 
 [lib]

{wse_server-2.0.2 → wse_server-2.0.7}/rust/src/wse/cluster.rs

@@ -1693,6 +1693,8 @@ pub(crate) async fn cluster_manager(
     // Dynamic peer tracking
     let known_peers: Arc<DashSet<String>> = Arc::new(DashSet::new());
     let connected_instances: Arc<DashMap<String, String>> = Arc::new(DashMap::new());
+    // Static peers have their own reconnect loop; never remove them from known_peers
+    let static_peers: std::collections::HashSet<String> = peers.iter().cloned().collect();
 
     // Channel for peer_reader to notify about newly discovered peers
     let (new_peer_tx, mut new_peer_rx) = mpsc::unbounded_channel::<String>();
@@ -1946,14 +1948,15 @@ pub(crate) async fn cluster_manager(
                             eprintln!("[WSE-Cluster] Ignoring invalid gossip address from cmd_rx: {addr}");
                             continue;
                         }
-                        // Forward PEER_ANNOUNCE to all peers except the source
-                        // Note: don't call write_framed() here -- peer_writer adds framing
-                        let mut frame_data = BytesMut::new();
-                        encode_peer_announce(&mut frame_data, &addr);
-                        let frame_bytes = frame_data.freeze();
-                        for (peer_addr, tx, _) in &peer_txs {
-                            if *peer_addr != exclude_peer {
-                                let _ = tx.try_send(frame_bytes.clone());
+                        // Forward only for newly discovered peers to prevent amplification
+                        if !known_peers.contains(&addr) {
+                            let mut frame_data = BytesMut::new();
+                            encode_peer_announce(&mut frame_data, &addr);
+                            let frame_bytes = frame_data.freeze();
+                            for (peer_addr, tx, _) in &peer_txs {
+                                if *peer_addr != exclude_peer {
+                                    let _ = tx.try_send(frame_bytes.clone());
+                                }
                             }
                         }
                     }
@@ -2036,15 +2039,17 @@ pub(crate) async fn cluster_manager(
                     if is_self {
                         continue;
                     }
-                    // Always forward PEER_ANNOUNCE to all peers except the source,
-                    // even for peers we already know about. Other peers may not know yet.
-                    // Note: don't call write_framed() here -- peer_writer adds framing
-                    let mut frame_data = BytesMut::new();
-                    encode_peer_announce(&mut frame_data, &addr);
-                    let frame_bytes = frame_data.freeze();
-                    for (peer_addr, tx, _) in &peer_txs {
-                        if *peer_addr != exclude_peer {
-                            let _ = tx.try_send(frame_bytes.clone());
+                    // Forward PEER_ANNOUNCE only for newly discovered peers
+                    // to prevent gossip amplification in dense clusters
+                    let is_new = !known_peers.contains(&addr);
+                    if is_new {
+                        let mut frame_data = BytesMut::new();
+                        encode_peer_announce(&mut frame_data, &addr);
+                        let frame_bytes = frame_data.freeze();
+                        for (peer_addr, tx, _) in &peer_txs {
+                            if *peer_addr != exclude_peer {
+                                let _ = tx.try_send(frame_bytes.clone());
+                            }
                         }
                     }
                     // Only spawn connection to this peer if we don't already know it
@@ -2082,6 +2087,10 @@ pub(crate) async fn cluster_manager(
                     }
                     Some(InterestUpdate::PeerDisconnected { peer_addr }) => {
                         remote_interest.remove(&peer_addr);
+                        // Allow gossip-discovered peers to be re-discovered
+                        if !static_peers.contains(&peer_addr) {
+                            known_peers.remove(&peer_addr);
+                        }
                     }
                     None => {
                         // Interest channel closed, continue running
@@ -2101,6 +2110,10 @@ pub(crate) async fn cluster_manager(
                     PeerRegistration::Deregister { peer_addr } => {
                         // Remove disconnected inbound peer
                         peer_txs.retain(|(a, _, _)| a != &peer_addr);
+                        // Allow gossip-discovered peers to be re-discovered
+                        if !static_peers.contains(&peer_addr) {
+                            known_peers.remove(&peer_addr);
+                        }
                     }
                 }
             }
@@ -2119,6 +2132,7 @@ pub(crate) async fn cluster_manager(
                                         let _ = tx.send(super::server::ServerCommand::BroadcastLocal {
                                             topic,
                                             data: msg,
+                                            skip_recovery: true,
                                         });
                                     }
                                 }
@@ -2130,6 +2144,7 @@ pub(crate) async fn cluster_manager(
                                         let _ = tx.send(super::server::ServerCommand::BroadcastLocal {
                                             topic,
                                             data: msg,
+                                            skip_recovery: true,
                                         });
                                     }
                                 }
@@ -2141,6 +2156,7 @@ pub(crate) async fn cluster_manager(
                                         let _ = tx.send(super::server::ServerCommand::BroadcastLocal {
                                             topic,
                                             data: msg,
+                                            skip_recovery: true,
                                         });
                                     }
                                 }

{wse_server-2.0.2 → wse_server-2.0.7}/rust/src/wse/presence.rs

@@ -142,7 +142,7 @@ impl PresenceManager {
             return TrackResult::AtCapacity;
         }
 
-        let (is_first, actually_inserted) = {
+        let (is_first, actually_inserted, has_remote_sentinel) = {
             let mut entry = topic_map
                 .entry(user_id.clone())
                 .or_insert_with(|| PresenceEntry {
@@ -151,7 +151,12 @@ impl PresenceManager {
                     connections: Vec::new(),
                     updated_at: now,
                 });
-            // Check for first LOCAL connection (ignore remote sentinels from cluster sync)
+            // Check for remote sentinel (from cluster sync) while holding the shard lock
+            let has_remote = entry
+                .connections
+                .iter()
+                .any(|c| c.conn_id.starts_with("__remote__"));
+            // Check for first LOCAL connection (ignore remote sentinels)
             let was_empty = !entry
                 .connections
                 .iter()
@@ -168,7 +173,7 @@ impl PresenceManager {
             };
             entry.data = data.clone();
             entry.updated_at = now;
-            (was_empty, inserted)
+            (was_empty, inserted, has_remote)
         };
 
         // Track which topics this conn has presence in
@@ -187,7 +192,11 @@ impl PresenceManager {
         }
 
         if is_first {
-            stats.num_users.fetch_add(1, Ordering::Relaxed);
+            // Only increment num_users if no remote sentinel exists
+            // (merge_remote_join already counted this user)
+            if !has_remote_sentinel {
+                stats.num_users.fetch_add(1, Ordering::Relaxed);
+            }
             TrackResult::FirstJoin {
                 user_id,
                 data: data.clone(),
@@ -381,8 +390,12 @@ impl PresenceManager {
             }
 
             for (user_id, data, _dead_conn_count) in dead_users {
-                // Collect dead conn_ids before removing the user entry
-                let actual_conn_count = if let Some((_, entry)) = topic_map.remove(&user_id) {
+                // Use remove_if to guard against concurrent track() adding a new conn
+                let actual_conn_count = if let Some((_, entry)) =
+                    topic_map.remove_if(&user_id, |_, e| {
+                        // Only remove if all connections are still dead
+                        e.connections.iter().all(|c| !is_conn_alive(&c.conn_id))
+                    }) {
                     let count = entry.connections.len();
                     for conn in &entry.connections {
                         self.conn_user_id.remove(&conn.conn_id);
@@ -390,21 +403,42 @@ impl PresenceManager {
                     }
                     count
                 } else {
+                    // Entry was either already removed or a live conn was added - skip
+                    // But still clean up any dead local connections from the entry
+                    if let Some(mut entry) = topic_map.get_mut(&user_id) {
+                        let before = entry.connections.len();
+                        entry.connections.retain(|c| {
+                            c.conn_id.starts_with("__remote__") || is_conn_alive(&c.conn_id)
+                        });
+                        let removed_count = before - entry.connections.len();
+                        if removed_count > 0
+                            && let Some(stats) = self.topic_presence_stats.get(&topic)
+                        {
+                            let _ = stats.num_connections.fetch_update(
+                                Ordering::Relaxed,
+                                Ordering::Relaxed,
+                                |current| Some(current.saturating_sub(removed_count)),
+                            );
+                        }
+                    }
                     0
                 };
-                if let Some(stats) = self.topic_presence_stats.get(&topic) {
-                    let _ = stats.num_users.fetch_update(
-                        Ordering::Relaxed,
-                        Ordering::Relaxed,
-                        |current| Some(current.saturating_sub(1)),
-                    );
-                    let _ = stats.num_connections.fetch_update(
-                        Ordering::Relaxed,
-                        Ordering::Relaxed,
-                        |current| Some(current.saturating_sub(actual_conn_count)),
-                    );
+                // Only decrement stats and emit leave if the entry was actually removed
+                if actual_conn_count > 0 {
+                    if let Some(stats) = self.topic_presence_stats.get(&topic) {
+                        let _ = stats.num_users.fetch_update(
+                            Ordering::Relaxed,
+                            Ordering::Relaxed,
+                            |current| Some(current.saturating_sub(1)),
+                        );
+                        let _ = stats.num_connections.fetch_update(
+                            Ordering::Relaxed,
+                            Ordering::Relaxed,
+                            |current| Some(current.saturating_sub(actual_conn_count)),
+                        );
+                    }
+                    removed.push((topic.clone(), user_id, data));
                 }
-                removed.push((topic.clone(), user_id, data));
             }
         }
 

{wse_server-2.0.2 → wse_server-2.0.7}/rust/src/wse/recovery.rs

@@ -164,9 +164,10 @@ impl TopicRecoveryBuffer {
             {
                 result.push(entry.data.clone());
             } else {
-                // Unreachable: the tail guard ensures every slot in [start, head_offset)
-                // was written after tail advanced past any previous occupant.
+                // Slot mismatch should be unreachable, but if it occurs in release,
+                // signal failure so the client re-subscribes rather than missing messages.
                 debug_assert!(false, "slot {} mismatch: expected offset {}", idx, offset);
+                return None;
             }
         }
 

{wse_server-2.0.2 → wse_server-2.0.7}/rust/src/wse/server.rs

@@ -170,14 +170,35 @@ const RATE_CAPACITY: f64 = 100_000.0;
 const RATE_REFILL: f64 = 10_000.0; // tokens per second
 
 pub(crate) enum ServerCommand {
-    SendText { conn_id: String, data: String },
-    SendBytes { conn_id: String, data: Vec<u8> },
-    SendPrebuilt { conn_id: String, message: Message },
-    BroadcastText { data: String },
-    BroadcastBytes { data: Vec<u8> },
-    BroadcastLocal { topic: String, data: String },
-    Disconnect { conn_id: String },
-    GetConnections { reply: oneshot::Sender<Vec<String>> },
+    SendText {
+        conn_id: String,
+        data: String,
+    },
+    SendBytes {
+        conn_id: String,
+        data: Vec<u8>,
+    },
+    SendPrebuilt {
+        conn_id: String,
+        message: Message,
+    },
+    BroadcastText {
+        data: String,
+    },
+    BroadcastBytes {
+        data: Vec<u8>,
+    },
+    BroadcastLocal {
+        topic: String,
+        data: String,
+        skip_recovery: bool,
+    },
+    Disconnect {
+        conn_id: String,
+    },
+    GetConnections {
+        reply: oneshot::Sender<Vec<String>>,
+    },
     Shutdown,
 }
 
@@ -1321,6 +1342,7 @@ async fn handle_connection(stream: TcpStream, addr: SocketAddr, state: Arc<Share
                     let _ = tx.send(ServerCommand::BroadcastLocal {
                         topic: topic.clone(),
                         data: msg,
+                        skip_recovery: true,
                     });
                 }
                 // Cluster sync: notify peers of presence leave
@@ -1605,10 +1627,14 @@ async fn process_commands(
                     }
                 }
             }
-            ServerCommand::BroadcastLocal { topic, data } => {
+            ServerCommand::BroadcastLocal {
+                topic,
+                data,
+                skip_recovery,
+            } => {
                 let preframed = pre_frame_text(&data);
-                // Store in recovery buffer (zero-copy: Bytes is Arc-refcounted)
-                if let Some(ref recovery) = state.recovery {
+                // Store in recovery buffer (skip for presence events)
+                if !skip_recovery && let Some(ref recovery) = state.recovery {
                     recovery.push(&topic, preframed.clone());
                 }
                 if state.conn_encryption.is_empty() {
@@ -1921,6 +1947,7 @@ impl RustWSEServer {
                                             ServerCommand::BroadcastLocal {
                                                 topic: topic.clone(),
                                                 data: msg,
+                                                skip_recovery: true,
                                             },
                                         );
                                     }
@@ -2029,6 +2056,7 @@ impl RustWSEServer {
                                                         let _ = tx.send(ServerCommand::BroadcastLocal {
                                                             topic: topic.clone(),
                                                             data: msg,
+                                                            skip_recovery: true,
                                                         });
                                                     }
                                                     if let Some(ref ctx) = cluster_tx {
@@ -2934,6 +2962,7 @@ impl RustWSEServer {
                         let _ = tx.send(ServerCommand::BroadcastLocal {
                             topic: topic.clone(),
                             data: msg,
+                            skip_recovery: true,
                         });
                     }
                     // Cluster sync: notify peers of presence join
@@ -3161,6 +3190,7 @@ impl RustWSEServer {
                         let _ = tx.send(ServerCommand::BroadcastLocal {
                             topic: topic.clone(),
                             data: msg,
+                            skip_recovery: true,
                         });
                     }
                     // Cluster sync: notify peers of presence leave
@@ -3206,6 +3236,7 @@ impl RustWSEServer {
         tx.send(ServerCommand::BroadcastLocal {
             topic: topic.to_owned(),
             data: data.to_owned(),
+            skip_recovery: false,
         })
         .map_err(|_| PyRuntimeError::new_err("Channel closed"))?;
         Ok(())
@@ -3218,6 +3249,7 @@ impl RustWSEServer {
             let _ = tx.send(ServerCommand::BroadcastLocal {
                 topic: topic.to_owned(),
                 data: data.to_owned(),
+                skip_recovery: false,
             });
         }
 
@@ -3393,6 +3425,7 @@ impl RustWSEServer {
                             let _ = tx.send(ServerCommand::BroadcastLocal {
                                 topic: topic.clone(),
                                 data: msg.clone(),
+                                skip_recovery: true,
                             });
                         }
                     }

{wse_server-2.0.2 → wse_server-2.0.7}/wse_server/__init__.py

@@ -26,7 +26,7 @@ from wse_server._wse_accel import (
 from wse_server.core.filters import EventFilter
 from wse_server.core.types import DeliveryGuarantee, EventPriority
 
-__version__ = "2.0.2"
+__version__ = "2.0.7"
 
 __all__ = [
     "RustWSEServer",