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 self.allocate_with_range(app_name, slot, 1, 65535)
33 }
34
35 pub fn allocate_with_range(
36 &mut self,
37 app_name: &str,
38 slot: &str,
39 port_range_start: u16,
40 port_range_end: u16,
41 ) -> Result<u16> {
42 let key = (app_name.to_string(), slot.to_string());
43
44 if let Some(&port) = self.app_slots.get(&key) {
45 return Ok(port);
46 }
47
48 let port = self.find_available_port_in_range(port_range_start, port_range_end)?;
49 self.used_ports.insert(
50 port,
51 PortAssignment {
52 app_name: app_name.to_string(),
53 slot: slot.to_string(),
54 port,
55 timestamp: chrono::Utc::now().to_rfc3339(),
56 },
57 );
58 self.app_slots.insert(key, port);
59
60 Ok(port)
61 }
62
63 fn find_available_port_in_range(
64 &self,
65 port_range_start: u16,
66 port_range_end: u16,
67 ) -> Result<u16> {
68 for port in port_range_start..=port_range_end {
69 if !self.used_ports.contains_key(&port) {
70 return Ok(port);
71 }
72 }
73 anyhow::bail!(
74 "No available ports found in range {}-{}",
75 port_range_start,
76 port_range_end
77 )
78 }
79
80 pub fn release(&mut self, app_name: &str, slot: &str) {
81 let key = (app_name.to_string(), slot.to_string());
82 if let Some(port) = self.app_slots.remove(&key) {
83 self.used_ports.remove(&port);
84 }
85 }
86
87 pub fn get_port(&self, app_name: &str, slot: &str) -> Option<u16> {
88 self.app_slots
89 .get(&(app_name.to_string(), slot.to_string()))
90 .copied()
91 }
92}
93
94pub struct PortManager {
95 allocator: Arc<Mutex<PortAllocator>>,
96 lock_file: PathBuf,
97}
98
99impl PortManager {
100 pub fn new(lock_dir: &str) -> Result<Self> {
101 let lock_dir = PathBuf::from(lock_dir);
102 fs::create_dir_all(&lock_dir)?;
103
104 let lock_file = lock_dir.join("ports.lock");
105
106 let allocator = Arc::new(Mutex::new(PortAllocator::new()));
107
108 Ok(Self {
109 allocator,
110 lock_file,
111 })
112 }
113
114 pub async fn allocate(&self, app_name: &str, slot: &str) -> Result<u16> {
115 self.allocate_with_range(app_name, slot, 1, 65535).await
116 }
117
118 pub async fn allocate_with_range(
119 &self,
120 app_name: &str,
121 slot: &str,
122 port_range_start: u16,
123 port_range_end: u16,
124 ) -> Result<u16> {
125 let port = {
126 let mut allocator = self.allocator.lock().await;
127 if let Some(port) = allocator.get_port(app_name, slot) {
128 return Ok(port);
129 }
130 allocator.allocate_with_range(app_name, slot, port_range_start, port_range_end)?
131 };
132 self.persist().await?;
133 Ok(port)
134 }
135
136 pub async fn release(&self, app_name: &str, slot: &str) {
137 {
138 let mut allocator = self.allocator.lock().await;
139 allocator.release(app_name, slot);
140 }
141 let _ = self.persist().await;
142 }
143
144 async fn persist(&self) -> Result<()> {
145 let allocator = self.allocator.lock().await;
146 let content = serde_json::to_string_pretty(&allocator.used_ports)?;
147 fs::write(&self.lock_file, content)?;
148 Ok(())
149 }
150
151 pub async fn load(&self) -> Result<()> {
152 if !self.lock_file.exists() {
153 return Ok(());
154 }
155
156 let content = fs::read_to_string(&self.lock_file)?;
157 let assignments: HashMap<u16, PortAssignment> = serde_json::from_str(&content)?;
158
159 let mut allocator = self.allocator.lock().await;
160 for (port, assignment) in assignments {
161 allocator.used_ports.insert(port, assignment.clone());
162 allocator
163 .app_slots
164 .insert((assignment.app_name, assignment.slot), port);
165 }
166
167 Ok(())
168 }
169
170 pub async fn get_port(&self, app_name: &str, slot: &str) -> Option<u16> {
171 self.allocator.lock().await.get_port(app_name, slot)
172 }
173
174 pub async fn get_app_name(&self, port: u16) -> Option<String> {
175 let allocator = self.allocator.lock().await;
176 allocator.used_ports.get(&port).map(|a| a.app_name.clone())
177 }
178}
179
180#[cfg(test)]
181mod tests {
182 use super::*;
183 use tempfile::TempDir;
184
185 #[tokio::test]
186 async fn test_port_allocation() {
187 let temp_dir = TempDir::new().unwrap();
188 let pm = PortManager::new(temp_dir.path().to_str().unwrap()).unwrap();
189
190 let port1 = pm.allocate("app1", "blue").await.unwrap();
191 let port2 = pm.allocate("app1", "green").await.unwrap();
192 let port3 = pm.allocate("app2", "blue").await.unwrap();
193
194 assert_ne!(port1, port2);
195 assert_ne!(port1, port3);
196 assert_ne!(port2, port3);
197
198 assert_eq!(pm.allocate("app1", "blue").await.unwrap(), port1);
199 }
200
201 #[tokio::test]
202 async fn test_port_release() {
203 let temp_dir = TempDir::new().unwrap();
204 let pm = PortManager::new(temp_dir.path().to_str().unwrap()).unwrap();
205
206 let port = pm.allocate("app1", "blue").await.unwrap();
207 pm.release("app1", "blue").await;
208
209 let new_port = pm.allocate("app1", "blue").await.unwrap();
210 assert_eq!(port, new_port);
211 }
212}