soli_proxy/app/
port_manager.rs

1use anyhow::Result;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::fs;
5use std::path::PathBuf;
6use std::sync::Arc;
7use tokio::sync::Mutex;
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct PortAssignment {
11    pub app_name: String,
12    pub slot: String,
13    pub port: u16,
14    pub timestamp: String,
15}
16
17#[derive(Default)]
18pub struct PortAllocator {
19    used_ports: HashMap<u16, PortAssignment>,
20    app_slots: HashMap<(String, String), u16>,
21}
22
23impl PortAllocator {
24    pub fn new() -> Self {
25        Self {
26            used_ports: HashMap::new(),
27            app_slots: HashMap::new(),
28        }
29    }
30
31    pub fn allocate(&mut self, app_name: &str, slot: &str) -> Result<u16> {
32        let key = (app_name.to_string(), slot.to_string());
33
34        if let Some(&port) = self.app_slots.get(&key) {
35            return Ok(port);
36        }
37
38        let port = self.find_available_port()?;
39        self.used_ports.insert(
40            port,
41            PortAssignment {
42                app_name: app_name.to_string(),
43                slot: slot.to_string(),
44                port,
45                timestamp: chrono::Utc::now().to_rfc3339(),
46            },
47        );
48        self.app_slots.insert(key, port);
49
50        Ok(port)
51    }
52
53    fn find_available_port(&self) -> Result<u16> {
54        if let Some(port) = portpicker::pick_unused_port() {
55            return Ok(port);
56        }
57        anyhow::bail!("No available ports found")
58    }
59
60    pub fn release(&mut self, app_name: &str, slot: &str) {
61        let key = (app_name.to_string(), slot.to_string());
62        if let Some(port) = self.app_slots.remove(&key) {
63            self.used_ports.remove(&port);
64        }
65    }
66
67    pub fn get_port(&self, app_name: &str, slot: &str) -> Option<u16> {
68        self.app_slots
69            .get(&(app_name.to_string(), slot.to_string()))
70            .copied()
71    }
72}
73
74pub struct PortManager {
75    allocator: Arc<Mutex<PortAllocator>>,
76    lock_file: PathBuf,
77}
78
79impl PortManager {
80    pub fn new(lock_dir: &str) -> Result<Self> {
81        let lock_dir = PathBuf::from(lock_dir);
82        fs::create_dir_all(&lock_dir)?;
83
84        let lock_file = lock_dir.join("ports.lock");
85
86        let allocator = Arc::new(Mutex::new(PortAllocator::new()));
87
88        Ok(Self {
89            allocator,
90            lock_file,
91        })
92    }
93
94    pub async fn allocate(&self, app_name: &str, slot: &str) -> Result<u16> {
95        let port = {
96            let mut allocator = self.allocator.lock().await;
97            if let Some(port) = allocator.get_port(app_name, slot) {
98                return Ok(port);
99            }
100            allocator.allocate(app_name, slot)?
101        };
102        self.persist().await?;
103        Ok(port)
104    }
105
106    pub async fn release(&self, app_name: &str, slot: &str) {
107        {
108            let mut allocator = self.allocator.lock().await;
109            allocator.release(app_name, slot);
110        }
111        let _ = self.persist().await;
112    }
113
114    async fn persist(&self) -> Result<()> {
115        let allocator = self.allocator.lock().await;
116        let content = serde_json::to_string_pretty(&allocator.used_ports)?;
117        fs::write(&self.lock_file, content)?;
118        Ok(())
119    }
120
121    pub async fn load(&self) -> Result<()> {
122        if !self.lock_file.exists() {
123            return Ok(());
124        }
125
126        let content = fs::read_to_string(&self.lock_file)?;
127        let assignments: HashMap<u16, PortAssignment> = serde_json::from_str(&content)?;
128
129        let mut allocator = self.allocator.lock().await;
130        for (port, assignment) in assignments {
131            allocator.used_ports.insert(port, assignment.clone());
132            allocator
133                .app_slots
134                .insert((assignment.app_name, assignment.slot), port);
135        }
136
137        Ok(())
138    }
139
140    pub async fn get_port(&self, app_name: &str, slot: &str) -> Option<u16> {
141        self.allocator.lock().await.get_port(app_name, slot)
142    }
143
144    pub async fn get_app_name(&self, port: u16) -> Option<String> {
145        let allocator = self.allocator.lock().await;
146        allocator.used_ports.get(&port).map(|a| a.app_name.clone())
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use tempfile::TempDir;
154
155    #[tokio::test]
156    async fn test_port_allocation() {
157        let temp_dir = TempDir::new().unwrap();
158        let pm = PortManager::new(temp_dir.path().to_str().unwrap()).unwrap();
159
160        let port1 = pm.allocate("app1", "blue").await.unwrap();
161        let port2 = pm.allocate("app1", "green").await.unwrap();
162        let port3 = pm.allocate("app2", "blue").await.unwrap();
163
164        assert_ne!(port1, port2);
165        assert_ne!(port1, port3);
166        assert_ne!(port2, port3);
167
168        assert_eq!(pm.allocate("app1", "blue").await.unwrap(), port1);
169    }
170
171    #[tokio::test]
172    async fn test_port_release() {
173        let temp_dir = TempDir::new().unwrap();
174        let pm = PortManager::new(temp_dir.path().to_str().unwrap()).unwrap();
175
176        let port = pm.allocate("app1", "blue").await.unwrap();
177        pm.release("app1", "blue").await;
178
179        let new_port = pm.allocate("app1", "blue").await.unwrap();
180        assert_ne!(port, new_port);
181    }
182}
soli_proxy/app/port_manager.rs

soli_proxy/app/
port_manager.rs
