fix(guard): add metrics

Author: MasterPtato

packages/common/logs/src/lib.rs

@@ -25,10 +25,12 @@ impl Logs {
 }
 
 impl Logs {
-	pub async fn start(self) -> Result<tokio::task::JoinHandle<()>> {
+	pub async fn start(mut self) -> Result<tokio::task::JoinHandle<()>> {
 		// Create logs dir if it does not exist
 		fs::create_dir_all(&self.path).await?;
 
+		self.rotate().await?;
+
 		Ok(tokio::spawn(self.run()))
 	}
 
@@ -112,10 +114,12 @@ impl Logs {
 }
 
 impl Logs {
-	pub fn start_sync(self) -> Result<std::thread::JoinHandle<()>> {
+	pub fn start_sync(mut self) -> Result<std::thread::JoinHandle<()>> {
 		// Create logs dir if it does not exist
 		std::fs::create_dir_all(&self.path)?;
 
+		self.rotate_sync()?;
+
 		Ok(std::thread::spawn(|| self.run_sync()))
 	}
 

packages/common/server-cli/src/commands/start.rs

@@ -54,7 +54,6 @@ impl Opts {
 			.and_then(|x| x.rivet.edge.as_ref())
 			.and_then(|x| x.redirect_logs_dir.as_ref())
 		{
-			std::fs::create_dir_all(logs_dir)?;
 			rivet_logs::Logs::new(logs_dir.clone(), LOGS_RETENTION)
 				.start()
 				.await?;

packages/core/services/cluster/src/workflows/server/install/install_scripts/mod.rs

@@ -227,6 +227,14 @@ pub async fn gen_initialize(
 			)?);
 			script.push(components::rivet::guard::fetch_tls(server_token)?);
 			script.push(components::rivet::guard::configure(config)?);
+
+			prometheus_targets.insert(
+				"guard".into(),
+				components::vector::PrometheusTarget {
+					endpoint: "http://127.0.0.1:8091".into(),
+					scrape_interval: 15,
+				},
+			);
 		}
 	}
 

packages/edge/infra/guard/core/src/metrics.rs

@@ -1,29 +1,67 @@
 use lazy_static::lazy_static;
-use rivet_metrics::prometheus::*;
+use rivet_metrics::{prometheus::*, REGISTRY};
 
 lazy_static! {
-	pub static ref ACTOR_REQUEST_TOTAL: IntCounterVec = register_int_counter_vec!(
-		"actor_request_total",
+	pub static ref ACTOR_REQUEST_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"guard_actor_request_total",
 		"Total number of requests to actor",
-		&["actor_id", "server_id", "method", "path"]
+		&["actor_id", "server_id", "method", "path"],
+		*REGISTRY,
 	)
 	.unwrap();
-	pub static ref ACTOR_REQUEST_PENDING: IntGaugeVec = register_int_gauge_vec!(
-		"actor_request_pending",
+	pub static ref ACTOR_REQUEST_PENDING: IntGaugeVec = register_int_gauge_vec_with_registry!(
+		"guard_actor_request_pending",
 		"Number of pending requests to actor",
-		&["actor_id", "server_id", "method", "path"]
+		&["actor_id", "server_id", "method", "path"],
+		*REGISTRY,
 	)
 	.unwrap();
-	pub static ref ACTOR_REQUEST_DURATION: HistogramVec = register_histogram_vec!(
-		"actor_request_duration_seconds",
+	pub static ref ACTOR_REQUEST_DURATION: HistogramVec = register_histogram_vec_with_registry!(
+		"guard_actor_request_duration_seconds",
 		"Request duration in seconds",
-		&["actor_id", "server_id", "status"]
+		&["actor_id", "server_id", "status"],
+		*REGISTRY,
 	)
 	.unwrap();
-	pub static ref ACTOR_REQUEST_ERRORS: IntCounterVec = register_int_counter_vec!(
-		"actor_request_errors_total",
+	pub static ref ACTOR_REQUEST_ERRORS: IntCounterVec = register_int_counter_vec_with_registry!(
+		"guard_actor_request_errors_total",
 		"Total number of errors when proxying requests to actor",
-		&["actor_id", "server_id", "error_type"]
+		&["actor_id", "server_id", "error_type"],
+		*REGISTRY,
+	)
+	.unwrap();
+	pub static ref ROUTE_CACHE_SIZE: IntGauge = register_int_gauge_with_registry!(
+		"guard_route_cache_size",
+		"Number of entries in the route cache",
+		*REGISTRY,
+	).unwrap();
+	pub static ref RATE_LIMITERS_COUNT: IntGauge = register_int_gauge_with_registry!(
+		"guard_rate_limiters_count",
+		"Number of active rate limiters",
+		*REGISTRY,
+	).unwrap();
+	pub static ref IN_FLIGHT_COUNTERS_COUNT: IntGauge = register_int_gauge_with_registry!(
+		"guard_in_flight_counters_count",
+		"Number of active in-flight counters",
+		*REGISTRY,
+	)
+	.unwrap();
+	pub static ref TCP_CONNECTIONS_TOTAL: IntGauge = register_int_gauge_with_registry!(
+		"guard_tcp_connections_total",
+		"Total number of open TCP connections",
+		*REGISTRY,
+	)
+	.unwrap();
+	pub static ref REQUEST_TOTAL: IntCounter = register_int_counter_with_registry!(
+		"guard_request_total",
+		"Total number of requests processed",
+		*REGISTRY,
+	)
+	.unwrap();
+	pub static ref REQUEST_PENDING: IntGauge = register_int_gauge_with_registry!(
+		"guard_request_pending",
+		"Total number of active requests being processed",
+		*REGISTRY,
 	)
 	.unwrap();
 }

packages/edge/infra/guard/core/src/proxy_service.rs

@@ -175,13 +175,17 @@ impl RouteCache {
 	#[tracing::instrument(skip_all)]
 	async fn insert(&self, hostname: String, path: String, result: RouteConfig) {
 		self.cache.upsert_async((hostname, path), result).await;
+
+		metrics::ROUTE_CACHE_SIZE.set(self.cache.len() as i64);
 	}
 
 	#[tracing::instrument(skip_all)]
 	async fn purge(&self, hostname: &str, path: &str) {
 		self.cache
 			.remove_async(&(hostname.to_owned(), path.to_owned()))
 			.await;
+
+		metrics::ROUTE_CACHE_SIZE.set(self.cache.len() as i64);
 	}
 }
 
@@ -253,8 +257,8 @@ pub struct ProxyState {
 	routing_fn: RoutingFn,
 	middleware_fn: MiddlewareFn,
 	route_cache: RouteCache,
-	rate_limiters: SccHashMap<Uuid, RateLimiter>,
-	in_flight_counters: SccHashMap<Uuid, InFlightCounter>,
+	rate_limiters: SccHashMap<std::net::IpAddr, RateLimiter>,
+	in_flight_counters: SccHashMap<std::net::IpAddr, InFlightCounter>,
 	port_type: PortType,
 }
 
@@ -447,79 +451,119 @@ impl ProxyState {
 	}
 
 	#[tracing::instrument(skip_all)]
-	async fn check_rate_limit(&self, actor_id: &Option<Uuid>) -> GlobalResult<bool> {
-		match actor_id {
-			Some(id) => {
-				let middleware_config = self.get_middleware_config(id).await?;
-
-				let entry = self
-					.rate_limiters
-					.entry_async(*id)
-					.instrument(tracing::info_span!("entry_async"))
-					.await;
-				if let scc::hash_map::Entry::Occupied(mut entry) = entry {
-					// Key exists, get and mutate existing RateLimiter
-					let write_guard = entry.get_mut();
-					Ok(write_guard.try_acquire())
-				} else {
-					// Key doesn't exist, insert a new RateLimiter
-					let mut limiter = RateLimiter::new(
-						middleware_config.rate_limit.requests,
-						middleware_config.rate_limit.period,
-					);
-					let result = limiter.try_acquire();
-					entry.insert_entry(limiter);
-					Ok(result)
-				}
-			}
+	async fn check_rate_limit(
+		&self,
+		ip_addr: std::net::IpAddr,
+		actor_id: &Option<Uuid>,
+	) -> GlobalResult<bool> {
+		// Use default middleware config, or actor-specific config if available
+		let middleware_config = match actor_id {
+			Some(id) => self.get_middleware_config(id).await?,
 			None => {
-				// No actor ID means no rate limiting
-				Ok(true)
+				// Default middleware config
+				MiddlewareConfig {
+					rate_limit: RateLimitConfig {
+						requests: 100, // 100 requests
+						period: 60,    // per 60 seconds
+					},
+					max_in_flight: MaxInFlightConfig {
+						amount: 20, // 20 concurrent requests
+					},
+					retry: RetryConfig {
+						max_attempts: 3,       // 3 retry attempts
+						initial_interval: 100, // 100ms initial interval
+					},
+					timeout: TimeoutConfig {
+						request_timeout: 30, // 30 seconds for requests
+					},
+				}
 			}
+		};
+
+		let entry = self
+			.rate_limiters
+			.entry_async(ip_addr)
+			.instrument(tracing::info_span!("entry_async"))
+			.await;
+		if let scc::hash_map::Entry::Occupied(mut entry) = entry {
+			// Key exists, get and mutate existing RateLimiter
+			let write_guard = entry.get_mut();
+			Ok(write_guard.try_acquire())
+		} else {
+			// Key doesn't exist, insert a new RateLimiter
+			let mut limiter = RateLimiter::new(
+				middleware_config.rate_limit.requests,
+				middleware_config.rate_limit.period,
+			);
+			let result = limiter.try_acquire();
+			entry.insert_entry(limiter);
+
+			metrics::RATE_LIMITERS_COUNT.set(self.rate_limiters.len() as i64);
+
+			Ok(result)
 		}
 	}
 
 	#[tracing::instrument(skip_all)]
-	async fn acquire_in_flight(&self, actor_id: &Option<Uuid>) -> GlobalResult<bool> {
-		match actor_id {
-			Some(id) => {
-				let middleware_config = self.get_middleware_config(id).await?;
-
-				let entry = self
-					.in_flight_counters
-					.entry_async(*id)
-					.instrument(tracing::info_span!("entry_async"))
-					.await;
-				if let scc::hash_map::Entry::Occupied(mut entry) = entry {
-					// Key exists, get and mutate existing InFlightCounter
-					let write_guard = entry.get_mut();
-					Ok(write_guard.try_acquire())
-				} else {
-					// Key doesn't exist, insert a new InFlightCounter
-					let mut counter = InFlightCounter::new(middleware_config.max_in_flight.amount);
-					let result = counter.try_acquire();
-					entry.insert_entry(counter);
-					Ok(result)
-				}
-			}
+	async fn acquire_in_flight(
+		&self,
+		ip_addr: std::net::IpAddr,
+		actor_id: &Option<Uuid>,
+	) -> GlobalResult<bool> {
+		// Use default middleware config, or actor-specific config if available
+		let middleware_config = match actor_id {
+			Some(id) => self.get_middleware_config(id).await?,
 			None => {
-				// No actor ID means no in-flight limiting
-				Ok(true)
+				// Default middleware config
+				MiddlewareConfig {
+					rate_limit: RateLimitConfig {
+						requests: 100, // 100 requests
+						period: 60,    // per 60 seconds
+					},
+					max_in_flight: MaxInFlightConfig {
+						amount: 20, // 20 concurrent requests
+					},
+					retry: RetryConfig {
+						max_attempts: 3,       // 3 retry attempts
+						initial_interval: 100, // 100ms initial interval
+					},
+					timeout: TimeoutConfig {
+						request_timeout: 30, // 30 seconds for requests
+					},
+				}
 			}
+		};
+
+		let entry = self
+			.in_flight_counters
+			.entry_async(ip_addr)
+			.instrument(tracing::info_span!("entry_async"))
+			.await;
+		if let scc::hash_map::Entry::Occupied(mut entry) = entry {
+			// Key exists, get and mutate existing InFlightCounter
+			let write_guard = entry.get_mut();
+			Ok(write_guard.try_acquire())
+		} else {
+			// Key doesn't exist, insert a new InFlightCounter
+			let mut counter = InFlightCounter::new(middleware_config.max_in_flight.amount);
+			let result = counter.try_acquire();
+			entry.insert_entry(counter);
+
+			metrics::IN_FLIGHT_COUNTERS_COUNT.set(self.in_flight_counters.len() as i64);
+
+			Ok(result)
 		}
 	}
 
 	#[tracing::instrument(skip_all)]
-	async fn release_in_flight(&self, actor_id: &Option<Uuid>) {
-		if let Some(id) = actor_id {
-			if let Some(mut counter) = self
-				.in_flight_counters
-				.get_async(id)
-				.instrument(tracing::info_span!("get_async"))
-				.await
-			{
-				counter.release();
-			}
+	async fn release_in_flight(&self, ip_addr: std::net::IpAddr) {
+		if let Some(mut counter) = self
+			.in_flight_counters
+			.get_async(&ip_addr)
+			.instrument(tracing::info_span!("get_async"))
+			.await
+		{
+			counter.release();
 		}
 	}
 }
@@ -607,8 +651,11 @@ impl ProxyService {
 		let actor_id_str = actor_id.map_or_else(|| "none".to_string(), |id| id.to_string());
 		let server_id_str = server_id.map_or_else(|| "none".to_string(), |id| id.to_string());
 
+		// Extract IP address from remote_addr
+		let client_ip = self.remote_addr.ip();
+
 		// Apply rate limiting
-		if !self.state.check_rate_limit(&actor_id).await? {
+		if !self.state.check_rate_limit(client_ip, &actor_id).await? {
 			metrics::ACTOR_REQUEST_ERRORS
 				.with_label_values(&[&actor_id_str, &server_id_str, "429"])
 				.inc();
@@ -619,7 +666,7 @@ impl ProxyService {
 		}
 
 		// Check in-flight limit
-		if !self.state.acquire_in_flight(&actor_id).await? {
+		if !self.state.acquire_in_flight(client_ip, &actor_id).await? {
 			metrics::ACTOR_REQUEST_ERRORS
 				.with_label_values(&[&actor_id_str, &server_id_str, "429"])
 				.inc();
@@ -646,10 +693,9 @@ impl ProxyService {
 
 		// Prepare to release in-flight counter when done
 		let state_clone = self.state.clone();
-		let actor_id_clone = actor_id;
 		crate::defer! {
 			tokio::spawn(async move {
-				state_clone.release_in_flight(&actor_id_clone).await;
+				state_clone.release_in_flight(client_ip).await;
 			}.instrument(tracing::info_span!("release_in_flight_task")));
 		}
 
@@ -1625,7 +1671,6 @@ impl ProxyService {
 					.observe(duration.as_secs_f64());
 
 				// Decrement pending metric at the end
-				info!("Decrementing pending metric");
 				metrics::ACTOR_REQUEST_PENDING
 					.with_label_values(&[&actor_id_str_clone, &server_id_str_clone, &method, &path])
 					.dec();