chore(tracing-utils): add tracing utils for custom instrument (#3293)

* chore(guard): add metadata to rate limit error

* chore(tracing-utils): add tracing utils for custom instrument

Author: NathanFlurry

Cargo.lock

@@ -374,6 +374,9 @@ path = "engine/packages/test-deps-docker"
 [workspace.dependencies.rivet-tracing-reconfigure]
 path = "engine/packages/tracing-reconfigure"
 
+[workspace.dependencies.rivet-tracing-utils]
+path = "engine/packages/tracing-utils"
+
 [workspace.dependencies.rivet-types]
 path = "engine/packages/types"
 

Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "rivet-tracing-utils"
+version.workspace = true
+authors.workspace = true
+license.workspace = true
+edition.workspace = true
+
+[dependencies]
+futures-util.workspace = true
+lazy_static.workspace = true
+rivet-metrics.workspace = true
+tracing.workspace = true

engine/packages/tracing-utils/Cargo.toml

@@ -0,0 +1,91 @@
+use std::{
+	future::Future,
+	pin::Pin,
+	task::{Context, Poll},
+	time::Instant,
+};
+
+use futures_util::future;
+use tracing::{Instrument, instrument::Instrumented};
+
+use rivet_metrics::KeyValue;
+
+/// Attempts to create a new future to select over a list of futures.
+/// Non-panicking version of [futures_util::future::select_all](https://docs.rs/futures/0.3.15/futures/future/fn.select_all.html).
+///
+/// If `iter` is empty, a `Pending` future is returned.
+pub async fn select_all_or_wait<I>(iter: I) -> <I::Item as Future>::Output
+where
+	I: IntoIterator,
+	I::Item: Future + Unpin,
+{
+	let futs = iter.into_iter().collect::<Vec<I::Item>>();
+
+	if !futs.is_empty() {
+		future::select_all(futs).await.0
+	} else {
+		std::future::pending().await
+	}
+}
+
+pub trait CustomInstrumentExt: Sized {
+	fn custom_instrument(self, span: tracing::Span) -> CustomInstrumented<Self> {
+		CustomInstrumented {
+			inner: self.instrument(span),
+			start: Instant::now(),
+		}
+	}
+}
+
+impl<F: Sized> CustomInstrumentExt for F {}
+
+pub struct CustomInstrumented<T> {
+	inner: Instrumented<T>,
+	start: Instant,
+}
+
+impl<T: Future> Future for CustomInstrumented<T> {
+	type Output = T::Output;
+
+	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+		let this = unsafe { self.get_unchecked_mut() };
+		let inner = unsafe { Pin::new_unchecked(&mut this.inner) };
+
+		let metadata = inner.span().metadata().clone();
+
+		match inner.poll(cx) {
+			Poll::Ready(val) => {
+				if let Some(metadata) = metadata {
+					if let (Some(file), Some(line)) = (metadata.file(), metadata.line()) {
+						metrics::INSTRUMENTED_FUTURE_DURATION.record(
+							this.start.elapsed().as_secs_f64(),
+							&[
+								KeyValue::new("location", format!("{file}:{line}")),
+								KeyValue::new("name", metadata.name()),
+							],
+						);
+					}
+				}
+				Poll::Ready(val)
+			}
+			Poll::Pending => Poll::Pending,
+		}
+	}
+}
+
+mod metrics {
+	use rivet_metrics::{
+		MICRO_BUCKETS,
+		otel::{global::*, metrics::*},
+	};
+
+	lazy_static::lazy_static! {
+		static ref METER: Meter = meter("rivet-util-core");
+
+		/// Expected attributes: "location", "name"
+		pub static ref INSTRUMENTED_FUTURE_DURATION: Histogram<f64> = METER.f64_histogram("rivet_instrumented_future_duration")
+			.with_description("Duration of a future.")
+			.with_boundaries(MICRO_BUCKETS.to_vec())
+			.build();
+	}
+}

engine/packages/tracing-utils/src/lib.rs

@@ -14,6 +14,7 @@ futures-util.workspace = true
 lazy_static.workspace = true
 rand.workspace = true
 rivet-metrics.workspace = true
+rivet-tracing-utils.workspace = true
 rocksdb.workspace = true
 serde.workspace = true
 thiserror.workspace = true

engine/packages/universaldb/Cargo.toml

@@ -2,6 +2,7 @@ use std::future::Future;
 
 use anyhow::{Context, Result, anyhow};
 use futures_util::FutureExt;
+use rivet_tracing_utils::CustomInstrumentExt;
 
 use crate::{
 	driver::{DatabaseDriverHandle, Erased},
@@ -20,6 +21,7 @@ impl Database {
 	}
 
 	/// Run a closure with automatic retry logic
+	#[tracing::instrument(skip_all)]
 	pub async fn run<'a, F, Fut, T>(&'a self, closure: F) -> Result<T>
 	where
 		F: Fn(RetryableTransaction) -> Fut + Send + Sync,
@@ -29,7 +31,9 @@ impl Database {
 		let closure = &closure;
 		self.driver
 			.run(Box::new(|tx| {
-				async move { closure(tx).await.map(|value| Box::new(value) as Erased) }.boxed()
+				async move { closure(tx).await.map(|value| Box::new(value) as Erased) }
+					.custom_instrument(tracing::info_span!("run_attempt"))
+					.boxed()
 			}))
 			.await
 			.and_then(|res| {

engine/packages/universaldb/src/database.rs

@@ -24,6 +24,7 @@ regex.workspace = true
 reqwest.workspace = true
 rivet-config.workspace = true
 rivet-metrics.workspace = true
+rivet-tracing-utils.workspace = true
 rivet-util-id.workspace = true
 serde.workspace = true
 serde_json.workspace = true

engine/packages/util/Cargo.toml

@@ -1,91 +1 @@
-use std::{
-	future::Future,
-	pin::Pin,
-	task::{Context, Poll},
-	time::Instant,
-};
-
-use futures_util::future;
-use tracing::{Instrument, instrument::Instrumented};
-
-use rivet_metrics::KeyValue;
-
-/// Attempts to create a new future to select over a list of futures.
-/// Non-panicking version of [futures_util::future::select_all](https://docs.rs/futures/0.3.15/futures/future/fn.select_all.html).
-///
-/// If `iter` is empty, a `Pending` future is returned.
-pub async fn select_all_or_wait<I>(iter: I) -> <I::Item as Future>::Output
-where
-	I: IntoIterator,
-	I::Item: Future + Unpin,
-{
-	let futs = iter.into_iter().collect::<Vec<I::Item>>();
-
-	if !futs.is_empty() {
-		future::select_all(futs).await.0
-	} else {
-		std::future::pending().await
-	}
-}
-
-pub trait CustomInstrumentExt: Sized {
-	fn custom_instrument(self, span: tracing::Span) -> CustomInstrumented<Self> {
-		CustomInstrumented {
-			inner: self.instrument(span),
-			start: Instant::now(),
-		}
-	}
-}
-
-impl<F: Sized> CustomInstrumentExt for F {}
-
-pub struct CustomInstrumented<T> {
-	inner: Instrumented<T>,
-	start: Instant,
-}
-
-impl<T: Future> Future for CustomInstrumented<T> {
-	type Output = T::Output;
-
-	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-		let this = unsafe { self.get_unchecked_mut() };
-		let inner = unsafe { Pin::new_unchecked(&mut this.inner) };
-
-		let metadata = inner.span().metadata().clone();
-
-		match inner.poll(cx) {
-			Poll::Ready(val) => {
-				if let Some(metadata) = metadata {
-					if let (Some(file), Some(line)) = (metadata.file(), metadata.line()) {
-						metrics::INSTRUMENTED_FUTURE_DURATION.record(
-							this.start.elapsed().as_secs_f64(),
-							&[
-								KeyValue::new("location", format!("{file}:{line}")),
-								KeyValue::new("name", metadata.name()),
-							],
-						);
-					}
-				}
-				Poll::Ready(val)
-			}
-			Poll::Pending => Poll::Pending,
-		}
-	}
-}
-
-mod metrics {
-	use rivet_metrics::{
-		MICRO_BUCKETS,
-		otel::{global::*, metrics::*},
-	};
-
-	lazy_static::lazy_static! {
-		static ref METER: Meter = meter("rivet-util-core");
-
-		/// Expected attributes: "location", "name"
-		pub static ref INSTRUMENTED_FUTURE_DURATION: Histogram<f64> = METER.f64_histogram("rivet_instrumented_future_duration")
-			.with_description("Duration of a future.")
-			.with_boundaries(MICRO_BUCKETS.to_vec())
-			.build();
-	}
-}
+pub use rivet_tracing_utils::*;