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