tatara_engine/domain/
port_allocator.rs

1//! Dynamic port allocation for tatara workloads.
2//!
3//! Assigns ports from a configurable range when tasks declare port 0.
4//! Tracks allocated ports per allocation for conflict detection.
5
6use serde::{Deserialize, Serialize};
7use std::collections::{HashMap, HashSet};
8use std::ops::RangeInclusive;
9use tokio::sync::RwLock;
10use uuid::Uuid;
11
12/// An allocated port for a task.
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct AllocatedPort {
15    /// Port label from the job spec (e.g., "http", "grpc").
16    pub label: String,
17
18    /// Static port from job spec (if explicitly specified).
19    pub static_port: Option<u16>,
20
21    /// The actual port assigned (dynamic or static).
22    pub assigned_port: u16,
23}
24
25/// Manages port allocation for a node.
26pub struct PortAllocator {
27    range: RangeInclusive<u16>,
28    /// (alloc_id, label) -> assigned port
29    allocated: RwLock<HashMap<(Uuid, String), u16>>,
30}
31
32impl PortAllocator {
33    /// Create a new allocator with the given port range.
34    pub fn new(start: u16, end: u16) -> Self {
35        Self {
36            range: start..=end,
37            allocated: RwLock::new(HashMap::new()),
38        }
39    }
40
41    /// Default range: 20000-32000.
42    pub fn default_range() -> Self {
43        Self::new(20000, 32000)
44    }
45
46    /// Allocate a port for a task. If `requested` is 0, assign dynamically.
47    /// If `requested` is non-zero, use it if available.
48    pub async fn allocate(
49        &self,
50        alloc_id: Uuid,
51        label: &str,
52        requested: u16,
53    ) -> Result<AllocatedPort, PortError> {
54        let mut allocated = self.allocated.write().await;
55
56        if requested != 0 {
57            // Static port — check for conflicts
58            let in_use = allocated.values().any(|&p| p == requested);
59            if in_use {
60                return Err(PortError::Conflict {
61                    port: requested,
62                    label: label.to_string(),
63                });
64            }
65            allocated.insert((alloc_id, label.to_string()), requested);
66            return Ok(AllocatedPort {
67                label: label.to_string(),
68                static_port: Some(requested),
69                assigned_port: requested,
70            });
71        }
72
73        // Dynamic allocation — find first available in range
74        let used: HashSet<u16> = allocated.values().copied().collect();
75        for port in self.range.clone() {
76            if !used.contains(&port) {
77                allocated.insert((alloc_id, label.to_string()), port);
78                return Ok(AllocatedPort {
79                    label: label.to_string(),
80                    static_port: None,
81                    assigned_port: port,
82                });
83            }
84        }
85
86        Err(PortError::Exhausted)
87    }
88
89    /// Release all ports for an allocation.
90    pub async fn release(&self, alloc_id: Uuid) {
91        let mut allocated = self.allocated.write().await;
92        allocated.retain(|(id, _), _| *id != alloc_id);
93    }
94
95    /// Check if a specific port is available on this node.
96    pub async fn is_available(&self, port: u16) -> bool {
97        let allocated = self.allocated.read().await;
98        !allocated.values().any(|&p| p == port)
99    }
100
101    /// Get all allocated ports for an allocation.
102    pub async fn get_ports(&self, alloc_id: Uuid) -> Vec<AllocatedPort> {
103        let allocated = self.allocated.read().await;
104        allocated
105            .iter()
106            .filter(|((id, _), _)| *id == alloc_id)
107            .map(|((_, label), &port)| AllocatedPort {
108                label: label.clone(),
109                static_port: None,
110                assigned_port: port,
111            })
112            .collect()
113    }
114}
115
116/// Port allocation errors.
117#[derive(Debug, thiserror::Error)]
118pub enum PortError {
119    #[error("port {port} already in use (requested for '{label}')")]
120    Conflict { port: u16, label: String },
121
122    #[error("all ports in range exhausted")]
123    Exhausted,
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[tokio::test]
131    async fn test_dynamic_allocation() {
132        let alloc = PortAllocator::new(30000, 30005);
133        let id = Uuid::new_v4();
134
135        let p1 = alloc.allocate(id, "http", 0).await.unwrap();
136        assert_eq!(p1.assigned_port, 30000);
137        assert!(p1.static_port.is_none());
138
139        let p2 = alloc.allocate(id, "grpc", 0).await.unwrap();
140        assert_eq!(p2.assigned_port, 30001);
141    }
142
143    #[tokio::test]
144    async fn test_static_allocation() {
145        let alloc = PortAllocator::new(30000, 30005);
146        let id = Uuid::new_v4();
147
148        let p = alloc.allocate(id, "http", 8080).await.unwrap();
149        assert_eq!(p.assigned_port, 8080);
150        assert_eq!(p.static_port, Some(8080));
151    }
152
153    #[tokio::test]
154    async fn test_conflict_detection() {
155        let alloc = PortAllocator::new(30000, 30005);
156        let id1 = Uuid::new_v4();
157        let id2 = Uuid::new_v4();
158
159        alloc.allocate(id1, "http", 8080).await.unwrap();
160        let err = alloc.allocate(id2, "http", 8080).await;
161        assert!(err.is_err());
162    }
163
164    #[tokio::test]
165    async fn test_release() {
166        let alloc = PortAllocator::new(30000, 30002);
167        let id = Uuid::new_v4();
168
169        alloc.allocate(id, "a", 0).await.unwrap();
170        alloc.allocate(id, "b", 0).await.unwrap();
171        alloc.allocate(id, "c", 0).await.unwrap();
172
173        // Range exhausted
174        let id2 = Uuid::new_v4();
175        assert!(alloc.allocate(id2, "d", 0).await.is_err());
176
177        // Release and try again
178        alloc.release(id).await;
179        assert!(alloc.allocate(id2, "d", 0).await.is_ok());
180    }
181}
tatara_engine/domain/port_allocator.rs

tatara_engine/domain/
port_allocator.rs
