1use crate::prelude::*;
2use compact_str::ToCompactString;
3use serde::{Deserialize, Serialize};
4use sqlx::{Row, postgres::PgRow};
5use std::{
6 collections::{BTreeMap, HashMap},
7 sync::LazyLock,
8};
9use utoipa::ToSchema;
10
11pub const CREATE_MAX_PORTS: usize = 65535;
12
13fn validate_filter_criteria(filter: &AllocationFilter, _context: &()) -> Result<(), garde::Error> {
14 if filter.is_empty() {
15 return Err(garde::Error::new(
16 "filter must narrow the selection by at least one criterion, use the all selector to target every allocation",
17 ));
18 }
19
20 Ok(())
21}
22
23#[derive(ToSchema, garde::Validate, Deserialize, Serialize, Default, Clone)]
24pub struct AllocationFilter {
25 #[garde(length(chars, min = 1, max = 128))]
26 #[schema(min_length = 1, max_length = 128)]
27 #[serde(
28 default,
29 deserialize_with = "crate::deserialize::deserialize_string_option"
30 )]
31 pub search: Option<compact_str::CompactString>,
32
33 #[garde(skip)]
34 #[schema(value_type = Option<String>)]
35 #[serde(default)]
36 pub ip: Option<std::net::IpAddr>,
37 #[garde(inner(range(min = 1, max = 65535)))]
38 #[schema(minimum = 1, maximum = 65535)]
39 #[serde(default)]
40 pub port_from: Option<i32>,
41 #[garde(inner(range(min = 1, max = 65535)))]
42 #[schema(minimum = 1, maximum = 65535)]
43 #[serde(default)]
44 pub port_to: Option<i32>,
45 #[garde(skip)]
46 #[serde(default)]
47 pub assigned: Option<bool>,
48}
49
50impl AllocationFilter {
51 #[inline]
52 pub fn is_empty(&self) -> bool {
53 self.search.is_none()
54 && self.ip.is_none()
55 && self.port_from.is_none()
56 && self.port_to.is_none()
57 && self.assigned.is_none()
58 }
59
60 #[inline]
61 pub fn bind<'q>(
62 &'q self,
63 query: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>,
64 ) -> sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments> {
65 query
66 .bind(self.search.as_deref())
67 .bind(self.ip.map(sqlx::types::ipnetwork::IpNetwork::from))
68 .bind(self.port_from)
69 .bind(self.port_to)
70 .bind(self.assigned)
71 }
72}
73
74#[derive(ToSchema, garde::Validate, Deserialize, Serialize)]
75#[serde(rename_all = "snake_case", tag = "type")]
76#[non_exhaustive]
77pub enum AllocationSelector {
78 Uuids {
79 #[garde(skip)]
80 uuids: Vec<uuid::Uuid>,
81 },
82 All,
83 Filter {
84 #[garde(dive, custom(validate_filter_criteria))]
85 filter: AllocationFilter,
86 },
87}
88
89impl AllocationSelector {
90 #[inline]
91 pub fn uuids(&self) -> Option<&[uuid::Uuid]> {
92 match self {
93 Self::Uuids { uuids } => Some(uuids),
94 _ => None,
95 }
96 }
97
98 #[inline]
99 pub fn filter(&self) -> Option<&AllocationFilter> {
100 match self {
101 Self::Filter { filter } => Some(filter),
102 _ => None,
103 }
104 }
105}
106
107#[derive(Serialize, Deserialize, Clone)]
108pub struct NodeAllocation {
109 pub uuid: uuid::Uuid,
110 pub server: Option<Fetchable<super::server::Server>>,
111
112 pub ip: sqlx::types::ipnetwork::IpNetwork,
113 pub ip_alias: Option<compact_str::CompactString>,
114 pub port: i32,
115
116 pub created: chrono::NaiveDateTime,
117
118 extension_data: super::ModelExtensionData,
119}
120
121impl BaseModel for NodeAllocation {
122 const NAME: &'static str = "node_allocation";
123
124 fn get_extension_list() -> &'static super::ModelExtensionList {
125 static EXTENSIONS: LazyLock<super::ModelExtensionList> =
126 LazyLock::new(|| parking_lot::RwLock::new(Vec::new()));
127
128 &EXTENSIONS
129 }
130
131 fn get_extension_data(&self) -> &super::ModelExtensionData {
132 &self.extension_data
133 }
134
135 #[inline]
136 fn base_columns(prefix: Option<&str>) -> BTreeMap<&'static str, compact_str::CompactString> {
137 let prefix = prefix.unwrap_or_default();
138
139 BTreeMap::from([
140 (
141 "node_allocations.uuid",
142 compact_str::format_compact!("{prefix}uuid"),
143 ),
144 (
145 "node_allocations.ip",
146 compact_str::format_compact!("{prefix}ip"),
147 ),
148 (
149 "node_allocations.ip_alias",
150 compact_str::format_compact!("{prefix}ip_alias"),
151 ),
152 (
153 "node_allocations.port",
154 compact_str::format_compact!("{prefix}port"),
155 ),
156 (
157 "node_allocations.created",
158 compact_str::format_compact!("{prefix}created"),
159 ),
160 ])
161 }
162
163 #[inline]
164 fn map(prefix: Option<&str>, row: &PgRow) -> Result<Self, crate::database::DatabaseError> {
165 let prefix = prefix.unwrap_or_default();
166
167 Ok(Self {
168 uuid: row.try_get(compact_str::format_compact!("{prefix}uuid").as_str())?,
169 server: if let Ok(server_uuid) = row.try_get::<uuid::Uuid, _>("server_uuid") {
170 Some(super::server::Server::get_fetchable(server_uuid))
171 } else {
172 None
173 },
174 ip: row.try_get(compact_str::format_compact!("{prefix}ip").as_str())?,
175 ip_alias: row.try_get(compact_str::format_compact!("{prefix}ip_alias").as_str())?,
176 port: row.try_get(compact_str::format_compact!("{prefix}port").as_str())?,
177 created: row.try_get(compact_str::format_compact!("{prefix}created").as_str())?,
178 extension_data: Self::map_extensions(prefix, row)?,
179 })
180 }
181}
182
183impl NodeAllocation {
184 pub async fn create(
185 database: &crate::database::Database,
186 node_uuid: uuid::Uuid,
187 ip: &sqlx::types::ipnetwork::IpNetwork,
188 ip_alias: Option<&str>,
189 port: i32,
190 ) -> Result<(), crate::database::DatabaseError> {
191 sqlx::query(
192 r#"
193 INSERT INTO node_allocations (node_uuid, ip, ip_alias, port)
194 VALUES ($1, $2, $3, $4)
195 "#,
196 )
197 .bind(node_uuid)
198 .bind(ip)
199 .bind(ip_alias)
200 .bind(port)
201 .execute(database.write())
202 .await?;
203
204 Ok(())
205 }
206
207 pub async fn create_many(
208 database: &crate::database::Database,
209 node_uuid: uuid::Uuid,
210 ip: &sqlx::types::ipnetwork::IpNetwork,
211 ip_alias: Option<&str>,
212 ports: &[i32],
213 ) -> Result<u64, crate::database::DatabaseError> {
214 let created = sqlx::query(
215 r#"
216 INSERT INTO node_allocations (node_uuid, ip, ip_alias, port)
217 SELECT DISTINCT $1, $2, $3, port
218 FROM UNNEST($4::int[]) AS port
219 ON CONFLICT (node_uuid, host(ip), port) DO NOTHING
220 "#,
221 )
222 .bind(node_uuid)
223 .bind(ip)
224 .bind(ip_alias)
225 .bind(ports)
226 .execute(database.write())
227 .await?
228 .rows_affected();
229
230 Ok(created)
231 }
232
233 pub async fn used_by_node(
234 state: &crate::State,
235 node: &super::node::Node,
236 ips: &[std::net::IpAddr],
237 ) -> Result<Vec<uuid::Uuid>, anyhow::Error> {
238 if ips.is_empty() {
239 return Ok(Vec::new());
240 }
241
242 let (used_ips, used_ports): (Vec<_>, Vec<_>) = node
243 .used_ports(state, ips)
244 .await?
245 .into_iter()
246 .flat_map(|(ip, ports)| {
247 ports
248 .into_iter()
249 .map(move |port| (sqlx::types::ipnetwork::IpNetwork::from(ip), port as i32))
250 })
251 .unzip();
252
253 if used_ips.is_empty() {
254 return Ok(Vec::new());
255 }
256
257 Ok(sqlx::query_scalar(
258 r#"
259 SELECT node_allocations.uuid
260 FROM node_allocations
261 JOIN UNNEST($2::inet[], $3::int[]) AS used(ip, port)
262 ON used.ip = node_allocations.ip AND used.port = node_allocations.port
263 WHERE node_allocations.node_uuid = $1
264 "#,
265 )
266 .bind(node.uuid)
267 .bind(&used_ips)
268 .bind(&used_ports)
269 .fetch_all(state.database.read())
270 .await?)
271 }
272
273 pub async fn used_by_node_any_ip(
274 state: &crate::State,
275 node: &super::node::Node,
276 ) -> Result<Vec<uuid::Uuid>, anyhow::Error> {
277 let ips: Vec<sqlx::types::ipnetwork::IpNetwork> = sqlx::query_scalar(
278 r#"
279 SELECT DISTINCT node_allocations.ip
280 FROM node_allocations
281 WHERE node_allocations.node_uuid = $1
282 "#,
283 )
284 .bind(node.uuid)
285 .fetch_all(state.database.read())
286 .await?;
287
288 Self::used_by_node(
289 state,
290 node,
291 &ips.iter().map(|ip| ip.ip()).collect::<Vec<_>>(),
292 )
293 .await
294 }
295
296 pub async fn get_random(
297 database: &crate::database::Database,
298 node_uuid: uuid::Uuid,
299 start_port: u16,
300 end_port: u16,
301 amount: i64,
302 exclude: &[uuid::Uuid],
303 ) -> Result<Vec<uuid::Uuid>, crate::database::DatabaseError> {
304 let rows = sqlx::query(
305 r#"
306 WITH eligible_ips AS (
307 SELECT node_allocations.ip
308 FROM node_allocations
309 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
310 WHERE
311 node_allocations.node_uuid = $1
312 AND node_allocations.port BETWEEN $2 AND $3
313 AND server_allocations.uuid IS NULL
314 AND NOT (node_allocations.uuid = ANY($5))
315 GROUP BY node_allocations.ip
316 HAVING COUNT(*) >= $4
317 ),
318 random_ip AS (
319 SELECT ip FROM eligible_ips ORDER BY RANDOM() LIMIT 1
320 )
321 SELECT node_allocations.uuid
322 FROM node_allocations
323 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
324 WHERE
325 node_allocations.node_uuid = $1
326 AND node_allocations.port BETWEEN $2 AND $3
327 AND server_allocations.uuid IS NULL
328 AND NOT (node_allocations.uuid = ANY($5))
329 AND node_allocations.ip = (SELECT ip FROM random_ip)
330 ORDER BY RANDOM()
331 LIMIT $4
332 "#,
333 )
334 .bind(node_uuid)
335 .bind(start_port as i32)
336 .bind(end_port as i32)
337 .bind(amount)
338 .bind(exclude)
339 .fetch_all(database.write())
340 .await?;
341
342 if rows.len() != amount as usize {
343 return Err(anyhow::anyhow!("only found {} available allocations", rows.len()).into());
344 }
345
346 Ok(rows
347 .into_iter()
348 .map(|row| row.get::<uuid::Uuid, _>("uuid"))
349 .collect())
350 }
351
352 pub async fn get_random_ip(
353 database: &crate::database::Database,
354 node_uuid: uuid::Uuid,
355 ip: &sqlx::types::ipnetwork::IpNetwork,
356 start_port: u16,
357 end_port: u16,
358 amount: i64,
359 exclude: &[uuid::Uuid],
360 ) -> Result<Vec<uuid::Uuid>, crate::database::DatabaseError> {
361 let rows = sqlx::query(
362 r#"
363 SELECT node_allocations.uuid
364 FROM node_allocations
365 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
366 WHERE
367 node_allocations.node_uuid = $1
368 AND node_allocations.ip = $2
369 AND node_allocations.port BETWEEN $3 AND $4
370 AND server_allocations.uuid IS NULL
371 AND NOT (node_allocations.uuid = ANY($6))
372 ORDER BY RANDOM()
373 LIMIT $5
374 "#,
375 )
376 .bind(node_uuid)
377 .bind(ip)
378 .bind(start_port as i32)
379 .bind(end_port as i32)
380 .bind(amount)
381 .bind(exclude)
382 .fetch_all(database.write())
383 .await?;
384
385 if rows.len() != amount as usize {
386 return Err(anyhow::anyhow!(
387 "only found {} available allocations on this IP",
388 rows.len()
389 )
390 .into());
391 }
392
393 Ok(rows
394 .into_iter()
395 .map(|row| row.get::<uuid::Uuid, _>("uuid"))
396 .collect())
397 }
398
399 pub async fn get_random_dedicated(
400 database: &crate::database::Database,
401 node_uuid: uuid::Uuid,
402 start_port: u16,
403 end_port: u16,
404 amount: i64,
405 exclude: &[uuid::Uuid],
406 ) -> Result<Vec<uuid::Uuid>, crate::database::DatabaseError> {
407 let rows = sqlx::query(
408 r#"
409 WITH eligible_ips AS (
410 SELECT node_allocations.ip
411 FROM node_allocations
412 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
413 WHERE node_allocations.node_uuid = $1
414 GROUP BY node_allocations.ip
415 HAVING
416 COUNT(server_allocations.uuid) = 0
417 AND COUNT(*) FILTER (WHERE node_allocations.uuid = ANY($5)) = 0
418 AND SUM(CASE WHEN node_allocations.port BETWEEN $2 AND $3 THEN 1 ELSE 0 END) >= $4
419 ),
420 random_ip AS (
421 SELECT ip FROM eligible_ips ORDER BY RANDOM() LIMIT 1
422 )
423 SELECT node_allocations.uuid
424 FROM node_allocations
425 WHERE
426 node_allocations.node_uuid = $1
427 AND node_allocations.port BETWEEN $2 AND $3
428 AND node_allocations.ip = (SELECT ip FROM random_ip)
429 ORDER BY RANDOM()
430 LIMIT $4
431 "#,
432 )
433 .bind(node_uuid)
434 .bind(start_port as i32)
435 .bind(end_port as i32)
436 .bind(amount)
437 .bind(exclude)
438 .fetch_all(database.write())
439 .await?;
440
441 if rows.len() != amount as usize {
442 return Err(anyhow::anyhow!(
443 "only found {} available dedicated allocations",
444 rows.len()
445 )
446 .into());
447 }
448
449 Ok(rows
450 .into_iter()
451 .map(|row| row.get::<uuid::Uuid, _>("uuid"))
452 .collect())
453 }
454
455 pub async fn get_preserved(
456 database: &crate::database::Database,
457 node_uuid: uuid::Uuid,
458 ports: &[i32],
459 exclude: &[uuid::Uuid],
460 ) -> Result<
461 Option<(sqlx::types::ipnetwork::IpNetwork, Vec<(uuid::Uuid, i32)>)>,
462 crate::database::DatabaseError,
463 > {
464 let rows = sqlx::query(
465 r#"
466 WITH best_ip AS (
467 SELECT node_allocations.ip
468 FROM node_allocations
469 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
470 WHERE
471 node_allocations.node_uuid = $1
472 AND node_allocations.port = ANY($2)
473 AND server_allocations.uuid IS NULL
474 AND NOT (node_allocations.uuid = ANY($3))
475 GROUP BY node_allocations.ip
476 ORDER BY COUNT(DISTINCT node_allocations.port) DESC, RANDOM()
477 LIMIT 1
478 )
479 SELECT node_allocations.uuid, node_allocations.ip, node_allocations.port
480 FROM node_allocations
481 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
482 WHERE
483 node_allocations.node_uuid = $1
484 AND node_allocations.port = ANY($2)
485 AND server_allocations.uuid IS NULL
486 AND NOT (node_allocations.uuid = ANY($3))
487 AND node_allocations.ip = (SELECT ip FROM best_ip)
488 "#,
489 )
490 .bind(node_uuid)
491 .bind(ports)
492 .bind(exclude)
493 .fetch_all(database.write())
494 .await?;
495
496 let Some(first) = rows.first() else {
497 return Ok(None);
498 };
499
500 let ip = first.get::<sqlx::types::ipnetwork::IpNetwork, _>("ip");
501
502 Ok(Some((
503 ip,
504 rows.into_iter()
505 .map(|row| (row.get::<uuid::Uuid, _>("uuid"), row.get::<i32, _>("port")))
506 .collect(),
507 )))
508 }
509
510 pub async fn by_node_uuid_ip_port_unused(
511 database: &crate::database::Database,
512 node_uuid: uuid::Uuid,
513 ip: &sqlx::types::ipnetwork::IpNetwork,
514 port: i32,
515 exclude: &[uuid::Uuid],
516 ) -> Result<Option<Self>, crate::database::DatabaseError> {
517 let row = sqlx::query(sqlx::AssertSqlSafe(format!(
518 r#"
519 SELECT {}
520 FROM node_allocations
521 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
522 WHERE node_allocations.node_uuid = $1 AND node_allocations.ip = $2 AND node_allocations.port = $3 AND server_allocations.uuid IS NULL
523 AND NOT (node_allocations.uuid = ANY($4))
524 "#,
525 Self::columns_sql(None)
526 )))
527 .bind(node_uuid)
528 .bind(ip)
529 .bind(port)
530 .bind(exclude)
531 .fetch_optional(database.read())
532 .await?;
533
534 row.try_map(|row| Self::map(None, &row))
535 }
536
537 pub async fn get_from_deployment<'a>(
538 state: &crate::State,
539 deployment: &'a super::egg_configuration::EggConfigAllocationsDeployment,
540 node: &super::node::Node,
541 variables: &mut HashMap<&'a str, compact_str::CompactString>,
542 ) -> Result<(Option<uuid::Uuid>, Vec<uuid::Uuid>), crate::database::DatabaseError> {
543 let database = &state.database;
544 let node_uuid = node.uuid;
545
546 let used = Self::used_by_node_any_ip(state, node).await?;
547
548 let mut primary = None;
549 let mut additional = Vec::new();
550
551 const MAX_ITER: usize = 100;
552
553 macro_rules! is_unused {
554 ($uuid:expr) => {
555 primary.as_ref().map_or(true, |p| p.uuid != $uuid) && !additional.contains(&$uuid)
556 };
557 }
558
559 macro_rules! get_random {
560 ($start_port:expr, $end_port:expr) => {{
561 let mut exclude = used.clone();
562 exclude.extend_from_slice(&additional);
563 if let Some(primary) = &primary {
564 exclude.push(primary.uuid);
565 Self::get_random_ip(
566 database,
567 node_uuid,
568 &primary.ip,
569 $start_port,
570 $end_port,
571 1,
572 &exclude,
573 )
574 .await?
575 } else {
576 Self::get_random(database, node_uuid, $start_port, $end_port, 1, &exclude)
577 .await?
578 }
579 }};
580 }
581
582 let mut success = false;
583
584 'primary: for i in 0..MAX_ITER {
585 if i != 0 && deployment.primary.is_none() {
586 break;
587 }
588
589 additional.clear();
590 variables.clear();
591
592 if let Some(primary_allocation) = &deployment.primary {
593 let random = if deployment.dedicated {
594 Self::get_random_dedicated(
595 database,
596 node_uuid,
597 primary_allocation.start_port,
598 primary_allocation.end_port,
599 1,
600 &used,
601 )
602 .await?
603 } else {
604 Self::get_random(
605 database,
606 node_uuid,
607 primary_allocation.start_port,
608 primary_allocation.end_port,
609 1,
610 &used,
611 )
612 .await?
613 };
614
615 let Some(allocation) = random.into_iter().next() else {
616 return Err(anyhow::anyhow!("no available primary allocation found").into());
617 };
618 let allocation =
619 match Self::by_node_uuid_uuid(database, node_uuid, allocation).await? {
620 Some(allocation) => allocation,
621 None => {
622 return Err(
623 anyhow::anyhow!("allocated primary allocation not found").into()
624 );
625 }
626 };
627
628 if let Some(variable_name) = &primary_allocation.assign_to_variable {
629 variables.insert(variable_name, allocation.port.to_compact_string());
630 }
631
632 primary = Some(allocation);
633 }
634
635 for additional_allocation in &deployment.additional {
636 match additional_allocation.mode {
637 super::egg_configuration::EggConfigAllocationDeploymentAdditionalAllocationMode::Random => {
638 let random = get_random!(1, u16::MAX);
639
640 let Some(allocation) = random.into_iter().next() else {
641 return Err(anyhow::anyhow!("no available additional allocation found").into());
642 };
643 additional.push(allocation);
644
645 if let Some(variable_name) = &additional_allocation.assign_to_variable {
646 let allocation = match Self::by_node_uuid_uuid(database, node_uuid, allocation).await? {
647 Some(allocation) => allocation,
648 None => {
649 return Err(
650 anyhow::anyhow!("allocated additional allocation not found").into()
651 );
652 }
653 };
654
655 variables.insert(variable_name, allocation.port.to_compact_string());
656 }
657 },
658 super::egg_configuration::EggConfigAllocationDeploymentAdditionalAllocationMode::Range {
659 start_port,
660 end_port,
661 } => {
662 let random = get_random!(start_port, end_port);
663
664 let Some(allocation) = random.into_iter().next() else {
665 return Err(anyhow::anyhow!("no available additional allocation found").into());
666 };
667 additional.push(allocation);
668
669 if let Some(variable_name) = &additional_allocation.assign_to_variable {
670 let allocation = match Self::by_node_uuid_uuid(database, node_uuid, allocation).await? {
671 Some(allocation) => allocation,
672 None => {
673 return Err(
674 anyhow::anyhow!("allocated additional allocation not found").into()
675 );
676 }
677 };
678
679 variables.insert(variable_name, allocation.port.to_compact_string());
680 }
681 }
682 super::egg_configuration::EggConfigAllocationDeploymentAdditionalAllocationMode::AddPrimary { value } => {
683 let primary = match &primary {
684 Some(primary) => primary,
685 None => {
686 return Err(anyhow::anyhow!("primary allocation is required for `add_primary` mode").into());
687 }
688 };
689
690 let allocation_port = primary.port + value as i32;
691
692 let allocation = match Self::by_node_uuid_ip_port_unused(database, node_uuid, &primary.ip, allocation_port, &used).await? {
693 Some(allocation) => allocation,
694 None => continue 'primary,
695 };
696 if !is_unused!(allocation.uuid) {
697 return Err(anyhow::anyhow!("allocated additional allocation is already in use").into());
698 }
699 additional.push(allocation.uuid);
700
701 if let Some(variable_name) = &additional_allocation.assign_to_variable {
702 variables.insert(variable_name, allocation.port.to_compact_string());
703 }
704 }
705 super::egg_configuration::EggConfigAllocationDeploymentAdditionalAllocationMode::SubtractPrimary { value } => {
706 let primary = match &primary {
707 Some(primary) => primary,
708 None => {
709 return Err(anyhow::anyhow!("primary allocation is required for `subtract_primary` mode").into());
710 }
711 };
712
713 let allocation_port = primary.port - value as i32;
714
715 let allocation = match Self::by_node_uuid_ip_port_unused(database, node_uuid, &primary.ip, allocation_port, &used).await? {
716 Some(allocation) => allocation,
717 None => continue 'primary,
718 };
719 if !is_unused!(allocation.uuid) {
720 return Err(anyhow::anyhow!("allocated additional allocation is already in use").into());
721 }
722 additional.push(allocation.uuid);
723
724 if let Some(variable_name) = &additional_allocation.assign_to_variable {
725 variables.insert(variable_name, allocation.port.to_compact_string());
726 }
727 }
728 super::egg_configuration::EggConfigAllocationDeploymentAdditionalAllocationMode::MultiplyPrimary { value } => {
729 let primary = match &primary {
730 Some(primary) => primary,
731 None => {
732 return Err(anyhow::anyhow!("primary allocation is required for `multiply_primary` mode").into());
733 }
734 };
735
736 let allocation_port = (primary.port as f64 * value) as i32;
737
738 let allocation = match Self::by_node_uuid_ip_port_unused(database, node_uuid, &primary.ip, allocation_port, &used).await? {
739 Some(allocation) => allocation,
740 None => continue 'primary,
741 };
742 if !is_unused!(allocation.uuid) {
743 return Err(anyhow::anyhow!("allocated additional allocation is already in use").into());
744 }
745 additional.push(allocation.uuid);
746
747 if let Some(variable_name) = &additional_allocation.assign_to_variable {
748 variables.insert(variable_name, allocation.port.to_compact_string());
749 }
750 }
751 super::egg_configuration::EggConfigAllocationDeploymentAdditionalAllocationMode::DividePrimary { value } => {
752 let primary = match &primary {
753 Some(primary) => primary,
754 None => {
755 return Err(anyhow::anyhow!("primary allocation is required for `divide_primary` mode").into());
756 }
757 };
758
759 let allocation_port = (primary.port as f64 / value) as i32;
760
761 let allocation = match Self::by_node_uuid_ip_port_unused(database, node_uuid, &primary.ip, allocation_port, &used).await? {
762 Some(allocation) => allocation,
763 None => continue 'primary,
764 };
765 if !is_unused!(allocation.uuid) {
766 return Err(anyhow::anyhow!("allocated additional allocation is already in use").into());
767 }
768 additional.push(allocation.uuid);
769
770 if let Some(variable_name) = &additional_allocation.assign_to_variable {
771 variables.insert(variable_name, allocation.port.to_compact_string());
772 }
773 }
774 }
775 }
776
777 success = true;
778 break;
779 }
780
781 if !success {
782 return Err(anyhow::anyhow!(
783 "could not satisfy all additional allocation rules after {MAX_ITER} attempts"
784 )
785 .into());
786 }
787
788 Ok((primary.map(|p| p.uuid), additional))
789 }
790
791 pub async fn by_node_uuid_uuid(
792 database: &crate::database::Database,
793 node_uuid: uuid::Uuid,
794 uuid: uuid::Uuid,
795 ) -> Result<Option<Self>, crate::database::DatabaseError> {
796 let row = sqlx::query(sqlx::AssertSqlSafe(format!(
797 r#"
798 SELECT {}
799 FROM node_allocations
800 WHERE node_allocations.node_uuid = $1 AND node_allocations.uuid = $2
801 "#,
802 Self::columns_sql(None)
803 )))
804 .bind(node_uuid)
805 .bind(uuid)
806 .fetch_optional(database.read())
807 .await?;
808
809 row.try_map(|row| Self::map(None, &row))
810 }
811
812 pub async fn available_by_node_uuid_with_pagination(
813 database: &crate::database::Database,
814 node_uuid: uuid::Uuid,
815 page: i64,
816 per_page: i64,
817 search: Option<&str>,
818 ) -> Result<super::Pagination<Self>, crate::database::DatabaseError> {
819 let offset = (page - 1) * per_page;
820
821 let rows = sqlx::query(sqlx::AssertSqlSafe(format!(
822 r#"
823 SELECT {}, COUNT(*) OVER() AS total_count
824 FROM node_allocations
825 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
826 WHERE node_allocations.node_uuid = $1 AND server_allocations.uuid IS NULL
827 AND (
828 $2 IS NULL
829 OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%'
830 OR (node_allocations.ip_alias IS NOT NULL AND node_allocations.ip_alias || ':' || node_allocations.port ILIKE '%' || $2 || '%')
831 )
832 ORDER BY node_allocations.ip, node_allocations.port
833 LIMIT $3 OFFSET $4
834 "#,
835 Self::columns_sql(None)
836 )))
837 .bind(node_uuid)
838 .bind(search)
839 .bind(per_page)
840 .bind(offset)
841 .fetch_all(database.read())
842 .await?;
843
844 Ok(super::Pagination {
845 total: rows
846 .first()
847 .map_or(Ok(0), |row| row.try_get("total_count"))?,
848 per_page,
849 page,
850 data: rows
851 .into_iter()
852 .map(|row| Self::map(None, &row))
853 .try_collect_vec()?,
854 })
855 }
856
857 pub async fn by_node_uuid_with_pagination(
858 database: &crate::database::Database,
859 node_uuid: uuid::Uuid,
860 filter: &AllocationFilter,
861 page: i64,
862 per_page: i64,
863 ) -> Result<super::Pagination<Self>, crate::database::DatabaseError> {
864 let offset = (page - 1) * per_page;
865
866 let query = sqlx::query(sqlx::AssertSqlSafe(format!(
867 r#"
868 SELECT {}, server_allocations.server_uuid, COUNT(*) OVER() AS total_count
869 FROM node_allocations
870 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
871 WHERE node_allocations.node_uuid = $1
872 AND ($2::text IS NULL
873 OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%'
874 OR (node_allocations.ip_alias IS NOT NULL AND node_allocations.ip_alias || ':' || node_allocations.port ILIKE '%' || $2 || '%')
875 OR server_allocations.notes ILIKE '%' || $2 || '%')
876 AND ($3::inet IS NULL OR host(node_allocations.ip) = host($3))
877 AND ($4::int IS NULL OR node_allocations.port >= $4)
878 AND ($5::int IS NULL OR node_allocations.port <= $5)
879 AND ($6::bool IS NULL OR (server_allocations.uuid IS NOT NULL) = $6)
880 ORDER BY node_allocations.ip, node_allocations.port
881 LIMIT $7 OFFSET $8
882 "#,
883 Self::columns_sql(None)
884 )))
885 .bind(node_uuid);
886
887 let rows = filter
888 .bind(query)
889 .bind(per_page)
890 .bind(offset)
891 .fetch_all(database.read())
892 .await?;
893
894 Ok(super::Pagination {
895 total: rows
896 .first()
897 .map_or(Ok(0), |row| row.try_get("total_count"))?,
898 per_page,
899 page,
900 data: rows
901 .into_iter()
902 .map(|row| Self::map(None, &row))
903 .try_collect_vec()?,
904 })
905 }
906
907 pub async fn distinct_ips_by_node_uuid(
908 database: &crate::database::Database,
909 node_uuid: uuid::Uuid,
910 ) -> Result<Vec<compact_str::CompactString>, crate::database::DatabaseError> {
911 let rows = sqlx::query(
912 r#"
913 SELECT host(node_allocations.ip) AS ip
914 FROM node_allocations
915 WHERE node_allocations.node_uuid = $1
916 GROUP BY host(node_allocations.ip)
917 ORDER BY MIN(node_allocations.ip)
918 "#,
919 )
920 .bind(node_uuid)
921 .fetch_all(database.read())
922 .await?;
923
924 rows.into_iter()
925 .map(|row| Ok(row.try_get::<String, _>("ip")?.into()))
926 .try_collect_vec()
927 }
928
929 pub async fn delete_by_selector(
930 database: &crate::database::Database,
931 node_uuid: uuid::Uuid,
932 selector: &AllocationSelector,
933 force: bool,
934 ) -> Result<(i64, i64), crate::database::DatabaseError> {
935 let filter = selector.filter().cloned().unwrap_or_default();
936
937 let query = sqlx::query(
938 r#"
939 WITH matched AS (
940 SELECT node_allocations.uuid, server_allocations.uuid IS NOT NULL AS in_use
941 FROM node_allocations
942 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
943 WHERE node_allocations.node_uuid = $1
944 AND ($7::uuid[] IS NULL OR node_allocations.uuid = ANY($7))
945 AND ($2::text IS NULL
946 OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%'
947 OR (node_allocations.ip_alias IS NOT NULL AND node_allocations.ip_alias || ':' || node_allocations.port ILIKE '%' || $2 || '%')
948 OR server_allocations.notes ILIKE '%' || $2 || '%')
949 AND ($3::inet IS NULL OR host(node_allocations.ip) = host($3))
950 AND ($4::int IS NULL OR node_allocations.port >= $4)
951 AND ($5::int IS NULL OR node_allocations.port <= $5)
952 AND ($6::bool IS NULL OR (server_allocations.uuid IS NOT NULL) = $6)
953 ), deleted AS (
954 DELETE FROM node_allocations
955 WHERE node_allocations.uuid IN (SELECT uuid FROM matched WHERE $8 OR NOT in_use)
956 RETURNING 1
957 )
958 SELECT (SELECT COUNT(*) FROM matched) AS matched_count,
959 (SELECT COUNT(*) FROM deleted) AS deleted_count
960 "#,
961 )
962 .bind(node_uuid);
963
964 let row = filter
965 .bind(query)
966 .bind(selector.uuids())
967 .bind(force)
968 .fetch_one(database.write())
969 .await?;
970
971 Ok((row.try_get("matched_count")?, row.try_get("deleted_count")?))
972 }
973
974 pub async fn update_by_selector(
975 database: &crate::database::Database,
976 node_uuid: uuid::Uuid,
977 selector: &AllocationSelector,
978 ip: &sqlx::types::ipnetwork::IpNetwork,
979 ip_alias: Option<Option<&str>>,
980 ) -> Result<(i64, i64), crate::database::DatabaseError> {
981 let filter = selector.filter().cloned().unwrap_or_default();
982
983 let query = sqlx::query(
984 r#"
985 WITH matched AS (
986 SELECT node_allocations.uuid, node_allocations.port
987 FROM node_allocations
988 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
989 WHERE node_allocations.node_uuid = $1
990 AND ($7::uuid[] IS NULL OR node_allocations.uuid = ANY($7))
991 AND ($2::text IS NULL
992 OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%'
993 OR (node_allocations.ip_alias IS NOT NULL AND node_allocations.ip_alias || ':' || node_allocations.port ILIKE '%' || $2 || '%')
994 OR server_allocations.notes ILIKE '%' || $2 || '%')
995 AND ($3::inet IS NULL OR host(node_allocations.ip) = host($3))
996 AND ($4::int IS NULL OR node_allocations.port >= $4)
997 AND ($5::int IS NULL OR node_allocations.port <= $5)
998 AND ($6::bool IS NULL OR (server_allocations.uuid IS NOT NULL) = $6)
999 ), winners AS (
1000 SELECT DISTINCT ON (matched.port) matched.uuid, matched.port
1001 FROM matched
1002 ORDER BY matched.port, matched.uuid
1003 ), eligible AS (
1004 SELECT winners.uuid
1005 FROM winners
1006 WHERE NOT EXISTS (
1007 SELECT 1 FROM node_allocations existing
1008 WHERE existing.node_uuid = $1
1009 AND host(existing.ip) = host($8::inet)
1010 AND existing.port = winners.port
1011 AND existing.uuid <> winners.uuid
1012 )
1013 ), updated AS (
1014 UPDATE node_allocations
1015 SET ip = $8,
1016 ip_alias = CASE WHEN $9 THEN $10::varchar ELSE node_allocations.ip_alias END
1017 WHERE node_allocations.uuid IN (SELECT uuid FROM eligible)
1018 RETURNING 1
1019 )
1020 SELECT (SELECT COUNT(*) FROM matched) AS matched_count,
1021 (SELECT COUNT(*) FROM updated) AS updated_count
1022 "#,
1023 )
1024 .bind(node_uuid);
1025
1026 let row = filter
1027 .bind(query)
1028 .bind(selector.uuids())
1029 .bind(ip)
1030 .bind(ip_alias.is_some())
1031 .bind(ip_alias.flatten())
1032 .fetch_one(database.write())
1033 .await?;
1034
1035 Ok((row.try_get("matched_count")?, row.try_get("updated_count")?))
1036 }
1037}
1038
1039#[async_trait::async_trait]
1040impl IntoAdminApiObject for NodeAllocation {
1041 type AdminApiObject = AdminApiNodeAllocation;
1042 type ExtraArgs<'a> = &'a crate::storage::StorageUrlRetriever<'a>;
1043
1044 async fn into_admin_api_object<'a>(
1045 self,
1046 state: &crate::State,
1047 storage_url_retriever: Self::ExtraArgs<'a>,
1048 ) -> Result<Self::AdminApiObject, crate::database::DatabaseError> {
1049 let api_object = AdminApiNodeAllocation::init_hooks(&self, state).await?;
1050
1051 let server = match self.server {
1052 Some(fetchable) => Some(
1053 fetchable
1054 .fetch_cached(&state.database)
1055 .await?
1056 .into_admin_api_object(state, storage_url_retriever)
1057 .await?,
1058 ),
1059 None => None,
1060 };
1061
1062 let api_object = finish_extendible!(
1063 AdminApiNodeAllocation {
1064 uuid: self.uuid,
1065 server,
1066 ip: self.ip.ip().to_compact_string(),
1067 ip_alias: self.ip_alias,
1068 port: self.port,
1069 created: self.created.and_utc(),
1070 },
1071 api_object,
1072 state
1073 )?;
1074
1075 Ok(api_object)
1076 }
1077}
1078
1079#[schema_extension_derive::extendible]
1080#[init_args(NodeAllocation, crate::State)]
1081#[hook_args(crate::State)]
1082#[derive(ToSchema, Serialize)]
1083#[schema(title = "NodeAllocation")]
1084pub struct AdminApiNodeAllocation {
1085 pub uuid: uuid::Uuid,
1086 pub server: Option<super::server::AdminApiServer>,
1087
1088 pub ip: compact_str::CompactString,
1089 pub ip_alias: Option<compact_str::CompactString>,
1090 pub port: i32,
1091
1092 pub created: chrono::DateTime<chrono::Utc>,
1093}