feox_server/
client_registry.rs

1use crate::network::Connection;
2use dashmap::DashMap;
3use std::net::SocketAddr;
4use std::sync::Arc;
5
6#[derive(Debug, Clone)]
7pub struct ClientInfo {
8    pub id: usize,
9    pub name: Option<String>,
10    pub addr: Option<SocketAddr>,
11    pub fd: i32,
12    pub connected_at: u64,
13    pub commands_processed: u64,
14    pub flags: Vec<String>,
15    pub thread_id: usize,
16    pub db: usize,
17}
18
19/// Global registry for all client connections
20pub struct ClientRegistry {
21    clients: Arc<DashMap<usize, ClientInfo>>,
22}
23
24impl ClientRegistry {
25    pub fn new() -> Self {
26        Self {
27            clients: Arc::new(DashMap::new()),
28        }
29    }
30
31    /// Register a new client connection
32    pub fn register(&self, conn: &Connection, thread_id: usize) {
33        let info = ClientInfo {
34            id: conn.connection_id,
35            name: conn.client_name.clone(),
36            addr: conn.client_addr,
37            fd: conn.fd(),
38            connected_at: conn.connected_at,
39            commands_processed: conn.commands_processed,
40            flags: conn.flags.clone(),
41            thread_id,
42            db: 0,
43        };
44        self.clients.insert(conn.connection_id, info);
45    }
46
47    /// Update client info
48    pub fn update(&self, conn: &Connection) {
49        if let Some(mut entry) = self.clients.get_mut(&conn.connection_id) {
50            entry.name = conn.client_name.clone();
51            entry.commands_processed = conn.commands_processed;
52            entry.flags = conn.flags.clone();
53        }
54    }
55
56    /// Unregister a client connection
57    pub fn unregister(&self, connection_id: usize) {
58        self.clients.remove(&connection_id);
59    }
60
61    /// Get all client information
62    pub fn get_all_clients(&self) -> Vec<ClientInfo> {
63        self.clients
64            .iter()
65            .map(|entry| entry.value().clone())
66            .collect()
67    }
68
69    /// Get client by ID
70    pub fn get_client(&self, connection_id: usize) -> Option<ClientInfo> {
71        self.clients.get(&connection_id).map(|e| e.clone())
72    }
73
74    /// Kill clients matching criteria
75    pub fn kill_clients(
76        &self,
77        filter_id: Option<usize>,
78        filter_addr: Option<&str>,
79        filter_type: Option<&str>,
80    ) -> Vec<usize> {
81        let mut to_kill = Vec::new();
82
83        for entry in self.clients.iter() {
84            let client = entry.value();
85            let mut should_kill = false;
86
87            if let Some(id) = filter_id {
88                if client.id == id {
89                    should_kill = true;
90                }
91            }
92
93            if let Some(addr_str) = filter_addr {
94                if let Some(addr) = client.addr {
95                    if addr.to_string() == addr_str {
96                        should_kill = true;
97                    }
98                }
99            }
100
101            if let Some(client_type) = filter_type {
102                match client_type {
103                    "normal" => {
104                        if !client.flags.contains(&"pubsub".to_string()) {
105                            should_kill = true;
106                        }
107                    }
108                    "pubsub" => {
109                        if client.flags.contains(&"pubsub".to_string()) {
110                            should_kill = true;
111                        }
112                    }
113                    _ => {}
114                }
115            }
116
117            if should_kill {
118                to_kill.push(client.id);
119            }
120        }
121
122        to_kill
123    }
124
125    /// Count total clients
126    pub fn client_count(&self) -> usize {
127        self.clients.len()
128    }
129}
130
131impl Default for ClientRegistry {
132    fn default() -> Self {
133        Self::new()
134    }
135}
feox_server/client_registry.rs

feox_server/
client_registry.rs
