feox_server/
client_registry.rs1use 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
20pub 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 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 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 pub fn unregister(&self, connection_id: usize) {
61 self.clients.remove(&connection_id);
62 }
63
64 pub fn get_all_clients(&self) -> Vec<ClientInfo> {
66 self.clients
67 .iter()
68 .map(|entry| entry.value().clone())
69 .collect()
70 }
71
72 pub fn get_client(&self, connection_id: usize) -> Option<ClientInfo> {
74 self.clients.get(&connection_id).map(|e| e.clone())
75 }
76
77 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 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}