batata_client/remote/
server_list.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3use crate::common::{parse_server_address, GRPC_PORT_OFFSET};
4use crate::error::{BatataError, Result};
5
6/// Server address information
7#[derive(Clone, Debug)]
8pub struct ServerAddress {
9    pub host: String,
10    pub port: u16,
11    pub grpc_port: u16,
12    pub tls_enabled: bool,
13}
14
15impl ServerAddress {
16    pub fn new(host: &str, port: u16) -> Self {
17        Self {
18            host: host.to_string(),
19            port,
20            grpc_port: port + GRPC_PORT_OFFSET,
21            tls_enabled: false,
22        }
23    }
24
25    /// Enable TLS for this server address
26    pub fn with_tls(mut self, enabled: bool) -> Self {
27        self.tls_enabled = enabled;
28        self
29    }
30
31    /// Get the host
32    pub fn host(&self) -> &str {
33        &self.host
34    }
35
36    /// Get the HTTP port
37    pub fn port(&self) -> u16 {
38        self.port
39    }
40
41    /// Get the address string (host:port)
42    pub fn address(&self) -> String {
43        format!("{}:{}", self.host, self.port)
44    }
45
46    pub fn grpc_endpoint(&self) -> String {
47        let scheme = if self.tls_enabled { "https" } else { "http" };
48        format!("{}://{}:{}", scheme, self.host, self.grpc_port)
49    }
50
51    pub fn http_endpoint(&self) -> String {
52        let scheme = if self.tls_enabled { "https" } else { "http" };
53        format!("{}://{}:{}", scheme, self.host, self.port)
54    }
55}
56
57/// Server list manager for load balancing
58pub struct ServerListManager {
59    servers: Vec<ServerAddress>,
60    current_index: AtomicUsize,
61}
62
63impl ServerListManager {
64    /// Create a new server list manager from address strings
65    pub fn new(addresses: Vec<String>) -> Result<Self> {
66        if addresses.is_empty() {
67            return Err(BatataError::InvalidParameter(
68                "Server addresses cannot be empty".to_string(),
69            ));
70        }
71
72        let servers: Vec<ServerAddress> = addresses
73            .iter()
74            .map(|addr| {
75                let (host, port) = parse_server_address(addr);
76                ServerAddress::new(&host, port)
77            })
78            .collect();
79
80        Ok(Self {
81            servers,
82            current_index: AtomicUsize::new(0),
83        })
84    }
85
86    /// Get the next server using round-robin
87    pub fn next_server(&self) -> &ServerAddress {
88        let index = self.current_index.fetch_add(1, Ordering::Relaxed) % self.servers.len();
89        &self.servers[index]
90    }
91
92    /// Get current server
93    pub fn current_server(&self) -> &ServerAddress {
94        let index = self.current_index.load(Ordering::Relaxed) % self.servers.len();
95        &self.servers[index]
96    }
97
98    /// Get all servers
99    pub fn all_servers(&self) -> &[ServerAddress] {
100        &self.servers
101    }
102
103    /// Get server count
104    pub fn server_count(&self) -> usize {
105        self.servers.len()
106    }
107
108    /// Mark current server as failed and move to next
109    pub fn fail_current(&self) -> &ServerAddress {
110        self.next_server()
111    }
112}
113
114impl Clone for ServerListManager {
115    fn clone(&self) -> Self {
116        Self {
117            servers: self.servers.clone(),
118            current_index: AtomicUsize::new(self.current_index.load(Ordering::Relaxed)),
119        }
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126
127    #[test]
128    fn test_server_address() {
129        let addr = ServerAddress::new("localhost", 8848);
130        assert_eq!(addr.grpc_port, 9848);
131        assert_eq!(addr.grpc_endpoint(), "http://localhost:9848");
132        assert_eq!(addr.http_endpoint(), "http://localhost:8848");
133    }
134
135    #[test]
136    fn test_server_list_manager() {
137        let manager = ServerListManager::new(vec![
138            "localhost:8848".to_string(),
139            "localhost:8849".to_string(),
140        ])
141        .unwrap();
142
143        assert_eq!(manager.server_count(), 2);
144
145        let first = manager.next_server().host.clone();
146        let second = manager.next_server().host.clone();
147        let third = manager.next_server().host.clone();
148
149        assert_eq!(first, "localhost");
150        assert_eq!(second, "localhost");
151        assert_eq!(third, "localhost");
152    }
153}
batata_client/remote/server_list.rs

batata_client/remote/
server_list.rs
