fix(guard, pb): add top level ray id, fix actor wf error handling (#3236)

Author: MasterPtato

Cargo.lock

@@ -12,7 +12,8 @@ axum-extra.workspace = true
 gas.workspace = true
 chrono.workspace = true
 http.workspace = true
-hyper = { workspace = true, features = ["full"] }
+# TODO: Make this use workspace version
+hyper = { version = "1.6.0", features = ["full"] }
 lazy_static.workspace = true
 opentelemetry.workspace = true
 rivet-cache.workspace = true

packages/common/api-builder/Cargo.toml

@@ -8,6 +8,7 @@ use axum::{
 	middleware::Next,
 	response::Response,
 };
+use hyper::header::HeaderName;
 use opentelemetry::trace::TraceContextExt;
 use rivet_metrics::KeyValue;
 use tower_http::trace::TraceLayer;
@@ -16,6 +17,8 @@ use tracing_opentelemetry::OpenTelemetrySpanExt;
 
 use crate::{ErrorExt, RequestIds, metrics};
 
+pub const X_RIVET_RAY_ID: HeaderName = HeaderName::from_static("x-rivet-ray-id");
+
 // TODO: Remove this since this is duplicate logs & traces, but this is just to see what Axum adds
 // natively vs our logging. We can add this once we're satisfied with our own logging.
 pub fn create_trace_layer()
@@ -46,11 +49,14 @@ pub async fn http_logging_middleware(
 		.span_context()
 		.clone();
 
-	// Generate request IDs
-	let request_ids = RequestIds::new(config.dc_label());
-
-	// Add request IDs to request extensions so they can be accessed by handlers
-	req.extensions_mut().insert(request_ids);
+	// Add request IDs to request extensions if not already added by guard so they can be accessed by handlers
+	let request_ids = if let Some(request_ids) = req.extensions().get::<RequestIds>() {
+		*request_ids
+	} else {
+		let request_ids = RequestIds::new(config.dc_label());
+		req.extensions_mut().insert(request_ids);
+		request_ids
+	};
 
 	// Create span for this request
 	let req_span = tracing::info_span!(
@@ -118,7 +124,7 @@ pub async fn http_logging_middleware(
 
 		// Add ray_id to response headers
 		if let Ok(ray_id_value) = request_ids.ray_id.to_string().parse() {
-			response.headers_mut().insert("rvt-ray-id", ray_id_value);
+			response.headers_mut().insert(X_RIVET_RAY_ID, ray_id_value);
 		}
 
 		let status = response.status();
@@ -175,8 +181,6 @@ pub async fn http_logging_middleware(
 		let duration = start.elapsed().as_secs_f64();
 
 		tracing::debug!(
-			ray_id = %request_ids.ray_id,
-			req_id = %request_ids.req_id,
 			%remote_addr,
 			%method,
 			%uri,

packages/common/api-builder/src/middleware.rs

@@ -1,7 +1,7 @@
 use anyhow::{Context, Result};
 use axum::{body::Body, response::Response};
 use futures_util::StreamExt;
-use rivet_api_builder::{ApiCtx, ErrorResponse, RawErrorResponse};
+use rivet_api_builder::{ApiCtx, ErrorResponse, RawErrorResponse, X_RIVET_RAY_ID};
 use serde::{Serialize, de::DeserializeOwned};
 use std::future::Future;
 
@@ -43,7 +43,7 @@ pub async fn request_remote_datacenter_raw(
 	let res = request
 		.send()
 		.await
-		.context("failed sending request to remote dc")?;
+		.context("failed parsing response from remote dc")?;
 	reqwest_to_axum_response(res)
 		.await
 		.context("failed parsing response from remote dc")
@@ -85,7 +85,7 @@ where
 	let res = request
 		.send()
 		.await
-		.context("failed sending request to remote dc")?;
+		.context("failed parsing response from remote dc")?;
 	parse_response::<T>(res)
 		.await
 		.context("failed parsing response from remote dc")
@@ -159,7 +159,7 @@ where
 	// Error only if all requests failed
 	if result_count == errors.len() {
 		if let Some(res) = errors.into_iter().next() {
-			return Err(res).context("all datacenter requests failed");
+			return Err(res).with_context(|| "all datacenter requests failed");
 		}
 	}
 
@@ -184,15 +184,25 @@ pub async fn reqwest_to_axum_response(reqwest_response: reqwest::Response) -> Re
 #[tracing::instrument(skip_all)]
 pub async fn parse_response<T: DeserializeOwned>(reqwest_response: reqwest::Response) -> Result<T> {
 	let status = reqwest_response.status();
+	let headers = reqwest_response.headers();
+	let ray_id = headers
+		.get(X_RIVET_RAY_ID)
+		.and_then(|v| v.to_str().ok())
+		.map(|x| x.to_string());
 	let response_text = reqwest_response.text().await?;
 
 	if status.is_success() {
 		serde_json::from_str::<T>(&response_text).map_err(Into::into)
 	} else {
-		Err(RawErrorResponse(
+		let error = RawErrorResponse(
 			status,
 			serde_json::from_str::<ErrorResponse>(&response_text)?,
-		)
-		.into())
+		);
+
+		if let Some(ray_id) = ray_id {
+			Err(error).with_context(|| format!("remote request failed (ray_id: {ray_id})"))
+		} else {
+			Err(error.into())
+		}
 	}
 }

packages/common/api-util/src/lib.rs

@@ -1,5 +1,4 @@
 use serde::{Deserialize, Serialize};
-use uuid::Uuid;
 
 // Properties not currently collected but should be added in future iterations:
 // - client_ssl_cipher: Requires TLS connection introspection
@@ -18,7 +17,9 @@ use uuid::Uuid;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct GuardHttpRequest {
-	pub request_id: Uuid,
+	// TODO:
+	// pub request_id: Id,
+	// pub ray_id: Id,
 	pub client_ip: String,
 	pub client_request_body_bytes: u64,
 	pub client_request_host: String,
@@ -31,9 +32,6 @@ pub struct GuardHttpRequest {
 	pub client_request_user_agent: String,
 	pub client_src_port: u16,
 	pub client_x_requested_with: String,
-	pub guard_datacenter_id: Uuid,
-	pub guard_cluster_id: Uuid,
-	pub guard_server_id: Uuid,
 	pub guard_end_timestamp: u64,
 	pub guard_response_body_bytes: u64,
 	pub guard_response_content_type: String,