chore: add tracing

README.md

@@ -48,7 +48,7 @@ Projects powering Rivet Engine:
 
 - **[Pegboard](packages/services/pegboard/)**: Actor orchestrator
 - **[Guard](packages/core/guard/)**: Proxy for routing traffic to Rivet Actors
-- **[Chirp](packages/common/chirp-workflow/)**: Core workflow engine that powers Rivet
+- **[Gasoline](packages/common/gasoline/)**: Core durable execution engine that powers Rivet
 
 ## Get Started
 

docker/dev/grafana/dashboards/gasoline.json

@@ -2580,4 +2580,4 @@
   "uid": "636d22f9-d18f-4086-8b45-7c50886a105c",
   "version": 1,
   "weekStart": ""
-}
+}

packages/common/universaldb/src/driver/postgres/database.rs

@@ -53,7 +53,7 @@ impl PostgresDatabaseDriver {
 			.context("failed to create btree_gist extension")?;
 
 		conn.execute(
-			"CREATE SEQUENCE global_version_seq START WITH 1 INCREMENT BY 1 MINVALUE 1",
+			"CREATE SEQUENCE IF NOT EXISTS global_version_seq START WITH 1 INCREMENT BY 1 MINVALUE 1",
 			&[],
 		)
 		.await

packages/core/actor-kv/Cargo.toml

@@ -8,6 +8,7 @@ edition.workspace = true
 [dependencies]
 anyhow.workspace = true
 futures-util.workspace = true
+gas.workspace = true
 rivet-runner-protocol.workspace = true
 rivet-util-id.workspace = true
 serde_bare.workspace = true

packages/core/actor-kv/src/lib.rs

@@ -3,9 +3,9 @@ use std::result::Result::{Err, Ok};
 use anyhow::*;
 use entry::{EntryBaseKey, EntryBuilder, EntryMetadataKey, EntryValueChunkKey};
 use futures_util::{StreamExt, TryStreamExt};
+use gas::prelude::*;
 use key::{KeyWrapper, ListKeyWrapper};
 use rivet_runner_protocol as rp;
-use rivet_util_id::Id;
 use universaldb::prelude::*;
 use universaldb::tuple::Subspace;
 use utils::{validate_entries, validate_keys};
@@ -27,6 +27,7 @@ fn subspace(actor_id: Id) -> universaldb::utils::Subspace {
 }
 
 /// Returns estimated size of the given subspace.
+#[tracing::instrument(skip_all)]
 pub async fn get_subspace_size(db: &universaldb::Database, subspace: &Subspace) -> Result<i64> {
 	let (start, end) = subspace.range();
 
@@ -38,6 +39,7 @@ pub async fn get_subspace_size(db: &universaldb::Database, subspace: &Subspace)
 }
 
 /// Gets keys from the KV store.
+#[tracing::instrument(skip_all)]
 pub async fn get(
 	db: &universaldb::Database,
 	actor_id: Id,
@@ -120,11 +122,13 @@ pub async fn get(
 			Ok((keys, values, metadata))
 		}
 	})
+	.custom_instrument(tracing::info_span!("kv_get_tx"))
 	.await
 	.map_err(Into::<anyhow::Error>::into)
 }
 
 /// Gets keys from the KV store.
+#[tracing::instrument(skip_all)]
 pub async fn list(
 	db: &universaldb::Database,
 	actor_id: Id,
@@ -210,11 +214,13 @@ pub async fn list(
 			Ok((keys, values, metadata))
 		}
 	})
+	.custom_instrument(tracing::info_span!("kv_list_tx"))
 	.await
 	.map_err(Into::<anyhow::Error>::into)
 }
 
 /// Puts keys into the KV store.
+#[tracing::instrument(skip_all)]
 pub async fn put(
 	db: &universaldb::Database,
 	actor_id: Id,
@@ -273,11 +279,13 @@ pub async fn put(
 				.await
 		}
 	})
+	.custom_instrument(tracing::info_span!("kv_put_tx"))
 	.await
 	.map_err(Into::into)
 }
 
 /// Deletes keys from the KV store. Cannot be undone.
+#[tracing::instrument(skip_all)]
 pub async fn delete(db: &universaldb::Database, actor_id: Id, keys: Vec<rp::KvKey>) -> Result<()> {
 	validate_keys(&keys)?;
 
@@ -293,16 +301,19 @@ pub async fn delete(db: &universaldb::Database, actor_id: Id, keys: Vec<rp::KvKe
 			Ok(())
 		}
 	})
+	.custom_instrument(tracing::info_span!("kv_delete_tx"))
 	.await
 	.map_err(Into::into)
 }
 
 /// Deletes all keys from the KV store. Cannot be undone.
+#[tracing::instrument(skip_all)]
 pub async fn delete_all(db: &universaldb::Database, actor_id: Id) -> Result<()> {
 	db.run(|tx| async move {
 		tx.clear_subspace_range(&subspace(actor_id));
 		Ok(())
 	})
+	.custom_instrument(tracing::info_span!("kv_delete_all_tx"))
 	.await
 	.map_err(Into::into)
 }

packages/core/pegboard-serverless/src/lib.rs

@@ -91,6 +91,7 @@ async fn tick(
 			.try_collect::<Vec<_>>()
 			.await
 		})
+		.custom_instrument(tracing::info_span!("tick_tx"))
 		.await?;
 
 	let runner_configs = ctx

packages/services/epoxy/src/ops/explicit_prepare.rs

@@ -36,13 +36,15 @@ pub async fn epoxy_explicit_prepare(
 	// Read config
 	let config = ctx
 		.udb()?
-		.run(move |tx| async move { utils::read_config(&tx, replica_id).await })
+		.run(|tx| async move { utils::read_config(&tx, replica_id).await })
+		.custom_instrument(tracing::info_span!("read_config_tx"))
 		.await?;
 
 	// EPaxos Step 25: Increment ballot number
 	let new_ballot = ctx
 		.udb()?
-		.run(move |tx| async move { replica::ballot::increment_ballot(&tx, replica_id).await })
+		.run(|tx| async move { replica::ballot::increment_ballot(&tx, replica_id).await })
+		.custom_instrument(tracing::info_span!("increment_ballot_tx"))
 		.await?;
 
 	// Get quorum members
@@ -248,6 +250,7 @@ fn compare_ballots(a: &protocol::Ballot, b: &protocol::Ballot) -> std::cmp::Orde
 	}
 }
 
+#[tracing::instrument(skip_all)]
 async fn send_prepares(
 	ctx: &OperationCtx,
 	config: &protocol::ClusterConfig,
@@ -300,6 +303,7 @@ async fn send_prepares(
 	Ok(responses)
 }
 
+#[tracing::instrument(skip_all)]
 async fn restart_phase1(
 	ctx: &OperationCtx,
 	_config: &protocol::ClusterConfig,

packages/services/epoxy/src/ops/kv/get_local.rs

@@ -41,6 +41,7 @@ pub async fn epoxy_kv_get_local(ctx: &OperationCtx, input: &Input) -> Result<Out
 				.await
 			}
 		})
+		.custom_instrument(tracing::info_span!("get_local_tx"))
 		.await?;
 
 	Ok(Output { value })

packages/services/epoxy/src/ops/kv/get_optimistic.rs

@@ -77,6 +77,7 @@ pub async fn epoxy_kv_get_optimistic(ctx: &OperationCtx, input: &Input) -> Resul
 				.await
 			}
 		})
+		.custom_instrument(tracing::info_span!("get_optimistic_tx"))
 		.await?;
 
 	if value.is_some() {
@@ -144,6 +145,7 @@ pub async fn epoxy_kv_get_optimistic(ctx: &OperationCtx, input: &Input) -> Resul
 						.await
 					}
 				})
+				.custom_instrument(tracing::info_span!("cache_value_tx"))
 				.await?;
 
 			return Ok(Output { value: Some(value) });

packages/services/epoxy/src/ops/propose.rs

@@ -33,16 +33,18 @@ pub async fn epoxy_propose(ctx: &OperationCtx, input: &Input) -> Result<Proposal
 	// Read config
 	let config = ctx
 		.udb()?
-		.run(move |tx| async move { utils::read_config(&tx, replica_id).await })
+		.run(|tx| async move { utils::read_config(&tx, replica_id).await })
+		.custom_instrument(tracing::info_span!("read_config_tx"))
 		.await?;
 
 	// Lead consensus
 	let payload = ctx
 		.udb()?
-		.run(move |tx| {
+		.run(|tx| {
 			let proposal = input.proposal.clone();
 			async move { replica::lead_consensus::lead_consensus(&*tx, replica_id, proposal).await }
 		})
+		.custom_instrument(tracing::info_span!("lead_consensus_tx"))
 		.await?;
 
 	// Get quorum members (only active replicas for voting)
@@ -55,11 +57,12 @@ pub async fn epoxy_propose(ctx: &OperationCtx, input: &Input) -> Result<Proposal
 	// Decide path
 	let path = ctx
 		.udb()?
-		.run(move |tx| {
+		.run(|tx| {
 			let pre_accept_oks = pre_accept_oks.clone();
 			let payload = payload.clone();
 			async move { replica::decide_path::decide_path(&*tx, pre_accept_oks, &payload) }
 		})
+		.custom_instrument(tracing::info_span!("decide_path_tx"))
 		.await?;
 
 	match path {
@@ -72,6 +75,7 @@ pub async fn epoxy_propose(ctx: &OperationCtx, input: &Input) -> Result<Proposal
 	}
 }
 
+#[tracing::instrument(skip_all)]
 pub async fn run_paxos_accept(
 	ctx: &OperationCtx,
 	config: &protocol::ClusterConfig,
@@ -84,10 +88,11 @@ pub async fn run_paxos_accept(
 
 	// Mark as accepted
 	ctx.udb()?
-		.run(move |tx| {
+		.run(|tx| {
 			let payload = payload.clone();
 			async move { replica::messages::accepted(&*tx, replica_id, payload).await }
 		})
+		.custom_instrument(tracing::info_span!("accept_tx"))
 		.await?;
 
 	// EPaxos Step 17
@@ -115,6 +120,7 @@ pub async fn run_paxos_accept(
 	}
 }
 
+#[tracing::instrument(skip_all)]
 pub async fn commit(
 	ctx: &OperationCtx,
 	config: &protocol::ClusterConfig,
@@ -130,14 +136,15 @@ pub async fn commit(
 	let cmd_err = {
 		let payload = payload.clone();
 		ctx.udb()?
-			.run(move |tx| {
+			.run(|tx| {
 				let payload = payload.clone();
 				async move {
 					let cmd_err = replica::messages::committed(&*tx, replica_id, &payload).await?;
 
 					Result::Ok(cmd_err)
 				}
 			})
+			.custom_instrument(tracing::info_span!("committed_tx"))
 			.await?
 	};
 
@@ -163,6 +170,7 @@ pub async fn commit(
 	}
 }
 
+#[tracing::instrument(skip_all)]
 async fn send_pre_accepts(
 	ctx: &OperationCtx,
 	config: &protocol::ClusterConfig,
@@ -207,6 +215,7 @@ async fn send_pre_accepts(
 	Ok(responses)
 }
 
+#[tracing::instrument(skip_all)]
 async fn send_accepts(
 	ctx: &OperationCtx,
 	config: &protocol::ClusterConfig,
@@ -252,6 +261,7 @@ async fn send_accepts(
 	Ok(responses.len() + 1)
 }
 
+#[tracing::instrument(skip_all)]
 async fn send_commits(
 	ctx: &OperationCtx,
 	config: &protocol::ClusterConfig,

packages/services/epoxy/src/ops/read_cluster_config.rs

@@ -22,6 +22,7 @@ pub async fn epoxy_read_cluster_config(ctx: &OperationCtx, input: &Input) -> Res
 			let replica_id = input.replica_id;
 			async move { utils::read_config(&tx, replica_id).await }
 		})
+		.custom_instrument(tracing::info_span!("read_cluster_config_tx"))
 		.await?;
 
 	Ok(Output { config })

packages/services/epoxy/src/replica/ballot.rs

@@ -6,6 +6,7 @@ use universaldb::utils::{FormalKey, IsolationLevel::*};
 use crate::keys;
 
 /// Get the current ballot for this replica
+#[tracing::instrument(skip_all)]
 pub async fn get_ballot(
 	tx: &Transaction,
 	replica_id: protocol::ReplicaId,
@@ -31,6 +32,7 @@ pub async fn get_ballot(
 }
 
 /// Increment the ballot number and return the new ballot
+#[tracing::instrument(skip_all)]
 pub async fn increment_ballot(
 	tx: &Transaction,
 	replica_id: protocol::ReplicaId,
@@ -67,6 +69,7 @@ pub fn compare_ballots(
 /// ballot if needed.
 ///
 /// Returns true if the ballot is valid (higher than previously seen).
+#[tracing::instrument(skip_all)]
 pub async fn validate_and_update_ballot_for_instance(
 	tx: &Transaction,
 	replica_id: protocol::ReplicaId,

packages/services/epoxy/src/replica/commit_kv.rs

@@ -6,6 +6,7 @@ use universaldb::prelude::*;
 use crate::{keys, ops::propose::CommandError, replica::utils};
 
 /// Commits a proposal to KV store.
+#[tracing::instrument(skip_all)]
 pub async fn commit_kv(
 	tx: &Transaction,
 	replica_id: ReplicaId,

packages/services/epoxy/src/replica/lead_consensus.rs

@@ -6,6 +6,7 @@ use universaldb::utils::{FormalKey, IsolationLevel::*};
 use crate::keys;
 use crate::replica::{ballot, messages, utils};
 
+#[tracing::instrument(skip_all)]
 pub async fn lead_consensus(
 	tx: &Transaction,
 	replica_id: protocol::ReplicaId,

packages/services/epoxy/src/replica/log.rs

@@ -15,6 +15,7 @@ pub fn state_order(state: &protocol::State) -> u8 {
 	}
 }
 
+#[tracing::instrument(skip_all)]
 pub async fn update_log(
 	tx: &Transaction,
 	replica_id: ReplicaId,