chore: clean up ip types

rivet-gg · Apr 18, 2024 · 6f118b4 · 6f118b4
1 parent c2e0023
commit 6f118b4
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 10 deletions.
diff --git a/svc/pkg/cluster/worker/src/workers/server_provision.rs b/svc/pkg/cluster/worker/src/workers/server_provision.rs
@@ -1,3 +1,5 @@
+use std::net::{IpAddr, Ipv4Addr};
+
 use chirp_worker::prelude::*;
 use futures_util::FutureExt;
 use proto::backend::{self, cluster::PoolType, pkg::*};
@@ -81,7 +83,7 @@ async fn inner(
 		WHERE server_id = $1
 		",
 		server_id,
-		&vlan_ip,
+		IpAddr::V4(vlan_ip),
 	)
 	.await?;
 
@@ -107,7 +109,7 @@ async fn inner(
 					provider_datacenter_id: datacenter.provider_datacenter_id.clone(),
 					hardware: Some(hardware.clone()),
 					pool_type: ctx.pool_type,
-					vlan_ip: vlan_ip.clone(),
+					vlan_ip: vlan_ip.to_string(),
 					tags: ctx.tags.clone(),
 				})
 				.await;
@@ -200,7 +202,7 @@ async fn get_vlan_ip(
 	_tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
 	server_id: Uuid,
 	pool_type: backend::cluster::PoolType,
-) -> GlobalResult<String> {
+) -> GlobalResult<Ipv4Addr> {
 	// Find next available vlan index
 	let mut vlan_addr_range = match pool_type {
 		PoolType::Job => util::net::job::vlan_addr_range(),
@@ -242,7 +244,7 @@ async fn get_vlan_ip(
 
 	let vlan_ip = unwrap!(vlan_addr_range.nth(network_idx as usize));
 
-	Ok(vlan_ip.to_string())
+	Ok(vlan_ip)
 }
 
 // This function is used to destroy leftovers from a failed partial provision.

diff --git a/svc/pkg/mm-config/ops/version-prepare/src/prewarm_ats.rs b/svc/pkg/mm-config/ops/version-prepare/src/prewarm_ats.rs
@@ -1,7 +1,11 @@
+use std::{
+	collections::{HashMap, HashSet},
+	net::IpAddr,
+};
+
 use proto::backend::{self, pkg::*};
 use rivet_operation::prelude::*;
 use serde_json::json;
-use std::collections::{HashMap, HashSet};
 
 /// Tracks which resources should be prewarmed in the ATS cache.
 #[derive(Debug)]
@@ -15,7 +19,7 @@ pub struct PrewarmAtsContext {
 #[derive(sqlx::FromRow)]
 struct VlanIp {
 	datacenter_id: Uuid,
-	vlan_ip: String,
+	vlan_ip: IpAddr,
 }
 
 /// Requests resources from the ATS cache to make sure any subsequent requests will be faster.

diff --git a/svc/pkg/mm/worker/src/workers/lobby_create/mod.rs b/svc/pkg/mm/worker/src/workers/lobby_create/mod.rs
@@ -1,6 +1,4 @@
-use std::collections::hash_map::DefaultHasher;
-use std::hash::Hasher;
-use std::ops::Deref;
+use std::{collections::hash_map::DefaultHasher, hash::Hasher, net::IpAddr, ops::Deref};
 
 use chirp_worker::prelude::*;
 use proto::backend::{self, pkg::*};
@@ -867,7 +865,7 @@ async fn resolve_image_artifact_url(
 			// prewarm_ats.rs @ prewarm_ats_cache
 			// Get vlan ip from build id hash for consistent routing
 			let (ats_vlan_ip,) = sql_fetch_one!(
-				[ctx, (String,)]
+				[ctx, (IpAddr,)]
 				"
 				WITH sel AS (
 					-- Select candidate vlan ips