From 0caa600717e25f8b8705eba86c04cab18d671723 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Thu, 18 Jun 2026 17:57:51 -0700 Subject: [PATCH 1/8] [SLOP(claude-opus-4-8)] fix(kitchen-sink): serve built frontend from server --- examples/kitchen-sink/src/server.ts | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/examples/kitchen-sink/src/server.ts b/examples/kitchen-sink/src/server.ts index f1e5ba9d70..bdac261b7c 100644 --- a/examples/kitchen-sink/src/server.ts +++ b/examples/kitchen-sink/src/server.ts @@ -1,6 +1,9 @@ +import { existsSync, readFileSync } from "node:fs"; import type { Server as HttpServer } from "node:http"; +import { resolve } from "node:path"; import * as v8 from "node:v8"; import { serve } from "@hono/node-server"; +import { serveStatic } from "@hono/node-server/serve-static"; import { Hono } from "hono"; import { registry } from "./index.ts"; import { resolveMode } from "./mode.ts"; @@ -164,6 +167,23 @@ if (mode === "serverful") { app.all("/api/rivet", (c) => registry.handler(c.req.raw)); } +// Serve the built frontend when it is present. The Vite build emits `dist/`, +// which only exists in production images, so dev runs skip this branch. +const distDir = "dist"; +const indexPath = resolve(process.cwd(), distDir, "index.html"); +if (existsSync(indexPath)) { + app.use("/*", serveStatic({ root: distDir })); + const indexHtml = readFileSync(indexPath, "utf8"); + app.get("/*", (c) => { + const path = new URL(c.req.url).pathname; + const last = path.slice(path.lastIndexOf("/") + 1); + // Fall through to 404 for asset-like paths so missing files do not + // resolve to the SPA shell. + if (last.includes(".")) return c.notFound(); + return c.html(indexHtml); + }); +} + const server = serve({ fetch: app.fetch, port }, () => { if (mode === "serverful") { console.log( From c40c936bcbc90e4cf5e9422eb3acc11fa3f785ec Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Fri, 19 Jun 2026 14:25:58 -0700 Subject: [PATCH 2/8] [SLOP(claude-opus-4-8-medium)] chore: add profiling cargo profile for heaptrack leak investigation --- Cargo.toml | 10 ++++ .../rivetkit-napi/src/napi_actor_events.rs | 47 +++++++++++++------ 2 files changed, 43 insertions(+), 14 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 3f544252f3..511aec200f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -672,6 +672,16 @@ lto = "fat" codegen-units = 1 opt-level = 3 +# Release-grade optimization with DWARF line info and no symbol stripping so +# heaptrack can resolve native allocation backtraces during leak investigation. +# Uses thin LTO and more codegen units to keep frames un-inlined and builds fast. +[profile.profiling] +inherits = "release" +debug = 1 +lto = "thin" +codegen-units = 16 +strip = false + [profile.quick] inherits = "dev" debug = false # no debug info → faster link, smaller binary diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs index bb2e91a6f1..9da9f0847e 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs @@ -173,21 +173,40 @@ async fn run_event_loop( dirty: &Arc, events: &mut ActorEvents, ) { - while let Some(event) = events.recv().await { + loop { pump_registered_tasks(tasks, registered_task_rx); - dispatch_event( - event, - bindings, - config, - ctx, - abort, - tasks, - registered_task_rx, - dirty, - ) - .await; - if ctx.has_end_reason() { - break; + + tokio::select! { + // Reap completed background tasks as they finish. A tokio JoinSet + // retains each finished task's allocation until it is joined, so + // without this the set grows for the entire actor lifetime and + // shows up as native (non-V8) RSS growth. + Some(result) = tasks.join_next(), if !tasks.is_empty() => { + if let Err(error) = result { + if !error.is_cancelled() { + tracing::error!(?error, "napi background task failed to join"); + } + } + } + event = events.recv() => { + let Some(event) = event else { + break; + }; + dispatch_event( + event, + bindings, + config, + ctx, + abort, + tasks, + registered_task_rx, + dirty, + ) + .await; + if ctx.has_end_reason() { + break; + } + } } } } From 9ff6358ec9e27f33e4c677ae4f21cfb0d25276d4 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Tue, 23 Jun 2026 12:49:51 -0700 Subject: [PATCH 3/8] Add observation and ser/de metrics for serde calls --- engine/packages/api-peer/src/internal.rs | 2 +- engine/packages/cache/src/req_config.rs | 6 +- engine/packages/depot/Cargo.toml | 6 +- engine/packages/gasoline-macros/src/lib.rs | 2 +- .../gasoline/src/builder/common/signal.rs | 2 +- .../gasoline/src/builder/common/workflow.rs | 4 +- .../gasoline/src/builder/workflow/lupe.rs | 8 +- .../gasoline/src/builder/workflow/message.rs | 2 +- .../gasoline/src/builder/workflow/signal.rs | 2 +- .../src/builder/workflow/sub_workflow.rs | 4 +- engine/packages/gasoline/src/ctx/message.rs | 7 +- engine/packages/gasoline/src/ctx/workflow.rs | 8 +- engine/packages/gasoline/src/db/mod.rs | 8 +- engine/packages/gasoline/src/history/event.rs | 7 +- engine/packages/gasoline/src/message.rs | 4 +- engine/packages/gasoline/src/registry.rs | 4 +- engine/packages/gasoline/src/signal.rs | 2 +- engine/packages/gasoline/src/workflow.rs | 4 +- engine/packages/ups-broadcast/src/lib.rs | 16 +- engine/packages/ups-broadcast/src/sim.rs | 2509 ----------------- engine/packages/util/src/lib.rs | 128 + engine/packages/util/src/metrics.rs | 34 + engine/packages/util/src/serde.rs | 106 + 23 files changed, 312 insertions(+), 2563 deletions(-) delete mode 100644 engine/packages/ups-broadcast/src/sim.rs create mode 100644 engine/packages/util/src/metrics.rs diff --git a/engine/packages/api-peer/src/internal.rs b/engine/packages/api-peer/src/internal.rs index 9b39644a08..eaa1e9ed34 100644 --- a/engine/packages/api-peer/src/internal.rs +++ b/engine/packages/api-peer/src/internal.rs @@ -62,7 +62,7 @@ pub async fn set_tracing_config( body: SetTracingConfigRequest, ) -> Result { // Broadcast message to all services via UPS - let message = serde_json::to_vec(&body)?; + let message = rivet_util::serde::json_to_vec!(&body)?; ctx.ups()? .publish(TracingConfigSubject, &message, PublishOpts::broadcast()) diff --git a/engine/packages/cache/src/req_config.rs b/engine/packages/cache/src/req_config.rs index c97df9c512..1b92618689 100644 --- a/engine/packages/cache/src/req_config.rs +++ b/engine/packages/cache/src/req_config.rs @@ -365,7 +365,7 @@ impl RequestConfig { keys: cache_keys.clone(), }; - let payload = serde_json::to_vec(&message)?; + let payload = rivet_util::serde::json_to_vec!(&message)?; if let Err(err) = ups .publish( @@ -495,12 +495,12 @@ impl RequestConfig { keys, getter, |value: &Value| -> Result> { - serde_json::to_vec(&value) + rivet_util::serde::json_to_vec!(&value) .map_err(Error::SerdeEncode) .map_err(Into::into) }, |value: &[u8]| -> Result { - serde_json::from_slice(value) + rivet_util::serde::json_from_slice!(value) .map_err(Error::SerdeDecode) .map_err(Into::into) }, diff --git a/engine/packages/depot/Cargo.toml b/engine/packages/depot/Cargo.toml index fbf16b652d..ff38b8cdd4 100644 --- a/engine/packages/depot/Cargo.toml +++ b/engine/packages/depot/Cargo.toml @@ -30,19 +30,19 @@ rivet-error.workspace = true rivet-metrics.workspace = true rivet-pools.workspace = true rivet-runtime.workspace = true +rivet-util.workspace = true +rusqlite.workspace = true scc.workspace = true -serde.workspace = true serde_bare.workspace = true serde_json.workspace = true +serde.workspace = true sha2.workspace = true -rusqlite.workspace = true tempfile.workspace = true tokio.workspace = true tokio-util.workspace = true tracing.workspace = true universaldb.workspace = true universalpubsub.workspace = true -util.workspace = true uuid.workspace = true vbare.workspace = true diff --git a/engine/packages/gasoline-macros/src/lib.rs b/engine/packages/gasoline-macros/src/lib.rs index ba241cb186..7b4810a118 100644 --- a/engine/packages/gasoline-macros/src/lib.rs +++ b/engine/packages/gasoline-macros/src/lib.rs @@ -403,7 +403,7 @@ pub fn signal(attr: TokenStream, item: TokenStream) -> TokenStream { } fn parse(_name: &str, body: &serde_json::value::RawValue) -> gas::prelude::WorkflowResult { - serde_json::from_str(body.get()).map_err(WorkflowError::DeserializeSignalBody) + rivet_util::serde::json_from_str!(body.get()).map_err(WorkflowError::DeserializeSignalBody) } } }; diff --git a/engine/packages/gasoline/src/builder/common/signal.rs b/engine/packages/gasoline/src/builder/common/signal.rs index 84a473c8dd..4a56c0a7dc 100644 --- a/engine/packages/gasoline/src/builder/common/signal.rs +++ b/engine/packages/gasoline/src/builder/common/signal.rs @@ -129,7 +129,7 @@ impl SignalBuilder { tracing::Span::current().record("signal_id", signal_id.to_string()); // Serialize input - let input_val = serde_json::value::to_raw_value(&self.body) + let input_val = rivet_util::serde::json_to_raw_value!(&self.body) .map_err(WorkflowError::SerializeSignalBody)?; match ( diff --git a/engine/packages/gasoline/src/builder/common/workflow.rs b/engine/packages/gasoline/src/builder/common/workflow.rs index da650d511b..f06ba4080e 100644 --- a/engine/packages/gasoline/src/builder/common/workflow.rs +++ b/engine/packages/gasoline/src/builder/common/workflow.rs @@ -81,7 +81,7 @@ where return self; } - match serde_json::to_value(&v) { + match rivet_util::serde::json_to_value!(&v) { Ok(v) => { self.tags.insert(k.to_string(), v); } @@ -125,7 +125,7 @@ where } // Serialize input - let input_val = serde_json::value::to_raw_value(&input) + let input_val = rivet_util::serde::json_to_raw_value!(&input) .map_err(WorkflowError::SerializeWorkflowInput)?; let actual_workflow_id = self diff --git a/engine/packages/gasoline/src/builder/workflow/lupe.rs b/engine/packages/gasoline/src/builder/workflow/lupe.rs index 8b45b79615..8dd72b34b8 100644 --- a/engine/packages/gasoline/src/builder/workflow/lupe.rs +++ b/engine/packages/gasoline/src/builder/workflow/lupe.rs @@ -72,7 +72,7 @@ impl<'a, S: Serialize + DeserializeOwned> LoopBuilder<'a, S> { (loop_event.iteration, state, output, None) } else { - let state_val = serde_json::value::to_raw_value(&state) + let state_val = rivet_util::serde::json_to_raw_value!(&state) .map_err(WorkflowError::SerializeLoopOutput)?; // Clone data to move into future @@ -219,7 +219,7 @@ impl<'a, S: Serialize + DeserializeOwned> LoopBuilder<'a, S> { if iteration % commit_interval.unwrap_or(DEFAULT_LOOP_COMMIT_INTERVAL) == 0 { - let state_val = serde_json::value::to_raw_value(&state) + let state_val = rivet_util::serde::json_to_raw_value!(&state) .map_err(WorkflowError::SerializeLoopOutput)?; // Clone data to move into future @@ -251,9 +251,9 @@ impl<'a, S: Serialize + DeserializeOwned> LoopBuilder<'a, S> { Loop::Break(res) => { iteration += 1; - let state_val = serde_json::value::to_raw_value(&state) + let state_val = rivet_util::serde::json_to_raw_value!(&state) .map_err(WorkflowError::SerializeLoopOutput)?; - let output_val = serde_json::value::to_raw_value(&res) + let output_val = rivet_util::serde::json_to_raw_value!(&res) .map_err(WorkflowError::SerializeLoopOutput)?; // Commit loop output and final state to db. Note that we don't defer this because diff --git a/engine/packages/gasoline/src/builder/workflow/message.rs b/engine/packages/gasoline/src/builder/workflow/message.rs index d21869956c..aebb3a08f2 100644 --- a/engine/packages/gasoline/src/builder/workflow/message.rs +++ b/engine/packages/gasoline/src/builder/workflow/message.rs @@ -76,7 +76,7 @@ impl<'a, M: Message> MessageBuilder<'a, M> { let start_instant = Instant::now(); // Serialize body - let body_val = serde_json::value::to_raw_value(&self.body) + let body_val = rivet_util::serde::json_to_raw_value!(&self.body) .map_err(WorkflowError::SerializeMessageBody)?; let topic = self.topic.unwrap_or_else(|| "*".to_string()); let tags = serde_json::Value::Object( diff --git a/engine/packages/gasoline/src/builder/workflow/signal.rs b/engine/packages/gasoline/src/builder/workflow/signal.rs index 34a7e8e63f..fb34607d72 100644 --- a/engine/packages/gasoline/src/builder/workflow/signal.rs +++ b/engine/packages/gasoline/src/builder/workflow/signal.rs @@ -145,7 +145,7 @@ impl<'a, T: Signal + Serialize> SignalBuilder<'a, T> { let db_write_duration; // Serialize input - let input_val = serde_json::value::to_raw_value(&self.body) + let input_val = rivet_util::serde::json_to_raw_value!(&self.body) .map_err(WorkflowError::SerializeSignalBody)?; match ( diff --git a/engine/packages/gasoline/src/builder/workflow/sub_workflow.rs b/engine/packages/gasoline/src/builder/workflow/sub_workflow.rs index 31c5ae8f34..d83a8e8c84 100644 --- a/engine/packages/gasoline/src/builder/workflow/sub_workflow.rs +++ b/engine/packages/gasoline/src/builder/workflow/sub_workflow.rs @@ -155,7 +155,7 @@ where } // Serialize input - let input_val = serde_json::value::to_raw_value(input) + let input_val = rivet_util::serde::json_to_raw_value!(input) .map_err(WorkflowError::SerializeWorkflowOutput)?; let actual_sub_workflow_id = ctx @@ -228,7 +228,7 @@ where // Err for version mismatch self.ctx.compare_version("sub workflow", self.version)?; - let input_val = serde_json::value::to_raw_value(&input) + let input_val = rivet_util::serde::json_to_raw_value!(&input) .map_err(WorkflowError::SerializeWorkflowInput)?; let mut branch = self .ctx diff --git a/engine/packages/gasoline/src/ctx/message.rs b/engine/packages/gasoline/src/ctx/message.rs index 181dfe0a28..a6edb2fe32 100644 --- a/engine/packages/gasoline/src/ctx/message.rs +++ b/engine/packages/gasoline/src/ctx/message.rs @@ -96,8 +96,8 @@ impl MessageCtx { let ts = duration_since_epoch.as_millis() as i64; // Serialize the body - let body_buf = - serde_json::to_string(&message_body).map_err(WorkflowError::SerializeMessage)?; + let body_buf = rivet_util::serde::json_to_string!(&message_body) + .map_err(WorkflowError::SerializeMessage)?; let body_buf_len = body_buf.len(); let body_buf = serde_json::value::RawValue::from_string(body_buf) .map_err(WorkflowError::SerializeMessage)?; @@ -111,7 +111,8 @@ impl MessageCtx { ts, body: &body_buf, }; - let message_buf = serde_json::to_vec(&message).map_err(WorkflowError::SerializeMessage)?; + let message_buf = + rivet_util::serde::json_to_vec!(&message).map_err(WorkflowError::SerializeMessage)?; tracing::debug!( %subject, diff --git a/engine/packages/gasoline/src/ctx/workflow.rs b/engine/packages/gasoline/src/ctx/workflow.rs index 7ed12c3e13..c89974fc62 100644 --- a/engine/packages/gasoline/src/ctx/workflow.rs +++ b/engine/packages/gasoline/src/ctx/workflow.rs @@ -302,9 +302,9 @@ impl WorkflowCtx { tracing::debug!("activity success"); // Write output - let input_val = serde_json::value::to_raw_value(input) + let input_val = rivet_util::serde::json_to_raw_value!(input) .map_err(WorkflowError::SerializeActivityInput)?; - let output_val = serde_json::value::to_raw_value(&output) + let output_val = rivet_util::serde::json_to_raw_value!(&output) .map_err(WorkflowError::SerializeActivityOutput)?; tokio::try_join!( @@ -346,7 +346,7 @@ impl WorkflowCtx { tracing::error!(?err, "activity error"); let err_str = err.to_string(); - let input_val = serde_json::value::to_raw_value(input) + let input_val = rivet_util::serde::json_to_raw_value!(input) .map_err(WorkflowError::SerializeActivityInput)?; // Write error (failed state) @@ -380,7 +380,7 @@ impl WorkflowCtx { tracing::debug!("activity timeout"); let err_str = err.to_string(); - let input_val = serde_json::value::to_raw_value(input) + let input_val = rivet_util::serde::json_to_raw_value!(input) .map_err(WorkflowError::SerializeActivityInput)?; self.db diff --git a/engine/packages/gasoline/src/db/mod.rs b/engine/packages/gasoline/src/db/mod.rs index 4f07de3758..5ac2af04c5 100644 --- a/engine/packages/gasoline/src/db/mod.rs +++ b/engine/packages/gasoline/src/db/mod.rs @@ -310,17 +310,19 @@ pub struct WorkflowData { impl WorkflowData { pub fn parse_input(&self) -> WorkflowResult { - serde_json::from_str(self.input.get()).map_err(WorkflowError::DeserializeWorkflowInput) + rivet_util::serde::json_from_str!(self.input.get()) + .map_err(WorkflowError::DeserializeWorkflowInput) } pub fn parse_state(&self) -> WorkflowResult { - serde_json::from_str(self.state.get()).map_err(WorkflowError::DeserializeWorkflowState) + rivet_util::serde::json_from_str!(self.state.get()) + .map_err(WorkflowError::DeserializeWorkflowState) } pub fn parse_output(&self) -> WorkflowResult> { self.output .as_ref() - .map(|x| serde_json::from_str(x.get())) + .map(|x| rivet_util::serde::json_from_str!(x.get())) .transpose() .map_err(WorkflowError::DeserializeWorkflowOutput) } diff --git a/engine/packages/gasoline/src/history/event.rs b/engine/packages/gasoline/src/history/event.rs index c7841f3140..0587f7442e 100644 --- a/engine/packages/gasoline/src/history/event.rs +++ b/engine/packages/gasoline/src/history/event.rs @@ -133,7 +133,7 @@ impl ActivityEvent { pub fn parse_output(&self) -> WorkflowResult> { self.output .as_ref() - .map(|x| serde_json::from_str(x.get())) + .map(|x| rivet_util::serde::json_from_str!(x.get())) .transpose() .map_err(WorkflowError::DeserializeActivityOutput) } @@ -166,13 +166,14 @@ pub struct LoopEvent { impl LoopEvent { pub fn parse_state(&self) -> WorkflowResult { - serde_json::from_str(self.state.get()).map_err(WorkflowError::DeserializeLoopState) + rivet_util::serde::json_from_str!(self.state.get()) + .map_err(WorkflowError::DeserializeLoopState) } pub fn parse_output(&self) -> WorkflowResult> { self.output .as_ref() - .map(|x| serde_json::from_str(x.get())) + .map(|x| rivet_util::serde::json_from_str!(x.get())) .transpose() .map_err(WorkflowError::DeserializeLoopOutput) } diff --git a/engine/packages/gasoline/src/message.rs b/engine/packages/gasoline/src/message.rs index 24eb9e983b..db30cc2b68 100644 --- a/engine/packages/gasoline/src/message.rs +++ b/engine/packages/gasoline/src/message.rs @@ -35,7 +35,7 @@ where wrapper: PubsubMessageWrapper<'_>, ) -> WorkflowResult { // Deserialize the body - let body = serde_json::from_str(wrapper.body.get()) + let body = rivet_util::serde::json_from_str!(wrapper.body.get()) .map_err(WorkflowError::DeserializeMessageBody)?; Ok(PubsubMessage { @@ -51,7 +51,7 @@ where pub(crate) fn deserialize_wrapper<'a>( buf: &'a [u8], ) -> WorkflowResult> { - serde_json::from_slice(buf).map_err(WorkflowError::DeserializeMessage) + rivet_util::serde::json_from_slice!(buf).map_err(WorkflowError::DeserializeMessage) } } diff --git a/engine/packages/gasoline/src/registry.rs b/engine/packages/gasoline/src/registry.rs index 73bad1ffb6..bcc9e77a40 100644 --- a/engine/packages/gasoline/src/registry.rs +++ b/engine/packages/gasoline/src/registry.rs @@ -62,7 +62,7 @@ impl Registry { run: |ctx| { async move { // Deserialize input - let input = serde_json::from_str(ctx.input().get()) + let input = rivet_util::serde::json_from_str!(ctx.input().get()) .map_err(WorkflowError::DeserializeWorkflowInput)?; // Run workflow @@ -79,7 +79,7 @@ impl Registry { }; // Serialize output - let output_val = serde_json::value::to_raw_value(&output) + let output_val = rivet_util::serde::json_to_raw_value!(&output) .map_err(WorkflowError::SerializeWorkflowOutput)?; Ok(output_val) diff --git a/engine/packages/gasoline/src/signal.rs b/engine/packages/gasoline/src/signal.rs index 47f2b27ac6..945a8c5344 100644 --- a/engine/packages/gasoline/src/signal.rs +++ b/engine/packages/gasoline/src/signal.rs @@ -73,7 +73,7 @@ macro_rules! join_signal { if name == <$types as gas::signal::Signal>::NAME { std::result::Result::Ok( Self::$names( - serde_json::from_str(body.get()) + rivet_util::serde::json_from_str!(body.get()) .map_err(WorkflowError::DeserializeSignalBody)? ) ) diff --git a/engine/packages/gasoline/src/workflow.rs b/engine/packages/gasoline/src/workflow.rs index e44c2e921b..bb16ded732 100644 --- a/engine/packages/gasoline/src/workflow.rs +++ b/engine/packages/gasoline/src/workflow.rs @@ -33,7 +33,7 @@ impl<'a, T: DeserializeOwned + Serialize> StateGuard<'a, T> { pub(crate) fn new( guard: MutexGuard<'a, (Box, bool)>, ) -> Result { - let value = serde_json::from_str::(guard.0.get())?; + let value = rivet_util::observe!(serde_json::from_str::(guard.0.get())?); Ok(Self { guard, @@ -60,7 +60,7 @@ impl<'a, T: DeserializeOwned + Serialize> std::ops::DerefMut for StateGuard<'a, impl<'a, T: DeserializeOwned + Serialize> Drop for StateGuard<'a, T> { fn drop(&mut self) { // TODO: Somehow don't panic when committing state back into mutex - self.guard.0 = serde_json::value::to_raw_value(&self.inner).expect("bad state"); + self.guard.0 = rivet_util::serde::json_to_raw_value!(&self.inner).expect("bad state"); } } diff --git a/engine/packages/ups-broadcast/src/lib.rs b/engine/packages/ups-broadcast/src/lib.rs index b3d0c15947..69cf170500 100644 --- a/engine/packages/ups-broadcast/src/lib.rs +++ b/engine/packages/ups-broadcast/src/lib.rs @@ -5,8 +5,6 @@ use universalpubsub::NextOutput; use universalpubsub::PublishOpts; use universalpubsub::Subject; -mod sim; - pub const BROADCAST_TOPIC: &str = "rivet.ups.broadcast"; pub struct BroadcastSubject; @@ -28,7 +26,7 @@ impl Subject for BroadcastSubject { } #[tracing::instrument(skip_all)] -pub async fn start(config: rivet_config::Config, pools: rivet_pools::Pools) -> Result<()> { +pub async fn start(_config: rivet_config::Config, pools: rivet_pools::Pools) -> Result<()> { let ups = pools.ups()?; let mut sub = ups.subscribe(BroadcastSubject).await?; @@ -38,18 +36,6 @@ pub async fn start(config: rivet_config::Config, pools: rivet_pools::Pools) -> R let handle = tokio::spawn(async move { while let Ok(NextOutput::Message(_)) = sub.next().await {} }); - if let Some(sim_config) = sim::Config::from_env()? { - let sim_udb = pools.udb().ok(); - let sim_ups = sim::pubsub_for_sim( - &config, - &ups, - sim_config.force_driver, - sim_config.disable_memory_optimization, - ) - .await?; - sim::spawn(sim_ups, sim_udb, sim_config); - } - loop { if let Err(err) = ups .publish(BroadcastSubject, &[], PublishOpts::broadcast()) diff --git a/engine/packages/ups-broadcast/src/sim.rs b/engine/packages/ups-broadcast/src/sim.rs deleted file mode 100644 index e901b14d8a..0000000000 --- a/engine/packages/ups-broadcast/src/sim.rs +++ /dev/null @@ -1,2509 +0,0 @@ -use std::{ - borrow::Cow, - env, fmt, hint, - sync::{ - atomic::{AtomicU64, Ordering}, - Arc, - }, - time::{Duration, Instant}, -}; - -use anyhow::{bail, Context, Result}; -use futures_util::{FutureExt, StreamExt}; -use gas::prelude::Id; -use rivet_pools::UdbPool; -use serde::Deserialize; -use universaldb::{ - prelude::{PackError, PackResult, TupleDepth, TuplePack, TupleUnpack, VersionstampOffset}, - utils::IsolationLevel::{Serializable, Snapshot}, - RangeOption, Subspace, -}; -use universalpubsub::{NextOutput, PubSub, PublishOpts, Subject, Subscriber}; - -const ENV_PREFIX: &str = "UPS_BROADCAST_SIM"; -const TICK: Duration = Duration::from_millis(10); -const PUBLISH_MAX_IN_FLIGHT: usize = 8_192; -const DEFAULT_TUNE_PATH: &str = "/tmp/ups-broadcast-sim-tune.json"; -const TUNE_POLL_INTERVAL: Duration = Duration::from_secs(1); -const TUNE_SUBJECT: &str = "rivet.ups.broadcast.sim.tune"; -const TUNE_SUBJECT_ROOT: &str = "rivet.ups.broadcast.sim.tune"; -const GATEWAY_MEMBERSHIP_PREFIX: &[u8] = b"rivet/ups-broadcast/sim/gateway-members"; -const GATEWAY_MEMBERSHIP_TX: &str = "ups_broadcast_sim_gateway_membership"; -const UDB_HOT_COUNTER_TX: &str = "ups_broadcast_sim_udb_hot_counter"; -const UDB_READ_SCAN_SEED_TX: &str = "ups_broadcast_sim_udb_read_scan_seed"; -const UDB_READ_SCAN_TX: &str = "ups_broadcast_sim_udb_read_scan"; -const UDB_CONFLICT_SEED_TX: &str = "ups_broadcast_sim_udb_conflict_seed"; -const UDB_CONFLICT_TX: &str = "ups_broadcast_sim_udb_conflict"; -const UDB_READ_SCAN_SEED_BATCH_SIZE: u64 = 500; -const UDB_CONFLICT_SEED_BATCH_SIZE: u64 = 500; -const READ_SCAN_KEY_ROOT: usize = 1; -const CONFLICT_KEY_ROOT: usize = 2; -static SUBJECT_SEQ: AtomicU64 = AtomicU64::new(0); -static HOT_COUNTER_SEQ: AtomicU64 = AtomicU64::new(0); -static READ_SCAN_SEQ: AtomicU64 = AtomicU64::new(0); -static CONFLICT_SEQ: AtomicU64 = AtomicU64::new(0); - -pub struct Config { - pub force_driver: bool, - pub disable_memory_optimization: bool, - tune_path: Option, - gateway_subjects: usize, - gateway_subscribers: usize, - gateway_publish_rps: f64, - gateway_payload_bytes: usize, - gateway_work_delay_ms: u64, - gateway_work_cpu_us: u64, - gateway_spread_replicas: usize, - gateway_spread_member_ttl_ms: u64, - envoy_subjects: usize, - envoy_responders: usize, - envoy_queue_group: Option, - envoy_request_unknown_root: bool, - envoy_request_rps: f64, - envoy_request_payload_bytes: usize, - envoy_request_timeout_ms: u64, - envoy_request_max_in_flight: usize, - envoy_work_delay_ms: u64, - envoy_work_cpu_us: u64, - envoy_eviction_subscribers: usize, - envoy_eviction_broadcast_rps: f64, - envoy_eviction_work_delay_ms: u64, - envoy_eviction_work_cpu_us: u64, - worker_bump_subscribers: usize, - worker_bump_broadcast_rps: f64, - worker_bump_work_delay_ms: u64, - worker_bump_work_cpu_us: u64, - serverless_subscribers: usize, - serverless_publish_rps: f64, - serverless_payload_bytes: usize, - serverless_work_delay_ms: u64, - serverless_work_cpu_us: u64, - cache_purge_subscribers: usize, - cache_purge_broadcast_rps: f64, - cache_purge_payload_bytes: usize, - cache_purge_work_delay_ms: u64, - cache_purge_work_cpu_us: u64, - tracing_config_subscribers: usize, - tracing_config_broadcast_rps: f64, - tracing_config_payload_bytes: usize, - tracing_config_work_delay_ms: u64, - tracing_config_work_cpu_us: u64, - route_stopped_subscribers: usize, - route_churn_rps: f64, - route_ephemeral_hold_ms: u64, - route_stopped_hold_ms: u64, - route_max_in_flight: usize, - route_work_delay_ms: u64, - route_work_cpu_us: u64, - workflow_signal_churn_rps: f64, - workflow_signal_hold_ms: u64, - workflow_signal_publish_rps: f64, - workflow_signal_work_delay_ms: u64, - workflow_signal_work_cpu_us: u64, - workflow_complete_publish_rps: f64, - udb_hot_counter_rps: f64, - udb_hot_counter_max_in_flight: usize, - udb_hot_counter_namespace_id: Id, - udb_hot_counter_actor_name: String, - udb_read_scan_rps: f64, - udb_read_scan_max_in_flight: usize, - udb_read_scan_seed_keys: u64, - udb_read_scan_keys_per_tx: usize, - udb_read_scan_value_bytes: usize, - udb_read_scan_unpack_keys: bool, - udb_conflict_rps: f64, - udb_conflict_max_in_flight: usize, - udb_conflict_keys: u64, -} - -impl Config { - pub fn from_env() -> Result> { - if !env_bool("ENABLED", false)? { - return Ok(None); - } - - let profile = env_string("PROFILE").with_context(|| { - format!("{ENV_PREFIX}_PROFILE must be set explicitly when {ENV_PREFIX}_ENABLED=true") - })?; - let mut config = match profile.as_str() { - "custom" => Self::custom(), - "staging_peak" => Self::staging_peak(), - other => bail!("unknown {ENV_PREFIX}_PROFILE: {other}"), - }; - - config.force_driver = env_bool("FORCE_DRIVER", config.force_driver)?; - config.disable_memory_optimization = env_bool( - "DISABLE_MEMORY_OPTIMIZATION", - config.disable_memory_optimization, - )?; - config.tune_path = env_string("TUNE_PATH") - .map(|x| if x.is_empty() { None } else { Some(x) }) - .unwrap_or(config.tune_path); - config.gateway_subjects = env_usize("GATEWAY_SUBJECTS", config.gateway_subjects)?; - config.gateway_subscribers = env_usize("GATEWAY_SUBSCRIBERS", config.gateway_subscribers)?; - config.gateway_publish_rps = env_f64("GATEWAY_PUBLISH_RPS", config.gateway_publish_rps)?; - config.gateway_payload_bytes = - env_usize("GATEWAY_PAYLOAD_BYTES", config.gateway_payload_bytes)?; - config.gateway_work_delay_ms = - env_u64("GATEWAY_WORK_DELAY_MS", config.gateway_work_delay_ms)?; - config.gateway_work_cpu_us = env_u64("GATEWAY_WORK_CPU_US", config.gateway_work_cpu_us)?; - config.gateway_spread_replicas = - env_usize("GATEWAY_SPREAD_REPLICAS", config.gateway_spread_replicas)?; - config.gateway_spread_member_ttl_ms = env_u64( - "GATEWAY_SPREAD_MEMBER_TTL_MS", - config.gateway_spread_member_ttl_ms, - )?; - config.envoy_subjects = env_usize("ENVOY_SUBJECTS", config.envoy_subjects)?; - config.envoy_responders = env_usize("ENVOY_RESPONDERS", config.envoy_responders)?; - config.envoy_queue_group = env_string("ENVOY_QUEUE_GROUP") - .map(|x| if x.is_empty() { None } else { Some(x) }) - .unwrap_or(config.envoy_queue_group); - config.envoy_request_unknown_root = env_bool( - "ENVOY_REQUEST_UNKNOWN_ROOT", - config.envoy_request_unknown_root, - )?; - config.envoy_request_rps = env_f64("ENVOY_REQUEST_RPS", config.envoy_request_rps)?; - config.envoy_request_payload_bytes = env_usize( - "ENVOY_REQUEST_PAYLOAD_BYTES", - config.envoy_request_payload_bytes, - )?; - config.envoy_request_timeout_ms = - env_u64("ENVOY_REQUEST_TIMEOUT_MS", config.envoy_request_timeout_ms)?; - config.envoy_request_max_in_flight = env_usize( - "ENVOY_REQUEST_MAX_IN_FLIGHT", - config.envoy_request_max_in_flight, - )?; - config.envoy_work_delay_ms = env_u64("ENVOY_WORK_DELAY_MS", config.envoy_work_delay_ms)?; - config.envoy_work_cpu_us = env_u64("ENVOY_WORK_CPU_US", config.envoy_work_cpu_us)?; - config.envoy_eviction_subscribers = env_usize( - "ENVOY_EVICTION_SUBSCRIBERS", - config.envoy_eviction_subscribers, - )?; - config.envoy_eviction_broadcast_rps = env_f64( - "ENVOY_EVICTION_BROADCAST_RPS", - config.envoy_eviction_broadcast_rps, - )?; - config.envoy_eviction_work_delay_ms = env_u64( - "ENVOY_EVICTION_WORK_DELAY_MS", - config.envoy_eviction_work_delay_ms, - )?; - config.envoy_eviction_work_cpu_us = env_u64( - "ENVOY_EVICTION_WORK_CPU_US", - config.envoy_eviction_work_cpu_us, - )?; - config.worker_bump_subscribers = - env_usize("WORKER_BUMP_SUBSCRIBERS", config.worker_bump_subscribers)?; - config.worker_bump_broadcast_rps = env_f64( - "WORKER_BUMP_BROADCAST_RPS", - config.worker_bump_broadcast_rps, - )?; - config.worker_bump_work_delay_ms = env_u64( - "WORKER_BUMP_WORK_DELAY_MS", - config.worker_bump_work_delay_ms, - )?; - config.worker_bump_work_cpu_us = - env_u64("WORKER_BUMP_WORK_CPU_US", config.worker_bump_work_cpu_us)?; - config.serverless_subscribers = - env_usize("SERVERLESS_SUBSCRIBERS", config.serverless_subscribers)?; - config.serverless_publish_rps = - env_f64("SERVERLESS_PUBLISH_RPS", config.serverless_publish_rps)?; - config.serverless_payload_bytes = - env_usize("SERVERLESS_PAYLOAD_BYTES", config.serverless_payload_bytes)?; - config.serverless_work_delay_ms = - env_u64("SERVERLESS_WORK_DELAY_MS", config.serverless_work_delay_ms)?; - config.serverless_work_cpu_us = - env_u64("SERVERLESS_WORK_CPU_US", config.serverless_work_cpu_us)?; - config.cache_purge_subscribers = - env_usize("CACHE_PURGE_SUBSCRIBERS", config.cache_purge_subscribers)?; - config.cache_purge_broadcast_rps = env_f64( - "CACHE_PURGE_BROADCAST_RPS", - config.cache_purge_broadcast_rps, - )?; - config.cache_purge_payload_bytes = env_usize( - "CACHE_PURGE_PAYLOAD_BYTES", - config.cache_purge_payload_bytes, - )?; - config.cache_purge_work_delay_ms = env_u64( - "CACHE_PURGE_WORK_DELAY_MS", - config.cache_purge_work_delay_ms, - )?; - config.cache_purge_work_cpu_us = - env_u64("CACHE_PURGE_WORK_CPU_US", config.cache_purge_work_cpu_us)?; - config.tracing_config_subscribers = env_usize( - "TRACING_CONFIG_SUBSCRIBERS", - config.tracing_config_subscribers, - )?; - config.tracing_config_broadcast_rps = env_f64( - "TRACING_CONFIG_BROADCAST_RPS", - config.tracing_config_broadcast_rps, - )?; - config.tracing_config_payload_bytes = env_usize( - "TRACING_CONFIG_PAYLOAD_BYTES", - config.tracing_config_payload_bytes, - )?; - config.tracing_config_work_delay_ms = env_u64( - "TRACING_CONFIG_WORK_DELAY_MS", - config.tracing_config_work_delay_ms, - )?; - config.tracing_config_work_cpu_us = env_u64( - "TRACING_CONFIG_WORK_CPU_US", - config.tracing_config_work_cpu_us, - )?; - config.route_stopped_subscribers = env_usize( - "ROUTE_STOPPED_SUBSCRIBERS", - config.route_stopped_subscribers, - )?; - config.route_churn_rps = env_f64("ROUTE_CHURN_RPS", config.route_churn_rps)?; - config.route_ephemeral_hold_ms = - env_u64("ROUTE_EPHEMERAL_HOLD_MS", config.route_ephemeral_hold_ms)?; - config.route_stopped_hold_ms = - env_u64("ROUTE_STOPPED_HOLD_MS", config.route_stopped_hold_ms)?; - config.route_max_in_flight = env_usize("ROUTE_MAX_IN_FLIGHT", config.route_max_in_flight)?; - config.route_work_delay_ms = env_u64("ROUTE_WORK_DELAY_MS", config.route_work_delay_ms)?; - config.route_work_cpu_us = env_u64("ROUTE_WORK_CPU_US", config.route_work_cpu_us)?; - config.workflow_signal_churn_rps = env_f64( - "WORKFLOW_SIGNAL_CHURN_RPS", - config.workflow_signal_churn_rps, - )?; - config.workflow_signal_hold_ms = - env_u64("WORKFLOW_SIGNAL_HOLD_MS", config.workflow_signal_hold_ms)?; - config.workflow_signal_publish_rps = env_f64( - "WORKFLOW_SIGNAL_PUBLISH_RPS", - config.workflow_signal_publish_rps, - )?; - config.workflow_signal_work_delay_ms = env_u64( - "WORKFLOW_SIGNAL_WORK_DELAY_MS", - config.workflow_signal_work_delay_ms, - )?; - config.workflow_signal_work_cpu_us = env_u64( - "WORKFLOW_SIGNAL_WORK_CPU_US", - config.workflow_signal_work_cpu_us, - )?; - config.workflow_complete_publish_rps = env_f64( - "WORKFLOW_COMPLETE_PUBLISH_RPS", - config.workflow_complete_publish_rps, - )?; - config.udb_hot_counter_rps = env_f64("UDB_HOT_COUNTER_RPS", config.udb_hot_counter_rps)?; - config.udb_hot_counter_max_in_flight = env_usize( - "UDB_HOT_COUNTER_MAX_IN_FLIGHT", - config.udb_hot_counter_max_in_flight, - )?; - config.udb_hot_counter_namespace_id = env_id( - "UDB_HOT_COUNTER_NAMESPACE_ID", - config.udb_hot_counter_namespace_id, - )?; - config.udb_hot_counter_actor_name = - env_string("UDB_HOT_COUNTER_ACTOR_NAME").unwrap_or(config.udb_hot_counter_actor_name); - config.udb_read_scan_rps = env_f64("UDB_READ_SCAN_RPS", config.udb_read_scan_rps)?; - config.udb_read_scan_max_in_flight = env_usize( - "UDB_READ_SCAN_MAX_IN_FLIGHT", - config.udb_read_scan_max_in_flight, - )?; - config.udb_read_scan_seed_keys = - env_u64("UDB_READ_SCAN_SEED_KEYS", config.udb_read_scan_seed_keys)?; - config.udb_read_scan_keys_per_tx = env_usize( - "UDB_READ_SCAN_KEYS_PER_TX", - config.udb_read_scan_keys_per_tx, - )?; - config.udb_read_scan_value_bytes = env_usize( - "UDB_READ_SCAN_VALUE_BYTES", - config.udb_read_scan_value_bytes, - )?; - config.udb_read_scan_unpack_keys = env_bool( - "UDB_READ_SCAN_UNPACK_KEYS", - config.udb_read_scan_unpack_keys, - )?; - config.udb_conflict_rps = env_f64("UDB_CONFLICT_RPS", config.udb_conflict_rps)?; - config.udb_conflict_max_in_flight = env_usize( - "UDB_CONFLICT_MAX_IN_FLIGHT", - config.udb_conflict_max_in_flight, - )?; - config.udb_conflict_keys = env_u64("UDB_CONFLICT_KEYS", config.udb_conflict_keys)?; - - validate_rate("GATEWAY_PUBLISH_RPS", config.gateway_publish_rps)?; - validate_rate("ENVOY_REQUEST_RPS", config.envoy_request_rps)?; - validate_rate( - "ENVOY_EVICTION_BROADCAST_RPS", - config.envoy_eviction_broadcast_rps, - )?; - validate_rate( - "WORKER_BUMP_BROADCAST_RPS", - config.worker_bump_broadcast_rps, - )?; - validate_rate("SERVERLESS_PUBLISH_RPS", config.serverless_publish_rps)?; - validate_rate( - "CACHE_PURGE_BROADCAST_RPS", - config.cache_purge_broadcast_rps, - )?; - validate_rate( - "TRACING_CONFIG_BROADCAST_RPS", - config.tracing_config_broadcast_rps, - )?; - validate_rate("ROUTE_CHURN_RPS", config.route_churn_rps)?; - validate_rate( - "WORKFLOW_SIGNAL_CHURN_RPS", - config.workflow_signal_churn_rps, - )?; - validate_rate( - "WORKFLOW_SIGNAL_PUBLISH_RPS", - config.workflow_signal_publish_rps, - )?; - validate_rate( - "WORKFLOW_COMPLETE_PUBLISH_RPS", - config.workflow_complete_publish_rps, - )?; - validate_rate("UDB_HOT_COUNTER_RPS", config.udb_hot_counter_rps)?; - validate_rate("UDB_READ_SCAN_RPS", config.udb_read_scan_rps)?; - validate_rate("UDB_CONFLICT_RPS", config.udb_conflict_rps)?; - - Ok(Some(config)) - } - - fn custom() -> Self { - Self { - force_driver: true, - disable_memory_optimization: false, - tune_path: Some(DEFAULT_TUNE_PATH.to_string()), - gateway_subjects: 0, - gateway_subscribers: 0, - gateway_publish_rps: 0.0, - gateway_payload_bytes: 192, - gateway_work_delay_ms: 0, - gateway_work_cpu_us: 0, - gateway_spread_replicas: 0, - gateway_spread_member_ttl_ms: 15_000, - envoy_subjects: 0, - envoy_responders: 0, - envoy_queue_group: None, - envoy_request_unknown_root: true, - envoy_request_rps: 0.0, - envoy_request_payload_bytes: 64, - envoy_request_timeout_ms: 30_000, - envoy_request_max_in_flight: 8_192, - envoy_work_delay_ms: 0, - envoy_work_cpu_us: 0, - envoy_eviction_subscribers: 0, - envoy_eviction_broadcast_rps: 0.0, - envoy_eviction_work_delay_ms: 0, - envoy_eviction_work_cpu_us: 0, - worker_bump_subscribers: 0, - worker_bump_broadcast_rps: 0.0, - worker_bump_work_delay_ms: 0, - worker_bump_work_cpu_us: 0, - serverless_subscribers: 0, - serverless_publish_rps: 0.0, - serverless_payload_bytes: 256, - serverless_work_delay_ms: 0, - serverless_work_cpu_us: 0, - cache_purge_subscribers: 0, - cache_purge_broadcast_rps: 0.0, - cache_purge_payload_bytes: 128, - cache_purge_work_delay_ms: 0, - cache_purge_work_cpu_us: 0, - tracing_config_subscribers: 0, - tracing_config_broadcast_rps: 0.0, - tracing_config_payload_bytes: 128, - tracing_config_work_delay_ms: 0, - tracing_config_work_cpu_us: 0, - route_stopped_subscribers: 0, - route_churn_rps: 0.0, - route_ephemeral_hold_ms: 25, - route_stopped_hold_ms: 7_500, - route_max_in_flight: 4_096, - route_work_delay_ms: 0, - route_work_cpu_us: 0, - workflow_signal_churn_rps: 0.0, - workflow_signal_hold_ms: 3_000, - workflow_signal_publish_rps: 0.0, - workflow_signal_work_delay_ms: 0, - workflow_signal_work_cpu_us: 0, - workflow_complete_publish_rps: 0.0, - udb_hot_counter_rps: 0.0, - udb_hot_counter_max_in_flight: 1_024, - udb_hot_counter_namespace_id: Id::nil(), - udb_hot_counter_actor_name: "sim-hot-namespace".to_string(), - udb_read_scan_rps: 0.0, - udb_read_scan_max_in_flight: 512, - udb_read_scan_seed_keys: 50_000, - udb_read_scan_keys_per_tx: 50, - udb_read_scan_value_bytes: 128, - udb_read_scan_unpack_keys: true, - udb_conflict_rps: 0.0, - udb_conflict_max_in_flight: 512, - udb_conflict_keys: 32, - } - } - - fn staging_peak() -> Self { - Self { - gateway_subjects: 20, - gateway_subscribers: 20, - gateway_publish_rps: 5_016.0, - gateway_payload_bytes: 192, - gateway_work_cpu_us: 100, - gateway_spread_replicas: 10, - envoy_subjects: 8, - envoy_responders: 8, - envoy_queue_group: Some("rivet-ups-broadcast-sim-envoy".to_string()), - envoy_request_unknown_root: true, - envoy_request_rps: 2_094.0, - envoy_request_payload_bytes: 64, - envoy_work_delay_ms: 10, - envoy_work_cpu_us: 250, - envoy_eviction_subscribers: 72, - envoy_eviction_work_delay_ms: 1, - worker_bump_subscribers: 10, - worker_bump_broadcast_rps: 141.0, - worker_bump_work_delay_ms: 225, - worker_bump_work_cpu_us: 2_000, - serverless_subscribers: 10, - serverless_work_delay_ms: 2, - serverless_work_cpu_us: 250, - cache_purge_subscribers: 20, - cache_purge_broadcast_rps: 0.62, - cache_purge_work_delay_ms: 2, - cache_purge_work_cpu_us: 250, - tracing_config_subscribers: 20, - tracing_config_work_delay_ms: 1, - route_stopped_subscribers: 2_100, - route_churn_rps: 6.5, - route_work_delay_ms: 5, - route_work_cpu_us: 500, - workflow_signal_churn_rps: 281.0, - workflow_signal_hold_ms: 2_900, - workflow_signal_publish_rps: 0.76, - workflow_signal_work_delay_ms: 15, - workflow_signal_work_cpu_us: 500, - workflow_complete_publish_rps: 0.04, - ..Self::custom() - } - } -} - -#[derive(Clone)] -struct Rate { - value: Arc, -} - -impl Rate { - fn new(value: f64) -> Self { - Self { - value: Arc::new(AtomicU64::new(value.to_bits())), - } - } - - fn load(&self) -> f64 { - f64::from_bits(self.value.load(Ordering::Relaxed)) - } - - fn store(&self, value: f64) { - self.value.store(value.to_bits(), Ordering::Relaxed); - } -} - -struct Rates { - gateway_publish_rps: Rate, - envoy_request_rps: Rate, - envoy_eviction_broadcast_rps: Rate, - worker_bump_broadcast_rps: Rate, - serverless_publish_rps: Rate, - cache_purge_broadcast_rps: Rate, - tracing_config_broadcast_rps: Rate, - route_churn_rps: Rate, - workflow_signal_churn_rps: Rate, - workflow_signal_publish_rps: Rate, - workflow_complete_publish_rps: Rate, - udb_hot_counter_rps: Rate, - udb_read_scan_rps: Rate, - udb_conflict_rps: Rate, -} - -impl Rates { - fn new(config: &Config) -> Self { - Self { - gateway_publish_rps: Rate::new(config.gateway_publish_rps), - envoy_request_rps: Rate::new(config.envoy_request_rps), - envoy_eviction_broadcast_rps: Rate::new(config.envoy_eviction_broadcast_rps), - worker_bump_broadcast_rps: Rate::new(config.worker_bump_broadcast_rps), - serverless_publish_rps: Rate::new(config.serverless_publish_rps), - cache_purge_broadcast_rps: Rate::new(config.cache_purge_broadcast_rps), - tracing_config_broadcast_rps: Rate::new(config.tracing_config_broadcast_rps), - route_churn_rps: Rate::new(config.route_churn_rps), - workflow_signal_churn_rps: Rate::new(config.workflow_signal_churn_rps), - workflow_signal_publish_rps: Rate::new(config.workflow_signal_publish_rps), - workflow_complete_publish_rps: Rate::new(config.workflow_complete_publish_rps), - udb_hot_counter_rps: Rate::new(config.udb_hot_counter_rps), - udb_read_scan_rps: Rate::new(config.udb_read_scan_rps), - udb_conflict_rps: Rate::new(config.udb_conflict_rps), - } - } -} - -#[derive(Clone)] -struct Workload { - delay_ms: Arc, - cpu_us: Arc, -} - -impl Workload { - fn new(delay_ms: u64, cpu_us: u64) -> Self { - Self { - delay_ms: Arc::new(AtomicU64::new(delay_ms)), - cpu_us: Arc::new(AtomicU64::new(cpu_us)), - } - } - - fn store_delay_ms(&self, value: u64) { - self.delay_ms.store(value, Ordering::Relaxed); - } - - fn store_cpu_us(&self, value: u64) { - self.cpu_us.store(value, Ordering::Relaxed); - } - - async fn run(&self) { - let cpu_us = self.cpu_us.load(Ordering::Relaxed); - if cpu_us > 0 { - burn_cpu(Duration::from_micros(cpu_us)); - } - - let delay_ms = self.delay_ms.load(Ordering::Relaxed); - if delay_ms > 0 { - tokio::time::sleep(Duration::from_millis(delay_ms)).await; - } - } -} - -struct Workloads { - gateway: Workload, - envoy: Workload, - envoy_eviction: Workload, - worker_bump: Workload, - serverless: Workload, - cache_purge: Workload, - tracing_config: Workload, - route: Workload, - workflow_signal: Workload, -} - -impl Workloads { - fn new(config: &Config) -> Self { - Self { - gateway: Workload::new(config.gateway_work_delay_ms, config.gateway_work_cpu_us), - envoy: Workload::new(config.envoy_work_delay_ms, config.envoy_work_cpu_us), - envoy_eviction: Workload::new( - config.envoy_eviction_work_delay_ms, - config.envoy_eviction_work_cpu_us, - ), - worker_bump: Workload::new( - config.worker_bump_work_delay_ms, - config.worker_bump_work_cpu_us, - ), - serverless: Workload::new( - config.serverless_work_delay_ms, - config.serverless_work_cpu_us, - ), - cache_purge: Workload::new( - config.cache_purge_work_delay_ms, - config.cache_purge_work_cpu_us, - ), - tracing_config: Workload::new( - config.tracing_config_work_delay_ms, - config.tracing_config_work_cpu_us, - ), - route: Workload::new(config.route_work_delay_ms, config.route_work_cpu_us), - workflow_signal: Workload::new( - config.workflow_signal_work_delay_ms, - config.workflow_signal_work_cpu_us, - ), - } - } -} - -#[derive(Debug, Deserialize)] -#[serde(deny_unknown_fields)] -struct TunePatch { - gateway_publish_rps: Option, - gateway_work_delay_ms: Option, - gateway_work_cpu_us: Option, - envoy_request_rps: Option, - envoy_work_delay_ms: Option, - envoy_work_cpu_us: Option, - envoy_eviction_broadcast_rps: Option, - envoy_eviction_work_delay_ms: Option, - envoy_eviction_work_cpu_us: Option, - worker_bump_broadcast_rps: Option, - worker_bump_work_delay_ms: Option, - worker_bump_work_cpu_us: Option, - serverless_publish_rps: Option, - serverless_work_delay_ms: Option, - serverless_work_cpu_us: Option, - cache_purge_broadcast_rps: Option, - cache_purge_work_delay_ms: Option, - cache_purge_work_cpu_us: Option, - tracing_config_broadcast_rps: Option, - tracing_config_work_delay_ms: Option, - tracing_config_work_cpu_us: Option, - route_churn_rps: Option, - route_work_delay_ms: Option, - route_work_cpu_us: Option, - workflow_signal_churn_rps: Option, - workflow_signal_publish_rps: Option, - workflow_signal_work_delay_ms: Option, - workflow_signal_work_cpu_us: Option, - workflow_complete_publish_rps: Option, - udb_hot_counter_rps: Option, - udb_read_scan_rps: Option, - udb_conflict_rps: Option, -} - -impl TunePatch { - fn apply(&self, rates: &Rates, workloads: &Workloads) -> Result<()> { - apply_rate( - "gateway_publish_rps", - self.gateway_publish_rps, - &rates.gateway_publish_rps, - )?; - apply_workload( - "gateway", - self.gateway_work_delay_ms, - self.gateway_work_cpu_us, - &workloads.gateway, - ); - apply_rate( - "envoy_request_rps", - self.envoy_request_rps, - &rates.envoy_request_rps, - )?; - apply_workload( - "envoy", - self.envoy_work_delay_ms, - self.envoy_work_cpu_us, - &workloads.envoy, - ); - apply_rate( - "envoy_eviction_broadcast_rps", - self.envoy_eviction_broadcast_rps, - &rates.envoy_eviction_broadcast_rps, - )?; - apply_workload( - "envoy_eviction", - self.envoy_eviction_work_delay_ms, - self.envoy_eviction_work_cpu_us, - &workloads.envoy_eviction, - ); - apply_rate( - "worker_bump_broadcast_rps", - self.worker_bump_broadcast_rps, - &rates.worker_bump_broadcast_rps, - )?; - apply_workload( - "worker_bump", - self.worker_bump_work_delay_ms, - self.worker_bump_work_cpu_us, - &workloads.worker_bump, - ); - apply_rate( - "serverless_publish_rps", - self.serverless_publish_rps, - &rates.serverless_publish_rps, - )?; - apply_workload( - "serverless", - self.serverless_work_delay_ms, - self.serverless_work_cpu_us, - &workloads.serverless, - ); - apply_rate( - "cache_purge_broadcast_rps", - self.cache_purge_broadcast_rps, - &rates.cache_purge_broadcast_rps, - )?; - apply_workload( - "cache_purge", - self.cache_purge_work_delay_ms, - self.cache_purge_work_cpu_us, - &workloads.cache_purge, - ); - apply_rate( - "tracing_config_broadcast_rps", - self.tracing_config_broadcast_rps, - &rates.tracing_config_broadcast_rps, - )?; - apply_workload( - "tracing_config", - self.tracing_config_work_delay_ms, - self.tracing_config_work_cpu_us, - &workloads.tracing_config, - ); - apply_rate( - "route_churn_rps", - self.route_churn_rps, - &rates.route_churn_rps, - )?; - apply_workload( - "route", - self.route_work_delay_ms, - self.route_work_cpu_us, - &workloads.route, - ); - apply_rate( - "workflow_signal_churn_rps", - self.workflow_signal_churn_rps, - &rates.workflow_signal_churn_rps, - )?; - apply_rate( - "workflow_signal_publish_rps", - self.workflow_signal_publish_rps, - &rates.workflow_signal_publish_rps, - )?; - apply_workload( - "workflow_signal", - self.workflow_signal_work_delay_ms, - self.workflow_signal_work_cpu_us, - &workloads.workflow_signal, - ); - apply_rate( - "workflow_complete_publish_rps", - self.workflow_complete_publish_rps, - &rates.workflow_complete_publish_rps, - )?; - apply_rate( - "udb_hot_counter_rps", - self.udb_hot_counter_rps, - &rates.udb_hot_counter_rps, - )?; - apply_rate( - "udb_read_scan_rps", - self.udb_read_scan_rps, - &rates.udb_read_scan_rps, - )?; - apply_rate( - "udb_conflict_rps", - self.udb_conflict_rps, - &rates.udb_conflict_rps, - )?; - - Ok(()) - } -} - -fn apply_rate(name: &'static str, value: Option, rate: &Rate) -> Result<()> { - if let Some(value) = value { - validate_rate(name, value)?; - rate.store(value); - } - - Ok(()) -} - -fn apply_workload( - name: &'static str, - delay_ms: Option, - cpu_us: Option, - workload: &Workload, -) { - if let Some(delay_ms) = delay_ms { - workload.store_delay_ms(delay_ms); - tracing::info!(name, delay_ms, "updated UPS simulation workload delay"); - } - if let Some(cpu_us) = cpu_us { - workload.store_cpu_us(cpu_us); - tracing::info!(name, cpu_us, "updated UPS simulation workload CPU"); - } -} - -pub async fn pubsub_for_sim( - config: &rivet_config::Config, - existing: &PubSub, - force_driver: bool, - disable_memory_optimization: bool, -) -> Result { - if !force_driver { - return Ok(existing.clone()); - } - - let mut root = (**config).clone(); - let mut pubsub = config.pubsub().clone(); - match &mut pubsub { - rivet_config::config::PubSub::Nats(nats) => { - nats.disable_memory_optimization = disable_memory_optimization; - } - rivet_config::config::PubSub::PostgresNotify(postgres) => { - postgres.disable_memory_optimization = disable_memory_optimization; - } - rivet_config::config::PubSub::Memory(memory) => { - memory.disable_memory_optimization = disable_memory_optimization; - } - } - root.pubsub = Some(pubsub); - - let sim_config = rivet_config::Config::from_root(root); - rivet_pools::db::ups::setup(&sim_config, "rivet-ups-broadcast-sim") - .await - .context("failed to create UPS simulation pubsub") -} - -pub fn spawn(ups: PubSub, udb: Option, config: Config) { - let rates = Arc::new(Rates::new(&config)); - let workloads = Arc::new(Workloads::new(&config)); - let workflow_signal_subjects = ActiveSubjects::default(); - - tracing::info!( - force_driver = config.force_driver, - disable_memory_optimization = config.disable_memory_optimization, - tune_path = ?config.tune_path, - gateway_publish_rps = config.gateway_publish_rps, - gateway_spread_replicas = config.gateway_spread_replicas, - envoy_request_rps = config.envoy_request_rps, - worker_bump_broadcast_rps = config.worker_bump_broadcast_rps, - worker_bump_work_delay_ms = config.worker_bump_work_delay_ms, - envoy_work_delay_ms = config.envoy_work_delay_ms, - workflow_signal_work_delay_ms = config.workflow_signal_work_delay_ms, - workflow_signal_churn_rps = config.workflow_signal_churn_rps, - udb_hot_counter_rps = config.udb_hot_counter_rps, - udb_read_scan_rps = config.udb_read_scan_rps, - udb_read_scan_seed_keys = config.udb_read_scan_seed_keys, - udb_read_scan_keys_per_tx = config.udb_read_scan_keys_per_tx, - udb_read_scan_unpack_keys = config.udb_read_scan_unpack_keys, - udb_conflict_rps = config.udb_conflict_rps, - udb_conflict_keys = config.udb_conflict_keys, - "starting UPS broadcast traffic simulator" - ); - - spawn_tuner( - ups.clone(), - rates.clone(), - workloads.clone(), - config.tune_path.clone(), - ); - - spawn_gateway_subscribers( - ups.clone(), - udb.clone(), - config.gateway_subjects, - config.gateway_subscribers, - config.gateway_spread_replicas, - Duration::from_millis(config.gateway_spread_member_ttl_ms), - workloads.gateway.clone(), - ); - let envoy_queue_group = config.envoy_queue_group.clone().map(Arc::new); - spawn_subject_subscribers( - ups.clone(), - "envoy", - "pegboard.envoy", - "pegboard.envoy.sim", - config.envoy_subjects, - config.envoy_responders, - envoy_queue_group, - Some(Arc::new(Vec::new())), - workloads.envoy.clone(), - ); - spawn_subject_subscribers( - ups.clone(), - "envoy eviction", - "pegboard.envoy.eviction", - "pegboard.envoy.eviction.sim", - config.envoy_subjects, - config.envoy_eviction_subscribers, - None, - None, - workloads.envoy_eviction.clone(), - ); - spawn_worker_bump_subscribers( - ups.clone(), - SimSubject::new("gasoline.worker.bump", "gasoline.worker.bump"), - config.worker_bump_subscribers, - workloads.worker_bump.clone(), - ); - spawn_same_subject_subscribers( - ups.clone(), - "serverless outbound", - SimSubject::new( - "pegboard.serverless.outbound", - "pegboard.serverless.outbound", - ), - config.serverless_subscribers, - None, - workloads.serverless.clone(), - ); - spawn_same_subject_subscribers( - ups.clone(), - "cache purge", - SimSubject::new("rivet.cache.purge", "rivet.cache.purge"), - config.cache_purge_subscribers, - None, - workloads.cache_purge.clone(), - ); - spawn_same_subject_subscribers( - ups.clone(), - "tracing config", - SimSubject::new("rivet.debug.tracing.config", "rivet.debug.tracing.config"), - config.tracing_config_subscribers, - None, - workloads.tracing_config.clone(), - ); - spawn_subject_subscribers( - ups.clone(), - "route stopped", - "gasoline.msg.pegboard_actor2_stopped", - "gasoline.msg.pegboard_actor2_stopped:actor", - config.route_stopped_subscribers, - config.route_stopped_subscribers, - None, - None, - Workload::new(0, 0), - ); - - let gateway_subjects = subjects( - "pegboard.gateway", - "pegboard.gateway.sim", - config.gateway_subjects, - ); - spawn_publish_rate( - ups.clone(), - "gateway publish", - gateway_subjects, - PublishOpts::one(), - rates.gateway_publish_rps.clone(), - Arc::new(payload(config.gateway_payload_bytes)), - ); - - if config.envoy_request_unknown_root { - spawn_request_rate( - ups.clone(), - raw_subjects("pegboard.envoy.sim", config.envoy_subjects), - rates.envoy_request_rps.clone(), - Arc::new(payload(config.envoy_request_payload_bytes)), - Duration::from_millis(config.envoy_request_timeout_ms), - config.envoy_request_max_in_flight, - ); - } else { - spawn_request_rate( - ups.clone(), - subjects( - "pegboard.envoy", - "pegboard.envoy.sim", - config.envoy_subjects, - ), - rates.envoy_request_rps.clone(), - Arc::new(payload(config.envoy_request_payload_bytes)), - Duration::from_millis(config.envoy_request_timeout_ms), - config.envoy_request_max_in_flight, - ); - } - - let envoy_eviction_subjects = subjects( - "pegboard.envoy.eviction", - "pegboard.envoy.eviction.sim", - config.envoy_subjects, - ); - spawn_publish_rate( - ups.clone(), - "envoy eviction broadcast", - envoy_eviction_subjects, - PublishOpts::broadcast(), - rates.envoy_eviction_broadcast_rps.clone(), - Arc::new(Vec::new()), - ); - spawn_publish_rate( - ups.clone(), - "worker bump broadcast", - vec![SimSubject::new( - "gasoline.worker.bump", - "gasoline.worker.bump", - )], - PublishOpts::broadcast(), - rates.worker_bump_broadcast_rps.clone(), - Arc::new(Vec::new()), - ); - spawn_publish_rate( - ups.clone(), - "serverless publish", - vec![SimSubject::new( - "pegboard.serverless.outbound", - "pegboard.serverless.outbound", - )], - PublishOpts::one(), - rates.serverless_publish_rps.clone(), - Arc::new(payload(config.serverless_payload_bytes)), - ); - spawn_publish_rate( - ups.clone(), - "cache purge broadcast", - vec![SimSubject::new("rivet.cache.purge", "rivet.cache.purge")], - PublishOpts::broadcast(), - rates.cache_purge_broadcast_rps.clone(), - Arc::new(payload(config.cache_purge_payload_bytes)), - ); - spawn_publish_rate( - ups.clone(), - "tracing config broadcast", - vec![SimSubject::new( - "rivet.debug.tracing.config", - "rivet.debug.tracing.config", - )], - PublishOpts::broadcast(), - rates.tracing_config_broadcast_rps.clone(), - Arc::new(payload(config.tracing_config_payload_bytes)), - ); - spawn_publish_active_rate( - ups.clone(), - "workflow signal broadcast", - workflow_signal_subjects.clone(), - PublishOpts::broadcast(), - rates.workflow_signal_publish_rps.clone(), - Arc::new(Vec::new()), - ); - spawn_publish_rate( - ups.clone(), - "workflow complete broadcast", - vec![unique_subject( - "gasoline.workflow.complete", - "gasoline.workflow.complete", - )], - PublishOpts::broadcast(), - rates.workflow_complete_publish_rps.clone(), - Arc::new(Vec::new()), - ); - - spawn_route_churn( - ups.clone(), - rates.route_churn_rps.clone(), - Duration::from_millis(config.route_ephemeral_hold_ms), - Duration::from_millis(config.route_stopped_hold_ms), - config.route_max_in_flight, - workloads.route.clone(), - ); - spawn_subscription_churn( - ups, - "workflow signal churn", - "gasoline.signal.for-workflow", - "gasoline.signal.for-workflow", - rates.workflow_signal_churn_rps.clone(), - Duration::from_millis(config.workflow_signal_hold_ms), - Some(workflow_signal_subjects), - workloads.workflow_signal.clone(), - ); - - spawn_udb_hot_counter( - udb.clone(), - rates.udb_hot_counter_rps.clone(), - config.udb_hot_counter_max_in_flight, - config.udb_hot_counter_namespace_id, - config.udb_hot_counter_actor_name, - ); - spawn_udb_read_scan( - udb.clone(), - rates.udb_read_scan_rps.clone(), - config.udb_read_scan_max_in_flight, - config.udb_read_scan_seed_keys, - config.udb_read_scan_keys_per_tx, - config.udb_read_scan_value_bytes, - config.udb_read_scan_unpack_keys, - ); - spawn_udb_conflict( - udb, - rates.udb_conflict_rps.clone(), - config.udb_conflict_max_in_flight, - config.udb_conflict_keys, - ); -} - -fn spawn_tuner( - ups: PubSub, - rates: Arc, - workloads: Arc, - tune_path: Option, -) { - let tune_subject = tune_subject(); - { - let ups = ups.clone(); - let rates = rates.clone(); - let workloads = workloads.clone(); - let tune_subject = tune_subject.clone(); - tokio::spawn(async move { - loop { - let mut sub = match ups.subscribe(tune_subject.clone()).await { - Ok(sub) => sub, - Err(err) => { - tracing::warn!(?err, "failed to subscribe to UPS simulation tune subject"); - tokio::time::sleep(Duration::from_secs(2)).await; - continue; - } - }; - - loop { - match sub.next().await { - Ok(NextOutput::Message(message)) => { - if let Err(err) = - apply_tune_patch_bytes(&message.payload, &rates, &workloads) - { - tracing::warn!(?err, "failed to apply UPS simulation tune message"); - } - } - Ok(NextOutput::Unsubscribed | NextOutput::NoResponders) => break, - Err(err) => { - tracing::warn!(?err, "UPS simulation tune subscriber failed"); - break; - } - } - } - } - }); - } - - if let Some(path) = tune_path { - tokio::spawn(async move { - let mut last_payload = None::>; - - loop { - match tokio::fs::read(&path).await { - Ok(payload) if Some(&payload) != last_payload.as_ref() => { - match apply_tune_patch_bytes(&payload, &rates, &workloads) { - Ok(()) => { - last_payload = Some(payload.clone()); - if let Err(err) = ups - .publish( - tune_subject.clone(), - &payload, - PublishOpts::broadcast(), - ) - .await - { - tracing::warn!( - ?err, - %path, - "failed to broadcast UPS simulation tune patch" - ); - } - } - Err(err) => { - tracing::warn!( - ?err, - %path, - "failed to apply UPS simulation tune file" - ); - } - } - } - Ok(_) => {} - Err(err) if err.kind() == std::io::ErrorKind::NotFound => {} - Err(err) => { - tracing::debug!(?err, %path, "failed to read UPS simulation tune file"); - } - } - - tokio::time::sleep(TUNE_POLL_INTERVAL).await; - } - }); - } -} - -fn apply_tune_patch_bytes(payload: &[u8], rates: &Rates, workloads: &Workloads) -> Result<()> { - let patch: TunePatch = - serde_json::from_slice(payload).context("failed to parse UPS simulation tune patch")?; - patch.apply(rates, workloads)?; - tracing::info!(?patch, "applied UPS simulation tune patch"); - Ok(()) -} - -fn tune_subject() -> SimSubject { - SimSubject::new(TUNE_SUBJECT, TUNE_SUBJECT_ROOT) -} - -fn spawn_subject_subscribers( - ups: PubSub, - label: &'static str, - root: &'static str, - prefix: &'static str, - subject_count: usize, - subscriber_count: usize, - queue_group: Option>, - reply_payload: Option>>, - workload: Workload, -) { - if subject_count == 0 || subscriber_count == 0 { - return; - } - - spawn_subject_subscribers_with_offset( - ups, - label, - root, - prefix, - subject_count, - subscriber_count, - 0, - queue_group, - reply_payload, - workload, - ); -} - -fn spawn_subject_subscribers_with_offset( - ups: PubSub, - label: &'static str, - root: &'static str, - prefix: &'static str, - subject_count: usize, - subscriber_count: usize, - subject_offset: usize, - queue_group: Option>, - reply_payload: Option>>, - workload: Workload, -) { - for idx in 0..subscriber_count { - let subject_idx = subject_offset.wrapping_add(idx) % subject_count; - let subject = SimSubject::new(format!("{prefix}.{subject_idx}"), root); - spawn_subscriber( - ups.clone(), - label, - subject, - queue_group.clone(), - reply_payload.clone(), - workload.clone(), - ); - } -} - -fn spawn_gateway_subscribers( - ups: PubSub, - udb: Option, - subject_count: usize, - subscriber_count: usize, - spread_replicas: usize, - member_ttl: Duration, - workload: Workload, -) { - if subject_count == 0 || subscriber_count == 0 { - return; - } - - let Some(udb) = udb.filter(|_| spread_replicas > 1 && subject_count > subscriber_count) else { - spawn_subject_subscribers( - ups, - "gateway", - "pegboard.gateway", - "pegboard.gateway.sim", - subject_count, - subscriber_count, - None, - None, - workload, - ); - return; - }; - - tokio::spawn(async move { - let member_id = gateway_member_id(); - - loop { - match gateway_subject_offset( - &udb, - &member_id, - spread_replicas, - subscriber_count, - member_ttl, - ) - .await - { - Ok(Some(subject_offset)) => { - spawn_gateway_membership_heartbeat(udb.clone(), member_id.clone(), member_ttl); - tracing::info!( - member_id, - subject_offset, - subject_count, - subscriber_count, - spread_replicas, - "starting spread gateway UPS simulation subscribers" - ); - spawn_subject_subscribers_with_offset( - ups, - "gateway", - "pegboard.gateway", - "pegboard.gateway.sim", - subject_count, - subscriber_count, - subject_offset, - None, - None, - workload, - ); - return; - } - Ok(None) => {} - Err(err) => { - tracing::warn!( - ?err, - member_id, - "failed to assign gateway UPS simulation subjects" - ); - } - } - - tokio::time::sleep(Duration::from_secs(2)).await; - } - }); -} - -fn spawn_gateway_membership_heartbeat(udb: UdbPool, member_id: String, member_ttl: Duration) { - let interval = (member_ttl / 3).max(Duration::from_secs(1)); - - tokio::spawn(async move { - loop { - if let Err(err) = gateway_refresh_member(&udb, &member_id).await { - tracing::warn!( - ?err, - member_id, - "failed to refresh gateway UPS simulation membership" - ); - } - - tokio::time::sleep(interval).await; - } - }); -} - -async fn gateway_refresh_member(udb: &UdbPool, member_id: &str) -> Result<()> { - let member_id = member_id.to_string(); - udb.txn(GATEWAY_MEMBERSHIP_TX, |tx| { - let member_id = member_id.clone(); - async move { - let now = now_ms(); - let group = gateway_member_group(&member_id); - let prefix = gateway_member_prefix(&group); - let member_key = gateway_member_key(&prefix, &member_id); - tx.informal().set(&member_key, &now.to_be_bytes()); - Ok(()) - } - }) - .await -} - -async fn gateway_subject_offset( - udb: &UdbPool, - member_id: &str, - expected_replicas: usize, - subscriber_count: usize, - member_ttl: Duration, -) -> Result> { - let member_id = member_id.to_string(); - let member_ttl_ms = duration_millis_u64(member_ttl); - let members = udb - .txn(GATEWAY_MEMBERSHIP_TX, |tx| { - let member_id = member_id.clone(); - async move { - let now = now_ms(); - let group = gateway_member_group(&member_id); - let prefix = gateway_member_prefix(&group); - let member_key = gateway_member_key(&prefix, &member_id); - tx.informal().set(&member_key, &now.to_be_bytes()); - - let mut end = prefix.clone(); - end.push(0xff); - let mut range: RangeOption<'static> = (prefix.clone()..end).into(); - range.limit = Some(expected_replicas.saturating_mul(4).max(32)); - - let min_fresh = now.saturating_sub(member_ttl_ms); - let informal = tx.informal(); - let mut stream = informal.get_ranges_keyvalues(range, Snapshot); - let mut members = Vec::new(); - while let Some(entry) = stream.next().await { - let entry = entry?; - let value = entry.value(); - if value.len() != 8 { - continue; - } - - let mut ts = [0; 8]; - ts.copy_from_slice(value); - if u64::from_be_bytes(ts) < min_fresh { - continue; - } - - if let Some(member) = gateway_member_from_key(&prefix, entry.key()) { - members.push(member); - } - } - - Ok(members) - } - }) - .await?; - - let mut members = members; - members.sort(); - members.dedup(); - - if members.len() < expected_replicas { - tracing::debug!( - member_id, - active_members = members.len(), - expected_replicas, - "waiting for stable gateway UPS simulation membership" - ); - return Ok(None); - } - - let Some(ordinal) = members.iter().position(|member| member == &member_id) else { - return Ok(None); - }; - - Ok(Some( - (ordinal % expected_replicas).saturating_mul(subscriber_count), - )) -} - -fn gateway_member_id() -> String { - env::var("HOSTNAME").unwrap_or_else(|_| format!("pid-{}", std::process::id())) -} - -fn gateway_member_group(member_id: &str) -> String { - member_id - .rsplit_once('-') - .map(|(group, _)| group) - .unwrap_or(member_id) - .to_string() -} - -fn gateway_member_prefix(group: &str) -> Vec { - let mut key = GATEWAY_MEMBERSHIP_PREFIX.to_vec(); - key.push(b'/'); - key.extend_from_slice(group.as_bytes()); - key.push(b'/'); - key -} - -fn gateway_member_key(prefix: &[u8], member_id: &str) -> Vec { - let mut key = prefix.to_vec(); - key.extend_from_slice(member_id.as_bytes()); - key -} - -fn gateway_member_from_key(prefix: &[u8], key: &[u8]) -> Option { - key.strip_prefix(prefix) - .and_then(|member| std::str::from_utf8(member).ok()) - .map(ToOwned::to_owned) -} - -fn spawn_same_subject_subscribers( - ups: PubSub, - label: &'static str, - subject: SimSubject, - subscriber_count: usize, - reply_payload: Option>>, - workload: Workload, -) { - for _ in 0..subscriber_count { - spawn_subscriber( - ups.clone(), - label, - subject.clone(), - None, - reply_payload.clone(), - workload.clone(), - ); - } -} - -fn spawn_worker_bump_subscribers( - ups: PubSub, - subject: SimSubject, - subscriber_count: usize, - workload: Workload, -) { - for _ in 0..subscriber_count { - let ups = ups.clone(); - let subject = subject.clone(); - let workload = workload.clone(); - tokio::spawn(async move { - loop { - let mut sub = match ups.subscribe(subject.clone()).await { - Ok(sub) => sub, - Err(err) => { - tracing::warn!( - ?err, - %subject, - "failed to subscribe for UPS worker bump simulation" - ); - tokio::time::sleep(Duration::from_secs(2)).await; - continue; - } - }; - - loop { - match sub.next().await { - Ok(NextOutput::Message(_)) => { - drain_ready_messages(&mut sub, "worker bump").await; - workload.run().await; - } - Ok(NextOutput::Unsubscribed | NextOutput::NoResponders) => break, - Err(err) => { - tracing::warn!( - ?err, - %subject, - "UPS worker bump simulation subscriber failed" - ); - break; - } - } - } - } - }); - } -} - -fn spawn_subscriber( - ups: PubSub, - label: &'static str, - subject: SimSubject, - queue_group: Option>, - reply_payload: Option>>, - workload: Workload, -) { - tokio::spawn(async move { - loop { - let sub_res = if let Some(queue_group) = queue_group.as_ref() { - ups.queue_subscribe(subject.clone(), queue_group.as_str()) - .await - } else { - ups.subscribe(subject.clone()).await - }; - let mut sub = match sub_res { - Ok(sub) => sub, - Err(err) => { - tracing::warn!(?err, %subject, label, "failed to subscribe for UPS simulation"); - tokio::time::sleep(Duration::from_secs(2)).await; - continue; - } - }; - - loop { - match sub.next().await { - Ok(NextOutput::Message(message)) => { - if let Some(reply_payload) = &reply_payload { - if let Err(err) = message.reply(reply_payload).await { - tracing::debug!( - ?err, - %subject, - label, - "failed to reply to UPS simulation message" - ); - } - } - workload.run().await; - } - Ok(NextOutput::Unsubscribed | NextOutput::NoResponders) => break, - Err(err) => { - tracing::warn!( - ?err, - %subject, - label, - "UPS simulation subscriber failed" - ); - break; - } - } - } - } - }); -} - -async fn drain_ready_messages(sub: &mut Subscriber, label: &'static str) { - for _ in 0..1023 { - match sub.next().now_or_never() { - Some(Ok(NextOutput::Message(_))) => {} - Some(Ok(NextOutput::Unsubscribed | NextOutput::NoResponders)) | None => break, - Some(Err(err)) => { - tracing::debug!(?err, label, "failed to drain UPS simulation messages"); - break; - } - } - } -} - -fn burn_cpu(duration: Duration) { - let start = Instant::now(); - let mut value = 0u64; - while start.elapsed() < duration { - value = value.wrapping_add(1); - hint::black_box(value); - } -} - -fn spawn_publish_rate( - ups: PubSub, - label: &'static str, - subjects: Vec, - opts: PublishOpts, - rate: Rate, - payload: Arc>, -) where - S: Subject + Clone + Send + Sync + 'static, -{ - if subjects.is_empty() { - return; - } - - tokio::spawn(async move { - let mut pacer = Pacer::new(); - let mut idx = 0usize; - let semaphore = Arc::new(tokio::sync::Semaphore::new(PUBLISH_MAX_IN_FLIGHT)); - - loop { - let count = pacer.next_count(rate.load()).await; - for _ in 0..count { - let Ok(permit) = semaphore.clone().try_acquire_owned() else { - continue; - }; - let subject = subjects[idx % subjects.len()].clone(); - idx = idx.wrapping_add(1); - let ups = ups.clone(); - let payload = payload.clone(); - tokio::spawn(async move { - let _permit = permit; - if let Err(err) = ups.publish(subject, &payload, opts).await { - tracing::warn!(?err, label, "UPS simulation publish failed"); - } - }); - } - } - }); -} - -#[derive(Clone, Default)] -struct ActiveSubjects { - subjects: Arc>>, -} - -impl ActiveSubjects { - async fn insert(&self, subject: SimSubject) { - self.subjects.write().await.push(subject); - } - - async fn remove(&self, subject: &SimSubject) { - self.subjects - .write() - .await - .retain(|existing| existing.subject != subject.subject); - } - - async fn get(&self, idx: usize) -> Option { - let subjects = self.subjects.read().await; - if subjects.is_empty() { - None - } else { - Some(subjects[idx % subjects.len()].clone()) - } - } -} - -fn spawn_publish_active_rate( - ups: PubSub, - label: &'static str, - subjects: ActiveSubjects, - opts: PublishOpts, - rate: Rate, - payload: Arc>, -) { - tokio::spawn(async move { - let mut pacer = Pacer::new(); - let mut idx = 0usize; - let semaphore = Arc::new(tokio::sync::Semaphore::new(PUBLISH_MAX_IN_FLIGHT)); - - loop { - let count = pacer.next_count(rate.load()).await; - for _ in 0..count { - let Some(subject) = subjects.get(idx).await else { - continue; - }; - idx = idx.wrapping_add(1); - - let Ok(permit) = semaphore.clone().try_acquire_owned() else { - continue; - }; - let ups = ups.clone(); - let payload = payload.clone(); - tokio::spawn(async move { - let _permit = permit; - if let Err(err) = ups.publish(subject, &payload, opts).await { - tracing::warn!(?err, label, "UPS simulation publish failed"); - } - }); - } - } - }); -} - -fn spawn_request_rate( - ups: PubSub, - subjects: Vec, - rate: Rate, - payload: Arc>, - timeout: Duration, - max_in_flight: usize, -) where - S: Subject + Clone + Send + Sync + 'static, -{ - if subjects.is_empty() || max_in_flight == 0 { - return; - } - - tokio::spawn(async move { - let mut pacer = Pacer::new(); - let mut idx = 0usize; - let semaphore = Arc::new(tokio::sync::Semaphore::new(max_in_flight)); - - loop { - let count = pacer.next_count(rate.load()).await; - for _ in 0..count { - let Ok(permit) = semaphore.clone().try_acquire_owned() else { - continue; - }; - let subject = subjects[idx % subjects.len()].clone(); - idx = idx.wrapping_add(1); - let ups = ups.clone(); - let payload = payload.clone(); - tokio::spawn(async move { - let _permit = permit; - if let Err(err) = ups.request_with_timeout(subject, &payload, timeout).await { - tracing::debug!(?err, "UPS simulation request failed"); - } - }); - } - } - }); -} - -fn spawn_udb_hot_counter( - udb: Option, - rate: Rate, - max_in_flight: usize, - namespace_id: Id, - actor_name: String, -) { - if max_in_flight == 0 { - return; - } - - let Some(udb) = udb else { - if rate.load() > 0.0 { - tracing::warn!("UPS simulation UDB hot counter enabled without a UDB pool"); - } - return; - }; - - tokio::spawn(async move { - let mut pacer = Pacer::new(); - let semaphore = Arc::new(tokio::sync::Semaphore::new(max_in_flight)); - - loop { - let count = pacer.next_count(rate.load()).await; - for _ in 0..count { - let Ok(permit) = semaphore.clone().try_acquire_owned() else { - continue; - }; - let udb = udb.clone(); - let actor_name = actor_name.clone(); - tokio::spawn(async move { - let _permit = permit; - let is_open = HOT_COUNTER_SEQ.fetch_add(1, Ordering::Relaxed) % 2 == 0; - let res = udb - .txn(UDB_HOT_COUNTER_TX, |tx| { - let actor_name = actor_name.clone(); - async move { - let tx = tx.with_subspace(namespace::keys::subspace()); - if is_open { - namespace::keys::metric::inc( - &tx, - namespace_id, - namespace::keys::metric::Metric::Requests( - actor_name.clone(), - "ws".to_string(), - ), - 1, - ); - namespace::keys::metric::inc( - &tx, - namespace_id, - namespace::keys::metric::Metric::ActiveRequests( - actor_name, - "ws".to_string(), - ), - 1, - ); - } else { - namespace::keys::metric::inc( - &tx, - namespace_id, - namespace::keys::metric::Metric::ActiveRequests( - actor_name, - "ws".to_string(), - ), - -1, - ); - } - - Ok(()) - } - }) - .await; - - if let Err(err) = res { - tracing::debug!(?err, "UPS simulation UDB hot counter transaction failed"); - } - }); - } - } - }); -} - -#[derive(Debug, Clone, Copy)] -struct ReadScanKey { - shard: u64, - index: u64, -} - -impl TuplePack for ReadScanKey { - fn pack( - &self, - w: &mut W, - tuple_depth: TupleDepth, - ) -> std::io::Result { - let t = (READ_SCAN_KEY_ROOT, self.shard, self.index); - t.pack(w, tuple_depth) - } -} - -impl<'de> TupleUnpack<'de> for ReadScanKey { - fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let (input, (root, shard, index)) = <(usize, u64, u64)>::unpack(input, tuple_depth)?; - if root != READ_SCAN_KEY_ROOT { - return Err(PackError::Message("expected READ_SCAN key root".into())); - } - - Ok((input, Self { shard, index })) - } -} - -#[derive(Debug, Clone, Copy)] -struct ConflictKey { - index: u64, -} - -impl TuplePack for ConflictKey { - fn pack( - &self, - w: &mut W, - tuple_depth: TupleDepth, - ) -> std::io::Result { - let t = (CONFLICT_KEY_ROOT, self.index); - t.pack(w, tuple_depth) - } -} - -impl<'de> TupleUnpack<'de> for ConflictKey { - fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let (input, (root, index)) = <(usize, u64)>::unpack(input, tuple_depth)?; - if root != CONFLICT_KEY_ROOT { - return Err(PackError::Message("expected CONFLICT key root".into())); - } - - Ok((input, Self { index })) - } -} - -fn sim_read_scan_subspace() -> Subspace { - Subspace::new(&("rivet", "ups-broadcast", "sim", "read-scan")) -} - -fn sim_conflict_subspace() -> Subspace { - Subspace::new(&("rivet", "ups-broadcast", "sim", "conflict")) -} - -fn read_scan_shard() -> u64 { - let member_id = gateway_member_id(); - let mut hash = 0xcbf2_9ce4_8422_2325u64; - for byte in member_id.as_bytes() { - hash ^= u64::from(*byte); - hash = hash.wrapping_mul(0x0000_0100_0000_01b3); - } - hash -} - -fn spawn_udb_read_scan( - udb: Option, - rate: Rate, - max_in_flight: usize, - seed_keys: u64, - keys_per_tx: usize, - value_bytes: usize, - unpack_keys: bool, -) { - if max_in_flight == 0 || keys_per_tx == 0 || seed_keys == 0 { - if rate.load() > 0.0 { - tracing::warn!( - max_in_flight, - seed_keys, - keys_per_tx, - "UPS simulation UDB read scan is enabled without enough configuration" - ); - } - return; - } - - let Some(udb) = udb else { - if rate.load() > 0.0 { - tracing::warn!("UPS simulation UDB read scan enabled without a UDB pool"); - } - return; - }; - - tokio::spawn(async move { - let shard = read_scan_shard(); - if let Err(err) = seed_udb_read_scan(&udb, shard, seed_keys, value_bytes).await { - tracing::warn!( - ?err, - shard, - "failed to seed UPS simulation UDB read scan keys" - ); - } - - let mut pacer = Pacer::new(); - let semaphore = Arc::new(tokio::sync::Semaphore::new(max_in_flight)); - - loop { - let count = pacer.next_count(rate.load()).await; - for _ in 0..count { - let Ok(permit) = semaphore.clone().try_acquire_owned() else { - continue; - }; - let udb = udb.clone(); - tokio::spawn(async move { - let _permit = permit; - if let Err(err) = - run_udb_read_scan(&udb, shard, seed_keys, keys_per_tx, unpack_keys).await - { - tracing::debug!(?err, "UPS simulation UDB read scan transaction failed"); - } - }); - } - } - }); -} - -async fn seed_udb_read_scan( - udb: &UdbPool, - shard: u64, - seed_keys: u64, - value_bytes: usize, -) -> Result<()> { - let value = Arc::new(payload(value_bytes)); - let mut start = 0; - - tracing::info!( - shard, - seed_keys, - value_bytes, - "seeding UPS simulation UDB read scan keys" - ); - - while start < seed_keys { - let end = start - .saturating_add(UDB_READ_SCAN_SEED_BATCH_SIZE) - .min(seed_keys); - let value = value.clone(); - udb.txn(UDB_READ_SCAN_SEED_TX, |tx| { - let value = value.clone(); - async move { - let tx = tx.with_subspace(sim_read_scan_subspace()); - for index in start..end { - let key = tx.pack(&ReadScanKey { shard, index }); - tx.set(&key, value.as_slice()); - } - - Ok(()) - } - }) - .await?; - start = end; - } - - tracing::info!(shard, seed_keys, "seeded UPS simulation UDB read scan keys"); - Ok(()) -} - -async fn run_udb_read_scan( - udb: &UdbPool, - shard: u64, - seed_keys: u64, - keys_per_tx: usize, - unpack_keys: bool, -) -> Result<()> { - let keys_per_tx_u64 = u64::try_from(keys_per_tx) - .unwrap_or(u64::MAX) - .min(seed_keys); - let start = READ_SCAN_SEQ.fetch_add(keys_per_tx_u64, Ordering::Relaxed) % seed_keys; - let end = start.saturating_add(keys_per_tx_u64).min(seed_keys); - let limit = usize::try_from(end.saturating_sub(start)).unwrap_or(keys_per_tx); - - udb.txn(UDB_READ_SCAN_TX, |tx| async move { - let tx = tx.with_subspace(sim_read_scan_subspace()); - let begin = tx.pack(&ReadScanKey { - shard, - index: start, - }); - let end = tx.pack(&ReadScanKey { shard, index: end }); - let mut range: RangeOption<'static> = (begin..end).into(); - range.limit = Some(limit); - - let informal = tx.informal(); - let mut stream = informal.get_ranges_keyvalues(range, Snapshot); - while let Some(entry) = stream.next().await { - let entry = entry?; - if unpack_keys { - let _ = tx.unpack::(entry.key())?; - } - hint::black_box(entry.value().len()); - } - - Ok(()) - }) - .await -} - -fn spawn_udb_conflict(udb: Option, rate: Rate, max_in_flight: usize, key_count: u64) { - if max_in_flight == 0 || key_count == 0 { - if rate.load() > 0.0 { - tracing::warn!( - max_in_flight, - key_count, - "UPS simulation UDB conflict load is enabled without enough configuration" - ); - } - return; - } - - let Some(udb) = udb else { - if rate.load() > 0.0 { - tracing::warn!("UPS simulation UDB conflict load enabled without a UDB pool"); - } - return; - }; - - tokio::spawn(async move { - if let Err(err) = seed_udb_conflict(&udb, key_count).await { - tracing::warn!( - ?err, - key_count, - "failed to seed UPS simulation UDB conflict keys" - ); - } - - let mut pacer = Pacer::new(); - let semaphore = Arc::new(tokio::sync::Semaphore::new(max_in_flight)); - - loop { - let count = pacer.next_count(rate.load()).await; - for _ in 0..count { - let Ok(permit) = semaphore.clone().try_acquire_owned() else { - continue; - }; - let udb = udb.clone(); - tokio::spawn(async move { - let _permit = permit; - let index = CONFLICT_SEQ.fetch_add(1, Ordering::Relaxed) % key_count; - if let Err(err) = run_udb_conflict(&udb, index).await { - tracing::debug!(?err, "UPS simulation UDB conflict transaction failed"); - } - }); - } - } - }); -} - -async fn seed_udb_conflict(udb: &UdbPool, key_count: u64) -> Result<()> { - let mut start = 0; - - tracing::info!(key_count, "seeding UPS simulation UDB conflict keys"); - - while start < key_count { - let end = start - .saturating_add(UDB_CONFLICT_SEED_BATCH_SIZE) - .min(key_count); - udb.txn(UDB_CONFLICT_SEED_TX, |tx| async move { - let tx = tx.with_subspace(sim_conflict_subspace()); - for index in start..end { - let key = tx.pack(&ConflictKey { index }); - tx.set(&key, &0u64.to_be_bytes()); - } - - Ok(()) - }) - .await?; - start = end; - } - - tracing::info!(key_count, "seeded UPS simulation UDB conflict keys"); - Ok(()) -} - -async fn run_udb_conflict(udb: &UdbPool, index: u64) -> Result<()> { - udb.txn(UDB_CONFLICT_TX, |tx| async move { - let tx = tx.with_subspace(sim_conflict_subspace()); - let key = tx.pack(&ConflictKey { index }); - let value = tx.get(&key, Serializable).await?; - let next = value - .as_ref() - .and_then(|value| value.as_slice().try_into().ok().map(u64::from_be_bytes)) - .unwrap_or(0) - .wrapping_add(1); - tx.set(&key, &next.to_be_bytes()); - Ok(()) - }) - .await -} - -fn spawn_route_churn( - ups: PubSub, - rate: Rate, - ephemeral_hold: Duration, - stopped_hold: Duration, - max_in_flight: usize, - workload: Workload, -) { - if max_in_flight == 0 { - return; - } - - tokio::spawn(async move { - let mut pacer = Pacer::new(); - let semaphore = Arc::new(tokio::sync::Semaphore::new(max_in_flight)); - - loop { - let count = pacer.next_count(rate.load()).await; - for _ in 0..count { - let Ok(permit) = semaphore.clone().try_acquire_owned() else { - continue; - }; - let ups = ups.clone(); - let workload = workload.clone(); - tokio::spawn(async move { - let _permit = permit; - let route_id = SUBJECT_SEQ.fetch_add(1, Ordering::Relaxed); - let mut ephemeral = Vec::new(); - for (root, prefix) in ROUTE_SUBJECTS { - let subject = - SimSubject::new(format!("{prefix}:actor_id:{route_id}"), *root); - match ups.subscribe(subject).await { - Ok(sub) => ephemeral.push(sub), - Err(err) => tracing::debug!( - ?err, - "failed to create UPS simulation route subscription" - ), - } - } - - let stopped = ups - .subscribe(SimSubject::new( - format!("gasoline.msg.pegboard_actor2_stopped:actor_id:{route_id}"), - "gasoline.msg.pegboard_actor2_stopped", - )) - .await - .ok(); - - workload.run().await; - tokio::time::sleep(ephemeral_hold).await; - drop(ephemeral); - tokio::time::sleep(stopped_hold).await; - drop(stopped); - }); - } - } - }); -} - -fn spawn_subscription_churn( - ups: PubSub, - label: &'static str, - root: &'static str, - prefix: &'static str, - rate: Rate, - hold: Duration, - active_subjects: Option, - workload: Workload, -) { - tokio::spawn(async move { - let mut pacer = Pacer::new(); - - loop { - let count = pacer.next_count(rate.load()).await; - for _ in 0..count { - let ups = ups.clone(); - let active_subjects = active_subjects.clone(); - let workload = workload.clone(); - tokio::spawn(async move { - let subject = unique_subject(root, prefix); - match ups.subscribe(subject.clone()).await { - Ok(mut sub) => { - if let Some(active_subjects) = active_subjects.as_ref() { - active_subjects.insert(subject.clone()).await; - } - - let deadline = tokio::time::Instant::now() + hold; - loop { - tokio::select! { - res = sub.next() => { - match res { - Ok(NextOutput::Message(_)) => workload.run().await, - Ok(NextOutput::Unsubscribed | NextOutput::NoResponders) => break, - Err(err) => { - tracing::debug!( - ?err, - %subject, - label, - "UPS simulation churn subscriber failed" - ); - break; - } - } - } - _ = tokio::time::sleep_until(deadline) => break, - } - } - - if let Some(active_subjects) = active_subjects.as_ref() { - active_subjects.remove(&subject).await; - } - drop(sub); - } - Err(err) => { - tracing::debug!( - ?err, - %subject, - label, - "failed to create UPS simulation churn subscription" - ); - } - } - }); - } - } - }); -} - -const ROUTE_SUBJECTS: &[(&str, &str)] = &[ - ( - "gasoline.msg.pegboard_actor_failed", - "gasoline.msg.pegboard_actor_failed", - ), - ( - "gasoline.msg.pegboard_actor_ready", - "gasoline.msg.pegboard_actor_ready", - ), - ( - "gasoline.msg.pegboard_actor_stopped", - "gasoline.msg.pegboard_actor_stopped", - ), - ( - "gasoline.msg.pegboard_actor_destroy_started", - "gasoline.msg.pegboard_actor_destroy_started", - ), - ( - "gasoline.msg.pegboard_actor_migrated_to_v2", - "gasoline.msg.pegboard_actor_migrated_to_v2", - ), - ( - "gasoline.msg.pegboard_actor2_ready", - "gasoline.msg.pegboard_actor2_ready", - ), - ( - "gasoline.msg.pegboard_actor2_stopped", - "gasoline.msg.pegboard_actor2_stopped", - ), - ( - "gasoline.msg.pegboard_actor2_failed", - "gasoline.msg.pegboard_actor2_failed", - ), - ( - "gasoline.msg.pegboard_actor2_destroy_started", - "gasoline.msg.pegboard_actor2_destroy_started", - ), -]; - -struct Pacer { - interval: tokio::time::Interval, - carry: f64, - last: Instant, -} - -impl Pacer { - fn new() -> Self { - let mut interval = tokio::time::interval(TICK); - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - Self { - interval, - carry: 0.0, - last: Instant::now(), - } - } - - async fn next_count(&mut self, rate_per_sec: f64) -> usize { - self.interval.tick().await; - let now = Instant::now(); - let elapsed = now.duration_since(self.last); - self.last = now; - self.carry += rate_per_sec * elapsed.as_secs_f64(); - let count = self.carry.floor() as usize; - self.carry -= count as f64; - count - } -} - -#[derive(Clone)] -struct SimSubject { - subject: String, - root: String, -} - -impl SimSubject { - fn new(subject: impl Into, root: impl Into) -> Self { - Self { - subject: subject.into(), - root: root.into(), - } - } -} - -impl fmt::Display for SimSubject { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - self.subject.fmt(f) - } -} - -impl Subject for SimSubject { - fn subject_root<'a>(&'a self) -> Option> { - Some(Cow::Borrowed(self.root.as_str())) - } - - fn as_str(&self) -> Option<&str> { - Some(self.subject.as_str()) - } -} - -#[derive(Clone)] -struct RawSubject { - subject: String, -} - -impl RawSubject { - fn new(subject: impl Into) -> Self { - Self { - subject: subject.into(), - } - } -} - -impl fmt::Display for RawSubject { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - self.subject.fmt(f) - } -} - -impl Subject for RawSubject { - fn as_str(&self) -> Option<&str> { - Some(self.subject.as_str()) - } -} - -fn subjects(root: &'static str, prefix: &'static str, count: usize) -> Vec { - (0..count) - .map(|idx| SimSubject::new(format!("{prefix}.{idx}"), root)) - .collect() -} - -fn raw_subjects(prefix: &'static str, count: usize) -> Vec { - (0..count) - .map(|idx| RawSubject::new(format!("{prefix}.{idx}"))) - .collect() -} - -fn unique_subject(root: &'static str, prefix: &'static str) -> SimSubject { - let idx = SUBJECT_SEQ.fetch_add(1, Ordering::Relaxed); - SimSubject::new(format!("{prefix}.{idx}"), root) -} - -fn payload(size: usize) -> Vec { - vec![b'x'; size] -} - -fn env_key(key: &str) -> String { - format!("{ENV_PREFIX}_{key}") -} - -fn env_string(key: &str) -> Option { - env::var(env_key(key)).ok() -} - -fn env_bool(key: &str, default: bool) -> Result { - let Some(value) = env_string(key) else { - return Ok(default); - }; - match value.to_ascii_lowercase().as_str() { - "1" | "true" | "yes" | "on" => Ok(true), - "0" | "false" | "no" | "off" => Ok(false), - _ => bail!("{ENV_PREFIX}_{key} must be a boolean"), - } -} - -fn env_usize(key: &str, default: usize) -> Result { - parse_env(key, default) -} - -fn env_u64(key: &str, default: u64) -> Result { - parse_env(key, default) -} - -fn env_f64(key: &str, default: f64) -> Result { - parse_env(key, default) -} - -fn env_id(key: &str, default: Id) -> Result { - let Some(value) = env_string(key) else { - return Ok(default); - }; - Id::parse(&value).with_context(|| format!("failed to parse {ENV_PREFIX}_{key}")) -} - -fn parse_env(key: &str, default: T) -> Result -where - T: std::str::FromStr, - T::Err: std::error::Error + Send + Sync + 'static, -{ - let Some(value) = env_string(key) else { - return Ok(default); - }; - value - .parse() - .with_context(|| format!("failed to parse {ENV_PREFIX}_{key}")) -} - -fn validate_rate(key: &str, rate: f64) -> Result<()> { - if rate.is_finite() && rate >= 0.0 { - Ok(()) - } else { - bail!("{ENV_PREFIX}_{key} must be a finite non-negative number") - } -} - -fn now_ms() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|duration| duration.as_millis() as u64) - .unwrap_or(0) -} - -fn duration_millis_u64(duration: Duration) -> u64 { - u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) -} diff --git a/engine/packages/util/src/lib.rs b/engine/packages/util/src/lib.rs index d899e608c3..9b77b4b487 100644 --- a/engine/packages/util/src/lib.rs +++ b/engine/packages/util/src/lib.rs @@ -1,3 +1,5 @@ +use std::fmt::Display; + pub use id::Id; pub use rivet_util_id as id; @@ -12,6 +14,7 @@ pub mod future; pub mod geo; pub mod math; pub mod metric; +pub mod metrics; pub mod req; pub mod serde; pub mod size; @@ -43,3 +46,128 @@ pub fn safe_slice(s: &str, start: usize, end: usize) -> &str { &s[new_start..=new_end] } + +/// Records the duration of the code inside the macro. +/// +/// ```rust +/// observe!(task()); +/// // or +/// observe!(long, long_task()); +/// ``` +/// +/// Supports async work. +/// Use `observe_with!` for callback. +/// ``` +#[macro_export] +macro_rules! observe { + (long, $($tt:tt)*) => {{ + let __start = std::time::Instant::now(); + + let __res = $($tt)*; + let __dt = __start.elapsed().as_secs_f64(); + + let __location = format!("{}:{}:{}", file!(), line!(), column!()); + $crate::metrics::LONG_OBSERVATION_DURATION.with_label_values(&[&__location]) + .observe(__dt); + + __res + }}; + ($($tt:tt)*) => {{ + let __start = std::time::Instant::now(); + + let __res = $($tt)*; + let __dt = __start.elapsed().as_secs_f64(); + + let __location = format!("{}:{}:{}", file!(), line!(), column!()); + $crate::metrics::OBSERVATION_DURATION.with_label_values(&[&__location]) + .observe(__dt); + + __res + }}; +} + +/// Records the duration of the code inside the macro and a callback macro. +/// +/// ```rust +/// observe_with!(task(), |dt, location| { +/// if dt > Duration::from_secs(10) { +/// tracing::warn!("long work at {location}"); +/// } +/// }); +/// // or +/// observe_with!(long, task(), |dt, location| { +/// if dt > Duration::from_secs(10) { +/// tracing::warn!("long work at {location}"); +/// } +/// }); +/// ``` +/// +/// Supports async work. +#[macro_export] +macro_rules! observe_with { + (long, $cb:expr, $($tt:tt)*) => {{ + let __start = std::time::Instant::now(); + + let __res = $($tt)*; + let __dt = __start.elapsed().as_secs_f64(); + + let __location = $crate::location!().to_string(); + + ($cb)(__dt, __location.as_str()); + + $crate::metrics::LONG_OBSERVATION_DURATION.with_label_values(&[__location.as_str()]) + .observe(__dt); + + __res + }}; + ($cb:expr, $($tt:tt)*) => {{ + let __start = std::time::Instant::now(); + + let __res = $($tt)*; + let __dt = __start.elapsed().as_secs_f64(); + + let __location = $crate::location!().to_string(); + + ($cb)(__dt, __location.as_str()); + + $crate::metrics::OBSERVATION_DURATION.with_label_values(&[__location.as_str()]) + .observe(__dt); + + __res + }}; +} + +#[derive(Debug)] +pub struct Location { + file: &'static str, + line: u32, + column: u32, +} + +impl Location { + pub fn new(file: &'static str, line: u32, column: u32) -> Self { + Location { file, line, column } + } +} + +impl Display for Location { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}:{}:{}", self.file, self.line, self.column) + } +} + +/// Constructs a `Location` object with the current file name, line number, and +/// column number. +/// +/// # Examples +/// +/// ``` +/// let loc = location!(); +/// println!("This code is at: {:?}", loc); +/// ``` +#[macro_export] +macro_rules! location { + () => { + $crate::Location::new(file!(), line!(), column!()) + }; +} diff --git a/engine/packages/util/src/metrics.rs b/engine/packages/util/src/metrics.rs new file mode 100644 index 0000000000..3529bebd88 --- /dev/null +++ b/engine/packages/util/src/metrics.rs @@ -0,0 +1,34 @@ +use rivet_metrics::{BUCKETS, MICRO_BUCKETS, REGISTRY, prometheus::*}; + +lazy_static::lazy_static! { + pub static ref OBSERVATION_DURATION: HistogramVec = register_histogram_vec_with_registry!( + "observation_duration", + "Duration of any code observation.", + &["location"], + MICRO_BUCKETS.to_vec(), + *REGISTRY + ).unwrap(); + pub static ref LONG_OBSERVATION_DURATION: HistogramVec = register_histogram_vec_with_registry!( + "long_observation_duration", + "Duration of any long code observation.", + &["location"], + BUCKETS.to_vec(), + *REGISTRY + ).unwrap(); + + pub static ref SERIALIZE_SIZE: HistogramVec = register_histogram_vec_with_registry!( + "serialize_size", + "Size in bytes for any serialization.", + &["format", "location"], + vec![16.0, 32.0, 64.0, 128.0, 256.0, 1024.0, 4096.0, 16384.0, 65536.0, 262144.0, 1048576.0, 4194304.0, 16777216.0], + *REGISTRY + ).unwrap(); + + pub static ref DESERIALIZE_SIZE: HistogramVec = register_histogram_vec_with_registry!( + "deserialize_size", + "Size in bytes for any deserialization.", + &["format", "location"], + vec![16.0, 32.0, 64.0, 128.0, 256.0, 1024.0, 4096.0, 16384.0, 65536.0, 262144.0, 1048576.0, 4194304.0, 16777216.0], + *REGISTRY + ).unwrap(); +} diff --git a/engine/packages/util/src/serde.rs b/engine/packages/util/src/serde.rs index 20c97419cb..9d2de3f4b0 100644 --- a/engine/packages/util/src/serde.rs +++ b/engine/packages/util/src/serde.rs @@ -1 +1,107 @@ pub use rivet_util_serde::*; + +/// Wraps `serde_json::to_vec` with observability. +#[macro_export] +macro_rules! json_to_vec { + ($value:expr) => {{ + let __res = $crate::observe!(serde_json::to_vec($value)); + if let std::result::Result::Ok(__res) = &__res { + $crate::metrics::SERIALIZE_SIZE + .with_label_values(&["json", $crate::location!().to_string().as_str()]) + .observe(__res.len() as f64); + } + __res + }}; +} +pub use json_to_vec; + +/// Wraps `serde_json::to_string` with observability. +#[macro_export] +macro_rules! json_to_string { + ($value:expr) => {{ + let __res = $crate::observe!(serde_json::to_string($value)); + if let std::result::Result::Ok(__res) = &__res { + $crate::metrics::SERIALIZE_SIZE + .with_label_values(&["json", $crate::location!().to_string().as_str()]) + .observe(__res.len() as f64); + } + __res + }}; +} +pub use json_to_string; + +/// Wraps `serde_json::to_value` with observability. +#[macro_export] +macro_rules! json_to_value { + ($value:expr) => {{ $crate::observe!(serde_json::to_value($value)) }}; +} +pub use json_to_value; + +/// Wraps `serde_json::value::to_raw_value` with observability. +#[macro_export] +macro_rules! json_to_raw_value { + ($value:expr) => {{ + let __res = $crate::observe!(serde_json::value::to_raw_value($value)); + if let std::result::Result::Ok(__res) = &__res { + $crate::metrics::SERIALIZE_SIZE + .with_label_values(&["json", $crate::location!().to_string().as_str()]) + .observe(__res.get().len() as f64); + } + __res + }}; +} +pub use json_to_raw_value; + +/// Wraps `serde_json::to_vec` with observability. +#[macro_export] +macro_rules! json_from_str { + ($value:expr) => {{ + let __bind = $value; + $crate::metrics::DESERIALIZE_SIZE + .with_label_values(&["json", $crate::location!().to_string().as_str()]) + .observe(__bind.len() as f64); + $crate::observe!(serde_json::from_str(__bind)) + }}; +} +pub use json_from_str; + +/// Wraps `serde_json::to_vec` with observability. +#[macro_export] +macro_rules! json_from_slice { + ($value:expr) => {{ + let __bind = $value; + $crate::metrics::DESERIALIZE_SIZE + .with_label_values(&["json", $crate::location!().to_string().as_str()]) + .observe(__bind.len() as f64); + $crate::observe!(serde_json::from_slice($value)) + }}; +} +pub use json_from_slice; + +/// Wraps `serde_bare::to_vec` with observability. +#[macro_export] +macro_rules! bare_to_vec { + ($value:expr) => {{ + let __res = $crate::observe!(serde_bare::to_vec($value)); + if let std::result::Result::Ok(__res) = &__res { + $crate::metrics::SERIALIZE_SIZE + .with_label_values(&["bare", $crate::location!().to_string().as_str()]) + .observe(__res.len() as f64); + } + __res + }}; +} +pub use bare_to_vec; + +/// Wraps `serde_bare::to_vec` with observability. +#[macro_export] +macro_rules! bare_from_slice { + ($value:expr) => {{ + let __bind = $value; + $crate::metrics::DESERIALIZE_SIZE + .with_label_values(&["bare", $crate::location!().to_string().as_str()]) + .observe(__bind.len() as f64); + $crate::observe!(serde_bare::from_slice($value)) + }}; +} +pub use bare_from_slice; From 885cd0312c74a291e4df5ef413adbc8c6e249029 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Tue, 23 Jun 2026 13:00:10 -0700 Subject: [PATCH 4/8] Add observation for serde_bare --- Cargo.lock | 3 +- .../depot/src/conveyer/types/branch.rs | 20 +-- .../depot/src/conveyer/types/compaction.rs | 4 +- .../depot/src/conveyer/types/history_pin.rs | 4 +- .../depot/src/conveyer/types/policy.rs | 8 +- .../src/conveyer/types/restore_points.rs | 4 +- .../depot/src/conveyer/types/storage.rs | 12 +- .../packages/engine/src/commands/udb/cli.rs | 22 +-- engine/packages/epoxy/src/http_client.rs | 5 +- engine/packages/epoxy/src/http_routes.rs | 2 +- engine/packages/epoxy/src/keys/keys.rs | 12 +- engine/packages/epoxy/src/keys/replica.rs | 4 +- engine/packages/gasoline/src/workflow.rs | 2 +- engine/packages/pegboard-runner/Cargo.toml | 1 + .../pegboard-runner/src/ws_to_tunnel_task.rs | 4 +- engine/packages/pegboard/src/keys/actor_kv.rs | 4 +- .../packages/runner-protocol/src/versioned.rs | 116 ++++++++----- engine/packages/util/src/serde.rs | 4 +- .../src/versioned/namespace_runner_config.rs | 48 ++++-- .../sdks/rust/depot-protocol/src/versioned.rs | 4 +- engine/sdks/rust/envoy-protocol/Cargo.toml | 2 +- .../rust/envoy-protocol/src/versioned/mod.rs | 153 +++++++++--------- .../sdks/rust/epoxy-protocol/src/versioned.rs | 44 +++-- .../sdks/rust/ups-protocol/src/versioned.rs | 18 ++- 24 files changed, 290 insertions(+), 210 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 12ad23aba3..78308fa999 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4147,6 +4147,7 @@ dependencies = [ "rivet-runner-protocol", "rivet-runtime", "rivet-types", + "rivet-util", "scc", "serde", "serde_bare", @@ -5516,7 +5517,7 @@ dependencies = [ "anyhow", "hex", "rand 0.8.5", - "rivet-util-serde", + "rivet-util", "rivet-vbare-compiler", "serde", "serde_bare", diff --git a/engine/packages/depot/src/conveyer/types/branch.rs b/engine/packages/depot/src/conveyer/types/branch.rs index f0c3faaf22..b0ee671ebf 100644 --- a/engine/packages/depot/src/conveyer/types/branch.rs +++ b/engine/packages/depot/src/conveyer/types/branch.rs @@ -77,14 +77,14 @@ impl OwnedVersionedData for VersionedDatabaseBranchRecord { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::Current(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::Current(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot DatabaseBranchRecord version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::Current(data) => serde_bare::to_vec(&data).map_err(Into::into), + Self::Current(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } } @@ -108,14 +108,14 @@ impl OwnedVersionedData for VersionedDatabasePointer { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot DatabasePointer version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + Self::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } } @@ -139,14 +139,14 @@ impl OwnedVersionedData for VersionedBucketBranchRecord { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot BucketBranchRecord version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + Self::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } } @@ -170,14 +170,14 @@ impl OwnedVersionedData for VersionedBucketPointer { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot BucketPointer version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + Self::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } } @@ -201,14 +201,14 @@ impl OwnedVersionedData for VersionedPointerSnapshot { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot PointerSnapshot version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + Self::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } } diff --git a/engine/packages/depot/src/conveyer/types/compaction.rs b/engine/packages/depot/src/conveyer/types/compaction.rs index 42b212bfb8..0d78e9b67a 100644 --- a/engine/packages/depot/src/conveyer/types/compaction.rs +++ b/engine/packages/depot/src/conveyer/types/compaction.rs @@ -61,14 +61,14 @@ macro_rules! impl_compaction_versioned_data { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot {} version: {version}", $name), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + Self::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } } diff --git a/engine/packages/depot/src/conveyer/types/history_pin.rs b/engine/packages/depot/src/conveyer/types/history_pin.rs index b1592c0eda..f6a07c57e1 100644 --- a/engine/packages/depot/src/conveyer/types/history_pin.rs +++ b/engine/packages/depot/src/conveyer/types/history_pin.rs @@ -43,14 +43,14 @@ impl OwnedVersionedData for VersionedDbHistoryPin { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot DbHistoryPin version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + Self::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } } diff --git a/engine/packages/depot/src/conveyer/types/policy.rs b/engine/packages/depot/src/conveyer/types/policy.rs index edbf5f2cda..a498dd78c3 100644 --- a/engine/packages/depot/src/conveyer/types/policy.rs +++ b/engine/packages/depot/src/conveyer/types/policy.rs @@ -59,14 +59,14 @@ impl OwnedVersionedData for VersionedPitrPolicy { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot PitrPolicy version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(policy) => serde_bare::to_vec(&policy).map_err(Into::into), + Self::V1(policy) => rivet_util::serde::bare_to_vec!(&policy).map_err(Into::into), } } } @@ -86,14 +86,14 @@ impl OwnedVersionedData for VersionedShardCachePolicy { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot ShardCachePolicy version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(policy) => serde_bare::to_vec(&policy).map_err(Into::into), + Self::V1(policy) => rivet_util::serde::bare_to_vec!(&policy).map_err(Into::into), } } } diff --git a/engine/packages/depot/src/conveyer/types/restore_points.rs b/engine/packages/depot/src/conveyer/types/restore_points.rs index 736365915d..f299e3cefa 100644 --- a/engine/packages/depot/src/conveyer/types/restore_points.rs +++ b/engine/packages/depot/src/conveyer/types/restore_points.rs @@ -161,14 +161,14 @@ impl OwnedVersionedData for VersionedRestorePointRecord { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot RestorePointRecord version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + Self::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } } diff --git a/engine/packages/depot/src/conveyer/types/storage.rs b/engine/packages/depot/src/conveyer/types/storage.rs index 64ce085c4c..749244b365 100644 --- a/engine/packages/depot/src/conveyer/types/storage.rs +++ b/engine/packages/depot/src/conveyer/types/storage.rs @@ -53,14 +53,14 @@ impl OwnedVersionedData for VersionedDBHead { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot DBHead version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + Self::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } } @@ -84,14 +84,14 @@ impl OwnedVersionedData for VersionedCommitRow { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot CommitRow version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + Self::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } } @@ -115,14 +115,14 @@ impl OwnedVersionedData for VersionedMetaCompact { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot MetaCompact version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + Self::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } } diff --git a/engine/packages/engine/src/commands/udb/cli.rs b/engine/packages/engine/src/commands/udb/cli.rs index a933ea0df4..c054c7eb72 100644 --- a/engine/packages/engine/src/commands/udb/cli.rs +++ b/engine/packages/engine/src/commands/udb/cli.rs @@ -929,21 +929,23 @@ impl SubCommand { // A v2 entry roundtrips byte-identically through the v2 // schema. v3 entries either fail to deserialize as v2 or // re-serialize to different bytes, so they are ignored. - let v2_entry: proto_v2::ChangelogEntry = - match serde_bare::from_slice(entry.value()) { - Ok(v) => v, - Err(_) => { - v3_count += 1; - continue; - } - }; - let reserialized = match serde_bare::to_vec(&v2_entry) { - Ok(b) => b, + let v2_entry: proto_v2::ChangelogEntry = match rivet_util::serde::bare_from_slice!( + entry.value() + ) { + Ok(v) => v, Err(_) => { v3_count += 1; continue; } }; + let reserialized = + match rivet_util::serde::bare_to_vec!(&v2_entry) { + Ok(b) => b, + Err(_) => { + v3_count += 1; + continue; + } + }; if reserialized != entry.value() { v3_count += 1; continue; diff --git a/engine/packages/epoxy/src/http_client.rs b/engine/packages/epoxy/src/http_client.rs index 539372d647..d495c3ed42 100644 --- a/engine/packages/epoxy/src/http_client.rs +++ b/engine/packages/epoxy/src/http_client.rs @@ -179,7 +179,8 @@ async fn send_request_to_address( let client = rivet_pools::reqwest::client().await?; // Create the request - let request = serde_bare::to_vec(&request).context("failed to serialize epoxy request")?; + let request = + rivet_util::serde::bare_to_vec!(&request).context("failed to serialize epoxy request")?; // Send the request let response_result = client @@ -223,7 +224,7 @@ async fn send_request_to_address( } let body = response.bytes().await?; - let response_body = serde_bare::from_slice(&body)?; + let response_body = rivet_util::serde::bare_from_slice!(&body)?; tracing::debug!( to_replica = to_replica_id, diff --git a/engine/packages/epoxy/src/http_routes.rs b/engine/packages/epoxy/src/http_routes.rs index 33c87586b2..865a860415 100644 --- a/engine/packages/epoxy/src/http_routes.rs +++ b/engine/packages/epoxy/src/http_routes.rs @@ -91,5 +91,5 @@ async fn handle_request(ctx: ApiCtx, request: protocol::Request) -> Result Result { - serde_bare::from_slice(raw).map_err(Into::into) + rivet_util::serde::bare_from_slice!(raw).map_err(Into::into) } fn serialize(&self, value: Self::Value) -> Result> { - serde_bare::to_vec(&value).map_err(Into::into) + rivet_util::serde::bare_to_vec!(&value).map_err(Into::into) } } @@ -203,11 +203,11 @@ impl FormalKey for KvAcceptedKey { type Value = KvAcceptedValue; fn deserialize(&self, raw: &[u8]) -> Result { - serde_bare::from_slice(raw).map_err(Into::into) + rivet_util::serde::bare_from_slice!(raw).map_err(Into::into) } fn serialize(&self, value: Self::Value) -> Result> { - serde_bare::to_vec(&value).map_err(Into::into) + rivet_util::serde::bare_to_vec!(&value).map_err(Into::into) } } @@ -368,11 +368,11 @@ impl FormalKey for ChangelogKey { // TODO: this is mistakenly not versioned. Transition to vbare so future // changes to ChangelogEntry don't require hand-rolled LegacyXxx fallbacks. fn deserialize(&self, raw: &[u8]) -> Result { - serde_bare::from_slice(raw).map_err(Into::into) + rivet_util::serde::bare_from_slice!(raw).map_err(Into::into) } fn serialize(&self, value: Self::Value) -> Result> { - serde_bare::to_vec(&value).map_err(Into::into) + rivet_util::serde::bare_to_vec!(&value).map_err(Into::into) } } diff --git a/engine/packages/epoxy/src/keys/replica.rs b/engine/packages/epoxy/src/keys/replica.rs index b0f7ce6988..beec8aa0d1 100644 --- a/engine/packages/epoxy/src/keys/replica.rs +++ b/engine/packages/epoxy/src/keys/replica.rs @@ -11,11 +11,11 @@ impl FormalKey for ConfigKey { // TODO: this is mistakenly not versioned. Transition to vbare so future // changes to ClusterConfig don't require hand-rolled LegacyXxx fallbacks. fn deserialize(&self, raw: &[u8]) -> Result { - serde_bare::from_slice(raw).map_err(Into::into) + rivet_util::serde::bare_from_slice!(raw).map_err(Into::into) } fn serialize(&self, value: Self::Value) -> Result> { - serde_bare::to_vec(&value).map_err(Into::into) + rivet_util::serde::bare_to_vec!(&value).map_err(Into::into) } } diff --git a/engine/packages/gasoline/src/workflow.rs b/engine/packages/gasoline/src/workflow.rs index bb16ded732..945f335177 100644 --- a/engine/packages/gasoline/src/workflow.rs +++ b/engine/packages/gasoline/src/workflow.rs @@ -33,7 +33,7 @@ impl<'a, T: DeserializeOwned + Serialize> StateGuard<'a, T> { pub(crate) fn new( guard: MutexGuard<'a, (Box, bool)>, ) -> Result { - let value = rivet_util::observe!(serde_json::from_str::(guard.0.get())?); + let value = rivet_util::serde::json_from_str!(guard.0.get())?; Ok(Self { guard, diff --git a/engine/packages/pegboard-runner/Cargo.toml b/engine/packages/pegboard-runner/Cargo.toml index d0d0e0e9eb..b23d7e0139 100644 --- a/engine/packages/pegboard-runner/Cargo.toml +++ b/engine/packages/pegboard-runner/Cargo.toml @@ -29,6 +29,7 @@ rivet-metrics.workspace = true rivet-runner-protocol.workspace = true rivet-runtime.workspace = true rivet-types.workspace = true +rivet-util.workspace = true scc.workspace = true serde_bare.workspace = true serde_json.workspace = true diff --git a/engine/packages/pegboard-runner/src/ws_to_tunnel_task.rs b/engine/packages/pegboard-runner/src/ws_to_tunnel_task.rs index a510172a75..483c53a598 100644 --- a/engine/packages/pegboard-runner/src/ws_to_tunnel_task.rs +++ b/engine/packages/pegboard-runner/src/ws_to_tunnel_task.rs @@ -1028,7 +1028,7 @@ async fn compat_ack_tunnel_message(conn: &Conn, payload: &[u8]) -> Result<()> { use rivet_runner_protocol::generated::v2 as protocol_v2; // Parse payload - let msg = serde_bare::from_slice::(&payload)?; + let msg: protocol_v2::ToServer = rivet_util::serde::bare_from_slice!(&payload)?; let protocol_v2::ToServer::ToServerTunnelMessage(msg) = msg else { return Ok(()); }; @@ -1036,7 +1036,7 @@ async fn compat_ack_tunnel_message(conn: &Conn, payload: &[u8]) -> Result<()> { tracing::debug!(?msg.request_id, ?msg.message_id, "sending v2 compat tunnel ack"); // Serialize response - let ack_msg = serde_bare::to_vec(&protocol_v2::ToClient::ToClientTunnelMessage( + let ack_msg = rivet_util::serde::bare_to_vec!(&protocol_v2::ToClient::ToClientTunnelMessage( protocol_v2::ToClientTunnelMessage { request_id: msg.request_id, message_id: msg.message_id, diff --git a/engine/packages/pegboard/src/keys/actor_kv.rs b/engine/packages/pegboard/src/keys/actor_kv.rs index c15e885813..82ce2072fd 100644 --- a/engine/packages/pegboard/src/keys/actor_kv.rs +++ b/engine/packages/pegboard/src/keys/actor_kv.rs @@ -150,11 +150,11 @@ impl FormalKey for EntryMetadataKey { // TODO: this is mistakenly not versioned. Transition to vbare so future // changes to KvMetadata don't require hand-rolled LegacyXxx fallbacks. fn deserialize(&self, raw: &[u8]) -> Result { - serde_bare::from_slice(raw).map_err(Into::into) + rivet_util::serde::bare_from_slice!(raw).map_err(Into::into) } fn serialize(&self, value: Self::Value) -> Result> { - serde_bare::to_vec(&value).map_err(Into::into) + rivet_util::serde::bare_to_vec!(&value).map_err(Into::into) } } diff --git a/engine/packages/runner-protocol/src/versioned.rs b/engine/packages/runner-protocol/src/versioned.rs index 1a94161592..92bc523a98 100644 --- a/engine/packages/runner-protocol/src/versioned.rs +++ b/engine/packages/runner-protocol/src/versioned.rs @@ -28,18 +28,24 @@ impl OwnedVersionedData for ToClientMk2 { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 4 => Ok(ToClientMk2::V4(serde_bare::from_slice(payload)?)), - 5 => Ok(ToClientMk2::V5(serde_bare::from_slice(payload)?)), - 6 | 7 => Ok(ToClientMk2::V7(serde_bare::from_slice(payload)?)), + 4 => Ok(ToClientMk2::V4(rivet_util::serde::bare_from_slice!( + payload + )?)), + 5 => Ok(ToClientMk2::V5(rivet_util::serde::bare_from_slice!( + payload + )?)), + 6 | 7 => Ok(ToClientMk2::V7(rivet_util::serde::bare_from_slice!( + payload + )?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - ToClientMk2::V4(data) => serde_bare::to_vec(&data).map_err(Into::into), - ToClientMk2::V5(data) => serde_bare::to_vec(&data).map_err(Into::into), - ToClientMk2::V7(data) => serde_bare::to_vec(&data).map_err(Into::into), + ToClientMk2::V4(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + ToClientMk2::V5(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + ToClientMk2::V7(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } @@ -421,19 +427,25 @@ impl OwnedVersionedData for ToServerMk2 { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 4 => Ok(ToServerMk2::V4(serde_bare::from_slice(payload)?)), + 4 => Ok(ToServerMk2::V4(rivet_util::serde::bare_from_slice!( + payload + )?)), // v5 and v6 have the same ToServer binary format - 5 | 6 => Ok(ToServerMk2::V6(serde_bare::from_slice(payload)?)), - 7 => Ok(ToServerMk2::V7(serde_bare::from_slice(payload)?)), + 5 | 6 => Ok(ToServerMk2::V6(rivet_util::serde::bare_from_slice!( + payload + )?)), + 7 => Ok(ToServerMk2::V7(rivet_util::serde::bare_from_slice!( + payload + )?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - ToServerMk2::V4(data) => serde_bare::to_vec(&data).map_err(Into::into), - ToServerMk2::V6(data) => serde_bare::to_vec(&data).map_err(Into::into), - ToServerMk2::V7(data) => serde_bare::to_vec(&data).map_err(Into::into), + ToServerMk2::V4(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + ToServerMk2::V6(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + ToServerMk2::V7(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } @@ -1008,16 +1020,20 @@ impl OwnedVersionedData for ToRunnerMk2 { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 4 => Ok(ToRunnerMk2::V4(serde_bare::from_slice(payload)?)), - 5 | 6 | 7 => Ok(ToRunnerMk2::V7(serde_bare::from_slice(payload)?)), + 4 => Ok(ToRunnerMk2::V4(rivet_util::serde::bare_from_slice!( + payload + )?)), + 5 | 6 | 7 => Ok(ToRunnerMk2::V7(rivet_util::serde::bare_from_slice!( + payload + )?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - ToRunnerMk2::V4(data) => serde_bare::to_vec(&data).map_err(Into::into), - ToRunnerMk2::V7(data) => serde_bare::to_vec(&data).map_err(Into::into), + ToRunnerMk2::V4(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + ToRunnerMk2::V7(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } @@ -1211,18 +1227,18 @@ impl OwnedVersionedData for ToClient { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(ToClient::V1(serde_bare::from_slice(payload)?)), - 2 => Ok(ToClient::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(ToClient::V3(serde_bare::from_slice(payload)?)), + 1 => Ok(ToClient::V1(rivet_util::serde::bare_from_slice!(payload)?)), + 2 => Ok(ToClient::V2(rivet_util::serde::bare_from_slice!(payload)?)), + 3 => Ok(ToClient::V3(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - ToClient::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), - ToClient::V2(data) => serde_bare::to_vec(&data).map_err(Into::into), - ToClient::V3(data) => serde_bare::to_vec(&data).map_err(Into::into), + ToClient::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + ToClient::V2(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + ToClient::V3(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } @@ -1561,18 +1577,18 @@ impl OwnedVersionedData for ToServer { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(ToServer::V1(serde_bare::from_slice(payload)?)), - 2 => Ok(ToServer::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(ToServer::V3(serde_bare::from_slice(payload)?)), + 1 => Ok(ToServer::V1(rivet_util::serde::bare_from_slice!(payload)?)), + 2 => Ok(ToServer::V2(rivet_util::serde::bare_from_slice!(payload)?)), + 3 => Ok(ToServer::V3(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - ToServer::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), - ToServer::V2(data) => serde_bare::to_vec(&data).map_err(Into::into), - ToServer::V3(data) => serde_bare::to_vec(&data).map_err(Into::into), + ToServer::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + ToServer::V2(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + ToServer::V3(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } @@ -1894,14 +1910,14 @@ impl OwnedVersionedData for ToRunner { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 | 2 | 3 => Ok(ToRunner::V3(serde_bare::from_slice(payload)?)), + 1 | 2 | 3 => Ok(ToRunner::V3(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - ToRunner::V3(data) => serde_bare::to_vec(&data).map_err(Into::into), + ToRunner::V3(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } @@ -1939,16 +1955,16 @@ impl OwnedVersionedData for ToGateway { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 | 2 | 3 => Ok(ToGateway::V3(serde_bare::from_slice(payload)?)), - 4 | 5 | 6 | 7 => Ok(ToGateway::V7(serde_bare::from_slice(payload)?)), + 1 | 2 | 3 => Ok(ToGateway::V3(rivet_util::serde::bare_from_slice!(payload)?)), + 4 | 5 | 6 | 7 => Ok(ToGateway::V7(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - ToGateway::V3(data) => serde_bare::to_vec(&data).map_err(Into::into), - ToGateway::V7(data) => serde_bare::to_vec(&data).map_err(Into::into), + ToGateway::V3(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + ToGateway::V7(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } @@ -2044,16 +2060,24 @@ impl OwnedVersionedData for ToServerlessServer { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 | 2 | 3 => Ok(ToServerlessServer::V3(serde_bare::from_slice(payload)?)), - 4 | 5 | 6 | 7 => Ok(ToServerlessServer::V7(serde_bare::from_slice(payload)?)), + 1 | 2 | 3 => Ok(ToServerlessServer::V3(rivet_util::serde::bare_from_slice!( + payload + )?)), + 4 | 5 | 6 | 7 => Ok(ToServerlessServer::V7(rivet_util::serde::bare_from_slice!( + payload + )?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - ToServerlessServer::V3(data) => serde_bare::to_vec(&data).map_err(Into::into), - ToServerlessServer::V7(data) => serde_bare::to_vec(&data).map_err(Into::into), + ToServerlessServer::V3(data) => { + rivet_util::serde::bare_to_vec!(&data).map_err(Into::into) + } + ToServerlessServer::V7(data) => { + rivet_util::serde::bare_to_vec!(&data).map_err(Into::into) + } } } @@ -2123,16 +2147,24 @@ impl OwnedVersionedData for ActorCommandKeyData { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 4 => Ok(ActorCommandKeyData::V4(serde_bare::from_slice(payload)?)), - 5 | 6 | 7 => Ok(ActorCommandKeyData::V7(serde_bare::from_slice(payload)?)), + 4 => Ok(ActorCommandKeyData::V4( + rivet_util::serde::bare_from_slice!(payload)?, + )), + 5 | 6 | 7 => Ok(ActorCommandKeyData::V7( + rivet_util::serde::bare_from_slice!(payload)?, + )), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - ActorCommandKeyData::V4(data) => serde_bare::to_vec(&data).map_err(Into::into), - ActorCommandKeyData::V7(data) => serde_bare::to_vec(&data).map_err(Into::into), + ActorCommandKeyData::V4(data) => { + rivet_util::serde::bare_to_vec!(&data).map_err(Into::into) + } + ActorCommandKeyData::V7(data) => { + rivet_util::serde::bare_to_vec!(&data).map_err(Into::into) + } } } diff --git a/engine/packages/util/src/serde.rs b/engine/packages/util/src/serde.rs index 9d2de3f4b0..18a9823171 100644 --- a/engine/packages/util/src/serde.rs +++ b/engine/packages/util/src/serde.rs @@ -78,7 +78,7 @@ macro_rules! json_from_slice { } pub use json_from_slice; -/// Wraps `serde_bare::to_vec` with observability. +/// Wraps `rivet_util::serde::bare_to_vec!` with observability. #[macro_export] macro_rules! bare_to_vec { ($value:expr) => {{ @@ -93,7 +93,7 @@ macro_rules! bare_to_vec { } pub use bare_to_vec; -/// Wraps `serde_bare::to_vec` with observability. +/// Wraps `rivet_util::serde::bare_to_vec!` with observability. #[macro_export] macro_rules! bare_from_slice { ($value:expr) => {{ diff --git a/engine/sdks/rust/data/src/versioned/namespace_runner_config.rs b/engine/sdks/rust/data/src/versioned/namespace_runner_config.rs index 6c4a40ad23..f887175e90 100644 --- a/engine/sdks/rust/data/src/versioned/namespace_runner_config.rs +++ b/engine/sdks/rust/data/src/versioned/namespace_runner_config.rs @@ -30,24 +30,48 @@ impl OwnedVersionedData for NamespaceRunnerConfig { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(NamespaceRunnerConfig::V1(serde_bare::from_slice(payload)?)), - 2 => Ok(NamespaceRunnerConfig::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(NamespaceRunnerConfig::V3(serde_bare::from_slice(payload)?)), - 4 => Ok(NamespaceRunnerConfig::V4(serde_bare::from_slice(payload)?)), - 5 => Ok(NamespaceRunnerConfig::V5(serde_bare::from_slice(payload)?)), - 6 => Ok(NamespaceRunnerConfig::V6(serde_bare::from_slice(payload)?)), + 1 => Ok(NamespaceRunnerConfig::V1( + rivet_util::serde::bare_from_slice!(payload)?, + )), + 2 => Ok(NamespaceRunnerConfig::V2( + rivet_util::serde::bare_from_slice!(payload)?, + )), + 3 => Ok(NamespaceRunnerConfig::V3( + rivet_util::serde::bare_from_slice!(payload)?, + )), + 4 => Ok(NamespaceRunnerConfig::V4( + rivet_util::serde::bare_from_slice!(payload)?, + )), + 5 => Ok(NamespaceRunnerConfig::V5( + rivet_util::serde::bare_from_slice!(payload)?, + )), + 6 => Ok(NamespaceRunnerConfig::V6( + rivet_util::serde::bare_from_slice!(payload)?, + )), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - NamespaceRunnerConfig::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), - NamespaceRunnerConfig::V2(data) => serde_bare::to_vec(&data).map_err(Into::into), - NamespaceRunnerConfig::V3(data) => serde_bare::to_vec(&data).map_err(Into::into), - NamespaceRunnerConfig::V4(data) => serde_bare::to_vec(&data).map_err(Into::into), - NamespaceRunnerConfig::V5(data) => serde_bare::to_vec(&data).map_err(Into::into), - NamespaceRunnerConfig::V6(data) => serde_bare::to_vec(&data).map_err(Into::into), + NamespaceRunnerConfig::V1(data) => { + rivet_util::serde::bare_to_vec!(&data).map_err(Into::into) + } + NamespaceRunnerConfig::V2(data) => { + rivet_util::serde::bare_to_vec!(&data).map_err(Into::into) + } + NamespaceRunnerConfig::V3(data) => { + rivet_util::serde::bare_to_vec!(&data).map_err(Into::into) + } + NamespaceRunnerConfig::V4(data) => { + rivet_util::serde::bare_to_vec!(&data).map_err(Into::into) + } + NamespaceRunnerConfig::V5(data) => { + rivet_util::serde::bare_to_vec!(&data).map_err(Into::into) + } + NamespaceRunnerConfig::V6(data) => { + rivet_util::serde::bare_to_vec!(&data).map_err(Into::into) + } } } diff --git a/engine/sdks/rust/depot-protocol/src/versioned.rs b/engine/sdks/rust/depot-protocol/src/versioned.rs index c8432da5ae..d90a617eec 100644 --- a/engine/sdks/rust/depot-protocol/src/versioned.rs +++ b/engine/sdks/rust/depot-protocol/src/versioned.rs @@ -22,14 +22,14 @@ impl OwnedVersionedData for DBHead { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid depot db head version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + Self::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } } diff --git a/engine/sdks/rust/envoy-protocol/Cargo.toml b/engine/sdks/rust/envoy-protocol/Cargo.toml index 4e3a222a97..907353804d 100644 --- a/engine/sdks/rust/envoy-protocol/Cargo.toml +++ b/engine/sdks/rust/envoy-protocol/Cargo.toml @@ -12,7 +12,7 @@ description = "Versioned Envoy protocol types for Rivet actor hosts" anyhow.workspace = true hex.workspace = true rand.workspace = true -rivet-util-serde.workspace = true +rivet-util.workspace = true serde_bare.workspace = true serde.workspace = true utoipa.workspace = true diff --git a/engine/sdks/rust/envoy-protocol/src/versioned/mod.rs b/engine/sdks/rust/envoy-protocol/src/versioned/mod.rs index ce618c0c9d..5357809472 100644 --- a/engine/sdks/rust/envoy-protocol/src/versioned/mod.rs +++ b/engine/sdks/rust/envoy-protocol/src/versioned/mod.rs @@ -129,24 +129,24 @@ impl OwnedVersionedData for ToEnvoy { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), - 2 => Ok(Self::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(Self::V3(serde_bare::from_slice(payload)?)), - 4 => Ok(Self::V4(serde_bare::from_slice(payload)?)), - 5 => Ok(Self::V5(serde_bare::from_slice(payload)?)), - 6 => Ok(Self::V6(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), + 2 => Ok(Self::V2(rivet_util::serde::bare_from_slice!(payload)?)), + 3 => Ok(Self::V3(rivet_util::serde::bare_from_slice!(payload)?)), + 4 => Ok(Self::V4(rivet_util::serde::bare_from_slice!(payload)?)), + 5 => Ok(Self::V5(rivet_util::serde::bare_from_slice!(payload)?)), + 6 => Ok(Self::V6(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V2(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V3(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V4(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V5(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V6(x) => serde_bare::to_vec(&x).map_err(Into::into), + Self::V1(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V2(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V3(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V4(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V5(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V6(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), } } @@ -261,24 +261,24 @@ impl OwnedVersionedData for ToRivet { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), - 2 => Ok(Self::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(Self::V3(serde_bare::from_slice(payload)?)), - 4 => Ok(Self::V4(serde_bare::from_slice(payload)?)), - 5 => Ok(Self::V5(serde_bare::from_slice(payload)?)), - 6 => Ok(Self::V6(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), + 2 => Ok(Self::V2(rivet_util::serde::bare_from_slice!(payload)?)), + 3 => Ok(Self::V3(rivet_util::serde::bare_from_slice!(payload)?)), + 4 => Ok(Self::V4(rivet_util::serde::bare_from_slice!(payload)?)), + 5 => Ok(Self::V5(rivet_util::serde::bare_from_slice!(payload)?)), + 6 => Ok(Self::V6(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V2(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V3(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V4(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V5(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V6(x) => serde_bare::to_vec(&x).map_err(Into::into), + Self::V1(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V2(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V3(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V4(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V5(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V6(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), } } @@ -393,24 +393,24 @@ impl OwnedVersionedData for ToEnvoyConn { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), - 2 => Ok(Self::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(Self::V3(serde_bare::from_slice(payload)?)), - 4 => Ok(Self::V4(serde_bare::from_slice(payload)?)), - 5 => Ok(Self::V5(serde_bare::from_slice(payload)?)), - 6 => Ok(Self::V6(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), + 2 => Ok(Self::V2(rivet_util::serde::bare_from_slice!(payload)?)), + 3 => Ok(Self::V3(rivet_util::serde::bare_from_slice!(payload)?)), + 4 => Ok(Self::V4(rivet_util::serde::bare_from_slice!(payload)?)), + 5 => Ok(Self::V5(rivet_util::serde::bare_from_slice!(payload)?)), + 6 => Ok(Self::V6(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V2(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V3(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V4(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V5(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V6(x) => serde_bare::to_vec(&x).map_err(Into::into), + Self::V1(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V2(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V3(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V4(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V5(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V6(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), } } @@ -525,24 +525,24 @@ impl OwnedVersionedData for ToGateway { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), - 2 => Ok(Self::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(Self::V3(serde_bare::from_slice(payload)?)), - 4 => Ok(Self::V4(serde_bare::from_slice(payload)?)), - 5 => Ok(Self::V5(serde_bare::from_slice(payload)?)), - 6 => Ok(Self::V6(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), + 2 => Ok(Self::V2(rivet_util::serde::bare_from_slice!(payload)?)), + 3 => Ok(Self::V3(rivet_util::serde::bare_from_slice!(payload)?)), + 4 => Ok(Self::V4(rivet_util::serde::bare_from_slice!(payload)?)), + 5 => Ok(Self::V5(rivet_util::serde::bare_from_slice!(payload)?)), + 6 => Ok(Self::V6(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V2(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V3(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V4(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V5(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V6(x) => serde_bare::to_vec(&x).map_err(Into::into), + Self::V1(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V2(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V3(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V4(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V5(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V6(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), } } @@ -657,24 +657,24 @@ impl OwnedVersionedData for ToOutbound { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), - 2 => Ok(Self::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(Self::V3(serde_bare::from_slice(payload)?)), - 4 => Ok(Self::V4(serde_bare::from_slice(payload)?)), - 5 => Ok(Self::V5(serde_bare::from_slice(payload)?)), - 6 => Ok(Self::V6(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), + 2 => Ok(Self::V2(rivet_util::serde::bare_from_slice!(payload)?)), + 3 => Ok(Self::V3(rivet_util::serde::bare_from_slice!(payload)?)), + 4 => Ok(Self::V4(rivet_util::serde::bare_from_slice!(payload)?)), + 5 => Ok(Self::V5(rivet_util::serde::bare_from_slice!(payload)?)), + 6 => Ok(Self::V6(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V2(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V3(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V4(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V5(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V6(x) => serde_bare::to_vec(&x).map_err(Into::into), + Self::V1(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V2(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V3(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V4(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V5(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V6(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), } } @@ -789,24 +789,24 @@ impl OwnedVersionedData for ActorCommandKeyData { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), - 2 => Ok(Self::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(Self::V3(serde_bare::from_slice(payload)?)), - 4 => Ok(Self::V4(serde_bare::from_slice(payload)?)), - 5 => Ok(Self::V5(serde_bare::from_slice(payload)?)), - 6 => Ok(Self::V6(serde_bare::from_slice(payload)?)), + 1 => Ok(Self::V1(rivet_util::serde::bare_from_slice!(payload)?)), + 2 => Ok(Self::V2(rivet_util::serde::bare_from_slice!(payload)?)), + 3 => Ok(Self::V3(rivet_util::serde::bare_from_slice!(payload)?)), + 4 => Ok(Self::V4(rivet_util::serde::bare_from_slice!(payload)?)), + 5 => Ok(Self::V5(rivet_util::serde::bare_from_slice!(payload)?)), + 6 => Ok(Self::V6(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Self::V1(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V2(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V3(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V4(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V5(x) => serde_bare::to_vec(&x).map_err(Into::into), - Self::V6(x) => serde_bare::to_vec(&x).map_err(Into::into), + Self::V1(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V2(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V3(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V4(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V5(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), + Self::V6(x) => rivet_util::serde::bare_to_vec!(&x).map_err(Into::into), } } @@ -934,8 +934,8 @@ mod tests { #[test] fn v1_start_command_deserializes_into_latest_without_sqlite_startup_data() -> Result<()> { - let payload = - serde_bare::to_vec(&v1::ToEnvoy::ToEnvoyCommands(vec![v1::CommandWrapper { + let payload = rivet_util::serde::bare_to_vec!(&v1::ToEnvoy::ToEnvoyCommands(vec![ + v1::CommandWrapper { checkpoint: v1::ActorCheckpoint { actor_id: "actor".into(), generation: 7, @@ -951,7 +951,8 @@ mod tests { hibernating_requests: Vec::new(), preloaded_kv: None, }), - }]))?; + } + ]))?; let decoded = ToEnvoy::deserialize(&payload, 1)?; let v6::ToEnvoy::ToEnvoyCommands(commands) = decoded else { @@ -969,7 +970,7 @@ mod tests { #[test] fn v2_sqlite_response_does_not_deserialize_to_stateless_protocol() -> Result<()> { - let payload = serde_bare::to_vec(&v2::ToEnvoy::ToEnvoySqliteCommitResponse( + let payload = rivet_util::serde::bare_to_vec!(&v2::ToEnvoy::ToEnvoySqliteCommitResponse( v2::ToEnvoySqliteCommitResponse { request_id: 1, data: v2::SqliteCommitResponse::SqliteErrorResponse(v2::SqliteErrorResponse { diff --git a/engine/sdks/rust/epoxy-protocol/src/versioned.rs b/engine/sdks/rust/epoxy-protocol/src/versioned.rs index eedebf7662..2dd501bbb8 100644 --- a/engine/sdks/rust/epoxy-protocol/src/versioned.rs +++ b/engine/sdks/rust/epoxy-protocol/src/versioned.rs @@ -26,16 +26,20 @@ impl OwnedVersionedData for CommittedValue { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 2 => Ok(CommittedValue::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(CommittedValue::V3(serde_bare::from_slice(payload)?)), + 2 => Ok(CommittedValue::V2(rivet_util::serde::bare_from_slice!( + payload + )?)), + 3 => Ok(CommittedValue::V3(rivet_util::serde::bare_from_slice!( + payload + )?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - CommittedValue::V2(data) => serde_bare::to_vec(&data).map_err(Into::into), - CommittedValue::V3(data) => serde_bare::to_vec(&data).map_err(Into::into), + CommittedValue::V2(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + CommittedValue::V3(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } @@ -89,16 +93,20 @@ impl OwnedVersionedData for CachedValue { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 2 => Ok(CachedValue::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(CachedValue::V3(serde_bare::from_slice(payload)?)), + 2 => Ok(CachedValue::V2(rivet_util::serde::bare_from_slice!( + payload + )?)), + 3 => Ok(CachedValue::V3(rivet_util::serde::bare_from_slice!( + payload + )?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - CachedValue::V2(data) => serde_bare::to_vec(&data).map_err(Into::into), - CachedValue::V3(data) => serde_bare::to_vec(&data).map_err(Into::into), + CachedValue::V2(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + CachedValue::V3(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } @@ -153,16 +161,20 @@ impl OwnedVersionedData for AcceptedValue { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 2 => Ok(AcceptedValue::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(AcceptedValue::V3(serde_bare::from_slice(payload)?)), + 2 => Ok(AcceptedValue::V2(rivet_util::serde::bare_from_slice!( + payload + )?)), + 3 => Ok(AcceptedValue::V3(rivet_util::serde::bare_from_slice!( + payload + )?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - AcceptedValue::V2(data) => serde_bare::to_vec(&data).map_err(Into::into), - AcceptedValue::V3(data) => serde_bare::to_vec(&data).map_err(Into::into), + AcceptedValue::V2(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + AcceptedValue::V3(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } @@ -216,16 +228,16 @@ impl OwnedVersionedData for Request { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 2 => Ok(Request::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(Request::V3(serde_bare::from_slice(payload)?)), + 2 => Ok(Request::V2(rivet_util::serde::bare_from_slice!(payload)?)), + 3 => Ok(Request::V3(rivet_util::serde::bare_from_slice!(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - Request::V2(data) => serde_bare::to_vec(&data).map_err(Into::into), - Request::V3(data) => serde_bare::to_vec(&data).map_err(Into::into), + Request::V2(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + Request::V3(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } diff --git a/engine/sdks/rust/ups-protocol/src/versioned.rs b/engine/sdks/rust/ups-protocol/src/versioned.rs index e94dff3890..ee756e64bc 100644 --- a/engine/sdks/rust/ups-protocol/src/versioned.rs +++ b/engine/sdks/rust/ups-protocol/src/versioned.rs @@ -26,18 +26,24 @@ impl OwnedVersionedData for UpsMessage { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(UpsMessage::V1(serde_bare::from_slice(payload)?)), - 2 => Ok(UpsMessage::V2(serde_bare::from_slice(payload)?)), - 3 => Ok(UpsMessage::V3(serde_bare::from_slice(payload)?)), + 1 => Ok(UpsMessage::V1(rivet_util::serde::bare_from_slice!( + payload + )?)), + 2 => Ok(UpsMessage::V2(rivet_util::serde::bare_from_slice!( + payload + )?)), + 3 => Ok(UpsMessage::V3(rivet_util::serde::bare_from_slice!( + payload + )?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - UpsMessage::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), - UpsMessage::V2(data) => serde_bare::to_vec(&data).map_err(Into::into), - UpsMessage::V3(data) => serde_bare::to_vec(&data).map_err(Into::into), + UpsMessage::V1(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + UpsMessage::V2(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), + UpsMessage::V3(data) => rivet_util::serde::bare_to_vec!(&data).map_err(Into::into), } } From 7717f5dc251f54d87be4e29fde3e88fb7236edf2 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Tue, 23 Jun 2026 14:24:25 -0700 Subject: [PATCH 5/8] [SLOP(claude-opus-4-8)] feat(rivetkit-core): add serde duration and size metrics --- .../rivetkit-core/src/actor/persist.rs | 16 +- .../packages/rivetkit-core/src/lib.rs | 1 + .../rivetkit-core/src/registry/http.rs | 100 +++++++---- .../rivetkit-core/src/serde_metrics.rs | 163 ++++++++++++++++++ 4 files changed, 237 insertions(+), 43 deletions(-) create mode 100644 rivetkit-rust/packages/rivetkit-core/src/serde_metrics.rs diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/persist.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/persist.rs index 52612c9dbf..ec29e45f8e 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/persist.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/persist.rs @@ -1,6 +1,8 @@ use anyhow::{Context, Result}; use vbare::OwnedVersionedData; +use crate::serde_metrics; + pub(crate) fn encode_latest_with_embedded_version( latest: T::Latest, version: u16, @@ -9,9 +11,11 @@ pub(crate) fn encode_latest_with_embedded_version( where T: OwnedVersionedData, { - T::wrap_latest(latest) - .serialize_with_embedded_version(version) - .with_context(|| format!("encode {label} versioned bare payload")) + serde_metrics::measure_serialize("bare", label, || { + T::wrap_latest(latest) + .serialize_with_embedded_version(version) + .with_context(|| format!("encode {label} versioned bare payload")) + }) } pub(crate) fn decode_latest_with_embedded_version( @@ -21,6 +25,8 @@ pub(crate) fn decode_latest_with_embedded_version( where T: OwnedVersionedData, { - ::deserialize_with_embedded_version(payload) - .with_context(|| format!("decode {label} versioned bare payload")) + serde_metrics::measure_deserialize("bare", label, payload.len(), || { + ::deserialize_with_embedded_version(payload) + .with_context(|| format!("decode {label} versioned bare payload")) + }) } diff --git a/rivetkit-rust/packages/rivetkit-core/src/lib.rs b/rivetkit-rust/packages/rivetkit-core/src/lib.rs index 7604aad0ab..8212027a1b 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/lib.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/lib.rs @@ -12,6 +12,7 @@ pub mod inspector_bundle; pub mod metrics_endpoint; pub mod registry; pub mod runtime; +pub(crate) mod serde_metrics; pub mod serverless; #[cfg(feature = "native-runtime")] pub mod serverless_http; diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs index 9839d79247..088debd34c 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs @@ -2,6 +2,7 @@ use super::dispatch::*; use super::inspector::*; use super::*; use crate::error::{ProtocolError, client_error_message, client_error_metadata}; +use crate::serde_metrics; use ::http; const HEADER_RIVET_ACTOR: &str = "x-rivet-actor"; @@ -773,6 +774,15 @@ pub(super) fn content_type_for_encoding(encoding: HttpResponseEncoding) -> &'sta } } +/// Bounded serde metric `format` label for the request/response encoding. +fn encoding_format_label(encoding: HttpResponseEncoding) -> &'static str { + match encoding { + HttpResponseEncoding::Json => "json", + HttpResponseEncoding::Cbor => "cbor", + HttpResponseEncoding::Bare => "bare", + } +} + pub(super) fn serialize_http_response_error( encoding: HttpResponseEncoding, group: &str, @@ -830,26 +840,31 @@ pub(super) fn decode_http_action_args( encoding: HttpResponseEncoding, body: &[u8], ) -> Result> { - match encoding { - HttpResponseEncoding::Json => { - let request: HttpActionRequestJson = - serde_json::from_slice(body).context("decode json HTTP action request")?; - let args = normalize_json_args(request.args); - encode_json_as_cbor(&args) - } - HttpResponseEncoding::Cbor => { - let request: HttpActionRequestJson = ciborium::from_reader(Cursor::new(body)) - .context("decode cbor HTTP action request")?; - let args = normalize_json_args(request.args); - encode_json_as_cbor(&args) - } - HttpResponseEncoding::Bare => { - let request = - ::deserialize_with_embedded_version(body) - .context("decode bare HTTP action request")?; - Ok(request.args) - } - } + serde_metrics::measure_deserialize( + encoding_format_label(encoding), + "http_action_request", + body.len(), + || match encoding { + HttpResponseEncoding::Json => { + let request: HttpActionRequestJson = + serde_json::from_slice(body).context("decode json HTTP action request")?; + let args = normalize_json_args(request.args); + encode_json_as_cbor(&args) + } + HttpResponseEncoding::Cbor => { + let request: HttpActionRequestJson = ciborium::from_reader(Cursor::new(body)) + .context("decode cbor HTTP action request")?; + let args = normalize_json_args(request.args); + encode_json_as_cbor(&args) + } + HttpResponseEncoding::Bare => { + let request = + ::deserialize_with_embedded_version(body) + .context("decode bare HTTP action request")?; + Ok(request.args) + } + }, + ) } fn normalize_json_args(args: JsonValue) -> Vec { @@ -900,25 +915,34 @@ pub(super) fn encode_http_action_response( encoding: HttpResponseEncoding, output: Vec, ) -> Result { - let body = match encoding { - HttpResponseEncoding::Json => serde_json::to_vec(&json!({ - "output": decode_cbor_json_or_null(&output), - }))?, - HttpResponseEncoding::Cbor => { - let mut out = Vec::new(); - ciborium::into_writer( - &json!({ + let body = serde_metrics::measure_serialize( + encoding_format_label(encoding), + "http_action_response", + || { + let body = match encoding { + HttpResponseEncoding::Json => serde_json::to_vec(&json!({ "output": decode_cbor_json_or_null(&output), - }), - &mut out, - )?; - out - } - HttpResponseEncoding::Bare => client_protocol::versioned::HttpActionResponse::wrap_latest( - client_protocol::HttpActionResponse { output }, - ) - .serialize_with_embedded_version(client_protocol::PROTOCOL_VERSION)?, - }; + }))?, + HttpResponseEncoding::Cbor => { + let mut out = Vec::new(); + ciborium::into_writer( + &json!({ + "output": decode_cbor_json_or_null(&output), + }), + &mut out, + )?; + out + } + HttpResponseEncoding::Bare => { + client_protocol::versioned::HttpActionResponse::wrap_latest( + client_protocol::HttpActionResponse { output }, + ) + .serialize_with_embedded_version(client_protocol::PROTOCOL_VERSION)? + } + }; + Ok(body) + }, + )?; Ok(HttpResponse { status: StatusCode::OK.as_u16(), headers: HashMap::from([( diff --git a/rivetkit-rust/packages/rivetkit-core/src/serde_metrics.rs b/rivetkit-rust/packages/rivetkit-core/src/serde_metrics.rs new file mode 100644 index 0000000000..b105c19f44 --- /dev/null +++ b/rivetkit-rust/packages/rivetkit-core/src/serde_metrics.rs @@ -0,0 +1,163 @@ +//! Duration and size metrics for serialization and deserialization hot paths. +//! +//! These mirror the engine-side serde observability but follow rivetkit's +//! metric conventions: `rivetkit_`-prefixed names registered through a +//! `LazyLock` collector struct, and `crate::time::Instant` so the same code +//! compiles for the wasm runtime. +//! +//! The `format` label is the wire format (`bare`, `json`, `cbor`). The +//! `location` label identifies the call site and must be a bounded, code-defined +//! string, never user input. + +use std::sync::LazyLock; +use std::time::Duration; + +use rivet_metrics::{ + MICRO_BUCKETS, + prometheus::{HistogramOpts, HistogramVec, Registry}, +}; + +use crate::time::Instant; + +const SERDE_LABELS: &[&str] = &["format", "location"]; + +/// Byte-size buckets shared by serialize and deserialize size histograms. +fn serde_size_buckets() -> Vec { + vec![ + 16.0, 32.0, 64.0, 128.0, 256.0, 1024.0, 4096.0, 16384.0, 65536.0, 262144.0, 1048576.0, + 4194304.0, 16777216.0, + ] +} + +struct SerdeMetricCollectors { + serialize_size: HistogramVec, + deserialize_size: HistogramVec, + serialize_duration_seconds: HistogramVec, + deserialize_duration_seconds: HistogramVec, +} + +static METRICS: LazyLock = LazyLock::new(SerdeMetricCollectors::new); + +impl SerdeMetricCollectors { + fn new() -> Self { + let serialize_size = HistogramVec::new( + HistogramOpts::new( + "rivetkit_serialize_size", + "size in bytes for any serialization", + ) + .buckets(serde_size_buckets()), + SERDE_LABELS, + ) + .expect("create rivetkit_serialize_size histogram"); + let deserialize_size = HistogramVec::new( + HistogramOpts::new( + "rivetkit_deserialize_size", + "size in bytes for any deserialization", + ) + .buckets(serde_size_buckets()), + SERDE_LABELS, + ) + .expect("create rivetkit_deserialize_size histogram"); + let serialize_duration_seconds = HistogramVec::new( + HistogramOpts::new( + "rivetkit_serialize_duration_seconds", + "duration in seconds for any serialization", + ) + .buckets(MICRO_BUCKETS.to_vec()), + SERDE_LABELS, + ) + .expect("create rivetkit_serialize_duration_seconds histogram"); + let deserialize_duration_seconds = HistogramVec::new( + HistogramOpts::new( + "rivetkit_deserialize_duration_seconds", + "duration in seconds for any deserialization", + ) + .buckets(MICRO_BUCKETS.to_vec()), + SERDE_LABELS, + ) + .expect("create rivetkit_deserialize_duration_seconds histogram"); + + register_metric(&rivet_metrics::REGISTRY, serialize_size.clone()); + register_metric(&rivet_metrics::REGISTRY, deserialize_size.clone()); + register_metric(&rivet_metrics::REGISTRY, serialize_duration_seconds.clone()); + register_metric( + &rivet_metrics::REGISTRY, + deserialize_duration_seconds.clone(), + ); + + Self { + serialize_size, + deserialize_size, + serialize_duration_seconds, + deserialize_duration_seconds, + } + } +} + +/// Records the duration and output size of a serialization producing `Vec`. +/// +/// The size is only recorded when the closure succeeds. +pub(crate) fn measure_serialize( + format: &str, + location: &str, + f: impl FnOnce() -> anyhow::Result>, +) -> anyhow::Result> { + let started = Instant::now(); + let result = f(); + observe( + &METRICS.serialize_duration_seconds, + format, + location, + started.elapsed(), + ); + if let Ok(bytes) = &result { + observe_size(&METRICS.serialize_size, format, location, bytes.len()); + } + result +} + +/// Records the duration and input size of a deserialization. +/// +/// The input size is recorded unconditionally because the bytes are available +/// regardless of whether decoding succeeds. +pub(crate) fn measure_deserialize( + format: &str, + location: &str, + input_len: usize, + f: impl FnOnce() -> anyhow::Result, +) -> anyhow::Result { + observe_size(&METRICS.deserialize_size, format, location, input_len); + let started = Instant::now(); + let result = f(); + observe( + &METRICS.deserialize_duration_seconds, + format, + location, + started.elapsed(), + ); + result +} + +fn observe(metric: &HistogramVec, format: &str, location: &str, elapsed: Duration) { + metric + .with_label_values(&[format, location]) + .observe(elapsed.as_secs_f64()); +} + +fn observe_size(metric: &HistogramVec, format: &str, location: &str, size: usize) { + metric + .with_label_values(&[format, location]) + .observe(size as f64); +} + +fn register_metric(registry: &Registry, metric: M) +where + M: rivet_metrics::prometheus::core::Collector + Clone + Send + Sync + 'static, +{ + if let Err(error) = registry.register(Box::new(metric)) { + tracing::warn!( + ?error, + "serde metric registration failed, using existing collector" + ); + } +} From c6eeb2a3fd5718849f2a3b8d1d241f5ccf08afb5 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Tue, 23 Jun 2026 17:28:30 -0700 Subject: [PATCH 6/8] [SLOP(claude-opus-4-8)] feat(util): add rate limiter primitive and ingress throttles for actor create and gateway websocket --- Cargo.lock | 1 + .../errors/actor.creation_rate_limit.json | 5 + engine/packages/config/src/config/pegboard.rs | 30 ++ engine/packages/gasoline/src/ctx/message.rs | 2 +- engine/packages/gasoline/src/error.rs | 2 +- .../packages/guard-core/src/proxy_service.rs | 16 +- engine/packages/guard-core/src/utils.rs | 44 +- engine/packages/pegboard-gateway2/src/lib.rs | 1 + .../pegboard-gateway2/src/shared_state.rs | 2 +- .../src/ws_to_tunnel_task.rs | 16 + engine/packages/pegboard/Cargo.toml | 1 + engine/packages/pegboard/src/errors.rs | 6 + .../packages/pegboard/src/ops/actor/create.rs | 34 ++ .../pegboard/src/workflows/actor/runtime.rs | 4 +- .../pegboard/src/workflows/actor2/runtime.rs | 2 +- .../workflows/runner_pool_metadata_poller.rs | 2 +- .../pegboard/src/workflows/serverless/conn.rs | 4 +- .../src/driver/postgres/mod.rs | 2 +- engine/packages/universalpubsub/src/pubsub.rs | 2 +- engine/packages/util/src/backoff.rs | 109 ---- engine/packages/util/src/lib.rs | 2 +- engine/packages/util/src/throttle.rs | 487 ++++++++++++++++++ 22 files changed, 606 insertions(+), 168 deletions(-) create mode 100644 engine/artifacts/errors/actor.creation_rate_limit.json delete mode 100644 engine/packages/util/src/backoff.rs create mode 100644 engine/packages/util/src/throttle.rs diff --git a/Cargo.lock b/Cargo.lock index 78308fa999..72af92dd43 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3938,6 +3938,7 @@ dependencies = [ "futures-util", "gasoline", "lazy_static", + "moka", "namespace", "nix 0.30.1", "portpicker", diff --git a/engine/artifacts/errors/actor.creation_rate_limit.json b/engine/artifacts/errors/actor.creation_rate_limit.json new file mode 100644 index 0000000000..ff4ee74181 --- /dev/null +++ b/engine/artifacts/errors/actor.creation_rate_limit.json @@ -0,0 +1,5 @@ +{ + "code": "creation_rate_limit", + "group": "actor", + "message": "Too many actors created at once. Try again later." +} \ No newline at end of file diff --git a/engine/packages/config/src/config/pegboard.rs b/engine/packages/config/src/config/pegboard.rs index 616af06a36..f05a98319f 100644 --- a/engine/packages/config/src/config/pegboard.rs +++ b/engine/packages/config/src/config/pegboard.rs @@ -117,6 +117,12 @@ pub struct Pegboard { pub gateway_hws_max_pending_size: Option, /// Max HTTP request body size in bytes for requests to actors. pub gateway_http_max_request_body_size: Option, + /// Max burst of inbound WebSocket messages on a single connection before throttling. + pub gateway_websocket_rate_limit_requests: Option, + /// Time to regain one inbound WebSocket message token on a single connection. + /// + /// Unit is in milliseconds. + pub gateway_websocket_rate_limit_drip_rate_ms: Option, // === Envoy Settings === /// How long to wait before considering an envoy lost and evicting all of its actors. @@ -162,6 +168,14 @@ pub struct Pegboard { /// /// Unit is in bytes. Default: 1,048,576 (1 MiB). pub preload_max_total_bytes: Option, + + // === Rate Limiting === + /// Max burst of actor creations per namespace before throttling. + pub actor_create_rate_limit_requests: Option, + /// Time to regain one actor creation token per namespace. + /// + /// Unit is in milliseconds. + pub actor_create_rate_limit_drip_rate_ms: Option, } impl Pegboard { @@ -369,6 +383,22 @@ impl Pegboard { self.serverless_drain_grace_period.unwrap_or(10_000) } + pub fn gateway_websocket_rate_limit_requests(&self) -> u64 { + self.gateway_websocket_rate_limit_requests.unwrap_or(2_000) + } + + pub fn gateway_websocket_rate_limit_drip_rate_ms(&self) -> u64 { + self.gateway_websocket_rate_limit_drip_rate_ms.unwrap_or(10) + } + + pub fn actor_create_rate_limit_requests(&self) -> u64 { + self.actor_create_rate_limit_requests.unwrap_or(500) + } + + pub fn actor_create_rate_limit_drip_rate_ms(&self) -> u64 { + self.actor_create_rate_limit_drip_rate_ms.unwrap_or(10) + } + pub fn preload_max_total_bytes(&self) -> u64 { self.preload_max_total_bytes.unwrap_or(1_048_576) } diff --git a/engine/packages/gasoline/src/ctx/message.rs b/engine/packages/gasoline/src/ctx/message.rs index a6edb2fe32..4c7fcf5e50 100644 --- a/engine/packages/gasoline/src/ctx/message.rs +++ b/engine/packages/gasoline/src/ctx/message.rs @@ -139,7 +139,7 @@ impl MessageCtx { M: Message, { // Infinite backoff since we want to wait until the service reboots. - let mut backoff = rivet_util::backoff::Backoff::default_infinite(); + let mut backoff = rivet_util::throttle::Backoff::default_infinite(); loop { // Ignore for infinite backoff backoff.tick().await; diff --git a/engine/packages/gasoline/src/error.rs b/engine/packages/gasoline/src/error.rs index 2fa7c686b9..1df67e17ad 100644 --- a/engine/packages/gasoline/src/error.rs +++ b/engine/packages/gasoline/src/error.rs @@ -193,7 +193,7 @@ impl WorkflowError { | WorkflowError::ActivityTimeout(_, error_count) | WorkflowError::OperationTimeout(_, error_count) => { // NOTE: Max retry is handled in `WorkflowCtx::activity` - let mut backoff = rivet_util::backoff::Backoff::new_at( + let mut backoff = rivet_util::throttle::Backoff::new_at( 8, None, RETRY_TIMEOUT_MS, diff --git a/engine/packages/guard-core/src/proxy_service.rs b/engine/packages/guard-core/src/proxy_service.rs index 8f53836ad0..9e94695509 100644 --- a/engine/packages/guard-core/src/proxy_service.rs +++ b/engine/packages/guard-core/src/proxy_service.rs @@ -33,7 +33,7 @@ use crate::RouteTarget; use crate::request_context::RequestContext; use crate::response_body::ResponseBody; use crate::route::{CacheKeyFn, ResolveRouteOutput, RouteCache, RoutingFn, RoutingOutput}; -use crate::utils::{InFlightCounter, RateLimiter}; +use crate::utils::InFlightCounter; use crate::{ WebSocketHandle, custom_serve::HibernationResult, errors, metrics, task_group::TaskGroup, utils, }; @@ -63,7 +63,7 @@ pub struct ProxyState { >, route_cache: RouteCache, // We use moka::Cache instead of scc::HashMap because it automatically handles TTL and capacity - rate_limiters: Cache>>, + rate_limiters: Cache>>, in_flight_counters: Cache>>, in_flight_requests: Cache, @@ -105,11 +105,11 @@ impl ProxyState { route_cache: RouteCache::new(route_cache_ttl), rate_limiters: Cache::builder() .max_capacity(10_000) - .time_to_live(PROXY_STATE_CACHE_TTL) + .time_to_idle(PROXY_STATE_CACHE_TTL) .build(), in_flight_counters: Cache::builder() .max_capacity(10_000) - .time_to_live(PROXY_STATE_CACHE_TTL) + .time_to_idle(PROXY_STATE_CACHE_TTL) .build(), in_flight_requests: Cache::builder().max_capacity(10_000_000).build(), tasks: TaskGroup::new(), @@ -224,9 +224,11 @@ impl ProxyState { if let Some(existing_limiter) = self.rate_limiters.get(&req_ctx.client_ip).await { existing_limiter } else { - let new_limiter = Arc::new(Mutex::new(RateLimiter::new( - req_ctx.rate_limit.requests, - req_ctx.rate_limit.period, + let new_limiter = Arc::new(Mutex::new(rivet_util::throttle::RateLimiter::new( + rivet_util::throttle::RateLimitMethod::FixedWindow { + requests: req_ctx.rate_limit.requests, + period: Duration::from_secs(req_ctx.rate_limit.period), + }, ))); self.rate_limiters .insert(req_ctx.client_ip, new_limiter.clone()) diff --git a/engine/packages/guard-core/src/utils.rs b/engine/packages/guard-core/src/utils.rs index 5503611dcc..7c5b3af382 100644 --- a/engine/packages/guard-core/src/utils.rs +++ b/engine/packages/guard-core/src/utils.rs @@ -7,7 +7,7 @@ use hyper::header::HeaderName; use rivet_api_builder::{ErrorResponse, RawErrorResponse}; use rivet_error::{INTERNAL_ERROR, RivetError}; use rivet_util::Id; -use std::time::{Duration, Instant}; +use std::time::Duration; use tokio_tungstenite::tungstenite::protocol::{CloseFrame, frame::coding::CloseCode}; use url::Url; @@ -19,7 +19,7 @@ const X_RIVET_TARGET: HeaderName = HeaderName::from_static("x-rivet-target"); const X_RIVET_ACTOR: HeaderName = HeaderName::from_static("x-rivet-actor"); const X_RIVET_TOKEN: HeaderName = HeaderName::from_static("x-rivet-token"); -// In-flight requests counter +// In-flight requests counter (semaphore) pub(crate) struct InFlightCounter { count: usize, max: usize, @@ -44,43 +44,6 @@ impl InFlightCounter { } } -// Rate limiter -pub(crate) struct RateLimiter { - requests_remaining: u64, - reset_time: Instant, - requests_limit: u64, - period: Duration, -} - -impl RateLimiter { - pub(crate) fn new(requests: u64, period_seconds: u64) -> Self { - Self { - requests_remaining: requests, - reset_time: Instant::now() + Duration::from_secs(period_seconds), - requests_limit: requests, - period: Duration::from_secs(period_seconds), - } - } - - pub(crate) fn try_acquire(&mut self) -> bool { - let now = Instant::now(); - - // Check if we need to reset the counter - if now >= self.reset_time { - self.requests_remaining = self.requests_limit; - self.reset_time = now + self.period; - } - - // Try to consume a request - if self.requests_remaining > 0 { - self.requests_remaining -= 1; - true - } else { - false - } - } -} - // Calculate backoff duration for a given retry attempt pub(crate) fn calculate_backoff(attempt: u32, initial_interval: u64) -> Duration { Duration::from_millis(initial_interval * 2u64.pow(attempt - 1)) @@ -177,7 +140,6 @@ pub(crate) fn err_into_response(err: anyhow::Error) -> Result StatusCode::BAD_GATEWAY, ("guard", "request_timeout") => StatusCode::GATEWAY_TIMEOUT, ("guard", "retry_attempts_exceeded") => StatusCode::BAD_GATEWAY, - ("actor", "not_found") => StatusCode::NOT_FOUND, ("guard", "service_unavailable") => StatusCode::SERVICE_UNAVAILABLE, ("guard", "actor_stopped_while_waiting") => StatusCode::SERVICE_UNAVAILABLE, ("guard", "tunnel_request_aborted") => StatusCode::SERVICE_UNAVAILABLE, @@ -188,6 +150,8 @@ pub(crate) fn err_into_response(err: anyhow::Error) -> Result StatusCode::NOT_FOUND, ("guard", "invalid_request_body") => StatusCode::PAYLOAD_TOO_LARGE, ("guard", "invalid_response_body") => StatusCode::BAD_GATEWAY, + ("actor", "creation_rate_limit") => StatusCode::TOO_MANY_REQUESTS, + ("actor", "not_found") => StatusCode::NOT_FOUND, _ => StatusCode::BAD_REQUEST, }; diff --git a/engine/packages/pegboard-gateway2/src/lib.rs b/engine/packages/pegboard-gateway2/src/lib.rs index 4bd33f15cd..e370270334 100644 --- a/engine/packages/pegboard-gateway2/src/lib.rs +++ b/engine/packages/pegboard-gateway2/src/lib.rs @@ -674,6 +674,7 @@ impl PegboardGateway2 { ); let ws_to_tunnel = tokio::spawn( ws_to_tunnel_task::task( + ctx.clone(), in_flight_req.clone(), ws_rx, ingress_bytes.clone(), diff --git a/engine/packages/pegboard-gateway2/src/shared_state.rs b/engine/packages/pegboard-gateway2/src/shared_state.rs index ff870d330d..67ad84a7b8 100644 --- a/engine/packages/pegboard-gateway2/src/shared_state.rs +++ b/engine/packages/pegboard-gateway2/src/shared_state.rs @@ -661,7 +661,7 @@ impl InFlightRequestHandle { // Cap retries so a permanently-gone receiver fails fast instead of pinning the // request forever. Worst-case backoff total is ~19s, which stays under the default // tunnel ping timeout (30s) so the ping path can take over if the receiver is truly lost. - let mut backoff = rivet_util::backoff::Backoff::new(6, Some(8), 100, 5); + let mut backoff = rivet_util::throttle::Backoff::new(6, Some(8), 100, 5); let first_attempt_at = Instant::now(); let mut attempt = 0; loop { diff --git a/engine/packages/pegboard-gateway2/src/ws_to_tunnel_task.rs b/engine/packages/pegboard-gateway2/src/ws_to_tunnel_task.rs index 1ebcff0e94..66ac4836d3 100644 --- a/engine/packages/pegboard-gateway2/src/ws_to_tunnel_task.rs +++ b/engine/packages/pegboard-gateway2/src/ws_to_tunnel_task.rs @@ -1,11 +1,13 @@ use anyhow::Result; use futures_util::TryStreamExt; +use gas::prelude::*; use rivet_envoy_protocol as protocol; use rivet_guard_core::websocket_handle::WebSocketReceiver; use std::sync::{ Arc, atomic::{AtomicU64, Ordering}, }; +use std::time::Duration; use tokio::sync::{Mutex, watch}; use tokio_tungstenite::tungstenite::Message; @@ -14,6 +16,7 @@ use crate::shared_state::{InFlightRequestHandle, display_id}; #[tracing::instrument(name = "ws_to_tunnel_task", skip_all)] pub async fn task( + ctx: StandaloneCtx, in_flight_req: InFlightRequestHandle, ws_rx: Arc>, ingress_bytes: Arc, @@ -21,7 +24,20 @@ pub async fn task( ) -> Result { let mut ws_rx = ws_rx.lock().await; + // Leaky bucket rate limit on consuming ws messages + let pegboard_config = ctx.config().pegboard(); + let mut rate_limit = rivet_util::throttle::RateLimiter::new( + rivet_util::throttle::RateLimitMethod::LeakyBucket { + requests: pegboard_config.gateway_websocket_rate_limit_requests(), + drip_rate: Duration::from_millis( + pegboard_config.gateway_websocket_rate_limit_drip_rate_ms(), + ), + }, + ); + loop { + rate_limit.acquire().await; + tokio::select! { res = ws_rx.try_next() => { if let Some(msg) = res? { diff --git a/engine/packages/pegboard/Cargo.toml b/engine/packages/pegboard/Cargo.toml index 3d878d6b1d..6e9701246f 100644 --- a/engine/packages/pegboard/Cargo.toml +++ b/engine/packages/pegboard/Cargo.toml @@ -17,6 +17,7 @@ foundationdb-tuple.workspace = true futures-util.workspace = true gas.workspace = true lazy_static.workspace = true +moka.workspace = true namespace.workspace = true nix.workspace = true rand.workspace = true diff --git a/engine/packages/pegboard/src/errors.rs b/engine/packages/pegboard/src/errors.rs index 13e21b55cb..45fb31fbd8 100644 --- a/engine/packages/pegboard/src/errors.rs +++ b/engine/packages/pegboard/src/errors.rs @@ -13,6 +13,12 @@ pub enum Actor { #[error("namespace_not_found", "The namespace does not exist.")] NamespaceNotFound, + #[error( + "creation_rate_limit", + "Too many actors created at once. Try again later." + )] + CreationRateLimit, + #[error( "input_too_large", "Actor input too large.", diff --git a/engine/packages/pegboard/src/ops/actor/create.rs b/engine/packages/pegboard/src/ops/actor/create.rs index c21e878e80..a0cb51bca0 100644 --- a/engine/packages/pegboard/src/ops/actor/create.rs +++ b/engine/packages/pegboard/src/ops/actor/create.rs @@ -1,7 +1,15 @@ use anyhow::{Context, Result}; use gas::prelude::*; +use moka::future::Cache; use rivet_api_util::{Method, request_remote_datacenter}; use rivet_types::actors::{Actor, CrashPolicy}; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; +use tokio::sync::Mutex; + +const RATE_LIMITER_CACHE_TTL: Duration = Duration::from_secs(60 * 60); +static RATE_LIMITERS: OnceLock>>> = + OnceLock::new(); #[derive(Debug)] pub struct Input { @@ -29,6 +37,32 @@ pub struct Output { #[operation] pub async fn pegboard_actor_create(ctx: &OperationCtx, input: &Input) -> Result { + let rate_limiter = RATE_LIMITERS + .get_or_init(|| { + Cache::builder() + .max_capacity(10_000) + .time_to_idle(RATE_LIMITER_CACHE_TTL) + .build() + }) + .entry(input.namespace_id) + .or_insert_with(async { + let pegboard_config = ctx.config().pegboard(); + Arc::new(Mutex::new(rivet_util::throttle::RateLimiter::new( + rivet_util::throttle::RateLimitMethod::LeakyBucket { + requests: pegboard_config.actor_create_rate_limit_requests(), + drip_rate: Duration::from_millis( + pegboard_config.actor_create_rate_limit_drip_rate_ms(), + ), + }, + ))) + }) + .await; + + // Limit actor creation per namespace id + if !rate_limiter.value().lock().await.try_acquire() { + return Err(crate::errors::Actor::CreationRateLimit.build()); + } + // Set up subscriptions before dispatching workflow let ( mut create_sub, diff --git a/engine/packages/pegboard/src/workflows/actor/runtime.rs b/engine/packages/pegboard/src/workflows/actor/runtime.rs index 9965b59d5d..505f0afa81 100644 --- a/engine/packages/pegboard/src/workflows/actor/runtime.rs +++ b/engine/packages/pegboard/src/workflows/actor/runtime.rs @@ -1307,8 +1307,8 @@ fn reschedule_backoff( retry_count: usize, base_retry_timeout: usize, max_exponent: usize, -) -> util::backoff::Backoff { - util::backoff::Backoff::new_at(max_exponent, None, base_retry_timeout, 500, retry_count) +) -> util::throttle::Backoff { + util::throttle::Backoff::new_at(max_exponent, None, base_retry_timeout, 500, retry_count) } #[derive(Debug, Serialize, Deserialize)] diff --git a/engine/packages/pegboard/src/workflows/actor2/runtime.rs b/engine/packages/pegboard/src/workflows/actor2/runtime.rs index e8e1d063e1..3eb8dd9a04 100644 --- a/engine/packages/pegboard/src/workflows/actor2/runtime.rs +++ b/engine/packages/pegboard/src/workflows/actor2/runtime.rs @@ -830,7 +830,7 @@ async fn compare_retry( if reset { state.reschedule_ts = None; } else { - let backoff = util::backoff::Backoff::new_at( + let backoff = util::throttle::Backoff::new_at( ctx.config().pegboard().reschedule_backoff_max_exponent(), None, ctx.config().pegboard().base_retry_timeout(), diff --git a/engine/packages/pegboard/src/workflows/runner_pool_metadata_poller.rs b/engine/packages/pegboard/src/workflows/runner_pool_metadata_poller.rs index bcb600f784..47f5e6e262 100644 --- a/engine/packages/pegboard/src/workflows/runner_pool_metadata_poller.rs +++ b/engine/packages/pegboard/src/workflows/runner_pool_metadata_poller.rs @@ -164,7 +164,7 @@ async fn poll_metadata(ctx: &ActivityCtx, input: &PollMetadataInput) -> Result

util::backoff::Backoff { - util::backoff::Backoff::new_at(max_exponent, None, base_retry_timeout, 500, retry_count) +) -> util::throttle::Backoff { + util::throttle::Backoff::new_at(max_exponent, None, base_retry_timeout, 500, retry_count) } /// Report an error to the error tracker workflow. diff --git a/engine/packages/universalpubsub/src/driver/postgres/mod.rs b/engine/packages/universalpubsub/src/driver/postgres/mod.rs index 3ef8f02beb..52bc219b2a 100644 --- a/engine/packages/universalpubsub/src/driver/postgres/mod.rs +++ b/engine/packages/universalpubsub/src/driver/postgres/mod.rs @@ -5,7 +5,7 @@ use base64::engine::general_purpose::STANDARD_NO_PAD as BASE64; use deadpool_postgres::{Config, ManagerConfig, Pool, PoolConfig, RecyclingMethod, Runtime}; use futures_util::future::poll_fn; use rivet_postgres_util::build_tls_config; -use rivet_util::backoff::Backoff; +use rivet_util::throttle::Backoff; use scc::HashMap; use std::hash::{DefaultHasher, Hash, Hasher}; use std::path::PathBuf; diff --git a/engine/packages/universalpubsub/src/pubsub.rs b/engine/packages/universalpubsub/src/pubsub.rs index 1640307e91..d1bc921af0 100644 --- a/engine/packages/universalpubsub/src/pubsub.rs +++ b/engine/packages/universalpubsub/src/pubsub.rs @@ -8,7 +8,7 @@ use scc::HashMap; use tokio::sync::broadcast; use uuid::Uuid; -use rivet_util::backoff::Backoff; +use rivet_util::throttle::Backoff; use crate::chunking::{ChunkTracker, FastPath, encode_chunk, split_payload_into_chunks}; use crate::driver::{PubSubDriverHandle, PublishOpts, SubscriberDriverHandle}; diff --git a/engine/packages/util/src/backoff.rs b/engine/packages/util/src/backoff.rs deleted file mode 100644 index 183f25042e..0000000000 --- a/engine/packages/util/src/backoff.rs +++ /dev/null @@ -1,109 +0,0 @@ -use rand::Rng; -use tokio::time::{Duration, Instant}; - -pub struct Backoff { - /// Maximum exponent for the backoff. - max_exponent: usize, - - /// Maximum amount of retries. - max_retries: Option, - - /// Base wait time in ms. - wait: usize, - - /// Maximum randomness. - randomness: usize, - - /// Iteration of the backoff. - i: usize, - - /// Timestamp to sleep until in ms. - sleep_until: Instant, -} - -impl Backoff { - pub fn new( - max_exponent: usize, - max_retries: Option, - wait: usize, - randomness: usize, - ) -> Backoff { - Backoff { - max_exponent, - max_retries, - wait, - randomness, - i: 0, - sleep_until: Instant::now(), - } - } - - pub fn new_at( - max_exponent: usize, - max_retries: Option, - wait: usize, - randomness: usize, - i: usize, - ) -> Backoff { - Backoff { - max_exponent, - max_retries, - wait, - randomness, - i, - sleep_until: Instant::now(), - } - } - - pub fn tick_index(&self) -> usize { - self.i - } - - /// Waits for the next backoff tick. - /// - /// Returns false if the index is greater than `max_retries`. - pub async fn tick(&mut self) -> bool { - if self.max_retries.map_or(false, |x| self.i > x) { - return false; - } - - tokio::time::sleep_until(self.sleep_until).await; - - let next_wait = self.current_duration() + rand::thread_rng().gen_range(0..self.randomness); - self.sleep_until += Duration::from_millis(next_wait as u64); - - self.i += 1; - - true - } - - /// Returns the instant of the next backoff tick. Does not wait. - /// - /// Returns None if the index is greater than `max_retries`. - pub fn step(&mut self) -> Option { - if self.max_retries.map_or(false, |x| self.i > x) { - return None; - } - - let next_wait = self.current_duration() + rand::thread_rng().gen_range(0..self.randomness); - self.sleep_until += Duration::from_millis(next_wait as u64); - - self.i += 1; - - Some(self.sleep_until) - } - - pub fn current_duration(&self) -> usize { - self.wait * 2usize.pow(self.i.min(self.max_exponent) as u32) - } - - pub fn default_infinite() -> Backoff { - Backoff::new(8, None, 1_000, 1_000) - } -} - -impl Default for Backoff { - fn default() -> Backoff { - Backoff::new(5, Some(16), 1_000, 1_000) - } -} diff --git a/engine/packages/util/src/lib.rs b/engine/packages/util/src/lib.rs index 9b77b4b487..0c088b9df8 100644 --- a/engine/packages/util/src/lib.rs +++ b/engine/packages/util/src/lib.rs @@ -4,7 +4,6 @@ pub use id::Id; pub use rivet_util_id as id; pub mod async_counter; -pub mod backoff; pub mod billing; pub mod build_meta; pub mod check; @@ -19,6 +18,7 @@ pub mod req; pub mod serde; pub mod size; pub mod sort; +pub mod throttle; pub mod timestamp; pub mod url; diff --git a/engine/packages/util/src/throttle.rs b/engine/packages/util/src/throttle.rs new file mode 100644 index 0000000000..38295f99f6 --- /dev/null +++ b/engine/packages/util/src/throttle.rs @@ -0,0 +1,487 @@ +use rand::Rng; +use tokio::time::{Duration, Instant}; + +pub struct Backoff { + /// Maximum exponent for the backoff. + max_exponent: usize, + + /// Maximum amount of retries. + max_retries: Option, + + /// Base wait time in ms. + wait: usize, + + /// Maximum randomness. + randomness: usize, + + /// Iteration of the backoff. + i: usize, + + /// Timestamp to sleep until in ms. + sleep_until: Instant, +} + +impl Backoff { + pub fn new( + max_exponent: usize, + max_retries: Option, + wait: usize, + randomness: usize, + ) -> Backoff { + Backoff { + max_exponent, + max_retries, + wait, + randomness, + i: 0, + sleep_until: Instant::now(), + } + } + + pub fn new_at( + max_exponent: usize, + max_retries: Option, + wait: usize, + randomness: usize, + i: usize, + ) -> Backoff { + Backoff { + max_exponent, + max_retries, + wait, + randomness, + i, + sleep_until: Instant::now(), + } + } + + pub fn tick_index(&self) -> usize { + self.i + } + + /// Waits for the next backoff tick. + /// + /// Returns false if the index is greater than `max_retries`. + pub async fn tick(&mut self) -> bool { + if self.max_retries.map_or(false, |x| self.i > x) { + return false; + } + + tokio::time::sleep_until(self.sleep_until).await; + + let next_wait = self.current_duration() + rand::thread_rng().gen_range(0..self.randomness); + self.sleep_until += Duration::from_millis(next_wait as u64); + + self.i += 1; + + true + } + + /// Returns the instant of the next backoff tick. Does not wait. + /// + /// Returns None if the index is greater than `max_retries`. + pub fn step(&mut self) -> Option { + if self.max_retries.map_or(false, |x| self.i > x) { + return None; + } + + let next_wait = self.current_duration() + rand::thread_rng().gen_range(0..self.randomness); + self.sleep_until += Duration::from_millis(next_wait as u64); + + self.i += 1; + + Some(self.sleep_until) + } + + pub fn current_duration(&self) -> usize { + self.wait * 2usize.pow(self.i.min(self.max_exponent) as u32) + } + + pub fn default_infinite() -> Backoff { + Backoff::new(8, None, 1_000, 1_000) + } +} + +impl Default for Backoff { + fn default() -> Backoff { + Backoff::new(5, Some(16), 1_000, 1_000) + } +} + +pub enum RateLimitMethod { + FixedWindow { + requests: u64, + period: Duration, + }, + LeakyBucket { + requests: u64, + /// How quickly to regain requests. 1 / drip_rate + drip_rate: Duration, + }, +} + +enum RateLimitState { + FixedWindow { + requests_remaining: u64, + requests_limit: u64, + reset_time: Instant, + period: Duration, + }, + LeakyBucket { + requests_remaining: u64, + requests_limit: u64, + last_acquire: Instant, + drip_rate: Duration, + accum_drip: f32, + }, +} + +pub struct RateLimiter { + state: RateLimitState, +} + +impl RateLimiter { + pub fn new(method: RateLimitMethod) -> Self { + Self { + state: match method { + RateLimitMethod::FixedWindow { requests, period } => RateLimitState::FixedWindow { + requests_remaining: requests, + requests_limit: requests, + reset_time: Instant::now() + period, + period, + }, + RateLimitMethod::LeakyBucket { + requests, + drip_rate, + } => RateLimitState::LeakyBucket { + requests_remaining: requests, + requests_limit: requests, + last_acquire: Instant::now(), + drip_rate: drip_rate, + accum_drip: 0.0, + }, + }, + } + } + + pub fn try_acquire(&mut self) -> bool { + match &mut self.state { + RateLimitState::FixedWindow { + requests_remaining, + requests_limit, + reset_time, + period, + } => { + let now = Instant::now(); + // Check if we need to reset the counter + if now >= *reset_time { + *requests_remaining = *requests_limit; + *reset_time = now + *period; + } + + // Try to consume a request + if *requests_remaining > 0 { + *requests_remaining -= 1; + true + } else { + false + } + } + RateLimitState::LeakyBucket { + requests_remaining, + requests_limit, + last_acquire, + drip_rate, + accum_drip, + } => { + let now = Instant::now(); + let dt = now - *last_acquire; + *last_acquire = now; + + // Drip bucket + if requests_remaining < requests_limit { + *accum_drip += dt.div_duration_f32(*drip_rate); + + *requests_remaining += + (*accum_drip as u64).min(*requests_limit - *requests_remaining); + + if *accum_drip >= 1.0 { + *accum_drip = accum_drip.fract(); + } + } + + if *requests_remaining > 0 { + *requests_remaining -= 1; + true + } else { + false + } + } + } + } + + pub async fn acquire(&mut self) { + match &mut self.state { + RateLimitState::FixedWindow { + requests_remaining, + requests_limit, + reset_time, + period, + } => { + let now = Instant::now(); + // Check if we need to reset the counter + if now >= *reset_time { + *requests_remaining = *requests_limit; + *reset_time = now + *period; + } + + // Try to consume a request + if *requests_remaining > 0 { + *requests_remaining -= 1; + } else { + tokio::time::sleep(*period).await; + + *requests_remaining = *requests_limit; + *reset_time = Instant::now() + *period; + } + } + RateLimitState::LeakyBucket { + requests_remaining, + requests_limit, + last_acquire, + drip_rate, + accum_drip, + } => { + let now = Instant::now(); + let dt = now - *last_acquire; + *last_acquire = now; + + // Drip bucket + if requests_remaining < requests_limit { + *accum_drip += dt.div_duration_f32(*drip_rate); + + *requests_remaining += + (*accum_drip as u64).min(*requests_limit - *requests_remaining); + + if *accum_drip >= 1.0 { + *accum_drip = accum_drip.fract(); + } + } + + if *requests_remaining > 0 { + *requests_remaining -= 1; + } else { + let deficit = 1.0 - *accum_drip; + tokio::time::sleep(drip_rate.mul_f32(deficit)).await; + + *last_acquire = Instant::now(); + *accum_drip = 0.0; + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::{RateLimitMethod, RateLimiter}; + use tokio::time::{Duration, Instant}; + + // MARK: FixedWindow / try_acquire + + #[tokio::test(start_paused = true)] + async fn fixed_window_allows_full_burst_then_blocks() { + let mut rl = RateLimiter::new(RateLimitMethod::FixedWindow { + requests: 3, + period: Duration::from_millis(100), + }); + + assert!(rl.try_acquire()); + assert!(rl.try_acquire()); + assert!(rl.try_acquire()); + // Limit reached within the window. + assert!(!rl.try_acquire()); + } + + #[tokio::test(start_paused = true)] + async fn fixed_window_does_not_refill_before_period() { + let mut rl = RateLimiter::new(RateLimitMethod::FixedWindow { + requests: 2, + period: Duration::from_millis(100), + }); + + assert!(rl.try_acquire()); + assert!(rl.try_acquire()); + assert!(!rl.try_acquire()); + + // Just shy of a full period: still no refill. The window is + // all-or-nothing, it does not drip partial credit. + tokio::time::advance(Duration::from_millis(99)).await; + assert!(!rl.try_acquire()); + } + + #[tokio::test(start_paused = true)] + async fn fixed_window_resets_to_full_after_period() { + let mut rl = RateLimiter::new(RateLimitMethod::FixedWindow { + requests: 2, + period: Duration::from_millis(100), + }); + + assert!(rl.try_acquire()); + assert!(rl.try_acquire()); + assert!(!rl.try_acquire()); + + // After a full period the window resets to its full allowance. + tokio::time::advance(Duration::from_millis(100)).await; + assert!(rl.try_acquire()); + assert!(rl.try_acquire()); + assert!(!rl.try_acquire()); + } + + // MARK: LeakyBucket / try_acquire + + #[tokio::test(start_paused = true)] + async fn leaky_bucket_allows_full_burst_then_blocks() { + let mut rl = RateLimiter::new(RateLimitMethod::LeakyBucket { + requests: 3, + drip_rate: Duration::from_millis(10), + }); + + assert!(rl.try_acquire()); + assert!(rl.try_acquire()); + assert!(rl.try_acquire()); + assert!(!rl.try_acquire()); + } + + #[tokio::test(start_paused = true)] + async fn leaky_bucket_drips_exactly_one_token_per_rate() { + let mut rl = RateLimiter::new(RateLimitMethod::LeakyBucket { + requests: 3, + drip_rate: Duration::from_millis(10), + }); + + // Drain the bucket. + for _ in 0..3 { + assert!(rl.try_acquire()); + } + assert!(!rl.try_acquire()); + + // Exactly one drip period yields exactly one token, no more. + tokio::time::advance(Duration::from_millis(10)).await; + assert!(rl.try_acquire()); + assert!(!rl.try_acquire()); + } + + #[tokio::test(start_paused = true)] + async fn leaky_bucket_refill_is_capped_at_capacity() { + let mut rl = RateLimiter::new(RateLimitMethod::LeakyBucket { + requests: 3, + drip_rate: Duration::from_millis(10), + }); + + for _ in 0..3 { + assert!(rl.try_acquire()); + } + assert!(!rl.try_acquire()); + + // Idle far longer than it takes to refill the whole bucket. Credit must + // not accumulate past capacity, so only `requests` tokens are available. + tokio::time::advance(Duration::from_millis(1_000)).await; + assert!(rl.try_acquire()); + assert!(rl.try_acquire()); + assert!(rl.try_acquire()); + assert!(!rl.try_acquire()); + } + + #[tokio::test(start_paused = true)] + async fn leaky_bucket_accumulates_fractional_drip_across_calls() { + let mut rl = RateLimiter::new(RateLimitMethod::LeakyBucket { + requests: 1, + drip_rate: Duration::from_millis(10), + }); + + // Consume the only token. + assert!(rl.try_acquire()); + assert!(!rl.try_acquire()); + + // Half a drip period: less than one whole token, still blocked. + tokio::time::advance(Duration::from_millis(5)).await; + assert!(!rl.try_acquire()); + + // Another half period: the fractional credit from the previous interval + // must carry over and complete one whole token. + tokio::time::advance(Duration::from_millis(5)).await; + assert!(rl.try_acquire()); + assert!(!rl.try_acquire()); + } + + // MARK: acquire (blocking) + + #[tokio::test(start_paused = true)] + async fn acquire_returns_immediately_while_tokens_remain() { + let mut rl = RateLimiter::new(RateLimitMethod::LeakyBucket { + requests: 3, + drip_rate: Duration::from_millis(10), + }); + + let start = Instant::now(); + rl.acquire().await; + rl.acquire().await; + rl.acquire().await; + // Burst is served without waiting. + assert_eq!(start.elapsed(), Duration::ZERO); + } + + #[tokio::test(start_paused = true)] + async fn acquire_blocks_until_a_token_is_available() { + let mut rl = RateLimiter::new(RateLimitMethod::LeakyBucket { + requests: 1, + drip_rate: Duration::from_millis(10), + }); + + // Drain the single token. + rl.acquire().await; + + // The next acquire must wait one full drip period for a token. + let start = Instant::now(); + rl.acquire().await; + assert!(start.elapsed() >= Duration::from_millis(10)); + } + + #[tokio::test(start_paused = true)] + async fn acquire_sustains_the_drip_rate_without_doubling() { + let mut rl = RateLimiter::new(RateLimitMethod::LeakyBucket { + requests: 1, + drip_rate: Duration::from_millis(10), + }); + + // Drain the initial burst token so every subsequent acquire starts empty. + rl.acquire().await; + + let start = Instant::now(); + // Five acquires, each starting from an empty bucket, must each cost one + // drip period, so the total is at least 5 * drip_rate. A limiter that + // admits the post-sleep request without debiting a token finishes in + // ~3 periods, effectively doubling the sustained rate. + for _ in 0..5 { + rl.acquire().await; + } + assert!(start.elapsed() >= Duration::from_millis(50)); + } + + #[tokio::test(start_paused = true)] + async fn fixed_window_acquire_blocks_until_window_resets() { + let mut rl = RateLimiter::new(RateLimitMethod::FixedWindow { + requests: 2, + period: Duration::from_millis(100), + }); + + rl.acquire().await; + rl.acquire().await; + + // The window is exhausted, so the next acquire must wait for the reset. + let start = Instant::now(); + rl.acquire().await; + assert!(start.elapsed() >= Duration::from_millis(100)); + } +} From 5c7f16e8623708127a4e1b9706745e0eb9a73d95 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Wed, 24 Jun 2026 12:00:17 -0700 Subject: [PATCH 7/8] [SLOP(claude-opus-4-8)] feat(pegboard-envoy): rate-limit envoy ws ingress and cap get_pages, trim unused metrics --- Cargo.lock | 1 + engine/artifacts/config-schema.json | 54 +++++++++++++++ engine/packages/config/src/config/pegboard.rs | 16 +++++ engine/packages/guard/src/routing/envoy.rs | 2 +- engine/packages/pegboard-envoy/Cargo.toml | 1 + engine/packages/pegboard-envoy/src/lib.rs | 8 +-- engine/packages/pegboard-envoy/src/metrics.rs | 69 ------------------- .../pegboard-envoy/src/ws_to_tunnel_task.rs | 35 +++++++++- .../packages/pegboard-gateway2/src/metrics.rs | 39 ----------- .../pegboard-gateway2/src/shared_state.rs | 1 - .../src/ws_to_tunnel_task.rs | 17 +++-- engine/packages/pegboard-outbound/src/lib.rs | 2 - .../packages/pegboard-outbound/src/metrics.rs | 47 ------------- .../packages/pegboard/src/ops/actor/create.rs | 4 +- 14 files changed, 122 insertions(+), 174 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 72af92dd43..b21a5d2859 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4011,6 +4011,7 @@ dependencies = [ "rivet-pools", "rivet-runtime", "rivet-types", + "rivet-util", "rusqlite", "scc", "serde", diff --git a/engine/artifacts/config-schema.json b/engine/artifacts/config-schema.json index ead08e9e6f..998771acc1 100644 --- a/engine/artifacts/config-schema.json +++ b/engine/artifacts/config-schema.json @@ -855,6 +855,24 @@ ], "format": "int64" }, + "actor_create_rate_limit_drip_rate_ms": { + "description": "Time to regain one actor creation token per namespace.\n\nUnit is in milliseconds.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0.0 + }, + "actor_create_rate_limit_requests": { + "description": "Max burst of actor creations per namespace before throttling.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0.0 + }, "actor_retry_duration_threshold": { "description": "How long to wait after starting to attempt to reallocate before before setting actor to sleep.\n\nUnit is in milliseconds.", "type": [ @@ -986,6 +1004,24 @@ "format": "uint64", "minimum": 0.0 }, + "envoy_websocket_rate_limit_drip_rate_us": { + "description": "Time to regain one inbound WebSocket message token on a single envoy connection.\n\nUnit is in microseconds. The envoy connection multiplexes every actor on a runner, so the sustained ceiling is far higher than the per-client gateway limit and needs sub-millisecond granularity to express.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0.0 + }, + "envoy_websocket_rate_limit_requests": { + "description": "Max burst of inbound WebSocket messages on a single envoy connection before throttling.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0.0 + }, "gateway_gc_interval_ms": { "description": "GC interval for in-flight requests in milliseconds.", "type": [ @@ -1057,6 +1093,24 @@ "format": "uint64", "minimum": 0.0 }, + "gateway_websocket_rate_limit_drip_rate_ms": { + "description": "Time to regain one inbound WebSocket message token on a single connection.\n\nUnit is in milliseconds.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0.0 + }, + "gateway_websocket_rate_limit_requests": { + "description": "Max burst of inbound WebSocket messages on a single connection before throttling.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0.0 + }, "hibernating_request_eligible_threshold": { "description": "How long after last ping before considering a hibernating request disconnected.\n\nUnit is in milliseconds.", "type": [ diff --git a/engine/packages/config/src/config/pegboard.rs b/engine/packages/config/src/config/pegboard.rs index f05a98319f..f0b660401e 100644 --- a/engine/packages/config/src/config/pegboard.rs +++ b/engine/packages/config/src/config/pegboard.rs @@ -149,6 +149,14 @@ pub struct Pegboard { pub envoy_expire_scheduler_max_concurrent_expires: Option, /// Maximum pending envoys tracked by the read-path envoy expire scheduler. pub envoy_expire_scheduler_max_pending: Option, + /// Max burst of inbound WebSocket messages on a single envoy connection before throttling. + pub envoy_websocket_rate_limit_requests: Option, + /// Time to regain one inbound WebSocket message token on a single envoy connection. + /// + /// Unit is in microseconds. The envoy connection multiplexes every actor on a runner, so the + /// sustained ceiling is far higher than the per-client gateway limit and needs sub-millisecond + /// granularity to express. + pub envoy_websocket_rate_limit_drip_rate_us: Option, // === Serverless Settings === /// **Deprecated** Configure the drain period in the runner config. @@ -391,6 +399,14 @@ impl Pegboard { self.gateway_websocket_rate_limit_drip_rate_ms.unwrap_or(10) } + pub fn envoy_websocket_rate_limit_requests(&self) -> u64 { + self.envoy_websocket_rate_limit_requests.unwrap_or(16_384) + } + + pub fn envoy_websocket_rate_limit_drip_rate_us(&self) -> u64 { + self.envoy_websocket_rate_limit_drip_rate_us.unwrap_or(200) + } + pub fn actor_create_rate_limit_requests(&self) -> u64 { self.actor_create_rate_limit_requests.unwrap_or(500) } diff --git a/engine/packages/guard/src/routing/envoy.rs b/engine/packages/guard/src/routing/envoy.rs index f2f6ae6ea7..c7895068bf 100644 --- a/engine/packages/guard/src/routing/envoy.rs +++ b/engine/packages/guard/src/routing/envoy.rs @@ -93,6 +93,6 @@ async fn route_envoy_internal( tracing::debug!("authenticated envoy connection"); } - let tunnel = pegboard_envoy::PegboardEnvoyWs::new(ctx.clone()); + let tunnel = pegboard_envoy::PegboardEnvoyWs::new(&ctx); Ok(RoutingOutput::CustomServe(Arc::new(tunnel))) } diff --git a/engine/packages/pegboard-envoy/Cargo.toml b/engine/packages/pegboard-envoy/Cargo.toml index 29ad9a85be..4d7f5f54ab 100644 --- a/engine/packages/pegboard-envoy/Cargo.toml +++ b/engine/packages/pegboard-envoy/Cargo.toml @@ -32,6 +32,7 @@ depot-client.workspace = true depot-client-embedded.workspace = true rivet-runtime.workspace = true rivet-types.workspace = true +rivet-util.workspace = true scc.workspace = true serde_bare.workspace = true serde_json.workspace = true diff --git a/engine/packages/pegboard-envoy/src/lib.rs b/engine/packages/pegboard-envoy/src/lib.rs index bcb5d2c77d..99ebe04cc7 100644 --- a/engine/packages/pegboard-envoy/src/lib.rs +++ b/engine/packages/pegboard-envoy/src/lib.rs @@ -45,12 +45,8 @@ pub struct PegboardEnvoyWs { } impl PegboardEnvoyWs { - pub fn new(ctx: StandaloneCtx) -> Self { - metrics::prepopulate(); - - let service = Self { ctx: ctx.clone() }; - - service + pub fn new(ctx: &StandaloneCtx) -> Self { + Self { ctx: ctx.clone() } } } diff --git a/engine/packages/pegboard-envoy/src/metrics.rs b/engine/packages/pegboard-envoy/src/metrics.rs index d44f177081..ad8303ac40 100644 --- a/engine/packages/pegboard-envoy/src/metrics.rs +++ b/engine/packages/pegboard-envoy/src/metrics.rs @@ -371,72 +371,3 @@ pub fn set_envoy_connection_state( (None, None) => {} } } - -pub fn prepopulate() { - ENVOY_CONNECTED.with_label_values(&["", ""]).set(0); - for state in EnvoyState::ALL { - ENVOY_CONNECTIONS_BY_STATE - .with_label_values(&["", "", "", state.as_str()]) - .set(0); - } - for (state, reasons) in [ - (EnvoyState::Starting, &["websocket_accepted"][..]), - (EnvoyState::Connected, &["init_complete"][..]), - (EnvoyState::Stopping, &["envoy_reported_stopping"][..]), - ( - EnvoyState::Disconnected, - &[ - "init_failed", - "websocket_closed", - "evicted", - "going_away", - "connection_error", - ][..], - ), - (EnvoyState::Lost, &["ping_timeout"][..]), - (EnvoyState::Stopped, &["graceful_shutdown_complete"][..]), - ] { - for reason in reasons { - ENVOY_STATE_TRANSITION_TOTAL - .with_label_values(&["", "", "", state.as_str(), reason]) - .inc_by(0); - } - } - let _ = ENVOY_LIFETIME_SECONDS.with_label_values(&["", ""]); - let _ = ENVOY_PING_LAG_SECONDS.with_label_values(&["", ""]); - for result in ["ok", "no_subscribers", "error"] { - TUNNEL_PUBLISH_TOTAL - .with_label_values(&["", "", result]) - .inc_by(0); - } - TUNNEL_TASKS_ACTIVE.with_label_values(&["", ""]).set(0); - WS_RESPONSES_IN_FLIGHT.set(0); - for task_kind in ["kv", "sqlite_page", "remote_sqlite", "tunnel_message"] { - ACTOR_TASKS_ACTIVE.with_label_values(&[task_kind]).set(0); - } - for branch in ["ws_msg", "completed_task"] { - let _ = WS_TO_TUNNEL_BRANCH_DURATION.with_label_values(&[branch]); - } - for result in ["ok", "error", "timeout"] { - let _ = ACTOR_WAKE_DURATION.with_label_values(&["", "", result]); - } - let _ = SQLITE_COMMIT_ENVOY_DISPATCH_DURATION.with_label_values(&["", ""]); - let _ = SQLITE_COMMIT_ENVOY_RESPONSE_DURATION.with_label_values(&["", ""]); - for request_type in ["get_pages", "commit", "exec", "execute"] { - for result in ["ok", "error"] { - SQLITE_REQUEST_TOTAL - .with_label_values(&["", "", request_type, result]) - .inc_by(0); - let _ = SQLITE_REQUEST_DURATION.with_label_values(&["", "", request_type, result]); - } - for direction in ["request", "response"] { - let _ = SQLITE_REQUEST_PAGES.with_label_values(&["", "", request_type, direction]); - } - let _ = SQLITE_REQUEST_DIRTY_PAGES.with_label_values(&["", "", request_type]); - for direction in ["request", "response"] { - SQLITE_REQUEST_PAYLOAD_BYTES - .with_label_values(&["", "", request_type, direction]) - .inc_by(0); - } - } -} diff --git a/engine/packages/pegboard-envoy/src/ws_to_tunnel_task.rs b/engine/packages/pegboard-envoy/src/ws_to_tunnel_task.rs index 785daf8c41..fcbe313453 100644 --- a/engine/packages/pegboard-envoy/src/ws_to_tunnel_task.rs +++ b/engine/packages/pegboard-envoy/src/ws_to_tunnel_task.rs @@ -45,6 +45,13 @@ use crate::{ const MAX_REMOTE_SQL_BIND_BYTES: usize = 128 * 1024; +/// Max number of pages a single `get_pages` request may ask for. Each requested page number is ~4 +/// bytes on the wire but forces the engine to fetch and materialize up to a full 4 KiB page inside +/// one UDB transaction, so an uncapped list is a large cost-asymmetry amplifier from an untrusted +/// runner. Sized well above observed production read batches (max ~1024 pages); the commit path has +/// its own lower cap (`MAX_COMMIT_DIRTY_PAGES`) since writes batch smaller than reads. +const MAX_GET_PAGES_PER_REQUEST: usize = 8192; + /// Wall-clock threshold above which a single handle_message invocation is logged as a head-of-line /// blocking risk. The ws_to_tunnel_task loop is strictly serial per envoy, so any handler that /// spends longer than this delays every subsequent WS message from the same envoy (including @@ -461,9 +468,25 @@ pub async fn task_inner( let mut term_signal = rivet_runtime::TermSignal::get(); let mut task_manager = TaskManager::new(ctx.clone(), conn.clone()); + // Leaky bucket rate limit on consuming envoy ws messages. The envoy connection multiplexes + // every actor on a runner, so this bounds the rate at which a single untrusted runner can drive + // engine work (task spawns, KV/SQLite ops). Reads are paused while empty, applying TCP + // backpressure to the runner rather than dropping protocol messages. + let mut rate_limit = rivet_util::throttle::RateLimiter::new( + rivet_util::throttle::RateLimitMethod::LeakyBucket { + requests: ctx.config().pegboard().envoy_websocket_rate_limit_requests(), + drip_rate: Duration::from_micros( + ctx.config().pegboard().envoy_websocket_rate_limit_drip_rate_us(), + ), + }, + ); + loop { tokio::select! { - recv = recv_msg(&mut ws_rx, &mut ws_to_tunnel_abort_rx, &mut term_signal) => { + recv = async { + rate_limit.acquire().await; + recv_msg(&mut ws_rx, &mut ws_to_tunnel_abort_rx, &mut term_signal).await + } => { let branch_start = Instant::now(); let branch_result: Result> = async { match recv? { @@ -1461,6 +1484,16 @@ async fn handle_sqlite_get_pages( conn: &Conn, request: protocol::SqliteGetPagesRequest, ) -> Result { + if request.pgnos.len() > MAX_GET_PAGES_PER_REQUEST { + return Ok(protocol::SqliteGetPagesResponse::SqliteErrorResponse( + sqlite_protocol_error_response(&format!( + "sqlite get_pages requested {} pages, exceeding limit {}", + request.pgnos.len(), + MAX_GET_PAGES_PER_REQUEST, + )), + )); + } + validate_sqlite_actor_for_request(ctx, conn, &request.actor_id, request.expected_generation) .await?; diff --git a/engine/packages/pegboard-gateway2/src/metrics.rs b/engine/packages/pegboard-gateway2/src/metrics.rs index e9c7a4acfe..244df7f882 100644 --- a/engine/packages/pegboard-gateway2/src/metrics.rs +++ b/engine/packages/pegboard-gateway2/src/metrics.rs @@ -66,42 +66,3 @@ lazy_static::lazy_static! { *REGISTRY ).unwrap(); } - -pub fn prepopulate() { - const RESULTS: &[&str] = &[ - "success", - "client_disconnect", - "actor_ready_timeout", - "request_timeout", - "envoy_error", - ]; - - for protocol in ["http", "websocket"] { - IN_FLIGHT.with_label_values(&["", "", protocol]).set(0); - IN_FLIGHT_DROPPED_TOTAL - .with_label_values(&["", "", protocol, "client_disconnect"]) - .inc_by(0); - TUNNEL_PING_DURATION.with_label_values(&["", "", protocol]); - LAST_PONG_AGE_SECONDS.with_label_values(&["", "", protocol]); - REQUEST_RETRIES_TOTAL.with_label_values(&["", "", protocol, "1"]); - - for result in RESULTS { - REQUEST_DURATION_SECONDS.with_label_values(&["", "", protocol, result]); - } - - for reason in [ - "server_close", - "client_close", - "abort", - "gc_timeout", - "shutdown", - ] { - CLOSE_SENT_TOTAL - .with_label_values(&["", "", protocol, reason]) - .inc_by(0); - } - } - for result in ["ok", "error", "timeout"] { - WEBSOCKET_OPEN_WAIT_SECONDS.with_label_values(&["", "", result]); - } -} diff --git a/engine/packages/pegboard-gateway2/src/shared_state.rs b/engine/packages/pegboard-gateway2/src/shared_state.rs index 67ad84a7b8..b3de6e5887 100644 --- a/engine/packages/pegboard-gateway2/src/shared_state.rs +++ b/engine/packages/pegboard-gateway2/src/shared_state.rs @@ -160,7 +160,6 @@ pub struct SharedState(Arc); impl SharedState { pub fn new(config: &rivet_config::Config, ups: PubSub) -> Self { - metrics::prepopulate(); init_slow_ping_threshold_from_env(); let gateway_id = protocol::util::generate_gateway_id(); diff --git a/engine/packages/pegboard-gateway2/src/ws_to_tunnel_task.rs b/engine/packages/pegboard-gateway2/src/ws_to_tunnel_task.rs index 66ac4836d3..27dec6914c 100644 --- a/engine/packages/pegboard-gateway2/src/ws_to_tunnel_task.rs +++ b/engine/packages/pegboard-gateway2/src/ws_to_tunnel_task.rs @@ -25,21 +25,26 @@ pub async fn task( let mut ws_rx = ws_rx.lock().await; // Leaky bucket rate limit on consuming ws messages - let pegboard_config = ctx.config().pegboard(); let mut rate_limit = rivet_util::throttle::RateLimiter::new( rivet_util::throttle::RateLimitMethod::LeakyBucket { - requests: pegboard_config.gateway_websocket_rate_limit_requests(), + requests: ctx + .config() + .pegboard() + .gateway_websocket_rate_limit_requests(), drip_rate: Duration::from_millis( - pegboard_config.gateway_websocket_rate_limit_drip_rate_ms(), + ctx.config() + .pegboard() + .gateway_websocket_rate_limit_drip_rate_ms(), ), }, ); loop { - rate_limit.acquire().await; - tokio::select! { - res = ws_rx.try_next() => { + res = async { + rate_limit.acquire().await; + ws_rx.try_next().await + } => { if let Some(msg) = res? { ingress_bytes.fetch_add(msg.len() as u64, Ordering::AcqRel); diff --git a/engine/packages/pegboard-outbound/src/lib.rs b/engine/packages/pegboard-outbound/src/lib.rs index 5d3f3a325d..00697ea7b0 100644 --- a/engine/packages/pegboard-outbound/src/lib.rs +++ b/engine/packages/pegboard-outbound/src/lib.rs @@ -25,8 +25,6 @@ const SSE_OPEN_WARN_THRESHOLD: Duration = Duration::from_secs(5); #[tracing::instrument(skip_all)] pub async fn start(config: rivet_config::Config, pools: rivet_pools::Pools) -> Result<()> { - metrics::prepopulate(); - let cache = rivet_cache::CacheInner::from_env(&config, pools.clone())?; let ctx = StandaloneCtx::new( db::DatabaseKv::new(config.clone(), pools.clone()).await?, diff --git a/engine/packages/pegboard-outbound/src/metrics.rs b/engine/packages/pegboard-outbound/src/metrics.rs index 3f74c822e7..a61cd4fefc 100644 --- a/engine/packages/pegboard-outbound/src/metrics.rs +++ b/engine/packages/pegboard-outbound/src/metrics.rs @@ -42,50 +42,3 @@ lazy_static::lazy_static! { *REGISTRY ).unwrap(); } - -pub fn prepopulate() { - const ERRORS: &[&str] = &[ - "http_error", - "connection_error", - "stream_ended_early", - "invalid_payload", - "downgrade", - "internal", - ]; - const STATUSES: &[&str] = &["429", "503", "5xx", "4xx", "2xx", "other", ""]; - const RESULTS: &[&str] = &[ - "success", - "error_http_429", - "error_http_503", - "error_http_5xx", - "error_http_4xx", - "error_http_other", - "error_connection", - "error_stream_ended", - "error_invalid_payload", - "error_downgrade", - "error_internal", - ]; - const DRAIN_REASONS: &[&str] = &[ - "lifespan_reached", - "going_away", - "actor_lost", - "connection_lost", - "term_signal", - "", - ]; - - for error in ERRORS { - for status in STATUSES { - REQ_ERROR_TOTAL - .with_label_values(&["", "", error, status]) - .inc_by(0); - } - } - - for result in RESULTS { - for drain_reason in DRAIN_REASONS { - REQ_DURATION_SECONDS.with_label_values(&["", "", result, drain_reason]); - } - } -} diff --git a/engine/packages/pegboard/src/ops/actor/create.rs b/engine/packages/pegboard/src/ops/actor/create.rs index a0cb51bca0..b2695c16d9 100644 --- a/engine/packages/pegboard/src/ops/actor/create.rs +++ b/engine/packages/pegboard/src/ops/actor/create.rs @@ -37,7 +37,7 @@ pub struct Output { #[operation] pub async fn pegboard_actor_create(ctx: &OperationCtx, input: &Input) -> Result { - let rate_limiter = RATE_LIMITERS + let rate_limit = RATE_LIMITERS .get_or_init(|| { Cache::builder() .max_capacity(10_000) @@ -59,7 +59,7 @@ pub async fn pegboard_actor_create(ctx: &OperationCtx, input: &Input) -> Result< .await; // Limit actor creation per namespace id - if !rate_limiter.value().lock().await.try_acquire() { + if !rate_limit.value().lock().await.try_acquire() { return Err(crate::errors::Actor::CreationRateLimit.build()); } From 3063bd1b112995653525d469b10b55be044de3d7 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Wed, 24 Jun 2026 15:50:19 -0700 Subject: [PATCH 8/8] [SLOP(claude-opus-4-8-high)] feat(universaldb): postgres leader-resolver driver overhaul --- Cargo.lock | 17 + Cargo.toml | 4 + .../packages/test-deps-docker/src/database.rs | 2 +- engine/packages/universaldb/Cargo.toml | 5 + ...onflict_tracker.rs => conflict_tracker.rs} | 20 +- .../universaldb/src/driver/postgres/codec.rs | 167 +++++++++ .../universaldb/src/driver/postgres/commit.rs | 185 ++++++++++ .../src/driver/postgres/database.rs | 241 +++++++------ .../src/driver/postgres/listener.rs | 228 ++++++++++++ .../universaldb/src/driver/postgres/mod.rs | 5 + .../src/driver/postgres/resolver/apply.rs | 139 ++++++++ .../src/driver/postgres/resolver/lease.rs | 82 +++++ .../src/driver/postgres/resolver/mod.rs | 336 ++++++++++++++++++ .../universaldb/src/driver/postgres/shared.rs | 161 +++++++++ .../src/driver/postgres/transaction.rs | 22 +- .../src/driver/postgres/transaction_task.rs | 283 +++------------ .../src/driver/rocksdb/database.rs | 6 +- .../universaldb/src/driver/rocksdb/mod.rs | 1 - .../src/driver/rocksdb/transaction.rs | 7 +- .../src/driver/rocksdb/transaction_task.rs | 6 +- engine/packages/universaldb/src/lib.rs | 1 + .../packages/universaldb/tests/integration.rs | 10 +- engine/sdks/rust/depot-protocol/Cargo.toml | 1 + .../sdks/rust/universaldb-commit/Cargo.toml | 16 + engine/sdks/rust/universaldb-commit/build.rs | 64 ++++ .../rust/universaldb-commit/src/generated.rs | 1 + .../sdks/rust/universaldb-commit/src/lib.rs | 6 + .../rust/universaldb-commit/src/versioned.rs | 38 ++ .../sdks/schemas/universaldb-commit/v1.bare | 70 ++++ scripts/run/postgres.sh | 2 +- scripts/run/restore-postgres.sh | 2 +- self-host/compose/dev-host/docker-compose.yml | 2 +- .../dev-multidc-multinode/docker-compose.yml | 6 +- .../compose/dev-multidc/docker-compose.yml | 6 +- .../compose/dev-multinode/docker-compose.yml | 2 +- self-host/compose/dev/docker-compose.yml | 2 +- .../compose/template/src/docker-compose.ts | 2 +- .../k8s/engine/12-postgres-statefulset.yaml | 2 +- .../docs/self-hosting/docker-compose.mdx | 2 +- .../docs/self-hosting/docker-container.mdx | 2 +- 40 files changed, 1747 insertions(+), 407 deletions(-) rename engine/packages/universaldb/src/{driver/rocksdb/transaction_conflict_tracker.rs => conflict_tracker.rs} (69%) create mode 100644 engine/packages/universaldb/src/driver/postgres/codec.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/commit.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/listener.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/resolver/apply.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/resolver/lease.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/resolver/mod.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/shared.rs create mode 100644 engine/sdks/rust/universaldb-commit/Cargo.toml create mode 100644 engine/sdks/rust/universaldb-commit/build.rs create mode 100644 engine/sdks/rust/universaldb-commit/src/generated.rs create mode 100644 engine/sdks/rust/universaldb-commit/src/lib.rs create mode 100644 engine/sdks/rust/universaldb-commit/src/versioned.rs create mode 100644 engine/sdks/schemas/universaldb-commit/v1.bare diff --git a/Cargo.lock b/Cargo.lock index b21a5d2859..7b0053c042 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5385,6 +5385,7 @@ name = "rivet-depot-protocol" version = "2.3.7" dependencies = [ "anyhow", + "rivet-util", "rivet-vbare-compiler", "serde", "serde_bare", @@ -5946,6 +5947,17 @@ dependencies = [ "vbare", ] +[[package]] +name = "rivet-universaldb-commit" +version = "2.3.2" +dependencies = [ + "anyhow", + "rivet-vbare-compiler", + "serde", + "serde_bare", + "vbare", +] + [[package]] name = "rivet-ups-broadcast" version = "0.1.0" @@ -8254,6 +8266,7 @@ version = "2.3.7" dependencies = [ "anyhow", "async-trait", + "base64 0.22.1", "deadpool-postgres", "foundationdb-tuple", "futures-util", @@ -8267,17 +8280,21 @@ dependencies = [ "rivet-postgres-util", "rivet-test-deps-docker", "rivet-tracing-utils", + "rivet-universaldb-commit", "rocksdb", + "scc", "serde", "tempfile", "thiserror 1.0.69", "tokio", "tokio-postgres", "tokio-postgres-rustls", + "tokio-util", "tracing", "tracing-subscriber", "url", "uuid", + "vbare", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 511aec200f..a9048c159f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,6 +67,7 @@ members = [ "engine/sdks/rust/depot-protocol", "engine/sdks/rust/test-envoy", "engine/sdks/rust/ups-protocol", + "engine/sdks/rust/universaldb-commit", "rivetkit-rust/packages/actor-persist", "rivetkit-rust/packages/client", "rivetkit-rust/packages/engine-process", @@ -655,6 +656,9 @@ members = [ [workspace.dependencies.rivet-ups-protocol] path = "engine/sdks/rust/ups-protocol" + [workspace.dependencies.rivet-universaldb-commit] + path = "engine/sdks/rust/universaldb-commit" + [profile.dev] overflow-checks = false # "line-tables-only" produces just the line-number DWARF needed for stack diff --git a/engine/packages/test-deps-docker/src/database.rs b/engine/packages/test-deps-docker/src/database.rs index 11532ebf25..955e4d7897 100644 --- a/engine/packages/test-deps-docker/src/database.rs +++ b/engine/packages/test-deps-docker/src/database.rs @@ -61,7 +61,7 @@ impl TestDatabase { }); let docker_config = DockerRunConfig { - image: "postgres:17".to_string(), + image: "postgres:18".to_string(), container_name: container_name.clone(), port_mapping: (port, 5432), env_vars: vec![ diff --git a/engine/packages/universaldb/Cargo.toml b/engine/packages/universaldb/Cargo.toml index 1fdc9b1c99..b2f6b5f045 100644 --- a/engine/packages/universaldb/Cargo.toml +++ b/engine/packages/universaldb/Cargo.toml @@ -9,6 +9,7 @@ edition.workspace = true [dependencies] anyhow.workspace = true async-trait.workspace = true +base64.workspace = true deadpool-postgres.workspace = true foundationdb-tuple.workspace = true futures-util.workspace = true @@ -18,16 +19,20 @@ rand.workspace = true rivet-metrics.workspace = true rivet-postgres-util.workspace = true rivet-tracing-utils.workspace = true +rivet-universaldb-commit.workspace = true rocksdb.workspace = true +scc.workspace = true serde.workspace = true tempfile.workspace = true thiserror.workspace = true tokio-postgres-rustls.workspace = true tokio-postgres.workspace = true +tokio-util.workspace = true tokio.workspace = true tracing.workspace = true url.workspace = true uuid.workspace = true +vbare.workspace = true [dev-dependencies] rivet-config.workspace = true diff --git a/engine/packages/universaldb/src/driver/rocksdb/transaction_conflict_tracker.rs b/engine/packages/universaldb/src/conflict_tracker.rs similarity index 69% rename from engine/packages/universaldb/src/driver/rocksdb/transaction_conflict_tracker.rs rename to engine/packages/universaldb/src/conflict_tracker.rs index 370240760a..127e2bf615 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/transaction_conflict_tracker.rs +++ b/engine/packages/universaldb/src/conflict_tracker.rs @@ -11,7 +11,7 @@ use tokio::sync::Mutex; use crate::options::ConflictRangeType; // Transactions cannot live longer than 5 seconds so we don't need to store transaction conflicts longer than -// that +// that. const TXN_CONFLICT_TTL: Duration = Duration::from_secs(10); #[derive(Debug)] @@ -22,6 +22,15 @@ struct PreviousTransaction { conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, } +/// In-process FoundationDB-style resolver. Holds the last `TXN_CONFLICT_TTL` of committed +/// transactions and rejects a committing transaction if any retained transaction has both an +/// overlapping version window and an overlapping conflict range of a differing type. +/// +/// Used by the rocksdb driver (single process) and by the postgres leader-resolver. The two +/// differ only in where the commit version comes from: rocksdb generates it from the in-process +/// `global_version` counter, while the postgres leader assigns it from the durable +/// `udb_version_seq` so it survives leader failover and matches the versionstamp. For that reason +/// `check_and_insert` takes the commit version from the caller instead of generating it. #[derive(Clone)] pub struct TransactionConflictTracker { // NOTE: We use a mutex because we need to lock reads across all active txns. This could be optimized to @@ -40,18 +49,23 @@ impl TransactionConflictTracker { } } - /// Each number returned is unique. + /// Each number returned is unique. Used by the in-process rocksdb driver to assign both start + /// and commit versions. The postgres leader does not use this; it assigns versions from the + /// durable Postgres sequence. pub fn next_global_version(&self) -> u64 { self.global_version.fetch_add(1, Ordering::SeqCst) } + /// Returns `true` on conflict (same polarity as the original rocksdb tracker). The caller + /// supplies `commit_version` (e.g. `nextval('udb_version_seq')` on the postgres leader, or + /// `next_global_version()` on rocksdb) so version assignment stays the caller's responsibility. pub async fn check_and_insert( &self, txn1_start_version: u64, + txn1_commit_version: u64, txn1_conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, ) -> bool { let mut txns = self.txns.lock().await; - let txn1_commit_version = self.next_global_version(); // Prune old entries txns.retain(|txn| txn.insert_instant.elapsed() < TXN_CONFLICT_TTL); diff --git a/engine/packages/universaldb/src/driver/postgres/codec.rs b/engine/packages/universaldb/src/driver/postgres/codec.rs new file mode 100644 index 0000000000..432ede2f34 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/codec.rs @@ -0,0 +1,167 @@ +use anyhow::Result; +use rivet_universaldb_commit::{self as proto, versioned}; +use vbare::OwnedVersionedData; + +use crate::{ + options::{ConflictRangeType, MutationType}, + tx_ops::Operation, +}; + +/// Decoded form of a `udb_commit_requests.payload` blob. +/// +/// `read_version` is intentionally omitted: it is also denormalized into the `read_version` column, +/// which is what the leader's drain reads, so decoding it here would be dead. +pub struct DecodedCommit { + pub conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, + pub operations: Vec, +} + +/// Encode a follower's commit request to the versioned BARE wire format with an embedded version +/// header so a leader running older or newer code can still decode it during a rolling deploy. +pub fn encode_commit_request( + read_version: u64, + conflict_ranges: &[(Vec, Vec, ConflictRangeType)], + operations: &[Operation], +) -> Result> { + let request = proto::CommitRequest { + read_version, + conflict_ranges: conflict_ranges + .iter() + .map(|(begin, end, kind)| proto::ConflictRange { + begin: begin.clone(), + end: end.clone(), + kind: conflict_range_type_to_proto(*kind), + }) + .collect(), + operations: operations.iter().map(operation_to_proto).collect(), + }; + + versioned::CommitRequest::wrap_latest(request) + .serialize_with_embedded_version(proto::PROTOCOL_VERSION) +} + +/// Decode a `udb_commit_requests.payload` blob produced by [`encode_commit_request`]. +pub fn decode_commit_request(payload: &[u8]) -> Result { + let request = versioned::CommitRequest::deserialize_with_embedded_version(payload)?; + + let conflict_ranges = request + .conflict_ranges + .into_iter() + .map(|range| { + ( + range.begin, + range.end, + conflict_range_type_from_proto(range.kind), + ) + }) + .collect(); + + let operations = request + .operations + .into_iter() + .map(operation_from_proto) + .collect(); + + Ok(DecodedCommit { + conflict_ranges, + operations, + }) +} + +fn conflict_range_type_to_proto(kind: ConflictRangeType) -> proto::ConflictRangeType { + match kind { + ConflictRangeType::Read => proto::ConflictRangeType::Read, + ConflictRangeType::Write => proto::ConflictRangeType::Write, + } +} + +fn conflict_range_type_from_proto(kind: proto::ConflictRangeType) -> ConflictRangeType { + match kind { + proto::ConflictRangeType::Read => ConflictRangeType::Read, + proto::ConflictRangeType::Write => ConflictRangeType::Write, + } +} + +fn operation_to_proto(op: &Operation) -> proto::Operation { + match op { + Operation::SetValue { key, value } => proto::Operation::SetValue(proto::SetValue { + key: key.clone(), + value: value.clone(), + }), + Operation::Clear { key } => proto::Operation::Clear(proto::Clear { key: key.clone() }), + Operation::ClearRange { begin, end } => proto::Operation::ClearRange(proto::ClearRange { + begin: begin.clone(), + end: end.clone(), + }), + Operation::AtomicOp { + key, + param, + op_type, + } => proto::Operation::AtomicOp(proto::AtomicOp { + key: key.clone(), + param: param.clone(), + op_type: mutation_type_to_proto(*op_type), + }), + } +} + +fn operation_from_proto(op: proto::Operation) -> Operation { + match op { + proto::Operation::SetValue(proto::SetValue { key, value }) => { + Operation::SetValue { key, value } + } + proto::Operation::Clear(proto::Clear { key }) => Operation::Clear { key }, + proto::Operation::ClearRange(proto::ClearRange { begin, end }) => { + Operation::ClearRange { begin, end } + } + proto::Operation::AtomicOp(proto::AtomicOp { + key, + param, + op_type, + }) => Operation::AtomicOp { + key, + param, + op_type: mutation_type_from_proto(op_type), + }, + } +} + +fn mutation_type_to_proto(op_type: MutationType) -> proto::MutationType { + match op_type { + MutationType::Add => proto::MutationType::Add, + MutationType::And => proto::MutationType::And, + MutationType::BitAnd => proto::MutationType::BitAnd, + MutationType::Or => proto::MutationType::Or, + MutationType::BitOr => proto::MutationType::BitOr, + MutationType::Xor => proto::MutationType::Xor, + MutationType::BitXor => proto::MutationType::BitXor, + MutationType::AppendIfFits => proto::MutationType::AppendIfFits, + MutationType::Max => proto::MutationType::Max, + MutationType::Min => proto::MutationType::Min, + MutationType::SetVersionstampedKey => proto::MutationType::SetVersionstampedKey, + MutationType::SetVersionstampedValue => proto::MutationType::SetVersionstampedValue, + MutationType::ByteMin => proto::MutationType::ByteMin, + MutationType::ByteMax => proto::MutationType::ByteMax, + MutationType::CompareAndClear => proto::MutationType::CompareAndClear, + } +} + +fn mutation_type_from_proto(op_type: proto::MutationType) -> MutationType { + match op_type { + proto::MutationType::Add => MutationType::Add, + proto::MutationType::And => MutationType::And, + proto::MutationType::BitAnd => MutationType::BitAnd, + proto::MutationType::Or => MutationType::Or, + proto::MutationType::BitOr => MutationType::BitOr, + proto::MutationType::Xor => MutationType::Xor, + proto::MutationType::BitXor => MutationType::BitXor, + proto::MutationType::AppendIfFits => MutationType::AppendIfFits, + proto::MutationType::Max => MutationType::Max, + proto::MutationType::Min => MutationType::Min, + proto::MutationType::SetVersionstampedKey => MutationType::SetVersionstampedKey, + proto::MutationType::SetVersionstampedValue => MutationType::SetVersionstampedValue, + proto::MutationType::ByteMin => MutationType::ByteMin, + proto::MutationType::ByteMax => MutationType::ByteMax, + proto::MutationType::CompareAndClear => MutationType::CompareAndClear, + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/commit.rs b/engine/packages/universaldb/src/driver/postgres/commit.rs new file mode 100644 index 0000000000..e61714dd0e --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/commit.rs @@ -0,0 +1,185 @@ +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result}; + +use crate::{error::DatabaseError, options::ConflictRangeType, tx_ops::Operation}; + +use super::{ + codec, + shared::{LeaseInfo, PostgresShared, commit_channel, reply_channel}, +}; + +/// How long to wait for a leader to be elected before giving up a submit as retryable. +const LEADER_WAIT_TIMEOUT: Duration = Duration::from_secs(5); +/// Backstop poll cadence while waiting for a commit result, in case a reply NOTIFY is missed. +const RESULT_POLL_INTERVAL: Duration = Duration::from_millis(250); + +/// Submit a follower transaction's commit to the leader and await the result. +/// +/// `read_version` is the watermark captured when this transaction opened its read snapshot. A pure +/// snapshot read-only transaction (no operations and no read conflict ranges) submits nothing. +pub async fn submit( + shared: &Arc, + read_version: i64, + operations: Vec, + conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, +) -> Result<()> { + // A transaction with no writes and no serializable read ranges has nothing to order or + // validate; it never needs the leader. + if operations.is_empty() + && conflict_ranges + .iter() + .all(|(_, _, kind)| matches!(kind, ConflictRangeType::Write)) + { + return Ok(()); + } + + let lease = wait_for_leader(shared).await?; + let payload = + codec::encode_commit_request(read_version.max(0) as u64, &conflict_ranges, &operations) + .context("failed to encode commit request")?; + let reply_channel = reply_channel(&shared.node_id); + + // Subscribe to our reply channel before inserting so we cannot miss the leader's NOTIFY. + let mut reply_rx = shared.listener.listen(&reply_channel).await; + + let conn = shared + .pool + .get() + .await + .context("failed to get connection for commit submit")?; + + let id: i64 = conn + .query_one( + "INSERT INTO udb_commit_requests (epoch, read_version, payload, reply_channel) + VALUES ($1, $2, $3, $4) + RETURNING id", + &[&lease.epoch, &read_version, &payload, &reply_channel], + ) + .await + .context("failed to enqueue commit request")? + .get(0); + + // Wake the leader's drain loop. + if let Err(err) = conn + .execute( + "SELECT pg_notify($1, $2)", + &[&commit_channel(&lease.leader_addr), &id.to_string()], + ) + .await + { + tracing::debug!( + ?err, + "failed to notify leader; relying on its poll backstop" + ); + } + + // Release the connection before waiting so a long wait does not pin a pool slot. The request + // row is durable, so await_result re-acquires a connection per poll. + drop(conn); + + await_result(shared, id, lease.epoch, &mut reply_rx).await +} + +/// Wait for a known leader, returning a retryable error if none is elected in time. +async fn wait_for_leader(shared: &Arc) -> Result { + let deadline = Instant::now() + LEADER_WAIT_TIMEOUT; + loop { + if let Some(lease) = shared.current_lease() { + return Ok(lease); + } + if Instant::now() >= deadline { + return Err(DatabaseError::NotCommitted.into()); + } + tokio::time::sleep(RESULT_POLL_INTERVAL).await; + } +} + +/// Poll the request row until it reaches a terminal status, woken by reply NOTIFYs with a polling +/// backstop. Bails as retryable if the leader epoch advances (our request is now orphaned and will +/// never be applied, so it is definitively not committed). +async fn await_result( + shared: &Arc, + id: i64, + submit_epoch: i64, + reply_rx: &mut tokio::sync::broadcast::Receiver, +) -> Result<()> { + loop { + // Re-acquire a connection per poll: the request row is durable, so a transient pool/query + // error just means we retry the poll rather than failing a possibly-applied commit. + match read_status(shared, id).await { + Ok(Some(Status::Committed)) => return Ok(()), + Ok(Some(Status::Conflict)) => return Err(DatabaseError::NotCommitted.into()), + Ok(Some(Status::Pending)) => {} + Ok(None) => { + // The row was GC'd before we observed a terminal status. Treat as not committed + // and let the retry loop resubmit. + return Err(DatabaseError::NotCommitted.into()); + } + Err(err) => { + tracing::debug!(?err, "transient error polling commit status, retrying"); + } + } + + // If a new leader took over, our old-epoch request will never be claimed. + if let Some(current) = shared.current_lease() { + if current.epoch != submit_epoch { + return Err(DatabaseError::NotCommitted.into()); + } + } + + tokio::select! { + res = reply_rx.recv() => { + match res { + Ok(_) | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + *reply_rx = shared + .listener + .listen(&reply_channel(&shared.node_id)) + .await; + } + } + } + _ = tokio::time::sleep(RESULT_POLL_INTERVAL) => {} + } + } +} + +enum Status { + Pending, + Committed, + Conflict, +} + +/// Read the current status of a commit request. `Ok(None)` means the row no longer exists. +async fn read_status(shared: &Arc, id: i64) -> Result> { + let conn = shared + .pool + .get() + .await + .context("failed to get connection for commit status poll")?; + + let row = conn + .query_opt( + "SELECT status FROM udb_commit_requests WHERE id = $1", + &[&id], + ) + .await + .context("failed to read commit request status")?; + + let Some(row) = row else { + return Ok(None); + }; + + let status: String = row.get(0); + let status = match status.as_str() { + "committed" => Status::Committed, + "conflict" => Status::Conflict, + // 'pending' or any in-flight state. + _ => Status::Pending, + }; + Ok(Some(status)) +} diff --git a/engine/packages/universaldb/src/driver/postgres/database.rs b/engine/packages/universaldb/src/driver/postgres/database.rs index 4d50d08ab3..6e24cf647d 100644 --- a/engine/packages/universaldb/src/driver/postgres/database.rs +++ b/engine/packages/universaldb/src/driver/postgres/database.rs @@ -13,6 +13,7 @@ use rivet_postgres_util::build_tls_config; use tokio::task::JoinHandle; use tokio_postgres_rustls::MakeRustlsConnect; use url::Url; +use uuid::Uuid; use crate::{ RetryableTransaction, Transaction, @@ -22,9 +23,15 @@ use crate::{ utils::{MaybeCommitted, calculate_tx_retry_backoff}, }; -use super::transaction::PostgresTransactionDriver; +use super::{ + listener::PgListener, resolver, shared::PostgresShared, transaction::PostgresTransactionDriver, +}; -const GC_INTERVAL: Duration = Duration::from_secs(5); +const GC_INTERVAL: Duration = Duration::from_secs(30); +/// Terminal and orphaned commit-request rows older than this are garbage collected. Must be well +/// beyond the longest a follower could spend awaiting a result, so a result is never deleted before +/// it is observed. +const COMMIT_ROW_MAX_AGE_SECS: i64 = 60; #[derive(Clone, Debug)] pub struct PostgresConfig { @@ -50,7 +57,7 @@ impl PostgresConfig { } pub struct PostgresDatabaseDriver { - pool: Pool, + shared: Arc, max_retries: AtomicI32, gc_handle: JoinHandle<()>, } @@ -63,7 +70,60 @@ impl PostgresDatabaseDriver { "creating PostgresDatabaseDriver" ); - // Create deadpool config from connection string + let ssl_disabled = if let Ok(url) = Url::parse(&config.connection_string) { + url.query_pairs() + .any(|(k, v)| k == "sslmode" && v == "disable") + } else { + false + }; + + let pool = Self::build_pool(&config, ssl_disabled)?; + + // Initialize the schema (idempotent). + { + let conn = pool + .get() + .await + .context("failed to get connection from postgres pool")?; + Self::init_schema(&conn).await?; + } + + // Unique per-process node id (no hyphens) used to name this node's NOTIFY channels. Kept + // short so `udb_commit_` stays within Postgres's 63-byte identifier limit. + let node_id = Uuid::new_v4().simple().to_string(); + + let listener = PgListener::new( + config.connection_string.clone(), + ssl_disabled, + config + .ssl_config + .as_ref() + .and_then(|c| c.ssl_root_cert_path.clone()), + config + .ssl_config + .as_ref() + .and_then(|c| c.ssl_client_cert_path.clone()), + config + .ssl_config + .as_ref() + .and_then(|c| c.ssl_client_key_path.clone()), + ); + + let shared = PostgresShared::new(pool, node_id, listener); + + // Every node runs the resolver; only the elected leader drains the commit queue. + resolver::spawn(shared.clone()); + + let gc_handle = Self::spawn_gc(shared.clone()); + + Ok(PostgresDatabaseDriver { + shared, + max_retries: AtomicI32::new(100), + gc_handle, + }) + } + + fn build_pool(config: &PostgresConfig, ssl_disabled: bool) -> Result { let mut pool_config = Config::new(); pool_config.url = Some(config.connection_string.clone()); pool_config.pool = Some(PoolConfig { @@ -74,21 +134,10 @@ impl PostgresDatabaseDriver { recycling_method: RecyclingMethod::Fast, }); - tracing::debug!("creating Postgres pool"); - - let ssl_disabled = if let Ok(url) = Url::parse(&config.connection_string) { - url.query_pairs() - .any(|(k, v)| k == "sslmode" && v == "disable") - } else { - false - }; - - let pool = if ssl_disabled { - let tls = tokio_postgres::NoTls; - + if ssl_disabled { pool_config - .create_pool(Some(Runtime::Tokio1), tls) - .context("failed to create postgres connection pool")? + .create_pool(Some(Runtime::Tokio1), tokio_postgres::NoTls) + .context("failed to create postgres connection pool") } else { let tls_config = build_tls_config( config @@ -104,139 +153,87 @@ impl PostgresDatabaseDriver { .as_ref() .and_then(|c| c.ssl_client_key_path.as_ref()), )?; - let tls = MakeRustlsConnect::new(tls_config); - pool_config - .create_pool(Some(Runtime::Tokio1), tls) - .context("failed to create postgres connection pool")? - }; - - tracing::debug!("Getting Postgres connection from pool"); - // Get a connection from the pool to create the table - let conn = pool - .get() - .await - .context("failed to get connection from postgres pool")?; - - // Enable btree gist - conn.execute("CREATE EXTENSION IF NOT EXISTS btree_gist", &[]) - .await - .context("failed to create btree_gist extension")?; - - conn.execute("CREATE UNLOGGED SEQUENCE IF NOT EXISTS global_version_seq START WITH 1 INCREMENT BY 1 MINVALUE 1", &[]) - .await - .context("failed to create global version sequence")?; + .create_pool(Some(Runtime::Tokio1), MakeRustlsConnect::new(tls_config)) + .context("failed to create postgres connection pool") + } + } - // Create the KV table if it doesn't exist - conn.execute( + async fn init_schema(conn: &deadpool_postgres::Client) -> Result<()> { + // Durable latest-value store. + conn.batch_execute( "CREATE TABLE IF NOT EXISTS kv ( key BYTEA PRIMARY KEY, value BYTEA NOT NULL - )", - &[], - ) - .await - .context("failed to create kv table")?; - - // Create range_type type if it doesn't exist - conn.execute( - "DO $$ BEGIN - CREATE TYPE range_type AS ENUM ('read', 'write'); - EXCEPTION - WHEN duplicate_object THEN null; - END $$", - &[], - ) - .await - .context("failed to create range_type enum")?; - - // Create bytearange type if it doesn't exist - conn.execute( - "DO $$ BEGIN - CREATE TYPE bytearange AS RANGE ( - SUBTYPE = bytea, - SUBTYPE_OPCLASS = bytea_ops - ); - EXCEPTION - WHEN duplicate_object THEN null; - END $$", - &[], + ); + + CREATE TABLE IF NOT EXISTS udb_lease ( + id INT PRIMARY KEY DEFAULT 1 CHECK (id = 1), + epoch BIGINT NOT NULL, + leader_addr TEXT NOT NULL, + durable_version BIGINT NOT NULL DEFAULT 0, + expires_at TIMESTAMPTZ NOT NULL + ); + + CREATE SEQUENCE IF NOT EXISTS udb_version_seq AS BIGINT + START WITH 1 INCREMENT BY 1 MINVALUE 1; + + CREATE TABLE IF NOT EXISTS udb_commit_requests ( + id BIGSERIAL PRIMARY KEY, + epoch BIGINT NOT NULL, + read_version BIGINT NOT NULL, + payload BYTEA NOT NULL, + reply_channel TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + commit_version BIGINT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + ); + + CREATE INDEX IF NOT EXISTS udb_commit_requests_pending + ON udb_commit_requests (id) WHERE status = 'pending';", ) .await - .context("failed to create bytearange type")?; - - // Create the conflict ranges table for non-snapshot reads - // This enforces consistent reads for ranges by preventing overlapping conflict ranges - conn.execute( - "CREATE UNLOGGED TABLE IF NOT EXISTS conflict_ranges ( - range_data BYTEARANGE NOT NULL, - conflict_type range_type NOT NULL, - start_version BIGINT NOT NULL, - commit_version BIGINT NOT NULL, - ts timestamp NOT NULL DEFAULT now(), - - EXCLUDE USING gist ( - -- Conflict if byte range overlaps... - range_data WITH &&, - -- And if conflict types are different... - conflict_type WITH <>, - -- And if the txn versions overlap... - int8range(start_version, commit_version, '[]') WITH &&, - -- But not if the start_version is the same (from the same txn) - start_version WITH <> - ) - )", - &[], - ) - .await - .context("failed to create conflict_ranges table")?; + .context("failed to initialize postgres schema")?; - // Create index on ts column for efficient garbage collection - conn.execute( - "CREATE INDEX IF NOT EXISTS idx_conflict_ranges_ts ON conflict_ranges (ts)", - &[], - ) - .await - .context("failed to create index on conflict_ranges ts column")?; + Ok(()) + } - let pool2 = pool.clone(); - let gc_handle = tokio::spawn(async move { + fn spawn_gc(shared: Arc) -> JoinHandle<()> { + tokio::spawn(async move { let mut interval = tokio::time::interval(GC_INTERVAL); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { interval.tick().await; - tracing::debug!(status=?pool2.status(), "postgres pool status"); + let conn = match shared.pool.get().await { + Ok(conn) => conn, + Err(err) => { + tracing::debug!(?err, "failed to get connection for commit gc"); + continue; + } + }; - // NOTE: Transactions have a max limit of 5 seconds, we delete after 10 seconds for extra padding - // Delete old conflict ranges if let Err(err) = conn .execute( - "DELETE FROM conflict_ranges where ts < now() - interval '10 seconds'", - &[], + "DELETE FROM udb_commit_requests + WHERE created_at < now() - ($1::bigint * interval '1 second')", + &[&COMMIT_ROW_MAX_AGE_SECS], ) .await { - tracing::error!(?err, "failed postgres gc task"); + tracing::error!(?err, "failed postgres commit-queue gc"); } } - }); - - Ok(PostgresDatabaseDriver { - pool, - max_retries: AtomicI32::new(100), - gc_handle, }) } } impl DatabaseDriver for PostgresDatabaseDriver { fn create_txn(&self) -> Result { - // Pass the connection pool and config to the transaction driver - Ok(Transaction::new(Arc::new( - PostgresTransactionDriver::with_config(self.pool.clone()), - ))) + Ok(Transaction::new(Arc::new(PostgresTransactionDriver::new( + self.shared.clone(), + )))) } fn run<'a>( diff --git a/engine/packages/universaldb/src/driver/postgres/listener.rs b/engine/packages/universaldb/src/driver/postgres/listener.rs new file mode 100644 index 0000000000..e5a7919dfe --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/listener.rs @@ -0,0 +1,228 @@ +use std::{path::PathBuf, sync::Arc, time::Duration}; + +use futures_util::future::poll_fn; +use rivet_postgres_util::build_tls_config; +use scc::HashMap; +use tokio::{ + io::{AsyncRead, AsyncWrite}, + sync::{Mutex, broadcast}, +}; +use tokio_postgres::AsyncMessage; +use tokio_postgres_rustls::MakeRustlsConnect; + +/// How long to wait between reconnect attempts for the dedicated LISTEN connection. +const RECONNECT_BACKOFF: Duration = Duration::from_secs(1); +/// Capacity of each channel's broadcast buffer. Notifications are wakeup signals with a polling +/// backstop, so a lagged receiver only delays a wake, never drops a durable commit. +const BROADCAST_CAPACITY: usize = 1024; + +struct Subscription { + tx: broadcast::Sender, +} + +/// Owns a single dedicated Postgres connection used exclusively for `LISTEN`. Demultiplexes +/// incoming `NOTIFY` payloads to per-channel broadcast senders and re-`LISTEN`s every registered +/// channel after a reconnect. +/// +/// This is separate from the deadpool pool because deadpool recycles connections and drops the +/// async notification stream; LISTEN requires owning the connection's message stream directly. +pub struct PgListener { + conn_str: String, + ssl_disabled: bool, + ssl_root_cert_path: Option, + ssl_client_cert_path: Option, + ssl_client_key_path: Option, + channels: Arc>, + client: Arc>>, +} + +impl PgListener { + pub fn new( + conn_str: String, + ssl_disabled: bool, + ssl_root_cert_path: Option, + ssl_client_cert_path: Option, + ssl_client_key_path: Option, + ) -> Self { + let channels: Arc> = Arc::new(HashMap::new()); + let client: Arc>> = Arc::new(Mutex::new(None)); + + tokio::spawn(Self::connection_lifecycle( + conn_str.clone(), + ssl_disabled, + ssl_root_cert_path.clone(), + ssl_client_cert_path.clone(), + ssl_client_key_path.clone(), + channels.clone(), + client.clone(), + )); + + Self { + conn_str, + ssl_disabled, + ssl_root_cert_path, + ssl_client_cert_path, + ssl_client_key_path, + channels, + client, + } + } + + /// Subscribe to a channel, registering a `LISTEN` if this is the first subscriber. Returns a + /// broadcast receiver of notification payloads. Idempotent per channel. + pub async fn listen(&self, channel: &str) -> broadcast::Receiver { + match self.channels.entry_async(channel.to_string()).await { + scc::hash_map::Entry::Occupied(entry) => entry.get().tx.subscribe(), + scc::hash_map::Entry::Vacant(entry) => { + let (tx, rx) = broadcast::channel(BROADCAST_CAPACITY); + entry.insert_entry(Subscription { tx }); + + // Best-effort immediate LISTEN; the lifecycle task re-LISTENs on reconnect. + if let Some(client) = &*self.client.lock().await { + if let Err(err) = client.execute(&format!("LISTEN \"{channel}\""), &[]).await { + tracing::warn!(?err, %channel, "failed to LISTEN, will retry on reconnect"); + } + } + + rx + } + } + } + + async fn connection_lifecycle( + conn_str: String, + ssl_disabled: bool, + ssl_root_cert_path: Option, + ssl_client_cert_path: Option, + ssl_client_key_path: Option, + channels: Arc>, + client: Arc>>, + ) { + loop { + let connected = if ssl_disabled { + Self::connect_and_run(&conn_str, tokio_postgres::NoTls, &channels, &client).await + } else { + match build_tls_config( + ssl_root_cert_path.as_ref(), + ssl_client_cert_path.as_ref(), + ssl_client_key_path.as_ref(), + ) { + Ok(tls_config) => { + Self::connect_and_run( + &conn_str, + MakeRustlsConnect::new(tls_config), + &channels, + &client, + ) + .await + } + Err(err) => { + tracing::error!(?err, "failed to build listener TLS config"); + false + } + } + }; + + if !connected { + tokio::time::sleep(RECONNECT_BACKOFF).await; + } + } + } + + /// Connects, re-LISTENs all channels, then drives the notification poll loop until the + /// connection closes. Returns `true` if a connection was successfully established (so the caller + /// can skip the reconnect backoff). + async fn connect_and_run( + conn_str: &str, + tls: T, + channels: &Arc>, + client: &Arc>>, + ) -> bool + where + T: tokio_postgres::tls::MakeTlsConnect, + T::Stream: AsyncRead + AsyncWrite + Unpin + Send + 'static, + T::TlsConnect: Send, + >::Future: Send, + { + let (new_client, connection) = match tokio_postgres::connect(conn_str, tls).await { + Ok(pair) => pair, + Err(err) => { + tracing::error!(?err, "failed to connect postgres listener"); + return false; + } + }; + + let channels_poll = channels.clone(); + let poll_handle = + tokio::spawn(async move { Self::poll_connection(connection, channels_poll).await }); + + // Re-LISTEN all registered channels on the fresh connection. + let mut registered = Vec::new(); + channels + .iter_async(|k, _| { + registered.push(k.clone()); + true + }) + .await; + for channel in ®istered { + if let Err(err) = new_client + .execute(&format!("LISTEN \"{channel}\""), &[]) + .await + { + tracing::error!(?err, %channel, "failed to re-LISTEN channel after reconnect"); + } + } + + *client.lock().await = Some(new_client); + + // Block until the poll loop ends (connection closed or errored). + let _ = poll_handle.await; + + *client.lock().await = None; + + true + } + + async fn poll_connection( + mut connection: tokio_postgres::Connection, + channels: Arc>, + ) where + S: AsyncRead + AsyncWrite + Unpin, + T: AsyncRead + AsyncWrite + Unpin, + { + loop { + match poll_fn(|cx| connection.poll_message(cx)).await { + Some(Ok(AsyncMessage::Notification(note))) => { + if let Some(sub) = channels.get_async(note.channel()).await { + // Ignore send errors: no active receiver just means no one is waiting + // right now; the polling backstop covers them. + let _ = sub.tx.send(note.payload().to_string()); + } + } + Some(Ok(_)) => {} + Some(Err(err)) => { + tracing::warn!(?err, "postgres listener connection error"); + break; + } + None => { + tracing::warn!("postgres listener connection closed"); + break; + } + } + } + } +} + +impl Clone for PgListener { + fn clone(&self) -> Self { + Self { + conn_str: self.conn_str.clone(), + ssl_disabled: self.ssl_disabled, + ssl_root_cert_path: self.ssl_root_cert_path.clone(), + ssl_client_cert_path: self.ssl_client_cert_path.clone(), + ssl_client_key_path: self.ssl_client_key_path.clone(), + channels: self.channels.clone(), + client: self.client.clone(), + } + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/mod.rs b/engine/packages/universaldb/src/driver/postgres/mod.rs index 1c24f9cb94..64f4bbd1bf 100644 --- a/engine/packages/universaldb/src/driver/postgres/mod.rs +++ b/engine/packages/universaldb/src/driver/postgres/mod.rs @@ -1,4 +1,9 @@ +mod codec; +mod commit; mod database; +mod listener; +mod resolver; +mod shared; mod transaction; mod transaction_task; diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs b/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs new file mode 100644 index 0000000000..d9729de884 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs @@ -0,0 +1,139 @@ +use anyhow::{Context, Result}; +use deadpool_postgres::Transaction; + +use crate::{ + atomic::apply_atomic_op, options::MutationType, tuple::Versionstamp, tx_ops::Operation, + versionstamp::substitute_raw_versionstamp, +}; + +/// Apply a winning transaction's operations to `kv` inside the leader's batch txn. +/// +/// `commit_version` is the Postgres-resolved version assigned to this commit (`nextval`). It is +/// substituted into the 8-byte committed-version slot of any versionstamped key/value so +/// versionstamps are globally monotonic with commit order across all follower processes. +pub async fn apply( + txn: &Transaction<'_>, + operations: Vec, + commit_version: u64, +) -> Result<()> { + // Distinguishes multiple versionstamped operations within a single commit so their 10-byte + // stamps stay unique (8-byte version shared, 2-byte counter incremented). + let mut versionstamp_counter: u16 = 0; + + for op in operations { + match op { + Operation::SetValue { key, value } => { + upsert(txn, &key, &value).await?; + } + Operation::Clear { key } => { + txn.execute("DELETE FROM kv WHERE key = $1", &[&key]) + .await + .context("failed to clear key")?; + } + Operation::ClearRange { begin, end } => { + txn.execute( + "DELETE FROM kv WHERE key >= $1 AND key < $2", + &[&begin, &end], + ) + .await + .context("failed to clear range")?; + } + Operation::AtomicOp { + key, + param, + op_type, + } => { + apply_atomic( + txn, + key, + param, + op_type, + commit_version, + &mut versionstamp_counter, + ) + .await?; + } + } + } + + Ok(()) +} + +async fn apply_atomic( + txn: &Transaction<'_>, + key: Vec, + param: Vec, + op_type: MutationType, + commit_version: u64, + versionstamp_counter: &mut u16, +) -> Result<()> { + match op_type { + MutationType::SetVersionstampedKey => { + let versionstamp = build_versionstamp(commit_version, versionstamp_counter); + let key = substitute_raw_versionstamp(key, &versionstamp) + .map_err(anyhow::Error::msg) + .context("failed substituting versionstamped key")?; + upsert(txn, &key, ¶m).await?; + } + MutationType::SetVersionstampedValue => { + let versionstamp = build_versionstamp(commit_version, versionstamp_counter); + let value = substitute_raw_versionstamp(param, &versionstamp) + .map_err(anyhow::Error::msg) + .context("failed substituting versionstamped value")?; + upsert(txn, &key, &value).await?; + } + // Read-modify-write atomics: the leader is the single writer, so reading the live value + // inside the apply txn and writing the result is serializable with no lost update. + MutationType::Add + | MutationType::And + | MutationType::BitAnd + | MutationType::Or + | MutationType::BitOr + | MutationType::Xor + | MutationType::BitXor + | MutationType::AppendIfFits + | MutationType::Max + | MutationType::Min + | MutationType::ByteMin + | MutationType::ByteMax + | MutationType::CompareAndClear => { + let current = txn + .query_opt("SELECT value FROM kv WHERE key = $1", &[&key]) + .await + .context("failed to read current value for atomic op")? + .map(|row| row.get::<_, Vec>(0)); + + let new_value = apply_atomic_op(current.as_deref(), ¶m, op_type); + + if let Some(new_value) = new_value { + upsert(txn, &key, &new_value).await?; + } else { + txn.execute("DELETE FROM kv WHERE key = $1", &[&key]) + .await + .context("failed to clear key after atomic op")?; + } + } + } + + Ok(()) +} + +async fn upsert(txn: &Transaction<'_>, key: &[u8], value: &[u8]) -> Result<()> { + txn.execute( + "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2", + &[&key, &value], + ) + .await + .context("failed to upsert kv")?; + Ok(()) +} + +/// Build a 10-byte versionstamp (plus the 2 user-version bytes the substitution helper ignores) +/// from the Postgres-resolved commit version and a per-commit counter. +fn build_versionstamp(commit_version: u64, counter: &mut u16) -> Versionstamp { + let mut bytes = [0u8; 12]; + bytes[0..8].copy_from_slice(&commit_version.to_be_bytes()); + bytes[8..10].copy_from_slice(&counter.to_be_bytes()); + *counter = counter.wrapping_add(1); + Versionstamp::from(bytes) +} diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs b/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs new file mode 100644 index 0000000000..a17b6c7759 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs @@ -0,0 +1,82 @@ +use anyhow::{Context, Result}; +use deadpool_postgres::Pool; + +use crate::driver::postgres::shared::LEASE_ID; + +/// Lease time-to-live. A leader renews well within this; a candidate may take over only after it +/// expires. +pub const LEASE_TTL_SECS: i64 = 10; + +/// Outcome of a leadership acquisition attempt. +pub struct Acquired { + pub epoch: i64, +} + +/// Attempt to acquire or take over the leader lease via an epoch CAS. Succeeds if there is no lease +/// row yet, or the existing lease has expired. Bumps `epoch` on every successful acquisition so a +/// superseded old leader is fenced out. +pub async fn try_acquire(pool: &Pool, node_id: &str) -> Result> { + let conn = pool + .get() + .await + .context("failed to get connection for lease acquire")?; + + // Take over an expired (or absent) lease. The INSERT seeds the singleton row on first ever + // election; thereafter the UPDATE path runs. + let row = conn + .query_opt( + "INSERT INTO udb_lease (id, epoch, leader_addr, durable_version, expires_at) + VALUES ($1, 1, $2, 0, now() + ($3 || ' seconds')::interval) + ON CONFLICT (id) DO UPDATE + SET epoch = udb_lease.epoch + 1, + leader_addr = EXCLUDED.leader_addr, + expires_at = now() + ($3 || ' seconds')::interval + WHERE udb_lease.expires_at < now() + RETURNING epoch", + &[&LEASE_ID, &node_id, &LEASE_TTL_SECS.to_string()], + ) + .await + .context("failed to run lease acquire query")?; + + Ok(row.map(|row| Acquired { epoch: row.get(0) })) +} + +/// Renew the lease, fenced on this leader's epoch. Returns `false` if the lease was lost (another +/// node took over, bumping the epoch), in which case the caller must step down. +pub async fn renew(pool: &Pool, node_id: &str, epoch: i64) -> Result { + let conn = pool + .get() + .await + .context("failed to get connection for lease renew")?; + + let updated = conn + .execute( + "UPDATE udb_lease + SET expires_at = now() + ($3 || ' seconds')::interval + WHERE id = $1 AND epoch = $2 AND leader_addr = $4", + &[&LEASE_ID, &epoch, &LEASE_TTL_SECS.to_string(), &node_id], + ) + .await + .context("failed to renew lease")?; + + Ok(updated == 1) +} + +/// Read the current durable version (`udb_lease.durable_version`). Used by a freshly elected leader +/// to learn the watermark floor it must continue from. +pub async fn current_durable_version(pool: &Pool) -> Result { + let conn = pool + .get() + .await + .context("failed to get connection for durable version read")?; + + let row = conn + .query_opt( + "SELECT durable_version FROM udb_lease WHERE id = $1", + &[&LEASE_ID], + ) + .await + .context("failed to read durable version")?; + + Ok(row.map(|row| row.get::<_, i64>(0)).unwrap_or(0)) +} diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs new file mode 100644 index 0000000000..79f1f1f48b --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs @@ -0,0 +1,336 @@ +mod apply; +mod lease; + +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result}; + +use crate::{conflict_tracker::TransactionConflictTracker, transaction::TXN_TIMEOUT}; + +use super::shared::{LEASE_ID, LeaseInfo, PostgresShared, WATERMARK_CHANNEL, commit_channel}; + +/// Max commits resolved+applied per batch (group commit). Amortizes the resolver, Postgres +/// round-trips, and fsync across the batch. +const DRAIN_BATCH_SIZE: i64 = 256; + +/// How often a leader renews its lease. Must be comfortably under `LEASE_TTL_SECS`. +const RENEW_INTERVAL: Duration = Duration::from_secs(3); + +/// Backstop poll cadence so a missed `udb_commit` NOTIFY cannot stall the drain indefinitely. +const POLL_BACKSTOP: Duration = Duration::from_millis(50); + +/// How long a candidate waits before retrying election when another node holds the lease. +const ELECTION_RETRY: Duration = Duration::from_secs(2); + +enum DrainOutcome { + /// Processed zero or more requests; still leader. + Drained, + /// Lost the lease (epoch bumped by a new leader). Step down. + LostLease, +} + +/// Spawn the per-process resolver task. Every node runs this; only the elected leader drains the +/// commit queue. +pub fn spawn(shared: Arc) { + tokio::spawn(run(shared)); +} + +async fn run(shared: Arc) { + loop { + match lease::try_acquire(&shared.pool, &shared.node_id).await { + Ok(Some(acquired)) => { + tracing::info!(epoch = acquired.epoch, node_id = %shared.node_id, "acquired udb leader lease"); + if let Err(err) = lead(&shared, acquired.epoch).await { + tracing::error!(?err, "udb leader loop errored, stepping down"); + } + tracing::info!(epoch = acquired.epoch, "stepped down from udb leader"); + } + Ok(None) => { + tokio::time::sleep(ELECTION_RETRY).await; + } + Err(err) => { + tracing::warn!(?err, "failed udb lease acquire attempt"); + tokio::time::sleep(ELECTION_RETRY).await; + } + } + } +} + +/// Leader main loop: hold the lease, drain the commit queue on wake or poll, and renew the lease. +async fn lead(shared: &Arc, epoch: i64) -> Result<()> { + // Publish our own lease into the cache immediately so our local commits route to us. + shared.set_lease(LeaseInfo { + epoch, + leader_addr: shared.node_id.clone(), + }); + + // The recovery floor: a freshly elected leader has a cold conflict window, so reject commits + // whose read_version predates the floor until the window warms (one TXN_TIMEOUT), forcing + // those followers to take a fresh read_version. + let recovery_version = recovery_floor(shared).await?; + let recovery_deadline = Instant::now() + TXN_TIMEOUT; + + let tracker = TransactionConflictTracker::new(); + + let mut wake_rx = shared + .listener + .listen(&commit_channel(&shared.node_id)) + .await; + + let mut renew_interval = tokio::time::interval(RENEW_INTERVAL); + renew_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut poll_interval = tokio::time::interval(POLL_BACKSTOP); + poll_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + // Drain anything already queued before our first wake. + if matches!( + drain(shared, epoch, &tracker, recovery_version, recovery_deadline).await?, + DrainOutcome::LostLease + ) { + return Ok(()); + } + + loop { + tokio::select! { + _ = renew_interval.tick() => { + if !lease::renew(&shared.pool, &shared.node_id, epoch).await? { + tracing::warn!(epoch, "lost udb lease on renew"); + return Ok(()); + } + } + res = wake_rx.recv() => { + match res { + Ok(_) | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + wake_rx = shared.listener.listen(&commit_channel(&shared.node_id)).await; + } + } + if matches!( + drain(shared, epoch, &tracker, recovery_version, recovery_deadline).await?, + DrainOutcome::LostLease + ) { + return Ok(()); + } + } + _ = poll_interval.tick() => { + if matches!( + drain(shared, epoch, &tracker, recovery_version, recovery_deadline).await?, + DrainOutcome::LostLease + ) { + return Ok(()); + } + } + } + } +} + +/// The version floor a freshly elected leader continues from: the higher of the durable watermark +/// and the sequence high-water. The LOGGED `udb_version_seq` is crash-safe, so this never regresses. +async fn recovery_floor(shared: &Arc) -> Result { + let durable = lease::current_durable_version(&shared.pool).await?; + + let conn = shared + .pool + .get() + .await + .context("failed to get connection for recovery floor")?; + let seq_high: i64 = conn + .query_one("SELECT last_value FROM udb_version_seq", &[]) + .await + .context("failed to read sequence high water")? + .get(0); + + Ok(durable.max(seq_high).max(0) as u64) +} + +/// Drain pending commit requests in id-ordered batches until none remain. Each batch resolves and +/// applies inside a single Postgres transaction (group commit), fenced on the leader's epoch. +async fn drain( + shared: &Arc, + epoch: i64, + tracker: &TransactionConflictTracker, + recovery_version: u64, + recovery_deadline: Instant, +) -> Result { + loop { + match drain_batch(shared, epoch, tracker, recovery_version, recovery_deadline).await? { + BatchOutcome::Empty => return Ok(DrainOutcome::Drained), + BatchOutcome::Processed => {} + BatchOutcome::LostLease => return Ok(DrainOutcome::LostLease), + } + } +} + +enum BatchOutcome { + Empty, + Processed, + LostLease, +} + +struct Reply { + channel: String, + id: i64, +} + +async fn drain_batch( + shared: &Arc, + epoch: i64, + tracker: &TransactionConflictTracker, + recovery_version: u64, + recovery_deadline: Instant, +) -> Result { + let mut conn = shared + .pool + .get() + .await + .context("failed to get connection for drain batch")?; + let txn = conn + .build_transaction() + .start() + .await + .context("failed to start drain batch txn")?; + + // Claim a batch in id order. FOR UPDATE SKIP LOCKED holds the rows for this txn so they are + // stamped terminal on COMMIT with no intermediate 'claimed' state to clean up. + let rows = txn + .query( + "SELECT id, read_version, payload, reply_channel + FROM udb_commit_requests + WHERE status = 'pending' AND epoch = $1 + ORDER BY id + LIMIT $2 + FOR UPDATE SKIP LOCKED", + &[&epoch, &DRAIN_BATCH_SIZE], + ) + .await + .context("failed to claim commit batch")?; + + if rows.is_empty() { + txn.rollback().await.ok(); + return Ok(BatchOutcome::Empty); + } + + let cold_window = Instant::now() < recovery_deadline; + let mut max_winner_cv: i64 = 0; + let mut replies = Vec::with_capacity(rows.len()); + + for row in &rows { + let id: i64 = row.get(0); + let read_version: i64 = row.get(1); + let payload: Vec = row.get(2); + let reply_channel: String = row.get(3); + + let decoded = super::codec::decode_commit_request(&payload) + .context("failed to decode commit payload")?; + + let commit_version: i64 = txn + .query_one("SELECT nextval('udb_version_seq')", &[]) + .await + .context("failed to get next commit version")? + .get(0); + + let start_version = read_version.max(0) as u64; + + // Cold-window guard: a commit whose read_version predates the recovery floor cannot be + // safely resolved against this leader's empty window. Reject it as retryable. + let conflicted = if cold_window && start_version < recovery_version { + true + } else { + tracker + .check_and_insert( + start_version, + commit_version.max(0) as u64, + decoded.conflict_ranges, + ) + .await + }; + + if conflicted { + txn.execute( + "UPDATE udb_commit_requests SET status = 'conflict' WHERE id = $1", + &[&id], + ) + .await + .context("failed to stamp conflict")?; + } else { + apply::apply(&txn, decoded.operations, commit_version.max(0) as u64) + .await + .context("failed to apply commit")?; + txn.execute( + "UPDATE udb_commit_requests SET status = 'committed', commit_version = $1 WHERE id = $2", + &[&commit_version, &id], + ) + .await + .context("failed to stamp committed")?; + max_winner_cv = max_winner_cv.max(commit_version); + } + + replies.push(Reply { + channel: reply_channel, + id, + }); + } + + // Advance the watermark, fenced on our epoch. A zombie old leader whose epoch was bumped sees + // zero rows updated and must step down before any of its writes become visible. + let new_durable: i64 = match txn + .query_opt( + "UPDATE udb_lease + SET durable_version = GREATEST(durable_version, $1) + WHERE id = $2 AND epoch = $3 + RETURNING durable_version", + &[&max_winner_cv, &LEASE_ID, &epoch], + ) + .await + .context("failed to advance watermark")? + { + Some(row) => row.get(0), + None => { + txn.rollback().await.ok(); + return Ok(BatchOutcome::LostLease); + } + }; + + txn.commit().await.context("failed to commit drain batch")?; + + // Watermark advances strictly after the apply txn is durably committed and visible, so a + // reader handed this read_version can never miss a write with commit_version <= read_version. + shared.advance_durable_version(new_durable); + + notify_after_commit(&conn, new_durable, &replies).await; + + Ok(BatchOutcome::Processed) +} + +/// Wake watermark listeners and the followers waiting on each processed request. Best-effort: a +/// missed NOTIFY is covered by the follower's polling backstop and the watermark refresh timer. +async fn notify_after_commit( + conn: &deadpool_postgres::Client, + new_durable: i64, + replies: &[Reply], +) { + if let Err(err) = conn + .execute( + "SELECT pg_notify($1, $2)", + &[&WATERMARK_CHANNEL, &new_durable.to_string()], + ) + .await + { + tracing::debug!(?err, "failed to notify watermark"); + } + + let channels: Vec<&str> = replies.iter().map(|r| r.channel.as_str()).collect(); + let ids: Vec = replies.iter().map(|r| r.id.to_string()).collect(); + if let Err(err) = conn + .execute( + "SELECT pg_notify(c, p) FROM unnest($1::text[], $2::text[]) AS t(c, p)", + &[&channels, &ids], + ) + .await + { + tracing::debug!(?err, "failed to notify commit replies"); + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/shared.rs b/engine/packages/universaldb/src/driver/postgres/shared.rs new file mode 100644 index 0000000000..0ab2963158 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/shared.rs @@ -0,0 +1,161 @@ +use std::{ + sync::{ + Arc, + atomic::{AtomicI64, Ordering}, + }, + time::Duration, +}; + +use deadpool_postgres::Pool; +use tokio::sync::{Notify, watch}; + +use super::listener::PgListener; + +/// The singleton row id of `udb_lease`. +pub const LEASE_ID: i32 = 1; + +/// How often the follower refreshes its cached lease row (epoch, leader channel, watermark) as a +/// backstop to the `udb_watermark` NOTIFY. A stale-but-older watermark only widens the conflict +/// window, so this can be loose. +const LEASE_REFRESH_INTERVAL: Duration = Duration::from_millis(500); + +/// Channel a follower NOTIFYs (and the leader LISTENs) to wake the leader's drain loop. +pub fn commit_channel(node_id: &str) -> String { + format!("udb_commit_{node_id}") +} + +/// Channel the leader NOTIFYs (and a follower LISTENs) to deliver a commit result. +pub fn reply_channel(node_id: &str) -> String { + format!("udb_reply_{node_id}") +} + +/// Channel the leader NOTIFYs on every watermark advance; all nodes LISTEN. +pub const WATERMARK_CHANNEL: &str = "udb_watermark"; + +/// Cached view of the current leader lease, as seen by a follower. +#[derive(Clone, Debug)] +pub struct LeaseInfo { + pub epoch: i64, + /// Node id of the current leader, used to build its commit channel. + pub leader_addr: String, +} + +/// Process-wide state shared by the follower transaction tasks and the leader resolver. Every node +/// is both a follower (it submits its own commits) and a candidate leader. +pub struct PostgresShared { + pub pool: Pool, + /// Unique per-process id used to name this node's NOTIFY channels. + pub node_id: String, + pub listener: PgListener, + /// Highest durable commit version (`udb_lease.durable_version`); the follower read version. + durable_version: AtomicI64, + /// Pinged whenever `durable_version` advances. + watermark_notify: Notify, + lease_tx: watch::Sender>, + lease_rx: watch::Receiver>, +} + +impl PostgresShared { + pub fn new(pool: Pool, node_id: String, listener: PgListener) -> Arc { + let (lease_tx, lease_rx) = watch::channel(None); + let shared = Arc::new(Self { + pool, + node_id, + listener, + durable_version: AtomicI64::new(0), + watermark_notify: Notify::new(), + lease_tx, + lease_rx, + }); + + tokio::spawn(Self::cache_refresh_task(shared.clone())); + + shared + } + + /// The cached follower read version (`durable_version`). + pub fn read_version(&self) -> i64 { + self.durable_version.load(Ordering::SeqCst) + } + + /// Advance the cached watermark monotonically and wake any waiters. + pub fn advance_durable_version(&self, version: i64) { + let prev = self.durable_version.fetch_max(version, Ordering::SeqCst); + if version > prev { + self.watermark_notify.notify_waiters(); + } + } + + /// Current cached lease, if known. + pub fn current_lease(&self) -> Option { + self.lease_rx.borrow().clone() + } + + /// Publish a freshly observed/elected lease into the cache. + pub fn set_lease(&self, lease: LeaseInfo) { + let _ = self.lease_tx.send(Some(lease)); + } + + /// Background task: keep `durable_version` and the cached lease fresh via the `udb_watermark` + /// NOTIFY plus a periodic poll of `udb_lease`. + async fn cache_refresh_task(shared: Arc) { + let mut watermark_rx = shared.listener.listen(WATERMARK_CHANNEL).await; + let mut interval = tokio::time::interval(LEASE_REFRESH_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + tokio::select! { + notify = watermark_rx.recv() => { + match notify { + Ok(payload) => { + if let Ok(version) = payload.parse::() { + shared.advance_durable_version(version); + } + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + // Re-subscribe; the listener recreates the channel on reconnect. + watermark_rx = shared.listener.listen(WATERMARK_CHANNEL).await; + } + } + } + _ = interval.tick() => { + shared.refresh_lease_row().await; + } + } + } + } + + async fn refresh_lease_row(&self) { + let conn = match self.pool.get().await { + Ok(conn) => conn, + Err(err) => { + tracing::debug!(?err, "failed to get connection for lease refresh"); + return; + } + }; + + let row = conn + .query_opt( + "SELECT epoch, leader_addr, durable_version FROM udb_lease WHERE id = $1", + &[&LEASE_ID], + ) + .await; + + match row { + Ok(Some(row)) => { + let epoch: i64 = row.get(0); + let leader_addr: String = row.get(1); + let durable_version: i64 = row.get(2); + self.advance_durable_version(durable_version); + self.set_lease(LeaseInfo { epoch, leader_addr }); + } + Ok(None) => { + // No lease row yet; no leader elected. + } + Err(err) => { + tracing::debug!(?err, "failed to refresh lease row"); + } + } + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/transaction.rs b/engine/packages/universaldb/src/driver/postgres/transaction.rs index f80cbba393..cc315feba4 100644 --- a/engine/packages/universaldb/src/driver/postgres/transaction.rs +++ b/engine/packages/universaldb/src/driver/postgres/transaction.rs @@ -1,11 +1,13 @@ use std::{ future::Future, pin::Pin, - sync::atomic::{AtomicBool, Ordering}, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, }; use anyhow::{Context, Result}; -use deadpool_postgres::Pool; use tokio::sync::{OnceCell, mpsc, oneshot}; use crate::{ @@ -18,33 +20,35 @@ use crate::{ value::{Slice, Value, Values}, }; -use super::transaction_task::{TransactionCommand, TransactionTask}; +use super::{ + shared::PostgresShared, + transaction_task::{TransactionCommand, TransactionTask}, +}; pub struct PostgresTransactionDriver { - pool: Pool, + shared: Arc, operations: TransactionOperations, committed: AtomicBool, tx_sender: OnceCell>, } impl PostgresTransactionDriver { - pub fn with_config(pool: Pool) -> Self { + pub fn new(shared: Arc) -> Self { PostgresTransactionDriver { - pool, + shared, operations: TransactionOperations::default(), committed: AtomicBool::new(false), tx_sender: OnceCell::new(), } } - /// Get or create the transaction task + /// Get or create the transaction task that owns this transaction's read snapshot. async fn ensure_transaction(&self) -> Result<&mpsc::UnboundedSender> { self.tx_sender .get_or_try_init(|| async { let (sender, receiver) = mpsc::unbounded_channel(); - // Spawn the transaction task with serializable isolation - let task = TransactionTask::new(self.pool.clone(), receiver); + let task = TransactionTask::new(self.shared.clone(), receiver); tokio::spawn(task.run()); anyhow::Ok(sender) diff --git a/engine/packages/universaldb/src/driver/postgres/transaction_task.rs b/engine/packages/universaldb/src/driver/postgres/transaction_task.rs index 297751f608..c4d3e3f9f6 100644 --- a/engine/packages/universaldb/src/driver/postgres/transaction_task.rs +++ b/engine/packages/universaldb/src/driver/postgres/transaction_task.rs @@ -1,17 +1,18 @@ -use anyhow::{Context, Result, anyhow, bail}; -use deadpool_postgres::{Pool, Transaction}; +use std::sync::Arc; + +use anyhow::{Result, anyhow, bail}; +use deadpool_postgres::Transaction; use tokio::sync::{mpsc, oneshot}; use tokio_postgres::IsolationLevel; use crate::{ - atomic::apply_atomic_op, - error::DatabaseError, - options::{ConflictRangeType, MutationType}, + options::ConflictRangeType, tx_ops::Operation, value::{KeyValue, Slice, Values}, - versionstamp::{generate_versionstamp, substitute_raw_versionstamp}, }; +use super::{commit, shared::PostgresShared}; + pub enum TransactionCommand { // Read operations Get { @@ -48,70 +49,57 @@ pub enum TransactionCommand { }, } -/// TransactionTask runs in a separate tokio task to manage a PostgreSQL transaction. -/// -/// This design is necessary because PostgreSQL transactions have lifetime constraints -/// that don't work well with the FoundationDB-style API. Specifically: -/// - The transaction must outlive all references to it -/// - We can't store the transaction in a mutex due to lifetime issues with the connection +/// TransactionTask runs in a separate tokio task to own a single pinned PostgreSQL `REPEATABLE READ` +/// snapshot connection for the lifetime of a follower transaction. /// -/// By running in a separate task and communicating via channels, we avoid these lifetime -/// issues while maintaining a single serializable transaction for all operations. +/// Reads go directly against this snapshot (they never involve the leader). Commits delegate to +/// [`commit::submit`], which enqueues the request on the leader and awaits the result. The +/// `read_version` is captured from the cached watermark before the snapshot is opened, so no write +/// with `commit_version <= read_version` can be invisible to the snapshot. pub struct TransactionTask { - pool: Pool, + shared: Arc, receiver: mpsc::UnboundedReceiver, } impl TransactionTask { - pub fn new(pool: Pool, receiver: mpsc::UnboundedReceiver) -> Self { - Self { pool, receiver } + pub fn new( + shared: Arc, + receiver: mpsc::UnboundedReceiver, + ) -> Self { + Self { shared, receiver } } pub async fn run(mut self) { - // Get connection from pool - let mut conn = match self.pool.get().await { + // Capture the read version BEFORE opening the snapshot so the snapshot reflects every write + // with commit_version <= read_version. + let read_version = self.shared.read_version(); + + let mut conn = match self.shared.pool.get().await { Ok(conn) => conn, Err(_) => { - // If we can't get a connection, respond to all pending commands with errors self.fail_receiver().await; return; } }; - // Start the read transaction let tx = match conn .build_transaction() .isolation_level(IsolationLevel::RepeatableRead) + .read_only(true) .start() .await { Ok(tx) => tx, Err(_) => { - // If we can't start a transaction, respond to all pending commands with errors self.fail_receiver().await; return; } }; - // TODO: Parallelize future - let start_version = match tx - .query_one("SELECT nextval('global_version_seq')", &[]) - .await - { - Ok(row) => row.get::<_, i64>(0), - Err(err) => { - tracing::error!(?err, "failed to get postgres txn start_version"); - self.fail_receiver().await; - return; - } - }; - - // Process commands while let Some(cmd) = self.receiver.recv().await { match cmd { TransactionCommand::Get { key, response } => { let result = self.handle_get(&tx, &key).await; - let _ = response.send(result); } TransactionCommand::GetKey { @@ -121,7 +109,6 @@ impl TransactionTask { response, } => { let result = self.handle_get_key(&tx, &key, or_equal, offset).await; - let _ = response.send(result); } TransactionCommand::GetRange { @@ -148,7 +135,6 @@ impl TransactionTask { reverse, ) .await; - let _ = response.send(result); } TransactionCommand::Commit { @@ -156,14 +142,12 @@ impl TransactionTask { conflict_ranges, response, } => { - let (_, result) = tokio::join!( - // Read-only txn, we don't care about the result - tx.commit(), - self.handle_commit(start_version, operations, conflict_ranges), - ); - + // The read snapshot is read-only; release it and submit the commit to the leader. + let _ = tx.commit().await; + let result = + commit::submit(&self.shared, read_version, operations, conflict_ranges) + .await; let _ = response.send(result); - // Exit after commit return; } TransactionCommand::GetEstimatedRangeSize { @@ -174,13 +158,12 @@ impl TransactionTask { let result = self .handle_get_estimated_range_size(&tx, &begin, &end) .await; - let _ = response.send(result); } } } - // If the channel is closed, the transaction will be rolled back when dropped + // If the channel is closed, the snapshot transaction is rolled back when dropped. } async fn handle_get(&mut self, tx: &Transaction<'_>, key: &[u8]) -> Result> { @@ -234,27 +217,18 @@ impl TransactionTask { reverse: bool, ) -> Result { // Determine SQL operators based on key selector types - // For begin selector: - // first_greater_or_equal: or_equal = false, offset = 1 -> ">=" - // first_greater_than: or_equal = true, offset = 1 -> ">" let begin_op = if begin_offset == 1 { if begin_or_equal { ">" } else { ">=" } } else { - // This shouldn't happen for begin in range queries ">=" }; - // For end selector: - // first_greater_than: or_equal = true, offset = 1 -> "<=" - // first_greater_or_equal: or_equal = false, offset = 1 -> "<" let end_op = if end_offset == 1 { if end_or_equal { "<=" } else { "<" } } else { - // This shouldn't happen for end in range queries "<" }; - // Build query with CTE that adds conflict range let query = if reverse { if let Some(limit) = limit { format!( @@ -301,22 +275,22 @@ impl TransactionTask { begin: &[u8], end: &[u8], ) -> Result { - // Sample's 1% of the range + // Sample 1% of the range. let query = " WITH range_stats AS ( - SELECT + SELECT COUNT(*) as estimated_count, COALESCE(SUM(pg_column_size(key) + pg_column_size(value)), 0) as sample_size - FROM kv TABLESAMPLE SYSTEM(1) + FROM kv TABLESAMPLE SYSTEM(1) WHERE key >= $1 AND key < $2 ), table_stats AS ( - SELECT reltuples::bigint as total_rows - FROM pg_class + SELECT reltuples::bigint as total_rows + FROM pg_class WHERE relname = 'kv' AND relkind = 'r' ) - SELECT - CASE + SELECT + CASE WHEN r.estimated_count = 0 THEN 0 ELSE (r.sample_size * 100)::bigint END as estimated_size @@ -329,165 +303,6 @@ impl TransactionTask { .map_err(map_postgres_error) } - async fn handle_commit( - &mut self, - start_version: i64, - operations: Vec, - conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, - ) -> Result<()> { - // Get connection from pool - let mut conn = self.pool.get().await?; - - // Start write transaction - let tx = conn - .build_transaction() - .isolation_level(IsolationLevel::ReadCommitted) - .start() - .await - .context("failed to start write txn")?; - - let mut begins = Vec::with_capacity(conflict_ranges.len()); - let mut ends = Vec::with_capacity(conflict_ranges.len()); - let mut conflict_types = Vec::with_capacity(conflict_ranges.len()); - - for (begin, end, conflict_type) in conflict_ranges { - let conflict_type = match conflict_type { - ConflictRangeType::Read => "read", - ConflictRangeType::Write => "write", - }; - - begins.push(begin); - ends.push(end); - conflict_types.push(conflict_type); - } - - let query = " - WITH data AS ( - SELECT nextval('global_version_seq') AS commit_version - ) - INSERT INTO conflict_ranges (range_data, conflict_type, start_version, commit_version) - SELECT - bytearange(begin_key, end_key, '[)'), - conflict_type::range_type, - $4, - data.commit_version - FROM UNNEST($1::bytea[], $2::bytea[], $3::text[]) AS t(begin_key, end_key, conflict_type), data"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - // Insert all conflict ranges at once - tx.execute(&stmt, &[&begins, &ends, &conflict_types, &start_version]) - .await - .map_err(map_postgres_error)?; - - let transaction_versionstamp = generate_versionstamp(0); - - for op in operations { - match op { - Operation::SetValue { key, value } => { - let query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key, &value]) - .await - .map_err(map_postgres_error)?; - } - Operation::Clear { key } => { - let query = "DELETE FROM kv WHERE key = $1"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key]) - .await - .map_err(map_postgres_error)?; - } - Operation::ClearRange { begin, end } => { - let query = "DELETE FROM kv WHERE key >= $1 AND key < $2"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&begin, &end]) - .await - .map_err(map_postgres_error)?; - } - Operation::AtomicOp { - key, - param, - op_type, - } => { - if matches!(op_type, MutationType::SetVersionstampedKey) { - let key = substitute_raw_versionstamp(key, &transaction_versionstamp) - .map_err(anyhow::Error::msg) - .context("failed substituting versionstamped key")?; - let query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key, ¶m]) - .await - .map_err(map_postgres_error)?; - continue; - } - - if matches!(op_type, MutationType::SetVersionstampedValue) { - let value = substitute_raw_versionstamp(param, &transaction_versionstamp) - .map_err(anyhow::Error::msg) - .context("failed substituting versionstamped value")?; - let query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key, &value]) - .await - .map_err(map_postgres_error)?; - continue; - } - - // TODO: All operations need to be done on the sql side, not in rust - - // Get current value from database - let current_query = "SELECT value FROM kv WHERE key = $1"; - let stmt = tx - .prepare_cached(current_query) - .await - .map_err(map_postgres_error)?; - - let current_row = tx - .query_opt(&stmt, &[&key]) - .await - .map_err(map_postgres_error)?; - - // Extract current value or use None if key doesn't exist - let current_value = current_row.map(|row| row.get::<_, Vec>(0)); - let current_slice = current_value.as_deref(); - - // Apply atomic operation - let new_value = apply_atomic_op(current_slice, ¶m, op_type); - - // Store the result - if let Some(new_value) = new_value { - let update_query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; - let stmt = tx - .prepare_cached(update_query) - .await - .map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key, &new_value]) - .await - .map_err(map_postgres_error)?; - } else { - let update_query = "DELETE FROM kv WHERE key = $1"; - let stmt = tx - .prepare_cached(update_query) - .await - .map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key]) - .await - .map_err(map_postgres_error)?; - } - } - } - } - - tx.commit().await.map_err(map_postgres_error) - } - async fn fail_receiver(&mut self) { while let Some(cmd) = self.receiver.recv().await { match cmd { @@ -511,31 +326,19 @@ impl TransactionTask { } } -/// Maps PostgreSQL error to DatabaseError +/// Maps a PostgreSQL error from the read path to a `DatabaseError` where appropriate. fn map_postgres_error(err: tokio_postgres::Error) -> anyhow::Error { - let error_str = if let Some(err) = err.as_db_error() { - err.to_string() - } else { - err.to_string() - }; + let error_str = err.to_string(); - if error_str.contains("exclusion_violation") - || error_str.contains("violates exclusion constraint") - { - // Retryable - another transaction has a conflicting range - DatabaseError::NotCommitted.into() - } else if error_str.contains("serialization failure") + if error_str.contains("serialization failure") || error_str.contains("could not serialize") || error_str.contains("deadlock detected") { - // Retryable - transaction conflict - DatabaseError::NotCommitted.into() + crate::error::DatabaseError::NotCommitted.into() } else if error_str.contains("current transaction is aborted") { - // Returned by the rest of the commands in a txn if it failed for exclusion reasons - DatabaseError::NotCommitted.into() + crate::error::DatabaseError::NotCommitted.into() } else { tracing::error!(%err, "postgres error"); - // Non-retryable error anyhow::Error::new(err) } } diff --git a/engine/packages/universaldb/src/driver/rocksdb/database.rs b/engine/packages/universaldb/src/driver/rocksdb/database.rs index ec0175ee5a..ba186d872e 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/database.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/database.rs @@ -17,9 +17,9 @@ use crate::{ utils::{MaybeCommitted, calculate_tx_retry_backoff}, }; -use super::{ - transaction::RocksDbTransactionDriver, transaction_conflict_tracker::TransactionConflictTracker, -}; +use crate::conflict_tracker::TransactionConflictTracker; + +use super::transaction::RocksDbTransactionDriver; pub struct RocksDbDatabaseDriver { db: Arc, diff --git a/engine/packages/universaldb/src/driver/rocksdb/mod.rs b/engine/packages/universaldb/src/driver/rocksdb/mod.rs index a24bd72603..b18e28edca 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/mod.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/mod.rs @@ -1,6 +1,5 @@ mod database; mod transaction; -mod transaction_conflict_tracker; mod transaction_task; pub use database::RocksDbDatabaseDriver; diff --git a/engine/packages/universaldb/src/driver/rocksdb/transaction.rs b/engine/packages/universaldb/src/driver/rocksdb/transaction.rs index e62c3eda79..85cf5edf5c 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/transaction.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/transaction.rs @@ -21,10 +21,9 @@ use crate::{ value::{Slice, Value, Values}, }; -use super::{ - transaction_conflict_tracker::TransactionConflictTracker, - transaction_task::{TransactionCommand, TransactionTask}, -}; +use crate::conflict_tracker::TransactionConflictTracker; + +use super::transaction_task::{TransactionCommand, TransactionTask}; pub struct RocksDbTransactionDriver { db: Arc, diff --git a/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs b/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs index 69835eb6a7..986e6ba2bd 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs @@ -6,9 +6,9 @@ use rocksdb::{ }; use tokio::sync::{mpsc, oneshot}; -use super::transaction_conflict_tracker::TransactionConflictTracker; use crate::{ atomic::apply_atomic_op, + conflict_tracker::TransactionConflictTracker, error::DatabaseError, key_selector::KeySelector, options::{ConflictRangeType, MutationType}, @@ -412,9 +412,11 @@ impl TransactionTask { } } + // rocksdb generates both start and commit versions from the in-process counter. + let commit_version = self.txn_conflict_tracker.next_global_version(); if self .txn_conflict_tracker - .check_and_insert(start_version, conflict_ranges) + .check_and_insert(start_version, commit_version, conflict_ranges) .await { return Err(DatabaseError::NotCommitted.into()); diff --git a/engine/packages/universaldb/src/lib.rs b/engine/packages/universaldb/src/lib.rs index 96260a177d..1b6fb78b87 100644 --- a/engine/packages/universaldb/src/lib.rs +++ b/engine/packages/universaldb/src/lib.rs @@ -1,4 +1,5 @@ pub(crate) mod atomic; +pub(crate) mod conflict_tracker; mod database; pub mod driver; pub mod error; diff --git a/engine/packages/universaldb/tests/integration.rs b/engine/packages/universaldb/tests/integration.rs index 52fa703116..45a5b3bdf7 100644 --- a/engine/packages/universaldb/tests/integration.rs +++ b/engine/packages/universaldb/tests/integration.rs @@ -137,11 +137,9 @@ async fn test_database_options(db: &Database) { use std::sync::Arc; use std::sync::atomic::{AtomicU32, Ordering}; use universaldb::error::DatabaseError; - use universaldb::options::DatabaseOption; // Test setting transaction retry limit - db.set_option(DatabaseOption::TransactionRetryLimit(5)) - .unwrap(); + db.txn_retry_limit(5).unwrap(); // Test that retry limit is respected by forcing conflicts let conflict_counter = Arc::new(AtomicU32::new(0)); @@ -172,8 +170,7 @@ async fn test_database_options(db: &Database) { assert_eq!(final_attempts, 3, "Should have taken 3 attempts"); // Now set a very low retry limit and verify it fails - db.set_option(DatabaseOption::TransactionRetryLimit(1)) - .unwrap(); + db.txn_retry_limit(1).unwrap(); let conflict_counter2 = Arc::new(AtomicU32::new(0)); let counter_clone2 = conflict_counter2.clone(); @@ -204,8 +201,7 @@ async fn test_database_options(db: &Database) { assert!(attempts <= 2, "Should not retry more than limit + 1"); // Reset to a reasonable retry limit - db.set_option(DatabaseOption::TransactionRetryLimit(100)) - .unwrap(); + db.txn_retry_limit(100).unwrap(); } async fn clear_test_namespace(db: &Database) -> Result<()> { diff --git a/engine/sdks/rust/depot-protocol/Cargo.toml b/engine/sdks/rust/depot-protocol/Cargo.toml index 274836220a..99d8117349 100644 --- a/engine/sdks/rust/depot-protocol/Cargo.toml +++ b/engine/sdks/rust/depot-protocol/Cargo.toml @@ -8,6 +8,7 @@ edition.workspace = true [dependencies] anyhow.workspace = true +rivet-util.workspace = true serde_bare.workspace = true serde.workspace = true vbare.workspace = true diff --git a/engine/sdks/rust/universaldb-commit/Cargo.toml b/engine/sdks/rust/universaldb-commit/Cargo.toml new file mode 100644 index 0000000000..126c5cb35e --- /dev/null +++ b/engine/sdks/rust/universaldb-commit/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "rivet-universaldb-commit" +publish = false +version.workspace = true +authors.workspace = true +license.workspace = true +edition.workspace = true + +[dependencies] +anyhow.workspace = true +serde_bare.workspace = true +serde.workspace = true +vbare.workspace = true + +[build-dependencies] +vbare-compiler.workspace = true diff --git a/engine/sdks/rust/universaldb-commit/build.rs b/engine/sdks/rust/universaldb-commit/build.rs new file mode 100644 index 0000000000..6400be6f2b --- /dev/null +++ b/engine/sdks/rust/universaldb-commit/build.rs @@ -0,0 +1,64 @@ +use std::{ + fs, + path::{Path, PathBuf}, +}; + +fn main() -> Result<(), Box> { + let manifest_dir = std::env::var("CARGO_MANIFEST_DIR")?; + let out_dir = PathBuf::from(std::env::var("OUT_DIR")?); + let workspace_root = Path::new(&manifest_dir) + .parent() + .and_then(|p| p.parent()) + .and_then(|p| p.parent()) + .ok_or("Failed to find workspace root")?; + + let schema_dir = workspace_root + .join("sdks") + .join("schemas") + .join("universaldb-commit"); + println!("cargo:rerun-if-changed={}", schema_dir.display()); + + let (highest_version, _) = find_highest_version(&schema_dir); + + let cfg = vbare_compiler::Config::default(); + vbare_compiler::process_schemas_with_config(&schema_dir, &cfg)?; + + // Append protocol version constant to generated file + let combined_imports_path = out_dir.join("combined_imports.rs"); + let mut combined = fs::read_to_string(&combined_imports_path)?; + combined.push_str(&format!( + "\npub const PROTOCOL_VERSION: u16 = {};\n", + highest_version + )); + fs::write(combined_imports_path, combined)?; + + Ok(()) +} + +fn find_highest_version(schema_dir: &Path) -> (u32, PathBuf) { + let mut highest_version = 0; + let mut highest_version_path = PathBuf::new(); + + for entry in fs::read_dir(schema_dir).unwrap().flatten() { + if !entry.path().is_dir() { + let path = entry.path(); + let bare_name = path + .file_name() + .unwrap() + .to_str() + .unwrap() + .split_once('.') + .unwrap() + .0; + + if let Ok(version) = bare_name[1..].parse::() { + if version > highest_version { + highest_version = version; + highest_version_path = path; + } + } + } + } + + (highest_version, highest_version_path) +} diff --git a/engine/sdks/rust/universaldb-commit/src/generated.rs b/engine/sdks/rust/universaldb-commit/src/generated.rs new file mode 100644 index 0000000000..84801af8dc --- /dev/null +++ b/engine/sdks/rust/universaldb-commit/src/generated.rs @@ -0,0 +1 @@ +include!(concat!(env!("OUT_DIR"), "/combined_imports.rs")); diff --git a/engine/sdks/rust/universaldb-commit/src/lib.rs b/engine/sdks/rust/universaldb-commit/src/lib.rs new file mode 100644 index 0000000000..41954bbe75 --- /dev/null +++ b/engine/sdks/rust/universaldb-commit/src/lib.rs @@ -0,0 +1,6 @@ +pub mod generated; +pub mod versioned; + +// Re-export latest +pub use generated::PROTOCOL_VERSION; +pub use generated::v1::*; diff --git a/engine/sdks/rust/universaldb-commit/src/versioned.rs b/engine/sdks/rust/universaldb-commit/src/versioned.rs new file mode 100644 index 0000000000..a9d637aa81 --- /dev/null +++ b/engine/sdks/rust/universaldb-commit/src/versioned.rs @@ -0,0 +1,38 @@ +use anyhow::{Ok, Result, bail}; +use vbare::OwnedVersionedData; + +use crate::generated::v1; + +// Only v1 exists today. When adding v2+, generate converters with +// `scripts/vbare-gen-converters` (see the envoy-protocol package for the +// resulting `versioned/` module layout) and wire them in here. +pub enum CommitRequest { + V1(v1::CommitRequest), +} + +impl OwnedVersionedData for CommitRequest { + type Latest = v1::CommitRequest; + + fn wrap_latest(latest: v1::CommitRequest) -> Self { + CommitRequest::V1(latest) + } + + fn unwrap_latest(self) -> Result { + match self { + CommitRequest::V1(data) => Ok(data), + } + } + + fn deserialize_version(payload: &[u8], version: u16) -> Result { + match version { + 1 => Ok(CommitRequest::V1(serde_bare::from_slice(payload)?)), + _ => bail!("invalid version: {version}"), + } + } + + fn serialize_version(self, _version: u16) -> Result> { + match self { + CommitRequest::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + } + } +} diff --git a/engine/sdks/schemas/universaldb-commit/v1.bare b/engine/sdks/schemas/universaldb-commit/v1.bare new file mode 100644 index 0000000000..2377312f84 --- /dev/null +++ b/engine/sdks/schemas/universaldb-commit/v1.bare @@ -0,0 +1,70 @@ +# Commit-queue wire format for the Postgres leader-resolver UDB driver. +# +# Followers encode a CommitRequest into the `payload` column of +# `udb_commit_requests`; the leader decodes it to resolve and apply. +# Rust-only (never leaves the engine), but versioned so rolling deploys can +# skew follower vs leader code. + +type ConflictRangeType enum { + READ + WRITE +} + +type ConflictRange struct { + begin: data + end: data + kind: ConflictRangeType +} + +# Order MUST match universaldb::options::MutationType declaration order so the +# enum tag round-trips. 15 variants today; never reorder, only append. +type MutationType enum { + ADD + AND + BIT_AND + OR + BIT_OR + XOR + BIT_XOR + APPEND_IF_FITS + MAX + MIN + SET_VERSIONSTAMPED_KEY + SET_VERSIONSTAMPED_VALUE + BYTE_MIN + BYTE_MAX + COMPARE_AND_CLEAR +} + +type SetValue struct { + key: data + value: data +} + +type Clear struct { + key: data +} + +type ClearRange struct { + begin: data + end: data +} + +type AtomicOp struct { + key: data + param: data + opType: MutationType +} + +type Operation union { + SetValue | + Clear | + ClearRange | + AtomicOp +} + +type CommitRequest struct { + readVersion: u64 + conflictRanges: list + operations: list +} diff --git a/scripts/run/postgres.sh b/scripts/run/postgres.sh index 30acd498f1..2ac817b598 100755 --- a/scripts/run/postgres.sh +++ b/scripts/run/postgres.sh @@ -2,7 +2,7 @@ set -euo pipefail CONTAINER_NAME="rivet-engine-postgres" -POSTGRES_IMAGE="postgres:17" +POSTGRES_IMAGE="postgres:18" if docker ps --all --format '{{.Names}}' | grep -qw "${CONTAINER_NAME}"; then if docker ps --format '{{.Names}}' | grep -qw "${CONTAINER_NAME}"; then diff --git a/scripts/run/restore-postgres.sh b/scripts/run/restore-postgres.sh index bd2b50e0b4..0a82410608 100755 --- a/scripts/run/restore-postgres.sh +++ b/scripts/run/restore-postgres.sh @@ -2,7 +2,7 @@ set -euo pipefail CONTAINER_NAME="rivet-engine-postgres" -POSTGRES_IMAGE="postgres:17" +POSTGRES_IMAGE="postgres:18" if [ $# -ne 1 ]; then echo "Usage: $0 " diff --git a/self-host/compose/dev-host/docker-compose.yml b/self-host/compose/dev-host/docker-compose.yml index 3b54702d63..96db5a57d0 100644 --- a/self-host/compose/dev-host/docker-compose.yml +++ b/self-host/compose/dev-host/docker-compose.yml @@ -72,7 +72,7 @@ services: network_mode: host postgres: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres diff --git a/self-host/compose/dev-multidc-multinode/docker-compose.yml b/self-host/compose/dev-multidc-multinode/docker-compose.yml index 75c500850b..7e421144c8 100644 --- a/self-host/compose/dev-multidc-multinode/docker-compose.yml +++ b/self-host/compose/dev-multidc-multinode/docker-compose.yml @@ -85,7 +85,7 @@ services: condition: service_healthy postgres-dc-a: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -355,7 +355,7 @@ services: - rivet-network-dc-a postgres-dc-b: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -617,7 +617,7 @@ services: - rivet-network-dc-b postgres-dc-c: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres diff --git a/self-host/compose/dev-multidc/docker-compose.yml b/self-host/compose/dev-multidc/docker-compose.yml index f20db41ff3..79cb619e0a 100644 --- a/self-host/compose/dev-multidc/docker-compose.yml +++ b/self-host/compose/dev-multidc/docker-compose.yml @@ -85,7 +85,7 @@ services: condition: service_healthy postgres-dc-a: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -235,7 +235,7 @@ services: - rivet-network-dc-a postgres-dc-b: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -377,7 +377,7 @@ services: - rivet-network-dc-b postgres-dc-c: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres diff --git a/self-host/compose/dev-multinode/docker-compose.yml b/self-host/compose/dev-multinode/docker-compose.yml index 381ece396d..a432743057 100644 --- a/self-host/compose/dev-multinode/docker-compose.yml +++ b/self-host/compose/dev-multinode/docker-compose.yml @@ -81,7 +81,7 @@ services: condition: service_healthy postgres: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres diff --git a/self-host/compose/dev/docker-compose.yml b/self-host/compose/dev/docker-compose.yml index 6c5dcc0ef2..91f82dce45 100644 --- a/self-host/compose/dev/docker-compose.yml +++ b/self-host/compose/dev/docker-compose.yml @@ -81,7 +81,7 @@ services: condition: service_healthy postgres: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres diff --git a/self-host/compose/template/src/docker-compose.ts b/self-host/compose/template/src/docker-compose.ts index ca2f6b1fc8..6db4a8ab9a 100644 --- a/self-host/compose/template/src/docker-compose.ts +++ b/self-host/compose/template/src/docker-compose.ts @@ -171,7 +171,7 @@ export function generateDockerCompose(context: TemplateContext) { ); services[postgresServiceName] = { restart: "unless-stopped", - image: "postgres:17-alpine", + image: "postgres:18-alpine", environment: [ "POSTGRES_USER=postgres", "POSTGRES_PASSWORD=postgres", diff --git a/self-host/k8s/engine/12-postgres-statefulset.yaml b/self-host/k8s/engine/12-postgres-statefulset.yaml index e1e14a4b75..c035d26611 100644 --- a/self-host/k8s/engine/12-postgres-statefulset.yaml +++ b/self-host/k8s/engine/12-postgres-statefulset.yaml @@ -20,7 +20,7 @@ spec: spec: containers: - name: postgres - image: postgres:17 + image: postgres:18 args: - postgres - -c diff --git a/website/src/content/docs/self-hosting/docker-compose.mdx b/website/src/content/docs/self-hosting/docker-compose.mdx index eae884ca99..daf5ab0096 100644 --- a/website/src/content/docs/self-hosting/docker-compose.mdx +++ b/website/src/content/docs/self-hosting/docker-compose.mdx @@ -192,7 +192,7 @@ PostgreSQL is the recommended backend for multi-node self-hosted deployments. It ```yaml services: postgres: - image: postgres:15 + image: postgres:18 environment: POSTGRES_DB: rivet POSTGRES_USER: rivet diff --git a/website/src/content/docs/self-hosting/docker-container.mdx b/website/src/content/docs/self-hosting/docker-container.mdx index e6d4366d68..46d544ddd7 100644 --- a/website/src/content/docs/self-hosting/docker-container.mdx +++ b/website/src/content/docs/self-hosting/docker-container.mdx @@ -161,7 +161,7 @@ docker run -d \ -e POSTGRES_USER=rivet \ -e POSTGRES_PASSWORD=rivet_password \ -v postgres-data:/var/lib/postgresql/data \ - postgres:15 + postgres:18 # Run Rivet Engine docker run -d \