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