From 0caa600717e25f8b8705eba86c04cab18d671723 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Thu, 18 Jun 2026 17:57:51 -0700 Subject: [PATCH 1/7] [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/7] [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/7] 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/7] 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/7] [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/7] [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/7] [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()); }