From 0caa600717e25f8b8705eba86c04cab18d671723 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Thu, 18 Jun 2026 17:57:51 -0700 Subject: [PATCH 01/16] [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 02/16] [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 03/16] 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 04/16] 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 05/16] [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 06/16] [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 07/16] [SLOP(claude-opus-4-8)] feat(pegboard-envoy): rate-limit envoy ws ingress and cap get_pages, trim unused metrics --- Cargo.lock | 1 + engine/artifacts/config-schema.json | 54 +++++++++++++++ engine/packages/config/src/config/pegboard.rs | 16 +++++ engine/packages/guard/src/routing/envoy.rs | 2 +- engine/packages/pegboard-envoy/Cargo.toml | 1 + engine/packages/pegboard-envoy/src/lib.rs | 8 +-- engine/packages/pegboard-envoy/src/metrics.rs | 69 ------------------- .../pegboard-envoy/src/ws_to_tunnel_task.rs | 35 +++++++++- .../packages/pegboard-gateway2/src/metrics.rs | 39 ----------- .../pegboard-gateway2/src/shared_state.rs | 1 - .../src/ws_to_tunnel_task.rs | 17 +++-- engine/packages/pegboard-outbound/src/lib.rs | 2 - .../packages/pegboard-outbound/src/metrics.rs | 47 ------------- .../packages/pegboard/src/ops/actor/create.rs | 4 +- 14 files changed, 122 insertions(+), 174 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 72af92dd43..b21a5d2859 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4011,6 +4011,7 @@ dependencies = [ "rivet-pools", "rivet-runtime", "rivet-types", + "rivet-util", "rusqlite", "scc", "serde", diff --git a/engine/artifacts/config-schema.json b/engine/artifacts/config-schema.json index ead08e9e6f..998771acc1 100644 --- a/engine/artifacts/config-schema.json +++ b/engine/artifacts/config-schema.json @@ -855,6 +855,24 @@ ], "format": "int64" }, + "actor_create_rate_limit_drip_rate_ms": { + "description": "Time to regain one actor creation token per namespace.\n\nUnit is in milliseconds.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0.0 + }, + "actor_create_rate_limit_requests": { + "description": "Max burst of actor creations per namespace before throttling.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0.0 + }, "actor_retry_duration_threshold": { "description": "How long to wait after starting to attempt to reallocate before before setting actor to sleep.\n\nUnit is in milliseconds.", "type": [ @@ -986,6 +1004,24 @@ "format": "uint64", "minimum": 0.0 }, + "envoy_websocket_rate_limit_drip_rate_us": { + "description": "Time to regain one inbound WebSocket message token on a single envoy connection.\n\nUnit is in microseconds. The envoy connection multiplexes every actor on a runner, so the sustained ceiling is far higher than the per-client gateway limit and needs sub-millisecond granularity to express.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0.0 + }, + "envoy_websocket_rate_limit_requests": { + "description": "Max burst of inbound WebSocket messages on a single envoy connection before throttling.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0.0 + }, "gateway_gc_interval_ms": { "description": "GC interval for in-flight requests in milliseconds.", "type": [ @@ -1057,6 +1093,24 @@ "format": "uint64", "minimum": 0.0 }, + "gateway_websocket_rate_limit_drip_rate_ms": { + "description": "Time to regain one inbound WebSocket message token on a single connection.\n\nUnit is in milliseconds.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0.0 + }, + "gateway_websocket_rate_limit_requests": { + "description": "Max burst of inbound WebSocket messages on a single connection before throttling.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0.0 + }, "hibernating_request_eligible_threshold": { "description": "How long after last ping before considering a hibernating request disconnected.\n\nUnit is in milliseconds.", "type": [ diff --git a/engine/packages/config/src/config/pegboard.rs b/engine/packages/config/src/config/pegboard.rs index f05a98319f..f0b660401e 100644 --- a/engine/packages/config/src/config/pegboard.rs +++ b/engine/packages/config/src/config/pegboard.rs @@ -149,6 +149,14 @@ pub struct Pegboard { pub envoy_expire_scheduler_max_concurrent_expires: Option, /// Maximum pending envoys tracked by the read-path envoy expire scheduler. pub envoy_expire_scheduler_max_pending: Option, + /// Max burst of inbound WebSocket messages on a single envoy connection before throttling. + pub envoy_websocket_rate_limit_requests: Option, + /// Time to regain one inbound WebSocket message token on a single envoy connection. + /// + /// Unit is in microseconds. The envoy connection multiplexes every actor on a runner, so the + /// sustained ceiling is far higher than the per-client gateway limit and needs sub-millisecond + /// granularity to express. + pub envoy_websocket_rate_limit_drip_rate_us: Option, // === Serverless Settings === /// **Deprecated** Configure the drain period in the runner config. @@ -391,6 +399,14 @@ impl Pegboard { self.gateway_websocket_rate_limit_drip_rate_ms.unwrap_or(10) } + pub fn envoy_websocket_rate_limit_requests(&self) -> u64 { + self.envoy_websocket_rate_limit_requests.unwrap_or(16_384) + } + + pub fn envoy_websocket_rate_limit_drip_rate_us(&self) -> u64 { + self.envoy_websocket_rate_limit_drip_rate_us.unwrap_or(200) + } + pub fn actor_create_rate_limit_requests(&self) -> u64 { self.actor_create_rate_limit_requests.unwrap_or(500) } diff --git a/engine/packages/guard/src/routing/envoy.rs b/engine/packages/guard/src/routing/envoy.rs index f2f6ae6ea7..c7895068bf 100644 --- a/engine/packages/guard/src/routing/envoy.rs +++ b/engine/packages/guard/src/routing/envoy.rs @@ -93,6 +93,6 @@ async fn route_envoy_internal( tracing::debug!("authenticated envoy connection"); } - let tunnel = pegboard_envoy::PegboardEnvoyWs::new(ctx.clone()); + let tunnel = pegboard_envoy::PegboardEnvoyWs::new(&ctx); Ok(RoutingOutput::CustomServe(Arc::new(tunnel))) } diff --git a/engine/packages/pegboard-envoy/Cargo.toml b/engine/packages/pegboard-envoy/Cargo.toml index 29ad9a85be..4d7f5f54ab 100644 --- a/engine/packages/pegboard-envoy/Cargo.toml +++ b/engine/packages/pegboard-envoy/Cargo.toml @@ -32,6 +32,7 @@ depot-client.workspace = true depot-client-embedded.workspace = true rivet-runtime.workspace = true rivet-types.workspace = true +rivet-util.workspace = true scc.workspace = true serde_bare.workspace = true serde_json.workspace = true diff --git a/engine/packages/pegboard-envoy/src/lib.rs b/engine/packages/pegboard-envoy/src/lib.rs index bcb5d2c77d..99ebe04cc7 100644 --- a/engine/packages/pegboard-envoy/src/lib.rs +++ b/engine/packages/pegboard-envoy/src/lib.rs @@ -45,12 +45,8 @@ pub struct PegboardEnvoyWs { } impl PegboardEnvoyWs { - pub fn new(ctx: StandaloneCtx) -> Self { - metrics::prepopulate(); - - let service = Self { ctx: ctx.clone() }; - - service + pub fn new(ctx: &StandaloneCtx) -> Self { + Self { ctx: ctx.clone() } } } diff --git a/engine/packages/pegboard-envoy/src/metrics.rs b/engine/packages/pegboard-envoy/src/metrics.rs index d44f177081..ad8303ac40 100644 --- a/engine/packages/pegboard-envoy/src/metrics.rs +++ b/engine/packages/pegboard-envoy/src/metrics.rs @@ -371,72 +371,3 @@ pub fn set_envoy_connection_state( (None, None) => {} } } - -pub fn prepopulate() { - ENVOY_CONNECTED.with_label_values(&["", ""]).set(0); - for state in EnvoyState::ALL { - ENVOY_CONNECTIONS_BY_STATE - .with_label_values(&["", "", "", state.as_str()]) - .set(0); - } - for (state, reasons) in [ - (EnvoyState::Starting, &["websocket_accepted"][..]), - (EnvoyState::Connected, &["init_complete"][..]), - (EnvoyState::Stopping, &["envoy_reported_stopping"][..]), - ( - EnvoyState::Disconnected, - &[ - "init_failed", - "websocket_closed", - "evicted", - "going_away", - "connection_error", - ][..], - ), - (EnvoyState::Lost, &["ping_timeout"][..]), - (EnvoyState::Stopped, &["graceful_shutdown_complete"][..]), - ] { - for reason in reasons { - ENVOY_STATE_TRANSITION_TOTAL - .with_label_values(&["", "", "", state.as_str(), reason]) - .inc_by(0); - } - } - let _ = ENVOY_LIFETIME_SECONDS.with_label_values(&["", ""]); - let _ = ENVOY_PING_LAG_SECONDS.with_label_values(&["", ""]); - for result in ["ok", "no_subscribers", "error"] { - TUNNEL_PUBLISH_TOTAL - .with_label_values(&["", "", result]) - .inc_by(0); - } - TUNNEL_TASKS_ACTIVE.with_label_values(&["", ""]).set(0); - WS_RESPONSES_IN_FLIGHT.set(0); - for task_kind in ["kv", "sqlite_page", "remote_sqlite", "tunnel_message"] { - ACTOR_TASKS_ACTIVE.with_label_values(&[task_kind]).set(0); - } - for branch in ["ws_msg", "completed_task"] { - let _ = WS_TO_TUNNEL_BRANCH_DURATION.with_label_values(&[branch]); - } - for result in ["ok", "error", "timeout"] { - let _ = ACTOR_WAKE_DURATION.with_label_values(&["", "", result]); - } - let _ = SQLITE_COMMIT_ENVOY_DISPATCH_DURATION.with_label_values(&["", ""]); - let _ = SQLITE_COMMIT_ENVOY_RESPONSE_DURATION.with_label_values(&["", ""]); - for request_type in ["get_pages", "commit", "exec", "execute"] { - for result in ["ok", "error"] { - SQLITE_REQUEST_TOTAL - .with_label_values(&["", "", request_type, result]) - .inc_by(0); - let _ = SQLITE_REQUEST_DURATION.with_label_values(&["", "", request_type, result]); - } - for direction in ["request", "response"] { - let _ = SQLITE_REQUEST_PAGES.with_label_values(&["", "", request_type, direction]); - } - let _ = SQLITE_REQUEST_DIRTY_PAGES.with_label_values(&["", "", request_type]); - for direction in ["request", "response"] { - SQLITE_REQUEST_PAYLOAD_BYTES - .with_label_values(&["", "", request_type, direction]) - .inc_by(0); - } - } -} diff --git a/engine/packages/pegboard-envoy/src/ws_to_tunnel_task.rs b/engine/packages/pegboard-envoy/src/ws_to_tunnel_task.rs index 785daf8c41..fcbe313453 100644 --- a/engine/packages/pegboard-envoy/src/ws_to_tunnel_task.rs +++ b/engine/packages/pegboard-envoy/src/ws_to_tunnel_task.rs @@ -45,6 +45,13 @@ use crate::{ const MAX_REMOTE_SQL_BIND_BYTES: usize = 128 * 1024; +/// Max number of pages a single `get_pages` request may ask for. Each requested page number is ~4 +/// bytes on the wire but forces the engine to fetch and materialize up to a full 4 KiB page inside +/// one UDB transaction, so an uncapped list is a large cost-asymmetry amplifier from an untrusted +/// runner. Sized well above observed production read batches (max ~1024 pages); the commit path has +/// its own lower cap (`MAX_COMMIT_DIRTY_PAGES`) since writes batch smaller than reads. +const MAX_GET_PAGES_PER_REQUEST: usize = 8192; + /// Wall-clock threshold above which a single handle_message invocation is logged as a head-of-line /// blocking risk. The ws_to_tunnel_task loop is strictly serial per envoy, so any handler that /// spends longer than this delays every subsequent WS message from the same envoy (including @@ -461,9 +468,25 @@ pub async fn task_inner( let mut term_signal = rivet_runtime::TermSignal::get(); let mut task_manager = TaskManager::new(ctx.clone(), conn.clone()); + // Leaky bucket rate limit on consuming envoy ws messages. The envoy connection multiplexes + // every actor on a runner, so this bounds the rate at which a single untrusted runner can drive + // engine work (task spawns, KV/SQLite ops). Reads are paused while empty, applying TCP + // backpressure to the runner rather than dropping protocol messages. + let mut rate_limit = rivet_util::throttle::RateLimiter::new( + rivet_util::throttle::RateLimitMethod::LeakyBucket { + requests: ctx.config().pegboard().envoy_websocket_rate_limit_requests(), + drip_rate: Duration::from_micros( + ctx.config().pegboard().envoy_websocket_rate_limit_drip_rate_us(), + ), + }, + ); + loop { tokio::select! { - recv = recv_msg(&mut ws_rx, &mut ws_to_tunnel_abort_rx, &mut term_signal) => { + recv = async { + rate_limit.acquire().await; + recv_msg(&mut ws_rx, &mut ws_to_tunnel_abort_rx, &mut term_signal).await + } => { let branch_start = Instant::now(); let branch_result: Result> = async { match recv? { @@ -1461,6 +1484,16 @@ async fn handle_sqlite_get_pages( conn: &Conn, request: protocol::SqliteGetPagesRequest, ) -> Result { + if request.pgnos.len() > MAX_GET_PAGES_PER_REQUEST { + return Ok(protocol::SqliteGetPagesResponse::SqliteErrorResponse( + sqlite_protocol_error_response(&format!( + "sqlite get_pages requested {} pages, exceeding limit {}", + request.pgnos.len(), + MAX_GET_PAGES_PER_REQUEST, + )), + )); + } + validate_sqlite_actor_for_request(ctx, conn, &request.actor_id, request.expected_generation) .await?; diff --git a/engine/packages/pegboard-gateway2/src/metrics.rs b/engine/packages/pegboard-gateway2/src/metrics.rs index e9c7a4acfe..244df7f882 100644 --- a/engine/packages/pegboard-gateway2/src/metrics.rs +++ b/engine/packages/pegboard-gateway2/src/metrics.rs @@ -66,42 +66,3 @@ lazy_static::lazy_static! { *REGISTRY ).unwrap(); } - -pub fn prepopulate() { - const RESULTS: &[&str] = &[ - "success", - "client_disconnect", - "actor_ready_timeout", - "request_timeout", - "envoy_error", - ]; - - for protocol in ["http", "websocket"] { - IN_FLIGHT.with_label_values(&["", "", protocol]).set(0); - IN_FLIGHT_DROPPED_TOTAL - .with_label_values(&["", "", protocol, "client_disconnect"]) - .inc_by(0); - TUNNEL_PING_DURATION.with_label_values(&["", "", protocol]); - LAST_PONG_AGE_SECONDS.with_label_values(&["", "", protocol]); - REQUEST_RETRIES_TOTAL.with_label_values(&["", "", protocol, "1"]); - - for result in RESULTS { - REQUEST_DURATION_SECONDS.with_label_values(&["", "", protocol, result]); - } - - for reason in [ - "server_close", - "client_close", - "abort", - "gc_timeout", - "shutdown", - ] { - CLOSE_SENT_TOTAL - .with_label_values(&["", "", protocol, reason]) - .inc_by(0); - } - } - for result in ["ok", "error", "timeout"] { - WEBSOCKET_OPEN_WAIT_SECONDS.with_label_values(&["", "", result]); - } -} diff --git a/engine/packages/pegboard-gateway2/src/shared_state.rs b/engine/packages/pegboard-gateway2/src/shared_state.rs index 67ad84a7b8..b3de6e5887 100644 --- a/engine/packages/pegboard-gateway2/src/shared_state.rs +++ b/engine/packages/pegboard-gateway2/src/shared_state.rs @@ -160,7 +160,6 @@ pub struct SharedState(Arc); impl SharedState { pub fn new(config: &rivet_config::Config, ups: PubSub) -> Self { - metrics::prepopulate(); init_slow_ping_threshold_from_env(); let gateway_id = protocol::util::generate_gateway_id(); diff --git a/engine/packages/pegboard-gateway2/src/ws_to_tunnel_task.rs b/engine/packages/pegboard-gateway2/src/ws_to_tunnel_task.rs index 66ac4836d3..27dec6914c 100644 --- a/engine/packages/pegboard-gateway2/src/ws_to_tunnel_task.rs +++ b/engine/packages/pegboard-gateway2/src/ws_to_tunnel_task.rs @@ -25,21 +25,26 @@ pub async fn task( let mut ws_rx = ws_rx.lock().await; // Leaky bucket rate limit on consuming ws messages - let pegboard_config = ctx.config().pegboard(); let mut rate_limit = rivet_util::throttle::RateLimiter::new( rivet_util::throttle::RateLimitMethod::LeakyBucket { - requests: pegboard_config.gateway_websocket_rate_limit_requests(), + requests: ctx + .config() + .pegboard() + .gateway_websocket_rate_limit_requests(), drip_rate: Duration::from_millis( - pegboard_config.gateway_websocket_rate_limit_drip_rate_ms(), + ctx.config() + .pegboard() + .gateway_websocket_rate_limit_drip_rate_ms(), ), }, ); loop { - rate_limit.acquire().await; - tokio::select! { - res = ws_rx.try_next() => { + res = async { + rate_limit.acquire().await; + ws_rx.try_next().await + } => { if let Some(msg) = res? { ingress_bytes.fetch_add(msg.len() as u64, Ordering::AcqRel); diff --git a/engine/packages/pegboard-outbound/src/lib.rs b/engine/packages/pegboard-outbound/src/lib.rs index 5d3f3a325d..00697ea7b0 100644 --- a/engine/packages/pegboard-outbound/src/lib.rs +++ b/engine/packages/pegboard-outbound/src/lib.rs @@ -25,8 +25,6 @@ const SSE_OPEN_WARN_THRESHOLD: Duration = Duration::from_secs(5); #[tracing::instrument(skip_all)] pub async fn start(config: rivet_config::Config, pools: rivet_pools::Pools) -> Result<()> { - metrics::prepopulate(); - let cache = rivet_cache::CacheInner::from_env(&config, pools.clone())?; let ctx = StandaloneCtx::new( db::DatabaseKv::new(config.clone(), pools.clone()).await?, diff --git a/engine/packages/pegboard-outbound/src/metrics.rs b/engine/packages/pegboard-outbound/src/metrics.rs index 3f74c822e7..a61cd4fefc 100644 --- a/engine/packages/pegboard-outbound/src/metrics.rs +++ b/engine/packages/pegboard-outbound/src/metrics.rs @@ -42,50 +42,3 @@ lazy_static::lazy_static! { *REGISTRY ).unwrap(); } - -pub fn prepopulate() { - const ERRORS: &[&str] = &[ - "http_error", - "connection_error", - "stream_ended_early", - "invalid_payload", - "downgrade", - "internal", - ]; - const STATUSES: &[&str] = &["429", "503", "5xx", "4xx", "2xx", "other", ""]; - const RESULTS: &[&str] = &[ - "success", - "error_http_429", - "error_http_503", - "error_http_5xx", - "error_http_4xx", - "error_http_other", - "error_connection", - "error_stream_ended", - "error_invalid_payload", - "error_downgrade", - "error_internal", - ]; - const DRAIN_REASONS: &[&str] = &[ - "lifespan_reached", - "going_away", - "actor_lost", - "connection_lost", - "term_signal", - "", - ]; - - for error in ERRORS { - for status in STATUSES { - REQ_ERROR_TOTAL - .with_label_values(&["", "", error, status]) - .inc_by(0); - } - } - - for result in RESULTS { - for drain_reason in DRAIN_REASONS { - REQ_DURATION_SECONDS.with_label_values(&["", "", result, drain_reason]); - } - } -} diff --git a/engine/packages/pegboard/src/ops/actor/create.rs b/engine/packages/pegboard/src/ops/actor/create.rs index a0cb51bca0..b2695c16d9 100644 --- a/engine/packages/pegboard/src/ops/actor/create.rs +++ b/engine/packages/pegboard/src/ops/actor/create.rs @@ -37,7 +37,7 @@ pub struct Output { #[operation] pub async fn pegboard_actor_create(ctx: &OperationCtx, input: &Input) -> Result { - let rate_limiter = RATE_LIMITERS + let rate_limit = RATE_LIMITERS .get_or_init(|| { Cache::builder() .max_capacity(10_000) @@ -59,7 +59,7 @@ pub async fn pegboard_actor_create(ctx: &OperationCtx, input: &Input) -> Result< .await; // Limit actor creation per namespace id - if !rate_limiter.value().lock().await.try_acquire() { + if !rate_limit.value().lock().await.try_acquire() { return Err(crate::errors::Actor::CreationRateLimit.build()); } From 3063bd1b112995653525d469b10b55be044de3d7 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Wed, 24 Jun 2026 15:50:19 -0700 Subject: [PATCH 08/16] [SLOP(claude-opus-4-8-high)] feat(universaldb): postgres leader-resolver driver overhaul --- Cargo.lock | 17 + Cargo.toml | 4 + .../packages/test-deps-docker/src/database.rs | 2 +- engine/packages/universaldb/Cargo.toml | 5 + ...onflict_tracker.rs => conflict_tracker.rs} | 20 +- .../universaldb/src/driver/postgres/codec.rs | 167 +++++++++ .../universaldb/src/driver/postgres/commit.rs | 185 ++++++++++ .../src/driver/postgres/database.rs | 241 +++++++------ .../src/driver/postgres/listener.rs | 228 ++++++++++++ .../universaldb/src/driver/postgres/mod.rs | 5 + .../src/driver/postgres/resolver/apply.rs | 139 ++++++++ .../src/driver/postgres/resolver/lease.rs | 82 +++++ .../src/driver/postgres/resolver/mod.rs | 336 ++++++++++++++++++ .../universaldb/src/driver/postgres/shared.rs | 161 +++++++++ .../src/driver/postgres/transaction.rs | 22 +- .../src/driver/postgres/transaction_task.rs | 283 +++------------ .../src/driver/rocksdb/database.rs | 6 +- .../universaldb/src/driver/rocksdb/mod.rs | 1 - .../src/driver/rocksdb/transaction.rs | 7 +- .../src/driver/rocksdb/transaction_task.rs | 6 +- engine/packages/universaldb/src/lib.rs | 1 + .../packages/universaldb/tests/integration.rs | 10 +- engine/sdks/rust/depot-protocol/Cargo.toml | 1 + .../sdks/rust/universaldb-commit/Cargo.toml | 16 + engine/sdks/rust/universaldb-commit/build.rs | 64 ++++ .../rust/universaldb-commit/src/generated.rs | 1 + .../sdks/rust/universaldb-commit/src/lib.rs | 6 + .../rust/universaldb-commit/src/versioned.rs | 38 ++ .../sdks/schemas/universaldb-commit/v1.bare | 70 ++++ scripts/run/postgres.sh | 2 +- scripts/run/restore-postgres.sh | 2 +- self-host/compose/dev-host/docker-compose.yml | 2 +- .../dev-multidc-multinode/docker-compose.yml | 6 +- .../compose/dev-multidc/docker-compose.yml | 6 +- .../compose/dev-multinode/docker-compose.yml | 2 +- self-host/compose/dev/docker-compose.yml | 2 +- .../compose/template/src/docker-compose.ts | 2 +- .../k8s/engine/12-postgres-statefulset.yaml | 2 +- .../docs/self-hosting/docker-compose.mdx | 2 +- .../docs/self-hosting/docker-container.mdx | 2 +- 40 files changed, 1747 insertions(+), 407 deletions(-) rename engine/packages/universaldb/src/{driver/rocksdb/transaction_conflict_tracker.rs => conflict_tracker.rs} (69%) create mode 100644 engine/packages/universaldb/src/driver/postgres/codec.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/commit.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/listener.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/resolver/apply.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/resolver/lease.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/resolver/mod.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/shared.rs create mode 100644 engine/sdks/rust/universaldb-commit/Cargo.toml create mode 100644 engine/sdks/rust/universaldb-commit/build.rs create mode 100644 engine/sdks/rust/universaldb-commit/src/generated.rs create mode 100644 engine/sdks/rust/universaldb-commit/src/lib.rs create mode 100644 engine/sdks/rust/universaldb-commit/src/versioned.rs create mode 100644 engine/sdks/schemas/universaldb-commit/v1.bare diff --git a/Cargo.lock b/Cargo.lock index b21a5d2859..7b0053c042 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5385,6 +5385,7 @@ name = "rivet-depot-protocol" version = "2.3.7" dependencies = [ "anyhow", + "rivet-util", "rivet-vbare-compiler", "serde", "serde_bare", @@ -5946,6 +5947,17 @@ dependencies = [ "vbare", ] +[[package]] +name = "rivet-universaldb-commit" +version = "2.3.2" +dependencies = [ + "anyhow", + "rivet-vbare-compiler", + "serde", + "serde_bare", + "vbare", +] + [[package]] name = "rivet-ups-broadcast" version = "0.1.0" @@ -8254,6 +8266,7 @@ version = "2.3.7" dependencies = [ "anyhow", "async-trait", + "base64 0.22.1", "deadpool-postgres", "foundationdb-tuple", "futures-util", @@ -8267,17 +8280,21 @@ dependencies = [ "rivet-postgres-util", "rivet-test-deps-docker", "rivet-tracing-utils", + "rivet-universaldb-commit", "rocksdb", + "scc", "serde", "tempfile", "thiserror 1.0.69", "tokio", "tokio-postgres", "tokio-postgres-rustls", + "tokio-util", "tracing", "tracing-subscriber", "url", "uuid", + "vbare", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 511aec200f..a9048c159f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,6 +67,7 @@ members = [ "engine/sdks/rust/depot-protocol", "engine/sdks/rust/test-envoy", "engine/sdks/rust/ups-protocol", + "engine/sdks/rust/universaldb-commit", "rivetkit-rust/packages/actor-persist", "rivetkit-rust/packages/client", "rivetkit-rust/packages/engine-process", @@ -655,6 +656,9 @@ members = [ [workspace.dependencies.rivet-ups-protocol] path = "engine/sdks/rust/ups-protocol" + [workspace.dependencies.rivet-universaldb-commit] + path = "engine/sdks/rust/universaldb-commit" + [profile.dev] overflow-checks = false # "line-tables-only" produces just the line-number DWARF needed for stack diff --git a/engine/packages/test-deps-docker/src/database.rs b/engine/packages/test-deps-docker/src/database.rs index 11532ebf25..955e4d7897 100644 --- a/engine/packages/test-deps-docker/src/database.rs +++ b/engine/packages/test-deps-docker/src/database.rs @@ -61,7 +61,7 @@ impl TestDatabase { }); let docker_config = DockerRunConfig { - image: "postgres:17".to_string(), + image: "postgres:18".to_string(), container_name: container_name.clone(), port_mapping: (port, 5432), env_vars: vec![ diff --git a/engine/packages/universaldb/Cargo.toml b/engine/packages/universaldb/Cargo.toml index 1fdc9b1c99..b2f6b5f045 100644 --- a/engine/packages/universaldb/Cargo.toml +++ b/engine/packages/universaldb/Cargo.toml @@ -9,6 +9,7 @@ edition.workspace = true [dependencies] anyhow.workspace = true async-trait.workspace = true +base64.workspace = true deadpool-postgres.workspace = true foundationdb-tuple.workspace = true futures-util.workspace = true @@ -18,16 +19,20 @@ rand.workspace = true rivet-metrics.workspace = true rivet-postgres-util.workspace = true rivet-tracing-utils.workspace = true +rivet-universaldb-commit.workspace = true rocksdb.workspace = true +scc.workspace = true serde.workspace = true tempfile.workspace = true thiserror.workspace = true tokio-postgres-rustls.workspace = true tokio-postgres.workspace = true +tokio-util.workspace = true tokio.workspace = true tracing.workspace = true url.workspace = true uuid.workspace = true +vbare.workspace = true [dev-dependencies] rivet-config.workspace = true diff --git a/engine/packages/universaldb/src/driver/rocksdb/transaction_conflict_tracker.rs b/engine/packages/universaldb/src/conflict_tracker.rs similarity index 69% rename from engine/packages/universaldb/src/driver/rocksdb/transaction_conflict_tracker.rs rename to engine/packages/universaldb/src/conflict_tracker.rs index 370240760a..127e2bf615 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/transaction_conflict_tracker.rs +++ b/engine/packages/universaldb/src/conflict_tracker.rs @@ -11,7 +11,7 @@ use tokio::sync::Mutex; use crate::options::ConflictRangeType; // Transactions cannot live longer than 5 seconds so we don't need to store transaction conflicts longer than -// that +// that. const TXN_CONFLICT_TTL: Duration = Duration::from_secs(10); #[derive(Debug)] @@ -22,6 +22,15 @@ struct PreviousTransaction { conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, } +/// In-process FoundationDB-style resolver. Holds the last `TXN_CONFLICT_TTL` of committed +/// transactions and rejects a committing transaction if any retained transaction has both an +/// overlapping version window and an overlapping conflict range of a differing type. +/// +/// Used by the rocksdb driver (single process) and by the postgres leader-resolver. The two +/// differ only in where the commit version comes from: rocksdb generates it from the in-process +/// `global_version` counter, while the postgres leader assigns it from the durable +/// `udb_version_seq` so it survives leader failover and matches the versionstamp. For that reason +/// `check_and_insert` takes the commit version from the caller instead of generating it. #[derive(Clone)] pub struct TransactionConflictTracker { // NOTE: We use a mutex because we need to lock reads across all active txns. This could be optimized to @@ -40,18 +49,23 @@ impl TransactionConflictTracker { } } - /// Each number returned is unique. + /// Each number returned is unique. Used by the in-process rocksdb driver to assign both start + /// and commit versions. The postgres leader does not use this; it assigns versions from the + /// durable Postgres sequence. pub fn next_global_version(&self) -> u64 { self.global_version.fetch_add(1, Ordering::SeqCst) } + /// Returns `true` on conflict (same polarity as the original rocksdb tracker). The caller + /// supplies `commit_version` (e.g. `nextval('udb_version_seq')` on the postgres leader, or + /// `next_global_version()` on rocksdb) so version assignment stays the caller's responsibility. pub async fn check_and_insert( &self, txn1_start_version: u64, + txn1_commit_version: u64, txn1_conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, ) -> bool { let mut txns = self.txns.lock().await; - let txn1_commit_version = self.next_global_version(); // Prune old entries txns.retain(|txn| txn.insert_instant.elapsed() < TXN_CONFLICT_TTL); diff --git a/engine/packages/universaldb/src/driver/postgres/codec.rs b/engine/packages/universaldb/src/driver/postgres/codec.rs new file mode 100644 index 0000000000..432ede2f34 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/codec.rs @@ -0,0 +1,167 @@ +use anyhow::Result; +use rivet_universaldb_commit::{self as proto, versioned}; +use vbare::OwnedVersionedData; + +use crate::{ + options::{ConflictRangeType, MutationType}, + tx_ops::Operation, +}; + +/// Decoded form of a `udb_commit_requests.payload` blob. +/// +/// `read_version` is intentionally omitted: it is also denormalized into the `read_version` column, +/// which is what the leader's drain reads, so decoding it here would be dead. +pub struct DecodedCommit { + pub conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, + pub operations: Vec, +} + +/// Encode a follower's commit request to the versioned BARE wire format with an embedded version +/// header so a leader running older or newer code can still decode it during a rolling deploy. +pub fn encode_commit_request( + read_version: u64, + conflict_ranges: &[(Vec, Vec, ConflictRangeType)], + operations: &[Operation], +) -> Result> { + let request = proto::CommitRequest { + read_version, + conflict_ranges: conflict_ranges + .iter() + .map(|(begin, end, kind)| proto::ConflictRange { + begin: begin.clone(), + end: end.clone(), + kind: conflict_range_type_to_proto(*kind), + }) + .collect(), + operations: operations.iter().map(operation_to_proto).collect(), + }; + + versioned::CommitRequest::wrap_latest(request) + .serialize_with_embedded_version(proto::PROTOCOL_VERSION) +} + +/// Decode a `udb_commit_requests.payload` blob produced by [`encode_commit_request`]. +pub fn decode_commit_request(payload: &[u8]) -> Result { + let request = versioned::CommitRequest::deserialize_with_embedded_version(payload)?; + + let conflict_ranges = request + .conflict_ranges + .into_iter() + .map(|range| { + ( + range.begin, + range.end, + conflict_range_type_from_proto(range.kind), + ) + }) + .collect(); + + let operations = request + .operations + .into_iter() + .map(operation_from_proto) + .collect(); + + Ok(DecodedCommit { + conflict_ranges, + operations, + }) +} + +fn conflict_range_type_to_proto(kind: ConflictRangeType) -> proto::ConflictRangeType { + match kind { + ConflictRangeType::Read => proto::ConflictRangeType::Read, + ConflictRangeType::Write => proto::ConflictRangeType::Write, + } +} + +fn conflict_range_type_from_proto(kind: proto::ConflictRangeType) -> ConflictRangeType { + match kind { + proto::ConflictRangeType::Read => ConflictRangeType::Read, + proto::ConflictRangeType::Write => ConflictRangeType::Write, + } +} + +fn operation_to_proto(op: &Operation) -> proto::Operation { + match op { + Operation::SetValue { key, value } => proto::Operation::SetValue(proto::SetValue { + key: key.clone(), + value: value.clone(), + }), + Operation::Clear { key } => proto::Operation::Clear(proto::Clear { key: key.clone() }), + Operation::ClearRange { begin, end } => proto::Operation::ClearRange(proto::ClearRange { + begin: begin.clone(), + end: end.clone(), + }), + Operation::AtomicOp { + key, + param, + op_type, + } => proto::Operation::AtomicOp(proto::AtomicOp { + key: key.clone(), + param: param.clone(), + op_type: mutation_type_to_proto(*op_type), + }), + } +} + +fn operation_from_proto(op: proto::Operation) -> Operation { + match op { + proto::Operation::SetValue(proto::SetValue { key, value }) => { + Operation::SetValue { key, value } + } + proto::Operation::Clear(proto::Clear { key }) => Operation::Clear { key }, + proto::Operation::ClearRange(proto::ClearRange { begin, end }) => { + Operation::ClearRange { begin, end } + } + proto::Operation::AtomicOp(proto::AtomicOp { + key, + param, + op_type, + }) => Operation::AtomicOp { + key, + param, + op_type: mutation_type_from_proto(op_type), + }, + } +} + +fn mutation_type_to_proto(op_type: MutationType) -> proto::MutationType { + match op_type { + MutationType::Add => proto::MutationType::Add, + MutationType::And => proto::MutationType::And, + MutationType::BitAnd => proto::MutationType::BitAnd, + MutationType::Or => proto::MutationType::Or, + MutationType::BitOr => proto::MutationType::BitOr, + MutationType::Xor => proto::MutationType::Xor, + MutationType::BitXor => proto::MutationType::BitXor, + MutationType::AppendIfFits => proto::MutationType::AppendIfFits, + MutationType::Max => proto::MutationType::Max, + MutationType::Min => proto::MutationType::Min, + MutationType::SetVersionstampedKey => proto::MutationType::SetVersionstampedKey, + MutationType::SetVersionstampedValue => proto::MutationType::SetVersionstampedValue, + MutationType::ByteMin => proto::MutationType::ByteMin, + MutationType::ByteMax => proto::MutationType::ByteMax, + MutationType::CompareAndClear => proto::MutationType::CompareAndClear, + } +} + +fn mutation_type_from_proto(op_type: proto::MutationType) -> MutationType { + match op_type { + proto::MutationType::Add => MutationType::Add, + proto::MutationType::And => MutationType::And, + proto::MutationType::BitAnd => MutationType::BitAnd, + proto::MutationType::Or => MutationType::Or, + proto::MutationType::BitOr => MutationType::BitOr, + proto::MutationType::Xor => MutationType::Xor, + proto::MutationType::BitXor => MutationType::BitXor, + proto::MutationType::AppendIfFits => MutationType::AppendIfFits, + proto::MutationType::Max => MutationType::Max, + proto::MutationType::Min => MutationType::Min, + proto::MutationType::SetVersionstampedKey => MutationType::SetVersionstampedKey, + proto::MutationType::SetVersionstampedValue => MutationType::SetVersionstampedValue, + proto::MutationType::ByteMin => MutationType::ByteMin, + proto::MutationType::ByteMax => MutationType::ByteMax, + proto::MutationType::CompareAndClear => MutationType::CompareAndClear, + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/commit.rs b/engine/packages/universaldb/src/driver/postgres/commit.rs new file mode 100644 index 0000000000..e61714dd0e --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/commit.rs @@ -0,0 +1,185 @@ +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result}; + +use crate::{error::DatabaseError, options::ConflictRangeType, tx_ops::Operation}; + +use super::{ + codec, + shared::{LeaseInfo, PostgresShared, commit_channel, reply_channel}, +}; + +/// How long to wait for a leader to be elected before giving up a submit as retryable. +const LEADER_WAIT_TIMEOUT: Duration = Duration::from_secs(5); +/// Backstop poll cadence while waiting for a commit result, in case a reply NOTIFY is missed. +const RESULT_POLL_INTERVAL: Duration = Duration::from_millis(250); + +/// Submit a follower transaction's commit to the leader and await the result. +/// +/// `read_version` is the watermark captured when this transaction opened its read snapshot. A pure +/// snapshot read-only transaction (no operations and no read conflict ranges) submits nothing. +pub async fn submit( + shared: &Arc, + read_version: i64, + operations: Vec, + conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, +) -> Result<()> { + // A transaction with no writes and no serializable read ranges has nothing to order or + // validate; it never needs the leader. + if operations.is_empty() + && conflict_ranges + .iter() + .all(|(_, _, kind)| matches!(kind, ConflictRangeType::Write)) + { + return Ok(()); + } + + let lease = wait_for_leader(shared).await?; + let payload = + codec::encode_commit_request(read_version.max(0) as u64, &conflict_ranges, &operations) + .context("failed to encode commit request")?; + let reply_channel = reply_channel(&shared.node_id); + + // Subscribe to our reply channel before inserting so we cannot miss the leader's NOTIFY. + let mut reply_rx = shared.listener.listen(&reply_channel).await; + + let conn = shared + .pool + .get() + .await + .context("failed to get connection for commit submit")?; + + let id: i64 = conn + .query_one( + "INSERT INTO udb_commit_requests (epoch, read_version, payload, reply_channel) + VALUES ($1, $2, $3, $4) + RETURNING id", + &[&lease.epoch, &read_version, &payload, &reply_channel], + ) + .await + .context("failed to enqueue commit request")? + .get(0); + + // Wake the leader's drain loop. + if let Err(err) = conn + .execute( + "SELECT pg_notify($1, $2)", + &[&commit_channel(&lease.leader_addr), &id.to_string()], + ) + .await + { + tracing::debug!( + ?err, + "failed to notify leader; relying on its poll backstop" + ); + } + + // Release the connection before waiting so a long wait does not pin a pool slot. The request + // row is durable, so await_result re-acquires a connection per poll. + drop(conn); + + await_result(shared, id, lease.epoch, &mut reply_rx).await +} + +/// Wait for a known leader, returning a retryable error if none is elected in time. +async fn wait_for_leader(shared: &Arc) -> Result { + let deadline = Instant::now() + LEADER_WAIT_TIMEOUT; + loop { + if let Some(lease) = shared.current_lease() { + return Ok(lease); + } + if Instant::now() >= deadline { + return Err(DatabaseError::NotCommitted.into()); + } + tokio::time::sleep(RESULT_POLL_INTERVAL).await; + } +} + +/// Poll the request row until it reaches a terminal status, woken by reply NOTIFYs with a polling +/// backstop. Bails as retryable if the leader epoch advances (our request is now orphaned and will +/// never be applied, so it is definitively not committed). +async fn await_result( + shared: &Arc, + id: i64, + submit_epoch: i64, + reply_rx: &mut tokio::sync::broadcast::Receiver, +) -> Result<()> { + loop { + // Re-acquire a connection per poll: the request row is durable, so a transient pool/query + // error just means we retry the poll rather than failing a possibly-applied commit. + match read_status(shared, id).await { + Ok(Some(Status::Committed)) => return Ok(()), + Ok(Some(Status::Conflict)) => return Err(DatabaseError::NotCommitted.into()), + Ok(Some(Status::Pending)) => {} + Ok(None) => { + // The row was GC'd before we observed a terminal status. Treat as not committed + // and let the retry loop resubmit. + return Err(DatabaseError::NotCommitted.into()); + } + Err(err) => { + tracing::debug!(?err, "transient error polling commit status, retrying"); + } + } + + // If a new leader took over, our old-epoch request will never be claimed. + if let Some(current) = shared.current_lease() { + if current.epoch != submit_epoch { + return Err(DatabaseError::NotCommitted.into()); + } + } + + tokio::select! { + res = reply_rx.recv() => { + match res { + Ok(_) | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + *reply_rx = shared + .listener + .listen(&reply_channel(&shared.node_id)) + .await; + } + } + } + _ = tokio::time::sleep(RESULT_POLL_INTERVAL) => {} + } + } +} + +enum Status { + Pending, + Committed, + Conflict, +} + +/// Read the current status of a commit request. `Ok(None)` means the row no longer exists. +async fn read_status(shared: &Arc, id: i64) -> Result> { + let conn = shared + .pool + .get() + .await + .context("failed to get connection for commit status poll")?; + + let row = conn + .query_opt( + "SELECT status FROM udb_commit_requests WHERE id = $1", + &[&id], + ) + .await + .context("failed to read commit request status")?; + + let Some(row) = row else { + return Ok(None); + }; + + let status: String = row.get(0); + let status = match status.as_str() { + "committed" => Status::Committed, + "conflict" => Status::Conflict, + // 'pending' or any in-flight state. + _ => Status::Pending, + }; + Ok(Some(status)) +} diff --git a/engine/packages/universaldb/src/driver/postgres/database.rs b/engine/packages/universaldb/src/driver/postgres/database.rs index 4d50d08ab3..6e24cf647d 100644 --- a/engine/packages/universaldb/src/driver/postgres/database.rs +++ b/engine/packages/universaldb/src/driver/postgres/database.rs @@ -13,6 +13,7 @@ use rivet_postgres_util::build_tls_config; use tokio::task::JoinHandle; use tokio_postgres_rustls::MakeRustlsConnect; use url::Url; +use uuid::Uuid; use crate::{ RetryableTransaction, Transaction, @@ -22,9 +23,15 @@ use crate::{ utils::{MaybeCommitted, calculate_tx_retry_backoff}, }; -use super::transaction::PostgresTransactionDriver; +use super::{ + listener::PgListener, resolver, shared::PostgresShared, transaction::PostgresTransactionDriver, +}; -const GC_INTERVAL: Duration = Duration::from_secs(5); +const GC_INTERVAL: Duration = Duration::from_secs(30); +/// Terminal and orphaned commit-request rows older than this are garbage collected. Must be well +/// beyond the longest a follower could spend awaiting a result, so a result is never deleted before +/// it is observed. +const COMMIT_ROW_MAX_AGE_SECS: i64 = 60; #[derive(Clone, Debug)] pub struct PostgresConfig { @@ -50,7 +57,7 @@ impl PostgresConfig { } pub struct PostgresDatabaseDriver { - pool: Pool, + shared: Arc, max_retries: AtomicI32, gc_handle: JoinHandle<()>, } @@ -63,7 +70,60 @@ impl PostgresDatabaseDriver { "creating PostgresDatabaseDriver" ); - // Create deadpool config from connection string + let ssl_disabled = if let Ok(url) = Url::parse(&config.connection_string) { + url.query_pairs() + .any(|(k, v)| k == "sslmode" && v == "disable") + } else { + false + }; + + let pool = Self::build_pool(&config, ssl_disabled)?; + + // Initialize the schema (idempotent). + { + let conn = pool + .get() + .await + .context("failed to get connection from postgres pool")?; + Self::init_schema(&conn).await?; + } + + // Unique per-process node id (no hyphens) used to name this node's NOTIFY channels. Kept + // short so `udb_commit_` stays within Postgres's 63-byte identifier limit. + let node_id = Uuid::new_v4().simple().to_string(); + + let listener = PgListener::new( + config.connection_string.clone(), + ssl_disabled, + config + .ssl_config + .as_ref() + .and_then(|c| c.ssl_root_cert_path.clone()), + config + .ssl_config + .as_ref() + .and_then(|c| c.ssl_client_cert_path.clone()), + config + .ssl_config + .as_ref() + .and_then(|c| c.ssl_client_key_path.clone()), + ); + + let shared = PostgresShared::new(pool, node_id, listener); + + // Every node runs the resolver; only the elected leader drains the commit queue. + resolver::spawn(shared.clone()); + + let gc_handle = Self::spawn_gc(shared.clone()); + + Ok(PostgresDatabaseDriver { + shared, + max_retries: AtomicI32::new(100), + gc_handle, + }) + } + + fn build_pool(config: &PostgresConfig, ssl_disabled: bool) -> Result { let mut pool_config = Config::new(); pool_config.url = Some(config.connection_string.clone()); pool_config.pool = Some(PoolConfig { @@ -74,21 +134,10 @@ impl PostgresDatabaseDriver { recycling_method: RecyclingMethod::Fast, }); - tracing::debug!("creating Postgres pool"); - - let ssl_disabled = if let Ok(url) = Url::parse(&config.connection_string) { - url.query_pairs() - .any(|(k, v)| k == "sslmode" && v == "disable") - } else { - false - }; - - let pool = if ssl_disabled { - let tls = tokio_postgres::NoTls; - + if ssl_disabled { pool_config - .create_pool(Some(Runtime::Tokio1), tls) - .context("failed to create postgres connection pool")? + .create_pool(Some(Runtime::Tokio1), tokio_postgres::NoTls) + .context("failed to create postgres connection pool") } else { let tls_config = build_tls_config( config @@ -104,139 +153,87 @@ impl PostgresDatabaseDriver { .as_ref() .and_then(|c| c.ssl_client_key_path.as_ref()), )?; - let tls = MakeRustlsConnect::new(tls_config); - pool_config - .create_pool(Some(Runtime::Tokio1), tls) - .context("failed to create postgres connection pool")? - }; - - tracing::debug!("Getting Postgres connection from pool"); - // Get a connection from the pool to create the table - let conn = pool - .get() - .await - .context("failed to get connection from postgres pool")?; - - // Enable btree gist - conn.execute("CREATE EXTENSION IF NOT EXISTS btree_gist", &[]) - .await - .context("failed to create btree_gist extension")?; - - conn.execute("CREATE UNLOGGED SEQUENCE IF NOT EXISTS global_version_seq START WITH 1 INCREMENT BY 1 MINVALUE 1", &[]) - .await - .context("failed to create global version sequence")?; + .create_pool(Some(Runtime::Tokio1), MakeRustlsConnect::new(tls_config)) + .context("failed to create postgres connection pool") + } + } - // Create the KV table if it doesn't exist - conn.execute( + async fn init_schema(conn: &deadpool_postgres::Client) -> Result<()> { + // Durable latest-value store. + conn.batch_execute( "CREATE TABLE IF NOT EXISTS kv ( key BYTEA PRIMARY KEY, value BYTEA NOT NULL - )", - &[], - ) - .await - .context("failed to create kv table")?; - - // Create range_type type if it doesn't exist - conn.execute( - "DO $$ BEGIN - CREATE TYPE range_type AS ENUM ('read', 'write'); - EXCEPTION - WHEN duplicate_object THEN null; - END $$", - &[], - ) - .await - .context("failed to create range_type enum")?; - - // Create bytearange type if it doesn't exist - conn.execute( - "DO $$ BEGIN - CREATE TYPE bytearange AS RANGE ( - SUBTYPE = bytea, - SUBTYPE_OPCLASS = bytea_ops - ); - EXCEPTION - WHEN duplicate_object THEN null; - END $$", - &[], + ); + + CREATE TABLE IF NOT EXISTS udb_lease ( + id INT PRIMARY KEY DEFAULT 1 CHECK (id = 1), + epoch BIGINT NOT NULL, + leader_addr TEXT NOT NULL, + durable_version BIGINT NOT NULL DEFAULT 0, + expires_at TIMESTAMPTZ NOT NULL + ); + + CREATE SEQUENCE IF NOT EXISTS udb_version_seq AS BIGINT + START WITH 1 INCREMENT BY 1 MINVALUE 1; + + CREATE TABLE IF NOT EXISTS udb_commit_requests ( + id BIGSERIAL PRIMARY KEY, + epoch BIGINT NOT NULL, + read_version BIGINT NOT NULL, + payload BYTEA NOT NULL, + reply_channel TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + commit_version BIGINT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + ); + + CREATE INDEX IF NOT EXISTS udb_commit_requests_pending + ON udb_commit_requests (id) WHERE status = 'pending';", ) .await - .context("failed to create bytearange type")?; - - // Create the conflict ranges table for non-snapshot reads - // This enforces consistent reads for ranges by preventing overlapping conflict ranges - conn.execute( - "CREATE UNLOGGED TABLE IF NOT EXISTS conflict_ranges ( - range_data BYTEARANGE NOT NULL, - conflict_type range_type NOT NULL, - start_version BIGINT NOT NULL, - commit_version BIGINT NOT NULL, - ts timestamp NOT NULL DEFAULT now(), - - EXCLUDE USING gist ( - -- Conflict if byte range overlaps... - range_data WITH &&, - -- And if conflict types are different... - conflict_type WITH <>, - -- And if the txn versions overlap... - int8range(start_version, commit_version, '[]') WITH &&, - -- But not if the start_version is the same (from the same txn) - start_version WITH <> - ) - )", - &[], - ) - .await - .context("failed to create conflict_ranges table")?; + .context("failed to initialize postgres schema")?; - // Create index on ts column for efficient garbage collection - conn.execute( - "CREATE INDEX IF NOT EXISTS idx_conflict_ranges_ts ON conflict_ranges (ts)", - &[], - ) - .await - .context("failed to create index on conflict_ranges ts column")?; + Ok(()) + } - let pool2 = pool.clone(); - let gc_handle = tokio::spawn(async move { + fn spawn_gc(shared: Arc) -> JoinHandle<()> { + tokio::spawn(async move { let mut interval = tokio::time::interval(GC_INTERVAL); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { interval.tick().await; - tracing::debug!(status=?pool2.status(), "postgres pool status"); + let conn = match shared.pool.get().await { + Ok(conn) => conn, + Err(err) => { + tracing::debug!(?err, "failed to get connection for commit gc"); + continue; + } + }; - // NOTE: Transactions have a max limit of 5 seconds, we delete after 10 seconds for extra padding - // Delete old conflict ranges if let Err(err) = conn .execute( - "DELETE FROM conflict_ranges where ts < now() - interval '10 seconds'", - &[], + "DELETE FROM udb_commit_requests + WHERE created_at < now() - ($1::bigint * interval '1 second')", + &[&COMMIT_ROW_MAX_AGE_SECS], ) .await { - tracing::error!(?err, "failed postgres gc task"); + tracing::error!(?err, "failed postgres commit-queue gc"); } } - }); - - Ok(PostgresDatabaseDriver { - pool, - max_retries: AtomicI32::new(100), - gc_handle, }) } } impl DatabaseDriver for PostgresDatabaseDriver { fn create_txn(&self) -> Result { - // Pass the connection pool and config to the transaction driver - Ok(Transaction::new(Arc::new( - PostgresTransactionDriver::with_config(self.pool.clone()), - ))) + Ok(Transaction::new(Arc::new(PostgresTransactionDriver::new( + self.shared.clone(), + )))) } fn run<'a>( diff --git a/engine/packages/universaldb/src/driver/postgres/listener.rs b/engine/packages/universaldb/src/driver/postgres/listener.rs new file mode 100644 index 0000000000..e5a7919dfe --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/listener.rs @@ -0,0 +1,228 @@ +use std::{path::PathBuf, sync::Arc, time::Duration}; + +use futures_util::future::poll_fn; +use rivet_postgres_util::build_tls_config; +use scc::HashMap; +use tokio::{ + io::{AsyncRead, AsyncWrite}, + sync::{Mutex, broadcast}, +}; +use tokio_postgres::AsyncMessage; +use tokio_postgres_rustls::MakeRustlsConnect; + +/// How long to wait between reconnect attempts for the dedicated LISTEN connection. +const RECONNECT_BACKOFF: Duration = Duration::from_secs(1); +/// Capacity of each channel's broadcast buffer. Notifications are wakeup signals with a polling +/// backstop, so a lagged receiver only delays a wake, never drops a durable commit. +const BROADCAST_CAPACITY: usize = 1024; + +struct Subscription { + tx: broadcast::Sender, +} + +/// Owns a single dedicated Postgres connection used exclusively for `LISTEN`. Demultiplexes +/// incoming `NOTIFY` payloads to per-channel broadcast senders and re-`LISTEN`s every registered +/// channel after a reconnect. +/// +/// This is separate from the deadpool pool because deadpool recycles connections and drops the +/// async notification stream; LISTEN requires owning the connection's message stream directly. +pub struct PgListener { + conn_str: String, + ssl_disabled: bool, + ssl_root_cert_path: Option, + ssl_client_cert_path: Option, + ssl_client_key_path: Option, + channels: Arc>, + client: Arc>>, +} + +impl PgListener { + pub fn new( + conn_str: String, + ssl_disabled: bool, + ssl_root_cert_path: Option, + ssl_client_cert_path: Option, + ssl_client_key_path: Option, + ) -> Self { + let channels: Arc> = Arc::new(HashMap::new()); + let client: Arc>> = Arc::new(Mutex::new(None)); + + tokio::spawn(Self::connection_lifecycle( + conn_str.clone(), + ssl_disabled, + ssl_root_cert_path.clone(), + ssl_client_cert_path.clone(), + ssl_client_key_path.clone(), + channels.clone(), + client.clone(), + )); + + Self { + conn_str, + ssl_disabled, + ssl_root_cert_path, + ssl_client_cert_path, + ssl_client_key_path, + channels, + client, + } + } + + /// Subscribe to a channel, registering a `LISTEN` if this is the first subscriber. Returns a + /// broadcast receiver of notification payloads. Idempotent per channel. + pub async fn listen(&self, channel: &str) -> broadcast::Receiver { + match self.channels.entry_async(channel.to_string()).await { + scc::hash_map::Entry::Occupied(entry) => entry.get().tx.subscribe(), + scc::hash_map::Entry::Vacant(entry) => { + let (tx, rx) = broadcast::channel(BROADCAST_CAPACITY); + entry.insert_entry(Subscription { tx }); + + // Best-effort immediate LISTEN; the lifecycle task re-LISTENs on reconnect. + if let Some(client) = &*self.client.lock().await { + if let Err(err) = client.execute(&format!("LISTEN \"{channel}\""), &[]).await { + tracing::warn!(?err, %channel, "failed to LISTEN, will retry on reconnect"); + } + } + + rx + } + } + } + + async fn connection_lifecycle( + conn_str: String, + ssl_disabled: bool, + ssl_root_cert_path: Option, + ssl_client_cert_path: Option, + ssl_client_key_path: Option, + channels: Arc>, + client: Arc>>, + ) { + loop { + let connected = if ssl_disabled { + Self::connect_and_run(&conn_str, tokio_postgres::NoTls, &channels, &client).await + } else { + match build_tls_config( + ssl_root_cert_path.as_ref(), + ssl_client_cert_path.as_ref(), + ssl_client_key_path.as_ref(), + ) { + Ok(tls_config) => { + Self::connect_and_run( + &conn_str, + MakeRustlsConnect::new(tls_config), + &channels, + &client, + ) + .await + } + Err(err) => { + tracing::error!(?err, "failed to build listener TLS config"); + false + } + } + }; + + if !connected { + tokio::time::sleep(RECONNECT_BACKOFF).await; + } + } + } + + /// Connects, re-LISTENs all channels, then drives the notification poll loop until the + /// connection closes. Returns `true` if a connection was successfully established (so the caller + /// can skip the reconnect backoff). + async fn connect_and_run( + conn_str: &str, + tls: T, + channels: &Arc>, + client: &Arc>>, + ) -> bool + where + T: tokio_postgres::tls::MakeTlsConnect, + T::Stream: AsyncRead + AsyncWrite + Unpin + Send + 'static, + T::TlsConnect: Send, + >::Future: Send, + { + let (new_client, connection) = match tokio_postgres::connect(conn_str, tls).await { + Ok(pair) => pair, + Err(err) => { + tracing::error!(?err, "failed to connect postgres listener"); + return false; + } + }; + + let channels_poll = channels.clone(); + let poll_handle = + tokio::spawn(async move { Self::poll_connection(connection, channels_poll).await }); + + // Re-LISTEN all registered channels on the fresh connection. + let mut registered = Vec::new(); + channels + .iter_async(|k, _| { + registered.push(k.clone()); + true + }) + .await; + for channel in ®istered { + if let Err(err) = new_client + .execute(&format!("LISTEN \"{channel}\""), &[]) + .await + { + tracing::error!(?err, %channel, "failed to re-LISTEN channel after reconnect"); + } + } + + *client.lock().await = Some(new_client); + + // Block until the poll loop ends (connection closed or errored). + let _ = poll_handle.await; + + *client.lock().await = None; + + true + } + + async fn poll_connection( + mut connection: tokio_postgres::Connection, + channels: Arc>, + ) where + S: AsyncRead + AsyncWrite + Unpin, + T: AsyncRead + AsyncWrite + Unpin, + { + loop { + match poll_fn(|cx| connection.poll_message(cx)).await { + Some(Ok(AsyncMessage::Notification(note))) => { + if let Some(sub) = channels.get_async(note.channel()).await { + // Ignore send errors: no active receiver just means no one is waiting + // right now; the polling backstop covers them. + let _ = sub.tx.send(note.payload().to_string()); + } + } + Some(Ok(_)) => {} + Some(Err(err)) => { + tracing::warn!(?err, "postgres listener connection error"); + break; + } + None => { + tracing::warn!("postgres listener connection closed"); + break; + } + } + } + } +} + +impl Clone for PgListener { + fn clone(&self) -> Self { + Self { + conn_str: self.conn_str.clone(), + ssl_disabled: self.ssl_disabled, + ssl_root_cert_path: self.ssl_root_cert_path.clone(), + ssl_client_cert_path: self.ssl_client_cert_path.clone(), + ssl_client_key_path: self.ssl_client_key_path.clone(), + channels: self.channels.clone(), + client: self.client.clone(), + } + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/mod.rs b/engine/packages/universaldb/src/driver/postgres/mod.rs index 1c24f9cb94..64f4bbd1bf 100644 --- a/engine/packages/universaldb/src/driver/postgres/mod.rs +++ b/engine/packages/universaldb/src/driver/postgres/mod.rs @@ -1,4 +1,9 @@ +mod codec; +mod commit; mod database; +mod listener; +mod resolver; +mod shared; mod transaction; mod transaction_task; diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs b/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs new file mode 100644 index 0000000000..d9729de884 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs @@ -0,0 +1,139 @@ +use anyhow::{Context, Result}; +use deadpool_postgres::Transaction; + +use crate::{ + atomic::apply_atomic_op, options::MutationType, tuple::Versionstamp, tx_ops::Operation, + versionstamp::substitute_raw_versionstamp, +}; + +/// Apply a winning transaction's operations to `kv` inside the leader's batch txn. +/// +/// `commit_version` is the Postgres-resolved version assigned to this commit (`nextval`). It is +/// substituted into the 8-byte committed-version slot of any versionstamped key/value so +/// versionstamps are globally monotonic with commit order across all follower processes. +pub async fn apply( + txn: &Transaction<'_>, + operations: Vec, + commit_version: u64, +) -> Result<()> { + // Distinguishes multiple versionstamped operations within a single commit so their 10-byte + // stamps stay unique (8-byte version shared, 2-byte counter incremented). + let mut versionstamp_counter: u16 = 0; + + for op in operations { + match op { + Operation::SetValue { key, value } => { + upsert(txn, &key, &value).await?; + } + Operation::Clear { key } => { + txn.execute("DELETE FROM kv WHERE key = $1", &[&key]) + .await + .context("failed to clear key")?; + } + Operation::ClearRange { begin, end } => { + txn.execute( + "DELETE FROM kv WHERE key >= $1 AND key < $2", + &[&begin, &end], + ) + .await + .context("failed to clear range")?; + } + Operation::AtomicOp { + key, + param, + op_type, + } => { + apply_atomic( + txn, + key, + param, + op_type, + commit_version, + &mut versionstamp_counter, + ) + .await?; + } + } + } + + Ok(()) +} + +async fn apply_atomic( + txn: &Transaction<'_>, + key: Vec, + param: Vec, + op_type: MutationType, + commit_version: u64, + versionstamp_counter: &mut u16, +) -> Result<()> { + match op_type { + MutationType::SetVersionstampedKey => { + let versionstamp = build_versionstamp(commit_version, versionstamp_counter); + let key = substitute_raw_versionstamp(key, &versionstamp) + .map_err(anyhow::Error::msg) + .context("failed substituting versionstamped key")?; + upsert(txn, &key, ¶m).await?; + } + MutationType::SetVersionstampedValue => { + let versionstamp = build_versionstamp(commit_version, versionstamp_counter); + let value = substitute_raw_versionstamp(param, &versionstamp) + .map_err(anyhow::Error::msg) + .context("failed substituting versionstamped value")?; + upsert(txn, &key, &value).await?; + } + // Read-modify-write atomics: the leader is the single writer, so reading the live value + // inside the apply txn and writing the result is serializable with no lost update. + MutationType::Add + | MutationType::And + | MutationType::BitAnd + | MutationType::Or + | MutationType::BitOr + | MutationType::Xor + | MutationType::BitXor + | MutationType::AppendIfFits + | MutationType::Max + | MutationType::Min + | MutationType::ByteMin + | MutationType::ByteMax + | MutationType::CompareAndClear => { + let current = txn + .query_opt("SELECT value FROM kv WHERE key = $1", &[&key]) + .await + .context("failed to read current value for atomic op")? + .map(|row| row.get::<_, Vec>(0)); + + let new_value = apply_atomic_op(current.as_deref(), ¶m, op_type); + + if let Some(new_value) = new_value { + upsert(txn, &key, &new_value).await?; + } else { + txn.execute("DELETE FROM kv WHERE key = $1", &[&key]) + .await + .context("failed to clear key after atomic op")?; + } + } + } + + Ok(()) +} + +async fn upsert(txn: &Transaction<'_>, key: &[u8], value: &[u8]) -> Result<()> { + txn.execute( + "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2", + &[&key, &value], + ) + .await + .context("failed to upsert kv")?; + Ok(()) +} + +/// Build a 10-byte versionstamp (plus the 2 user-version bytes the substitution helper ignores) +/// from the Postgres-resolved commit version and a per-commit counter. +fn build_versionstamp(commit_version: u64, counter: &mut u16) -> Versionstamp { + let mut bytes = [0u8; 12]; + bytes[0..8].copy_from_slice(&commit_version.to_be_bytes()); + bytes[8..10].copy_from_slice(&counter.to_be_bytes()); + *counter = counter.wrapping_add(1); + Versionstamp::from(bytes) +} diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs b/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs new file mode 100644 index 0000000000..a17b6c7759 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs @@ -0,0 +1,82 @@ +use anyhow::{Context, Result}; +use deadpool_postgres::Pool; + +use crate::driver::postgres::shared::LEASE_ID; + +/// Lease time-to-live. A leader renews well within this; a candidate may take over only after it +/// expires. +pub const LEASE_TTL_SECS: i64 = 10; + +/// Outcome of a leadership acquisition attempt. +pub struct Acquired { + pub epoch: i64, +} + +/// Attempt to acquire or take over the leader lease via an epoch CAS. Succeeds if there is no lease +/// row yet, or the existing lease has expired. Bumps `epoch` on every successful acquisition so a +/// superseded old leader is fenced out. +pub async fn try_acquire(pool: &Pool, node_id: &str) -> Result> { + let conn = pool + .get() + .await + .context("failed to get connection for lease acquire")?; + + // Take over an expired (or absent) lease. The INSERT seeds the singleton row on first ever + // election; thereafter the UPDATE path runs. + let row = conn + .query_opt( + "INSERT INTO udb_lease (id, epoch, leader_addr, durable_version, expires_at) + VALUES ($1, 1, $2, 0, now() + ($3 || ' seconds')::interval) + ON CONFLICT (id) DO UPDATE + SET epoch = udb_lease.epoch + 1, + leader_addr = EXCLUDED.leader_addr, + expires_at = now() + ($3 || ' seconds')::interval + WHERE udb_lease.expires_at < now() + RETURNING epoch", + &[&LEASE_ID, &node_id, &LEASE_TTL_SECS.to_string()], + ) + .await + .context("failed to run lease acquire query")?; + + Ok(row.map(|row| Acquired { epoch: row.get(0) })) +} + +/// Renew the lease, fenced on this leader's epoch. Returns `false` if the lease was lost (another +/// node took over, bumping the epoch), in which case the caller must step down. +pub async fn renew(pool: &Pool, node_id: &str, epoch: i64) -> Result { + let conn = pool + .get() + .await + .context("failed to get connection for lease renew")?; + + let updated = conn + .execute( + "UPDATE udb_lease + SET expires_at = now() + ($3 || ' seconds')::interval + WHERE id = $1 AND epoch = $2 AND leader_addr = $4", + &[&LEASE_ID, &epoch, &LEASE_TTL_SECS.to_string(), &node_id], + ) + .await + .context("failed to renew lease")?; + + Ok(updated == 1) +} + +/// Read the current durable version (`udb_lease.durable_version`). Used by a freshly elected leader +/// to learn the watermark floor it must continue from. +pub async fn current_durable_version(pool: &Pool) -> Result { + let conn = pool + .get() + .await + .context("failed to get connection for durable version read")?; + + let row = conn + .query_opt( + "SELECT durable_version FROM udb_lease WHERE id = $1", + &[&LEASE_ID], + ) + .await + .context("failed to read durable version")?; + + Ok(row.map(|row| row.get::<_, i64>(0)).unwrap_or(0)) +} diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs new file mode 100644 index 0000000000..79f1f1f48b --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs @@ -0,0 +1,336 @@ +mod apply; +mod lease; + +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result}; + +use crate::{conflict_tracker::TransactionConflictTracker, transaction::TXN_TIMEOUT}; + +use super::shared::{LEASE_ID, LeaseInfo, PostgresShared, WATERMARK_CHANNEL, commit_channel}; + +/// Max commits resolved+applied per batch (group commit). Amortizes the resolver, Postgres +/// round-trips, and fsync across the batch. +const DRAIN_BATCH_SIZE: i64 = 256; + +/// How often a leader renews its lease. Must be comfortably under `LEASE_TTL_SECS`. +const RENEW_INTERVAL: Duration = Duration::from_secs(3); + +/// Backstop poll cadence so a missed `udb_commit` NOTIFY cannot stall the drain indefinitely. +const POLL_BACKSTOP: Duration = Duration::from_millis(50); + +/// How long a candidate waits before retrying election when another node holds the lease. +const ELECTION_RETRY: Duration = Duration::from_secs(2); + +enum DrainOutcome { + /// Processed zero or more requests; still leader. + Drained, + /// Lost the lease (epoch bumped by a new leader). Step down. + LostLease, +} + +/// Spawn the per-process resolver task. Every node runs this; only the elected leader drains the +/// commit queue. +pub fn spawn(shared: Arc) { + tokio::spawn(run(shared)); +} + +async fn run(shared: Arc) { + loop { + match lease::try_acquire(&shared.pool, &shared.node_id).await { + Ok(Some(acquired)) => { + tracing::info!(epoch = acquired.epoch, node_id = %shared.node_id, "acquired udb leader lease"); + if let Err(err) = lead(&shared, acquired.epoch).await { + tracing::error!(?err, "udb leader loop errored, stepping down"); + } + tracing::info!(epoch = acquired.epoch, "stepped down from udb leader"); + } + Ok(None) => { + tokio::time::sleep(ELECTION_RETRY).await; + } + Err(err) => { + tracing::warn!(?err, "failed udb lease acquire attempt"); + tokio::time::sleep(ELECTION_RETRY).await; + } + } + } +} + +/// Leader main loop: hold the lease, drain the commit queue on wake or poll, and renew the lease. +async fn lead(shared: &Arc, epoch: i64) -> Result<()> { + // Publish our own lease into the cache immediately so our local commits route to us. + shared.set_lease(LeaseInfo { + epoch, + leader_addr: shared.node_id.clone(), + }); + + // The recovery floor: a freshly elected leader has a cold conflict window, so reject commits + // whose read_version predates the floor until the window warms (one TXN_TIMEOUT), forcing + // those followers to take a fresh read_version. + let recovery_version = recovery_floor(shared).await?; + let recovery_deadline = Instant::now() + TXN_TIMEOUT; + + let tracker = TransactionConflictTracker::new(); + + let mut wake_rx = shared + .listener + .listen(&commit_channel(&shared.node_id)) + .await; + + let mut renew_interval = tokio::time::interval(RENEW_INTERVAL); + renew_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut poll_interval = tokio::time::interval(POLL_BACKSTOP); + poll_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + // Drain anything already queued before our first wake. + if matches!( + drain(shared, epoch, &tracker, recovery_version, recovery_deadline).await?, + DrainOutcome::LostLease + ) { + return Ok(()); + } + + loop { + tokio::select! { + _ = renew_interval.tick() => { + if !lease::renew(&shared.pool, &shared.node_id, epoch).await? { + tracing::warn!(epoch, "lost udb lease on renew"); + return Ok(()); + } + } + res = wake_rx.recv() => { + match res { + Ok(_) | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + wake_rx = shared.listener.listen(&commit_channel(&shared.node_id)).await; + } + } + if matches!( + drain(shared, epoch, &tracker, recovery_version, recovery_deadline).await?, + DrainOutcome::LostLease + ) { + return Ok(()); + } + } + _ = poll_interval.tick() => { + if matches!( + drain(shared, epoch, &tracker, recovery_version, recovery_deadline).await?, + DrainOutcome::LostLease + ) { + return Ok(()); + } + } + } + } +} + +/// The version floor a freshly elected leader continues from: the higher of the durable watermark +/// and the sequence high-water. The LOGGED `udb_version_seq` is crash-safe, so this never regresses. +async fn recovery_floor(shared: &Arc) -> Result { + let durable = lease::current_durable_version(&shared.pool).await?; + + let conn = shared + .pool + .get() + .await + .context("failed to get connection for recovery floor")?; + let seq_high: i64 = conn + .query_one("SELECT last_value FROM udb_version_seq", &[]) + .await + .context("failed to read sequence high water")? + .get(0); + + Ok(durable.max(seq_high).max(0) as u64) +} + +/// Drain pending commit requests in id-ordered batches until none remain. Each batch resolves and +/// applies inside a single Postgres transaction (group commit), fenced on the leader's epoch. +async fn drain( + shared: &Arc, + epoch: i64, + tracker: &TransactionConflictTracker, + recovery_version: u64, + recovery_deadline: Instant, +) -> Result { + loop { + match drain_batch(shared, epoch, tracker, recovery_version, recovery_deadline).await? { + BatchOutcome::Empty => return Ok(DrainOutcome::Drained), + BatchOutcome::Processed => {} + BatchOutcome::LostLease => return Ok(DrainOutcome::LostLease), + } + } +} + +enum BatchOutcome { + Empty, + Processed, + LostLease, +} + +struct Reply { + channel: String, + id: i64, +} + +async fn drain_batch( + shared: &Arc, + epoch: i64, + tracker: &TransactionConflictTracker, + recovery_version: u64, + recovery_deadline: Instant, +) -> Result { + let mut conn = shared + .pool + .get() + .await + .context("failed to get connection for drain batch")?; + let txn = conn + .build_transaction() + .start() + .await + .context("failed to start drain batch txn")?; + + // Claim a batch in id order. FOR UPDATE SKIP LOCKED holds the rows for this txn so they are + // stamped terminal on COMMIT with no intermediate 'claimed' state to clean up. + let rows = txn + .query( + "SELECT id, read_version, payload, reply_channel + FROM udb_commit_requests + WHERE status = 'pending' AND epoch = $1 + ORDER BY id + LIMIT $2 + FOR UPDATE SKIP LOCKED", + &[&epoch, &DRAIN_BATCH_SIZE], + ) + .await + .context("failed to claim commit batch")?; + + if rows.is_empty() { + txn.rollback().await.ok(); + return Ok(BatchOutcome::Empty); + } + + let cold_window = Instant::now() < recovery_deadline; + let mut max_winner_cv: i64 = 0; + let mut replies = Vec::with_capacity(rows.len()); + + for row in &rows { + let id: i64 = row.get(0); + let read_version: i64 = row.get(1); + let payload: Vec = row.get(2); + let reply_channel: String = row.get(3); + + let decoded = super::codec::decode_commit_request(&payload) + .context("failed to decode commit payload")?; + + let commit_version: i64 = txn + .query_one("SELECT nextval('udb_version_seq')", &[]) + .await + .context("failed to get next commit version")? + .get(0); + + let start_version = read_version.max(0) as u64; + + // Cold-window guard: a commit whose read_version predates the recovery floor cannot be + // safely resolved against this leader's empty window. Reject it as retryable. + let conflicted = if cold_window && start_version < recovery_version { + true + } else { + tracker + .check_and_insert( + start_version, + commit_version.max(0) as u64, + decoded.conflict_ranges, + ) + .await + }; + + if conflicted { + txn.execute( + "UPDATE udb_commit_requests SET status = 'conflict' WHERE id = $1", + &[&id], + ) + .await + .context("failed to stamp conflict")?; + } else { + apply::apply(&txn, decoded.operations, commit_version.max(0) as u64) + .await + .context("failed to apply commit")?; + txn.execute( + "UPDATE udb_commit_requests SET status = 'committed', commit_version = $1 WHERE id = $2", + &[&commit_version, &id], + ) + .await + .context("failed to stamp committed")?; + max_winner_cv = max_winner_cv.max(commit_version); + } + + replies.push(Reply { + channel: reply_channel, + id, + }); + } + + // Advance the watermark, fenced on our epoch. A zombie old leader whose epoch was bumped sees + // zero rows updated and must step down before any of its writes become visible. + let new_durable: i64 = match txn + .query_opt( + "UPDATE udb_lease + SET durable_version = GREATEST(durable_version, $1) + WHERE id = $2 AND epoch = $3 + RETURNING durable_version", + &[&max_winner_cv, &LEASE_ID, &epoch], + ) + .await + .context("failed to advance watermark")? + { + Some(row) => row.get(0), + None => { + txn.rollback().await.ok(); + return Ok(BatchOutcome::LostLease); + } + }; + + txn.commit().await.context("failed to commit drain batch")?; + + // Watermark advances strictly after the apply txn is durably committed and visible, so a + // reader handed this read_version can never miss a write with commit_version <= read_version. + shared.advance_durable_version(new_durable); + + notify_after_commit(&conn, new_durable, &replies).await; + + Ok(BatchOutcome::Processed) +} + +/// Wake watermark listeners and the followers waiting on each processed request. Best-effort: a +/// missed NOTIFY is covered by the follower's polling backstop and the watermark refresh timer. +async fn notify_after_commit( + conn: &deadpool_postgres::Client, + new_durable: i64, + replies: &[Reply], +) { + if let Err(err) = conn + .execute( + "SELECT pg_notify($1, $2)", + &[&WATERMARK_CHANNEL, &new_durable.to_string()], + ) + .await + { + tracing::debug!(?err, "failed to notify watermark"); + } + + let channels: Vec<&str> = replies.iter().map(|r| r.channel.as_str()).collect(); + let ids: Vec = replies.iter().map(|r| r.id.to_string()).collect(); + if let Err(err) = conn + .execute( + "SELECT pg_notify(c, p) FROM unnest($1::text[], $2::text[]) AS t(c, p)", + &[&channels, &ids], + ) + .await + { + tracing::debug!(?err, "failed to notify commit replies"); + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/shared.rs b/engine/packages/universaldb/src/driver/postgres/shared.rs new file mode 100644 index 0000000000..0ab2963158 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/shared.rs @@ -0,0 +1,161 @@ +use std::{ + sync::{ + Arc, + atomic::{AtomicI64, Ordering}, + }, + time::Duration, +}; + +use deadpool_postgres::Pool; +use tokio::sync::{Notify, watch}; + +use super::listener::PgListener; + +/// The singleton row id of `udb_lease`. +pub const LEASE_ID: i32 = 1; + +/// How often the follower refreshes its cached lease row (epoch, leader channel, watermark) as a +/// backstop to the `udb_watermark` NOTIFY. A stale-but-older watermark only widens the conflict +/// window, so this can be loose. +const LEASE_REFRESH_INTERVAL: Duration = Duration::from_millis(500); + +/// Channel a follower NOTIFYs (and the leader LISTENs) to wake the leader's drain loop. +pub fn commit_channel(node_id: &str) -> String { + format!("udb_commit_{node_id}") +} + +/// Channel the leader NOTIFYs (and a follower LISTENs) to deliver a commit result. +pub fn reply_channel(node_id: &str) -> String { + format!("udb_reply_{node_id}") +} + +/// Channel the leader NOTIFYs on every watermark advance; all nodes LISTEN. +pub const WATERMARK_CHANNEL: &str = "udb_watermark"; + +/// Cached view of the current leader lease, as seen by a follower. +#[derive(Clone, Debug)] +pub struct LeaseInfo { + pub epoch: i64, + /// Node id of the current leader, used to build its commit channel. + pub leader_addr: String, +} + +/// Process-wide state shared by the follower transaction tasks and the leader resolver. Every node +/// is both a follower (it submits its own commits) and a candidate leader. +pub struct PostgresShared { + pub pool: Pool, + /// Unique per-process id used to name this node's NOTIFY channels. + pub node_id: String, + pub listener: PgListener, + /// Highest durable commit version (`udb_lease.durable_version`); the follower read version. + durable_version: AtomicI64, + /// Pinged whenever `durable_version` advances. + watermark_notify: Notify, + lease_tx: watch::Sender>, + lease_rx: watch::Receiver>, +} + +impl PostgresShared { + pub fn new(pool: Pool, node_id: String, listener: PgListener) -> Arc { + let (lease_tx, lease_rx) = watch::channel(None); + let shared = Arc::new(Self { + pool, + node_id, + listener, + durable_version: AtomicI64::new(0), + watermark_notify: Notify::new(), + lease_tx, + lease_rx, + }); + + tokio::spawn(Self::cache_refresh_task(shared.clone())); + + shared + } + + /// The cached follower read version (`durable_version`). + pub fn read_version(&self) -> i64 { + self.durable_version.load(Ordering::SeqCst) + } + + /// Advance the cached watermark monotonically and wake any waiters. + pub fn advance_durable_version(&self, version: i64) { + let prev = self.durable_version.fetch_max(version, Ordering::SeqCst); + if version > prev { + self.watermark_notify.notify_waiters(); + } + } + + /// Current cached lease, if known. + pub fn current_lease(&self) -> Option { + self.lease_rx.borrow().clone() + } + + /// Publish a freshly observed/elected lease into the cache. + pub fn set_lease(&self, lease: LeaseInfo) { + let _ = self.lease_tx.send(Some(lease)); + } + + /// Background task: keep `durable_version` and the cached lease fresh via the `udb_watermark` + /// NOTIFY plus a periodic poll of `udb_lease`. + async fn cache_refresh_task(shared: Arc) { + let mut watermark_rx = shared.listener.listen(WATERMARK_CHANNEL).await; + let mut interval = tokio::time::interval(LEASE_REFRESH_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + tokio::select! { + notify = watermark_rx.recv() => { + match notify { + Ok(payload) => { + if let Ok(version) = payload.parse::() { + shared.advance_durable_version(version); + } + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + // Re-subscribe; the listener recreates the channel on reconnect. + watermark_rx = shared.listener.listen(WATERMARK_CHANNEL).await; + } + } + } + _ = interval.tick() => { + shared.refresh_lease_row().await; + } + } + } + } + + async fn refresh_lease_row(&self) { + let conn = match self.pool.get().await { + Ok(conn) => conn, + Err(err) => { + tracing::debug!(?err, "failed to get connection for lease refresh"); + return; + } + }; + + let row = conn + .query_opt( + "SELECT epoch, leader_addr, durable_version FROM udb_lease WHERE id = $1", + &[&LEASE_ID], + ) + .await; + + match row { + Ok(Some(row)) => { + let epoch: i64 = row.get(0); + let leader_addr: String = row.get(1); + let durable_version: i64 = row.get(2); + self.advance_durable_version(durable_version); + self.set_lease(LeaseInfo { epoch, leader_addr }); + } + Ok(None) => { + // No lease row yet; no leader elected. + } + Err(err) => { + tracing::debug!(?err, "failed to refresh lease row"); + } + } + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/transaction.rs b/engine/packages/universaldb/src/driver/postgres/transaction.rs index f80cbba393..cc315feba4 100644 --- a/engine/packages/universaldb/src/driver/postgres/transaction.rs +++ b/engine/packages/universaldb/src/driver/postgres/transaction.rs @@ -1,11 +1,13 @@ use std::{ future::Future, pin::Pin, - sync::atomic::{AtomicBool, Ordering}, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, }; use anyhow::{Context, Result}; -use deadpool_postgres::Pool; use tokio::sync::{OnceCell, mpsc, oneshot}; use crate::{ @@ -18,33 +20,35 @@ use crate::{ value::{Slice, Value, Values}, }; -use super::transaction_task::{TransactionCommand, TransactionTask}; +use super::{ + shared::PostgresShared, + transaction_task::{TransactionCommand, TransactionTask}, +}; pub struct PostgresTransactionDriver { - pool: Pool, + shared: Arc, operations: TransactionOperations, committed: AtomicBool, tx_sender: OnceCell>, } impl PostgresTransactionDriver { - pub fn with_config(pool: Pool) -> Self { + pub fn new(shared: Arc) -> Self { PostgresTransactionDriver { - pool, + shared, operations: TransactionOperations::default(), committed: AtomicBool::new(false), tx_sender: OnceCell::new(), } } - /// Get or create the transaction task + /// Get or create the transaction task that owns this transaction's read snapshot. async fn ensure_transaction(&self) -> Result<&mpsc::UnboundedSender> { self.tx_sender .get_or_try_init(|| async { let (sender, receiver) = mpsc::unbounded_channel(); - // Spawn the transaction task with serializable isolation - let task = TransactionTask::new(self.pool.clone(), receiver); + let task = TransactionTask::new(self.shared.clone(), receiver); tokio::spawn(task.run()); anyhow::Ok(sender) diff --git a/engine/packages/universaldb/src/driver/postgres/transaction_task.rs b/engine/packages/universaldb/src/driver/postgres/transaction_task.rs index 297751f608..c4d3e3f9f6 100644 --- a/engine/packages/universaldb/src/driver/postgres/transaction_task.rs +++ b/engine/packages/universaldb/src/driver/postgres/transaction_task.rs @@ -1,17 +1,18 @@ -use anyhow::{Context, Result, anyhow, bail}; -use deadpool_postgres::{Pool, Transaction}; +use std::sync::Arc; + +use anyhow::{Result, anyhow, bail}; +use deadpool_postgres::Transaction; use tokio::sync::{mpsc, oneshot}; use tokio_postgres::IsolationLevel; use crate::{ - atomic::apply_atomic_op, - error::DatabaseError, - options::{ConflictRangeType, MutationType}, + options::ConflictRangeType, tx_ops::Operation, value::{KeyValue, Slice, Values}, - versionstamp::{generate_versionstamp, substitute_raw_versionstamp}, }; +use super::{commit, shared::PostgresShared}; + pub enum TransactionCommand { // Read operations Get { @@ -48,70 +49,57 @@ pub enum TransactionCommand { }, } -/// TransactionTask runs in a separate tokio task to manage a PostgreSQL transaction. -/// -/// This design is necessary because PostgreSQL transactions have lifetime constraints -/// that don't work well with the FoundationDB-style API. Specifically: -/// - The transaction must outlive all references to it -/// - We can't store the transaction in a mutex due to lifetime issues with the connection +/// TransactionTask runs in a separate tokio task to own a single pinned PostgreSQL `REPEATABLE READ` +/// snapshot connection for the lifetime of a follower transaction. /// -/// By running in a separate task and communicating via channels, we avoid these lifetime -/// issues while maintaining a single serializable transaction for all operations. +/// Reads go directly against this snapshot (they never involve the leader). Commits delegate to +/// [`commit::submit`], which enqueues the request on the leader and awaits the result. The +/// `read_version` is captured from the cached watermark before the snapshot is opened, so no write +/// with `commit_version <= read_version` can be invisible to the snapshot. pub struct TransactionTask { - pool: Pool, + shared: Arc, receiver: mpsc::UnboundedReceiver, } impl TransactionTask { - pub fn new(pool: Pool, receiver: mpsc::UnboundedReceiver) -> Self { - Self { pool, receiver } + pub fn new( + shared: Arc, + receiver: mpsc::UnboundedReceiver, + ) -> Self { + Self { shared, receiver } } pub async fn run(mut self) { - // Get connection from pool - let mut conn = match self.pool.get().await { + // Capture the read version BEFORE opening the snapshot so the snapshot reflects every write + // with commit_version <= read_version. + let read_version = self.shared.read_version(); + + let mut conn = match self.shared.pool.get().await { Ok(conn) => conn, Err(_) => { - // If we can't get a connection, respond to all pending commands with errors self.fail_receiver().await; return; } }; - // Start the read transaction let tx = match conn .build_transaction() .isolation_level(IsolationLevel::RepeatableRead) + .read_only(true) .start() .await { Ok(tx) => tx, Err(_) => { - // If we can't start a transaction, respond to all pending commands with errors self.fail_receiver().await; return; } }; - // TODO: Parallelize future - let start_version = match tx - .query_one("SELECT nextval('global_version_seq')", &[]) - .await - { - Ok(row) => row.get::<_, i64>(0), - Err(err) => { - tracing::error!(?err, "failed to get postgres txn start_version"); - self.fail_receiver().await; - return; - } - }; - - // Process commands while let Some(cmd) = self.receiver.recv().await { match cmd { TransactionCommand::Get { key, response } => { let result = self.handle_get(&tx, &key).await; - let _ = response.send(result); } TransactionCommand::GetKey { @@ -121,7 +109,6 @@ impl TransactionTask { response, } => { let result = self.handle_get_key(&tx, &key, or_equal, offset).await; - let _ = response.send(result); } TransactionCommand::GetRange { @@ -148,7 +135,6 @@ impl TransactionTask { reverse, ) .await; - let _ = response.send(result); } TransactionCommand::Commit { @@ -156,14 +142,12 @@ impl TransactionTask { conflict_ranges, response, } => { - let (_, result) = tokio::join!( - // Read-only txn, we don't care about the result - tx.commit(), - self.handle_commit(start_version, operations, conflict_ranges), - ); - + // The read snapshot is read-only; release it and submit the commit to the leader. + let _ = tx.commit().await; + let result = + commit::submit(&self.shared, read_version, operations, conflict_ranges) + .await; let _ = response.send(result); - // Exit after commit return; } TransactionCommand::GetEstimatedRangeSize { @@ -174,13 +158,12 @@ impl TransactionTask { let result = self .handle_get_estimated_range_size(&tx, &begin, &end) .await; - let _ = response.send(result); } } } - // If the channel is closed, the transaction will be rolled back when dropped + // If the channel is closed, the snapshot transaction is rolled back when dropped. } async fn handle_get(&mut self, tx: &Transaction<'_>, key: &[u8]) -> Result> { @@ -234,27 +217,18 @@ impl TransactionTask { reverse: bool, ) -> Result { // Determine SQL operators based on key selector types - // For begin selector: - // first_greater_or_equal: or_equal = false, offset = 1 -> ">=" - // first_greater_than: or_equal = true, offset = 1 -> ">" let begin_op = if begin_offset == 1 { if begin_or_equal { ">" } else { ">=" } } else { - // This shouldn't happen for begin in range queries ">=" }; - // For end selector: - // first_greater_than: or_equal = true, offset = 1 -> "<=" - // first_greater_or_equal: or_equal = false, offset = 1 -> "<" let end_op = if end_offset == 1 { if end_or_equal { "<=" } else { "<" } } else { - // This shouldn't happen for end in range queries "<" }; - // Build query with CTE that adds conflict range let query = if reverse { if let Some(limit) = limit { format!( @@ -301,22 +275,22 @@ impl TransactionTask { begin: &[u8], end: &[u8], ) -> Result { - // Sample's 1% of the range + // Sample 1% of the range. let query = " WITH range_stats AS ( - SELECT + SELECT COUNT(*) as estimated_count, COALESCE(SUM(pg_column_size(key) + pg_column_size(value)), 0) as sample_size - FROM kv TABLESAMPLE SYSTEM(1) + FROM kv TABLESAMPLE SYSTEM(1) WHERE key >= $1 AND key < $2 ), table_stats AS ( - SELECT reltuples::bigint as total_rows - FROM pg_class + SELECT reltuples::bigint as total_rows + FROM pg_class WHERE relname = 'kv' AND relkind = 'r' ) - SELECT - CASE + SELECT + CASE WHEN r.estimated_count = 0 THEN 0 ELSE (r.sample_size * 100)::bigint END as estimated_size @@ -329,165 +303,6 @@ impl TransactionTask { .map_err(map_postgres_error) } - async fn handle_commit( - &mut self, - start_version: i64, - operations: Vec, - conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, - ) -> Result<()> { - // Get connection from pool - let mut conn = self.pool.get().await?; - - // Start write transaction - let tx = conn - .build_transaction() - .isolation_level(IsolationLevel::ReadCommitted) - .start() - .await - .context("failed to start write txn")?; - - let mut begins = Vec::with_capacity(conflict_ranges.len()); - let mut ends = Vec::with_capacity(conflict_ranges.len()); - let mut conflict_types = Vec::with_capacity(conflict_ranges.len()); - - for (begin, end, conflict_type) in conflict_ranges { - let conflict_type = match conflict_type { - ConflictRangeType::Read => "read", - ConflictRangeType::Write => "write", - }; - - begins.push(begin); - ends.push(end); - conflict_types.push(conflict_type); - } - - let query = " - WITH data AS ( - SELECT nextval('global_version_seq') AS commit_version - ) - INSERT INTO conflict_ranges (range_data, conflict_type, start_version, commit_version) - SELECT - bytearange(begin_key, end_key, '[)'), - conflict_type::range_type, - $4, - data.commit_version - FROM UNNEST($1::bytea[], $2::bytea[], $3::text[]) AS t(begin_key, end_key, conflict_type), data"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - // Insert all conflict ranges at once - tx.execute(&stmt, &[&begins, &ends, &conflict_types, &start_version]) - .await - .map_err(map_postgres_error)?; - - let transaction_versionstamp = generate_versionstamp(0); - - for op in operations { - match op { - Operation::SetValue { key, value } => { - let query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key, &value]) - .await - .map_err(map_postgres_error)?; - } - Operation::Clear { key } => { - let query = "DELETE FROM kv WHERE key = $1"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key]) - .await - .map_err(map_postgres_error)?; - } - Operation::ClearRange { begin, end } => { - let query = "DELETE FROM kv WHERE key >= $1 AND key < $2"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&begin, &end]) - .await - .map_err(map_postgres_error)?; - } - Operation::AtomicOp { - key, - param, - op_type, - } => { - if matches!(op_type, MutationType::SetVersionstampedKey) { - let key = substitute_raw_versionstamp(key, &transaction_versionstamp) - .map_err(anyhow::Error::msg) - .context("failed substituting versionstamped key")?; - let query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key, ¶m]) - .await - .map_err(map_postgres_error)?; - continue; - } - - if matches!(op_type, MutationType::SetVersionstampedValue) { - let value = substitute_raw_versionstamp(param, &transaction_versionstamp) - .map_err(anyhow::Error::msg) - .context("failed substituting versionstamped value")?; - let query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key, &value]) - .await - .map_err(map_postgres_error)?; - continue; - } - - // TODO: All operations need to be done on the sql side, not in rust - - // Get current value from database - let current_query = "SELECT value FROM kv WHERE key = $1"; - let stmt = tx - .prepare_cached(current_query) - .await - .map_err(map_postgres_error)?; - - let current_row = tx - .query_opt(&stmt, &[&key]) - .await - .map_err(map_postgres_error)?; - - // Extract current value or use None if key doesn't exist - let current_value = current_row.map(|row| row.get::<_, Vec>(0)); - let current_slice = current_value.as_deref(); - - // Apply atomic operation - let new_value = apply_atomic_op(current_slice, ¶m, op_type); - - // Store the result - if let Some(new_value) = new_value { - let update_query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; - let stmt = tx - .prepare_cached(update_query) - .await - .map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key, &new_value]) - .await - .map_err(map_postgres_error)?; - } else { - let update_query = "DELETE FROM kv WHERE key = $1"; - let stmt = tx - .prepare_cached(update_query) - .await - .map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key]) - .await - .map_err(map_postgres_error)?; - } - } - } - } - - tx.commit().await.map_err(map_postgres_error) - } - async fn fail_receiver(&mut self) { while let Some(cmd) = self.receiver.recv().await { match cmd { @@ -511,31 +326,19 @@ impl TransactionTask { } } -/// Maps PostgreSQL error to DatabaseError +/// Maps a PostgreSQL error from the read path to a `DatabaseError` where appropriate. fn map_postgres_error(err: tokio_postgres::Error) -> anyhow::Error { - let error_str = if let Some(err) = err.as_db_error() { - err.to_string() - } else { - err.to_string() - }; + let error_str = err.to_string(); - if error_str.contains("exclusion_violation") - || error_str.contains("violates exclusion constraint") - { - // Retryable - another transaction has a conflicting range - DatabaseError::NotCommitted.into() - } else if error_str.contains("serialization failure") + if error_str.contains("serialization failure") || error_str.contains("could not serialize") || error_str.contains("deadlock detected") { - // Retryable - transaction conflict - DatabaseError::NotCommitted.into() + crate::error::DatabaseError::NotCommitted.into() } else if error_str.contains("current transaction is aborted") { - // Returned by the rest of the commands in a txn if it failed for exclusion reasons - DatabaseError::NotCommitted.into() + crate::error::DatabaseError::NotCommitted.into() } else { tracing::error!(%err, "postgres error"); - // Non-retryable error anyhow::Error::new(err) } } diff --git a/engine/packages/universaldb/src/driver/rocksdb/database.rs b/engine/packages/universaldb/src/driver/rocksdb/database.rs index ec0175ee5a..ba186d872e 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/database.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/database.rs @@ -17,9 +17,9 @@ use crate::{ utils::{MaybeCommitted, calculate_tx_retry_backoff}, }; -use super::{ - transaction::RocksDbTransactionDriver, transaction_conflict_tracker::TransactionConflictTracker, -}; +use crate::conflict_tracker::TransactionConflictTracker; + +use super::transaction::RocksDbTransactionDriver; pub struct RocksDbDatabaseDriver { db: Arc, diff --git a/engine/packages/universaldb/src/driver/rocksdb/mod.rs b/engine/packages/universaldb/src/driver/rocksdb/mod.rs index a24bd72603..b18e28edca 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/mod.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/mod.rs @@ -1,6 +1,5 @@ mod database; mod transaction; -mod transaction_conflict_tracker; mod transaction_task; pub use database::RocksDbDatabaseDriver; diff --git a/engine/packages/universaldb/src/driver/rocksdb/transaction.rs b/engine/packages/universaldb/src/driver/rocksdb/transaction.rs index e62c3eda79..85cf5edf5c 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/transaction.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/transaction.rs @@ -21,10 +21,9 @@ use crate::{ value::{Slice, Value, Values}, }; -use super::{ - transaction_conflict_tracker::TransactionConflictTracker, - transaction_task::{TransactionCommand, TransactionTask}, -}; +use crate::conflict_tracker::TransactionConflictTracker; + +use super::transaction_task::{TransactionCommand, TransactionTask}; pub struct RocksDbTransactionDriver { db: Arc, diff --git a/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs b/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs index 69835eb6a7..986e6ba2bd 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs @@ -6,9 +6,9 @@ use rocksdb::{ }; use tokio::sync::{mpsc, oneshot}; -use super::transaction_conflict_tracker::TransactionConflictTracker; use crate::{ atomic::apply_atomic_op, + conflict_tracker::TransactionConflictTracker, error::DatabaseError, key_selector::KeySelector, options::{ConflictRangeType, MutationType}, @@ -412,9 +412,11 @@ impl TransactionTask { } } + // rocksdb generates both start and commit versions from the in-process counter. + let commit_version = self.txn_conflict_tracker.next_global_version(); if self .txn_conflict_tracker - .check_and_insert(start_version, conflict_ranges) + .check_and_insert(start_version, commit_version, conflict_ranges) .await { return Err(DatabaseError::NotCommitted.into()); diff --git a/engine/packages/universaldb/src/lib.rs b/engine/packages/universaldb/src/lib.rs index 96260a177d..1b6fb78b87 100644 --- a/engine/packages/universaldb/src/lib.rs +++ b/engine/packages/universaldb/src/lib.rs @@ -1,4 +1,5 @@ pub(crate) mod atomic; +pub(crate) mod conflict_tracker; mod database; pub mod driver; pub mod error; diff --git a/engine/packages/universaldb/tests/integration.rs b/engine/packages/universaldb/tests/integration.rs index 52fa703116..45a5b3bdf7 100644 --- a/engine/packages/universaldb/tests/integration.rs +++ b/engine/packages/universaldb/tests/integration.rs @@ -137,11 +137,9 @@ async fn test_database_options(db: &Database) { use std::sync::Arc; use std::sync::atomic::{AtomicU32, Ordering}; use universaldb::error::DatabaseError; - use universaldb::options::DatabaseOption; // Test setting transaction retry limit - db.set_option(DatabaseOption::TransactionRetryLimit(5)) - .unwrap(); + db.txn_retry_limit(5).unwrap(); // Test that retry limit is respected by forcing conflicts let conflict_counter = Arc::new(AtomicU32::new(0)); @@ -172,8 +170,7 @@ async fn test_database_options(db: &Database) { assert_eq!(final_attempts, 3, "Should have taken 3 attempts"); // Now set a very low retry limit and verify it fails - db.set_option(DatabaseOption::TransactionRetryLimit(1)) - .unwrap(); + db.txn_retry_limit(1).unwrap(); let conflict_counter2 = Arc::new(AtomicU32::new(0)); let counter_clone2 = conflict_counter2.clone(); @@ -204,8 +201,7 @@ async fn test_database_options(db: &Database) { assert!(attempts <= 2, "Should not retry more than limit + 1"); // Reset to a reasonable retry limit - db.set_option(DatabaseOption::TransactionRetryLimit(100)) - .unwrap(); + db.txn_retry_limit(100).unwrap(); } async fn clear_test_namespace(db: &Database) -> Result<()> { diff --git a/engine/sdks/rust/depot-protocol/Cargo.toml b/engine/sdks/rust/depot-protocol/Cargo.toml index 274836220a..99d8117349 100644 --- a/engine/sdks/rust/depot-protocol/Cargo.toml +++ b/engine/sdks/rust/depot-protocol/Cargo.toml @@ -8,6 +8,7 @@ edition.workspace = true [dependencies] anyhow.workspace = true +rivet-util.workspace = true serde_bare.workspace = true serde.workspace = true vbare.workspace = true diff --git a/engine/sdks/rust/universaldb-commit/Cargo.toml b/engine/sdks/rust/universaldb-commit/Cargo.toml new file mode 100644 index 0000000000..126c5cb35e --- /dev/null +++ b/engine/sdks/rust/universaldb-commit/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "rivet-universaldb-commit" +publish = false +version.workspace = true +authors.workspace = true +license.workspace = true +edition.workspace = true + +[dependencies] +anyhow.workspace = true +serde_bare.workspace = true +serde.workspace = true +vbare.workspace = true + +[build-dependencies] +vbare-compiler.workspace = true diff --git a/engine/sdks/rust/universaldb-commit/build.rs b/engine/sdks/rust/universaldb-commit/build.rs new file mode 100644 index 0000000000..6400be6f2b --- /dev/null +++ b/engine/sdks/rust/universaldb-commit/build.rs @@ -0,0 +1,64 @@ +use std::{ + fs, + path::{Path, PathBuf}, +}; + +fn main() -> Result<(), Box> { + let manifest_dir = std::env::var("CARGO_MANIFEST_DIR")?; + let out_dir = PathBuf::from(std::env::var("OUT_DIR")?); + let workspace_root = Path::new(&manifest_dir) + .parent() + .and_then(|p| p.parent()) + .and_then(|p| p.parent()) + .ok_or("Failed to find workspace root")?; + + let schema_dir = workspace_root + .join("sdks") + .join("schemas") + .join("universaldb-commit"); + println!("cargo:rerun-if-changed={}", schema_dir.display()); + + let (highest_version, _) = find_highest_version(&schema_dir); + + let cfg = vbare_compiler::Config::default(); + vbare_compiler::process_schemas_with_config(&schema_dir, &cfg)?; + + // Append protocol version constant to generated file + let combined_imports_path = out_dir.join("combined_imports.rs"); + let mut combined = fs::read_to_string(&combined_imports_path)?; + combined.push_str(&format!( + "\npub const PROTOCOL_VERSION: u16 = {};\n", + highest_version + )); + fs::write(combined_imports_path, combined)?; + + Ok(()) +} + +fn find_highest_version(schema_dir: &Path) -> (u32, PathBuf) { + let mut highest_version = 0; + let mut highest_version_path = PathBuf::new(); + + for entry in fs::read_dir(schema_dir).unwrap().flatten() { + if !entry.path().is_dir() { + let path = entry.path(); + let bare_name = path + .file_name() + .unwrap() + .to_str() + .unwrap() + .split_once('.') + .unwrap() + .0; + + if let Ok(version) = bare_name[1..].parse::() { + if version > highest_version { + highest_version = version; + highest_version_path = path; + } + } + } + } + + (highest_version, highest_version_path) +} diff --git a/engine/sdks/rust/universaldb-commit/src/generated.rs b/engine/sdks/rust/universaldb-commit/src/generated.rs new file mode 100644 index 0000000000..84801af8dc --- /dev/null +++ b/engine/sdks/rust/universaldb-commit/src/generated.rs @@ -0,0 +1 @@ +include!(concat!(env!("OUT_DIR"), "/combined_imports.rs")); diff --git a/engine/sdks/rust/universaldb-commit/src/lib.rs b/engine/sdks/rust/universaldb-commit/src/lib.rs new file mode 100644 index 0000000000..41954bbe75 --- /dev/null +++ b/engine/sdks/rust/universaldb-commit/src/lib.rs @@ -0,0 +1,6 @@ +pub mod generated; +pub mod versioned; + +// Re-export latest +pub use generated::PROTOCOL_VERSION; +pub use generated::v1::*; diff --git a/engine/sdks/rust/universaldb-commit/src/versioned.rs b/engine/sdks/rust/universaldb-commit/src/versioned.rs new file mode 100644 index 0000000000..a9d637aa81 --- /dev/null +++ b/engine/sdks/rust/universaldb-commit/src/versioned.rs @@ -0,0 +1,38 @@ +use anyhow::{Ok, Result, bail}; +use vbare::OwnedVersionedData; + +use crate::generated::v1; + +// Only v1 exists today. When adding v2+, generate converters with +// `scripts/vbare-gen-converters` (see the envoy-protocol package for the +// resulting `versioned/` module layout) and wire them in here. +pub enum CommitRequest { + V1(v1::CommitRequest), +} + +impl OwnedVersionedData for CommitRequest { + type Latest = v1::CommitRequest; + + fn wrap_latest(latest: v1::CommitRequest) -> Self { + CommitRequest::V1(latest) + } + + fn unwrap_latest(self) -> Result { + match self { + CommitRequest::V1(data) => Ok(data), + } + } + + fn deserialize_version(payload: &[u8], version: u16) -> Result { + match version { + 1 => Ok(CommitRequest::V1(serde_bare::from_slice(payload)?)), + _ => bail!("invalid version: {version}"), + } + } + + fn serialize_version(self, _version: u16) -> Result> { + match self { + CommitRequest::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + } + } +} diff --git a/engine/sdks/schemas/universaldb-commit/v1.bare b/engine/sdks/schemas/universaldb-commit/v1.bare new file mode 100644 index 0000000000..2377312f84 --- /dev/null +++ b/engine/sdks/schemas/universaldb-commit/v1.bare @@ -0,0 +1,70 @@ +# Commit-queue wire format for the Postgres leader-resolver UDB driver. +# +# Followers encode a CommitRequest into the `payload` column of +# `udb_commit_requests`; the leader decodes it to resolve and apply. +# Rust-only (never leaves the engine), but versioned so rolling deploys can +# skew follower vs leader code. + +type ConflictRangeType enum { + READ + WRITE +} + +type ConflictRange struct { + begin: data + end: data + kind: ConflictRangeType +} + +# Order MUST match universaldb::options::MutationType declaration order so the +# enum tag round-trips. 15 variants today; never reorder, only append. +type MutationType enum { + ADD + AND + BIT_AND + OR + BIT_OR + XOR + BIT_XOR + APPEND_IF_FITS + MAX + MIN + SET_VERSIONSTAMPED_KEY + SET_VERSIONSTAMPED_VALUE + BYTE_MIN + BYTE_MAX + COMPARE_AND_CLEAR +} + +type SetValue struct { + key: data + value: data +} + +type Clear struct { + key: data +} + +type ClearRange struct { + begin: data + end: data +} + +type AtomicOp struct { + key: data + param: data + opType: MutationType +} + +type Operation union { + SetValue | + Clear | + ClearRange | + AtomicOp +} + +type CommitRequest struct { + readVersion: u64 + conflictRanges: list + operations: list +} diff --git a/scripts/run/postgres.sh b/scripts/run/postgres.sh index 30acd498f1..2ac817b598 100755 --- a/scripts/run/postgres.sh +++ b/scripts/run/postgres.sh @@ -2,7 +2,7 @@ set -euo pipefail CONTAINER_NAME="rivet-engine-postgres" -POSTGRES_IMAGE="postgres:17" +POSTGRES_IMAGE="postgres:18" if docker ps --all --format '{{.Names}}' | grep -qw "${CONTAINER_NAME}"; then if docker ps --format '{{.Names}}' | grep -qw "${CONTAINER_NAME}"; then diff --git a/scripts/run/restore-postgres.sh b/scripts/run/restore-postgres.sh index bd2b50e0b4..0a82410608 100755 --- a/scripts/run/restore-postgres.sh +++ b/scripts/run/restore-postgres.sh @@ -2,7 +2,7 @@ set -euo pipefail CONTAINER_NAME="rivet-engine-postgres" -POSTGRES_IMAGE="postgres:17" +POSTGRES_IMAGE="postgres:18" if [ $# -ne 1 ]; then echo "Usage: $0 " diff --git a/self-host/compose/dev-host/docker-compose.yml b/self-host/compose/dev-host/docker-compose.yml index 3b54702d63..96db5a57d0 100644 --- a/self-host/compose/dev-host/docker-compose.yml +++ b/self-host/compose/dev-host/docker-compose.yml @@ -72,7 +72,7 @@ services: network_mode: host postgres: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres diff --git a/self-host/compose/dev-multidc-multinode/docker-compose.yml b/self-host/compose/dev-multidc-multinode/docker-compose.yml index 75c500850b..7e421144c8 100644 --- a/self-host/compose/dev-multidc-multinode/docker-compose.yml +++ b/self-host/compose/dev-multidc-multinode/docker-compose.yml @@ -85,7 +85,7 @@ services: condition: service_healthy postgres-dc-a: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -355,7 +355,7 @@ services: - rivet-network-dc-a postgres-dc-b: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -617,7 +617,7 @@ services: - rivet-network-dc-b postgres-dc-c: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres diff --git a/self-host/compose/dev-multidc/docker-compose.yml b/self-host/compose/dev-multidc/docker-compose.yml index f20db41ff3..79cb619e0a 100644 --- a/self-host/compose/dev-multidc/docker-compose.yml +++ b/self-host/compose/dev-multidc/docker-compose.yml @@ -85,7 +85,7 @@ services: condition: service_healthy postgres-dc-a: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -235,7 +235,7 @@ services: - rivet-network-dc-a postgres-dc-b: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -377,7 +377,7 @@ services: - rivet-network-dc-b postgres-dc-c: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres diff --git a/self-host/compose/dev-multinode/docker-compose.yml b/self-host/compose/dev-multinode/docker-compose.yml index 381ece396d..a432743057 100644 --- a/self-host/compose/dev-multinode/docker-compose.yml +++ b/self-host/compose/dev-multinode/docker-compose.yml @@ -81,7 +81,7 @@ services: condition: service_healthy postgres: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres diff --git a/self-host/compose/dev/docker-compose.yml b/self-host/compose/dev/docker-compose.yml index 6c5dcc0ef2..91f82dce45 100644 --- a/self-host/compose/dev/docker-compose.yml +++ b/self-host/compose/dev/docker-compose.yml @@ -81,7 +81,7 @@ services: condition: service_healthy postgres: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres diff --git a/self-host/compose/template/src/docker-compose.ts b/self-host/compose/template/src/docker-compose.ts index ca2f6b1fc8..6db4a8ab9a 100644 --- a/self-host/compose/template/src/docker-compose.ts +++ b/self-host/compose/template/src/docker-compose.ts @@ -171,7 +171,7 @@ export function generateDockerCompose(context: TemplateContext) { ); services[postgresServiceName] = { restart: "unless-stopped", - image: "postgres:17-alpine", + image: "postgres:18-alpine", environment: [ "POSTGRES_USER=postgres", "POSTGRES_PASSWORD=postgres", diff --git a/self-host/k8s/engine/12-postgres-statefulset.yaml b/self-host/k8s/engine/12-postgres-statefulset.yaml index e1e14a4b75..c035d26611 100644 --- a/self-host/k8s/engine/12-postgres-statefulset.yaml +++ b/self-host/k8s/engine/12-postgres-statefulset.yaml @@ -20,7 +20,7 @@ spec: spec: containers: - name: postgres - image: postgres:17 + image: postgres:18 args: - postgres - -c diff --git a/website/src/content/docs/self-hosting/docker-compose.mdx b/website/src/content/docs/self-hosting/docker-compose.mdx index eae884ca99..daf5ab0096 100644 --- a/website/src/content/docs/self-hosting/docker-compose.mdx +++ b/website/src/content/docs/self-hosting/docker-compose.mdx @@ -192,7 +192,7 @@ PostgreSQL is the recommended backend for multi-node self-hosted deployments. It ```yaml services: postgres: - image: postgres:15 + image: postgres:18 environment: POSTGRES_DB: rivet POSTGRES_USER: rivet diff --git a/website/src/content/docs/self-hosting/docker-container.mdx b/website/src/content/docs/self-hosting/docker-container.mdx index e6d4366d68..46d544ddd7 100644 --- a/website/src/content/docs/self-hosting/docker-container.mdx +++ b/website/src/content/docs/self-hosting/docker-container.mdx @@ -161,7 +161,7 @@ docker run -d \ -e POSTGRES_USER=rivet \ -e POSTGRES_PASSWORD=rivet_password \ -v postgres-data:/var/lib/postgresql/data \ - postgres:15 + postgres:18 # Run Rivet Engine docker run -d \ From 2534d5c84a46ec35446378cdb9f6a26e764070d6 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Thu, 25 Jun 2026 12:00:32 -0700 Subject: [PATCH 09/16] [SLOP(claude-opus-4-8-high)] feat(ups): table-backed postgres transport with coalesced doorbell --- .../src/driver/postgres/doorbell.rs | 133 ++++ .../src/driver/postgres/mod.rs | 692 +++++++++--------- 2 files changed, 477 insertions(+), 348 deletions(-) create mode 100644 engine/packages/universalpubsub/src/driver/postgres/doorbell.rs diff --git a/engine/packages/universalpubsub/src/driver/postgres/doorbell.rs b/engine/packages/universalpubsub/src/driver/postgres/doorbell.rs new file mode 100644 index 0000000000..f08ff7bbf3 --- /dev/null +++ b/engine/packages/universalpubsub/src/driver/postgres/doorbell.rs @@ -0,0 +1,133 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use deadpool_postgres::Pool; +use tokio::sync::Notify; +use tokio::time::Instant; + +/// Number of doorbell shards. A subject maps to a shard via `hash(subject_hash) % K`. +/// Subscribers LISTEN their subject's shard channel; publishers wake the local +/// doorbell task which NOTIFYs the shard. +pub const DOORBELL_SHARD_COUNT: usize = 32; + +/// Debounce window. Caps each (process, shard) NOTIFY rate at one per window, which +/// bounds how many backends are woken per shard over time. +const DOORBELL_WINDOW: Duration = Duration::from_millis(5); + +/// Returns the NOTIFY channel name for a doorbell shard. +pub fn shard_channel(shard: usize) -> String { + format!("ups_db_{shard}") +} + +/// Returns the doorbell shard for a subject hash. +pub fn shard_for(subject_hash: &str) -> usize { + use std::hash::{DefaultHasher, Hash, Hasher}; + let mut hasher = DefaultHasher::new(); + subject_hash.hash(&mut hasher); + (hasher.finish() as usize) % DOORBELL_SHARD_COUNT +} + +/// Coalesced, payload-free NOTIFY doorbell. +/// +/// Publishers call [`Doorbell::mark_dirty`] after committing a row. A single +/// per-process task drains dirty shards and emits at most one NOTIFY per shard per +/// debounce window using leading-edge fire plus a trailing-edge flush. The doorbell +/// is a latency optimization only. Correctness comes from the table plus the +/// subscriber poll backstop, so a dropped or failed NOTIFY only adds latency. +pub struct Doorbell { + dirty: [AtomicBool; DOORBELL_SHARD_COUNT], + notify: Notify, + pool: Arc, +} + +impl Doorbell { + pub fn new(pool: Arc) -> Arc { + let doorbell = Arc::new(Self { + dirty: std::array::from_fn(|_| AtomicBool::new(false)), + notify: Notify::new(), + pool, + }); + + let task_doorbell = doorbell.clone(); + tokio::spawn(async move { task_doorbell.run().await }); + + doorbell + } + + /// Marks a shard dirty and wakes the doorbell task. Never blocks. + pub fn mark_dirty(&self, shard: usize) { + self.dirty[shard].store(true, Ordering::Release); + self.notify.notify_one(); + } + + async fn run(self: Arc) { + // Per-shard timestamp of the last NOTIFY emitted by this process. + let mut last_notify: [Option; DOORBELL_SHARD_COUNT] = [None; DOORBELL_SHARD_COUNT]; + // Per-shard deadline for a pending trailing-edge NOTIFY, if any. + let mut trailing: [Option; DOORBELL_SHARD_COUNT] = [None; DOORBELL_SHARD_COUNT]; + + loop { + // Arm on the next pending trailing deadline so the trailing edge fires + // even with no further publishes. Wait on the notify permit otherwise. + let next_deadline = trailing.iter().filter_map(|x| *x).min(); + match next_deadline { + Some(deadline) => { + tokio::select! { + _ = self.notify.notified() => {} + _ = tokio::time::sleep_until(deadline) => {} + } + } + None => { + self.notify.notified().await; + } + } + + let now = Instant::now(); + for shard in 0..DOORBELL_SHARD_COUNT { + let is_dirty = self.dirty[shard].swap(false, Ordering::AcqRel); + if is_dirty { + match last_notify[shard] { + Some(last) if now.duration_since(last) < DOORBELL_WINDOW => { + // Within the window. Defer to a trailing-edge NOTIFY at + // window end so at most one NOTIFY fires per shard per W. + if trailing[shard].is_none() { + trailing[shard] = Some(last + DOORBELL_WINDOW); + } + } + _ => { + // Leading edge. Fire immediately for low idle latency. + self.notify_shard(shard).await; + last_notify[shard] = Some(now); + trailing[shard] = None; + } + } + } + + // Flush a trailing-edge NOTIFY whose window has elapsed. + if let Some(deadline) = trailing[shard] { + if now >= deadline { + self.notify_shard(shard).await; + last_notify[shard] = Some(now); + trailing[shard] = None; + } + } + } + } + } + + async fn notify_shard(&self, shard: usize) { + let channel = shard_channel(shard); + match self.pool.get().await { + Ok(conn) => { + // Payload-free doorbell. The payload lives in the table. + if let Err(err) = conn.execute("SELECT pg_notify($1, '')", &[&channel]).await { + tracing::warn!(?err, %channel, "failed to emit doorbell notify"); + } + } + Err(err) => { + tracing::warn!(?err, %channel, "failed to get connection for doorbell notify"); + } + } + } +} diff --git a/engine/packages/universalpubsub/src/driver/postgres/mod.rs b/engine/packages/universalpubsub/src/driver/postgres/mod.rs index 52bc219b2a..d361eb1149 100644 --- a/engine/packages/universalpubsub/src/driver/postgres/mod.rs +++ b/engine/packages/universalpubsub/src/driver/postgres/mod.rs @@ -1,12 +1,11 @@ -use anyhow::{Context, Result, anyhow}; +use anyhow::{Context, Result}; use async_trait::async_trait; -use base64::Engine; -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::throttle::Backoff; use scc::HashMap; +use std::collections::VecDeque; use std::hash::{DefaultHasher, Hash, Hasher}; use std::path::PathBuf; use std::sync::Arc; @@ -21,37 +20,37 @@ use crate::driver::{PubSubDriver, SubscriberDriver, SubscriberDriverHandle}; use crate::metrics; use crate::pubsub::DriverOutput; -#[derive(Clone)] -struct Subscription { - // Channel to send messages to this subscription - tx: broadcast::Sender>, -} +mod doorbell; -impl Subscription { - fn new(tx: broadcast::Sender>) -> Self { - Self { tx } - } -} +use doorbell::{Doorbell, shard_channel, shard_for}; -/// > In the default configuration it must be shorter than 8000 bytes -/// -/// https://www.postgresql.org/docs/17/sql-notify.html -const MAX_NOTIFY_LENGTH: usize = 8000; +/// The transport is the table, not the NOTIFY payload, so there is no per-message +/// size cap from the 8000-byte NOTIFY limit. Match the NATS ceiling so chunking +/// behaves identically across drivers. +pub const POSTGRES_MAX_MESSAGE_SIZE: usize = 1024 * 1024; -/// Base64 encoding ratio -const BYTES_PER_BLOCK: usize = 3; -const CHARS_PER_BLOCK: usize = 4; +/// Poll backstop interval. Every subscriber reads its table on this interval +/// regardless of doorbell wakeups. This is the correctness floor that makes delivery +/// independent of any NOTIFY arriving. +const POLL_INTERVAL: Duration = Duration::from_secs(1); -/// Calculate max message size if encoded as base64 -/// -/// We need to remove BYTES_PER_BLOCK since there might be a tail on the base64-encoded data that -/// would bump it over the limit. -pub const POSTGRES_MAX_MESSAGE_SIZE: usize = - (MAX_NOTIFY_LENGTH * BYTES_PER_BLOCK) / CHARS_PER_BLOCK - BYTES_PER_BLOCK; +/// Idle-in-transaction timeout applied to the LISTEN connection. A wedged listener +/// holding a transaction open would otherwise fill the shared notify queue and fail +/// NOTIFY cluster-wide. Bounding it keeps a stuck listener degrading to added latency +/// rather than a cluster outage. +const LISTEN_IDLE_IN_TRANSACTION_TIMEOUT_MS: i64 = 30_000; const QUEUE_SUB_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10); /// How long a queue subscriber's heartbeat must be within to be considered active. const QUEUE_SUB_TTL_SECS: i64 = 30; + +/// How often to GC expired broadcast messages. +const MESSAGE_GC_INTERVAL: Duration = Duration::from_secs(5); +/// Max age before a broadcast message row is garbage collected. Must exceed the poll +/// interval plus the reconnect gap. A subscriber that falls behind this misses +/// messages, matching NATS-core at-most-once semantics for slow consumers. +const MESSAGE_MAX_AGE_SECS: i64 = 10; + /// How often to GC orphaned queue messages. const QUEUE_MESSAGE_GC_INTERVAL: Duration = Duration::from_secs(300); /// Max age before an unconsumed queue message is garbage collected. @@ -61,9 +60,11 @@ const QUEUE_MESSAGE_MAX_AGE_SECS: i64 = 3600; pub struct PostgresDriver { pool: Arc, client: Arc>>, - subscriptions: Arc>, - /// Wakeup channels for queue subscriptions, keyed by queue channel name. - queue_subscriptions: Arc>, + /// Wakeup channels keyed by doorbell shard channel name. Shared by broadcast and + /// queue subscribers whose subjects map to the same shard. Carries empty wakeups + /// only; payload lives in the table. + shard_subscriptions: Arc>>, + doorbell: Arc, client_ready: tokio::sync::watch::Receiver, } @@ -103,8 +104,9 @@ impl PostgresDriver { .context("failed to create postgres pool")?; tracing::debug!("postgres pool created successfully"); - let subscriptions: Arc> = Arc::new(HashMap::new()); - let queue_subscriptions: Arc> = Arc::new(HashMap::new()); + let pool = Arc::new(pool); + let shard_subscriptions: Arc>> = + Arc::new(HashMap::new()); let client: Arc>> = Arc::new(Mutex::new(None)); // Create channel for client ready notifications @@ -113,8 +115,7 @@ impl PostgresDriver { // Spawn connection lifecycle task tokio::spawn(Self::spawn_connection_lifecycle( conn_str.clone(), - subscriptions.clone(), - queue_subscriptions.clone(), + shard_subscriptions.clone(), client.clone(), ready_tx, ssl_root_cert_path.clone(), @@ -122,26 +123,41 @@ impl PostgresDriver { ssl_client_key_path.clone(), )); + let doorbell = Doorbell::new(pool.clone()); + let driver = Self { - pool: Arc::new(pool), + pool, client, - subscriptions, - queue_subscriptions, + shard_subscriptions, + doorbell, client_ready, }; // Wait for initial connection to be established driver.wait_for_client().await?; - // Create queue tables eagerly so they exist before any publish or subscribe + // Create tables eagerly so they exist before any publish or subscribe. { let conn = driver .pool .get() .await - .context("failed to get connection for queue table creation")?; + .context("failed to get connection for table creation")?; conn.batch_execute( - "CREATE TABLE IF NOT EXISTS ups_queue_subs ( \ + // Broadcast transport table. UNLOGGED gives at-most-once across a + // crash, matching NATS-core semantics, and avoids WAL fsync on every + // publish. The real subject is stored so receivers can verify it and + // reject DefaultHasher subject-hash collisions. + "CREATE UNLOGGED TABLE IF NOT EXISTS ups_messages ( \ + id BIGSERIAL PRIMARY KEY, \ + subject_hash TEXT NOT NULL, \ + subject TEXT NOT NULL, \ + payload BYTEA NOT NULL, \ + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() \ + ); \ + CREATE INDEX IF NOT EXISTS ups_messages_subject_id \ + ON ups_messages (subject_hash, id); \ + CREATE TABLE IF NOT EXISTS ups_queue_subs ( \ id TEXT PRIMARY KEY, \ subject_hash TEXT NOT NULL, \ queue_hash TEXT NOT NULL, \ @@ -149,7 +165,7 @@ impl PostgresDriver { ); \ CREATE INDEX IF NOT EXISTS ups_queue_subs_subject_queue \ ON ups_queue_subs (subject_hash, queue_hash); \ - CREATE TABLE IF NOT EXISTS ups_queue_messages ( \ + CREATE UNLOGGED TABLE IF NOT EXISTS ups_queue_messages ( \ id BIGSERIAL PRIMARY KEY, \ subject_hash TEXT NOT NULL, \ queue_hash TEXT NOT NULL, \ @@ -160,10 +176,33 @@ impl PostgresDriver { ON ups_queue_messages (subject_hash, queue_hash, id);", ) .await - .context("failed to create queue tables")?; - tracing::debug!("queue tables ready"); + .context("failed to create tables")?; + tracing::debug!("tables ready"); } + // Spawn GC task for expired broadcast messages + let message_gc_driver = driver.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(MESSAGE_GC_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + interval.tick().await; + if let Ok(conn) = message_gc_driver.pool.get().await { + let result = conn + .execute( + "DELETE FROM ups_messages \ + WHERE created_at < NOW() - ($1::bigint * INTERVAL '1 second')", + &[&MESSAGE_MAX_AGE_SECS], + ) + .await; + if let Err(e) = result { + tracing::warn!(?e, "failed to gc broadcast messages"); + } + } + } + }); + // Spawn GC task for orphaned queue messages let gc_driver = driver.clone(); tokio::spawn(async move { @@ -193,8 +232,7 @@ impl PostgresDriver { /// Manages the connection lifecycle with automatic reconnection async fn spawn_connection_lifecycle( conn_str: String, - subscriptions: Arc>, - queue_subscriptions: Arc>, + shard_subscriptions: Arc>>, client: Arc>>, ready_tx: tokio::sync::watch::Sender, ssl_root_cert_path: Option, @@ -227,41 +265,42 @@ impl PostgresDriver { // Spawn the polling task immediately // This must be done before any operations on the client - let subscriptions_clone = subscriptions.clone(); - let queue_subscriptions_clone = queue_subscriptions.clone(); + let shard_subscriptions_clone = shard_subscriptions.clone(); let poll_handle = tokio::spawn(async move { - Self::poll_connection(conn, subscriptions_clone, queue_subscriptions_clone) - .await; + Self::poll_connection(conn, shard_subscriptions_clone).await; }); - // Get regular channels to re-subscribe to + // Bound a stuck listener so it cannot wedge the shared notify queue. + if let Result::Err(e) = new_client + .execute( + &format!( + "SET idle_in_transaction_session_timeout = '{}'", + LISTEN_IDLE_IN_TRANSACTION_TIMEOUT_MS + ), + &[], + ) + .await + { + tracing::warn!(?e, "failed to set idle_in_transaction_session_timeout"); + } + + // Get shard channels to re-subscribe to let mut channels = Vec::new(); - subscriptions + shard_subscriptions .iter_async(|k, _| { channels.push(k.clone()); true }) .await; - // Get queue wakeup channels to re-subscribe to - let mut queue_channels = Vec::new(); - queue_subscriptions - .iter_async(|k, _| { - queue_channels.push(k.clone()); - true - }) - .await; - - let needs_resubscribe = !channels.is_empty() || !queue_channels.is_empty(); - if needs_resubscribe { + if !channels.is_empty() { tracing::debug!( - regular_channels = channels.len(), - queue_channels = queue_channels.len(), - "re-subscribing to channels after reconnection" + channels = channels.len(), + "re-subscribing to doorbell shards after reconnection" ); } - for channel in channels.iter().chain(queue_channels.iter()) { + for channel in channels.iter() { tracing::debug!(?channel, "re-subscribing to channel"); if let Result::Err(e) = new_client .execute(&format!("LISTEN \"{}\"", channel), &[]) @@ -298,30 +337,20 @@ impl PostgresDriver { /// Polls the connection for notifications until it closes or errors async fn poll_connection( mut conn: tokio_postgres::Connection, - subscriptions: Arc>, - queue_subscriptions: Arc>, + shard_subscriptions: Arc>>, ) where T: tokio_postgres::tls::TlsStream + Unpin, { loop { match poll_fn(|cx| conn.poll_message(cx)).await { Some(std::result::Result::Ok(AsyncMessage::Notification(note))) => { - tracing::trace!(channel = %note.channel(), "received notification"); - if let Some(sub) = subscriptions.get_async(note.channel()).await { - let bytes = match BASE64.decode(note.payload()) { - std::result::Result::Ok(b) => b, - std::result::Result::Err(err) => { - tracing::error!(?err, "failed decoding base64"); - continue; - } - }; - tracing::trace!(channel = %note.channel(), bytes_len = bytes.len(), "sending to broadcast channel"); - let _ = sub.tx.send(bytes); - } else if let Some(sub) = queue_subscriptions.get_async(note.channel()).await { - // Queue notifications are wakeup signals only; payload lives in the table - let _ = sub.tx.send(Vec::new()); + tracing::trace!(channel = %note.channel(), "received doorbell wakeup"); + // Doorbell notifications are payload-free wakeup signals only. + // Subscribers read their payload from the table. + if let Some(sub) = shard_subscriptions.get_async(note.channel()).await { + let _ = sub.send(()); } else { - tracing::warn!(channel = %note.channel(), "received notification for unknown channel"); + tracing::trace!(channel = %note.channel(), "wakeup for unknown shard"); } } Some(std::result::Result::Ok(_)) => { @@ -361,8 +390,9 @@ impl PostgresDriver { } fn hash_subject(&self, subject: &str) -> String { - // Postgres channel names have a 64 character limit - // Hash the subject to ensure it fits + // Postgres channel names have a 64 character limit, but this hash is also the + // table index key. Collisions are possible and resolved by verifying the real + // subject stored alongside each row. let mut hasher = DefaultHasher::new(); subject.hash(&mut hasher); format!("ups_{:x}", hasher.finish()) @@ -374,57 +404,68 @@ impl PostgresDriver { format!("{:x}", hasher.finish()) } - /// Returns the NOTIFY channel name for a (subject, queue) pair. - fn queue_channel(&self, subject_hash: &str, queue_hash: &str) -> String { - // Max length: "ups_q_" (6) + 16 + "_" (1) + 16 = 39 chars, well within 64 - format!("ups_q_{}_{}", subject_hash, queue_hash) - } - - /// Inserts messages into the queue table and notifies active queue subscribers. - async fn publish_to_queues(&self, subject: &str, payload: &[u8]) -> Result<()> { - let subject_hash = self.hash_subject(subject); - + /// Returns the current max broadcast message id, used as a subscriber's starting + /// cursor so it only sees future messages (NATS at-most-once, no replay). + async fn current_max_id(&self) -> Result { let conn = self .pool .get() .await - .context("failed to get connection for queue publish")?; - - // Find active queue groups for this subject - let rows = conn - .query( - "SELECT DISTINCT queue_hash FROM ups_queue_subs \ - WHERE subject_hash = $1 \ - AND heartbeat_at > NOW() - ($2::bigint * INTERVAL '1 second')", - &[&subject_hash, &QUEUE_SUB_TTL_SECS], - ) + .context("failed to get connection for cursor init")?; + let row = conn + .query_one("SELECT COALESCE(MAX(id), 0) FROM ups_messages", &[]) .await - .context("failed to query active queue subs")?; + .context("failed to read current max id")?; + Ok(row.get(0)) + } - for row in rows { - let queue_hash: String = row.get(0); - let channel = self.queue_channel(&subject_hash, &queue_hash); + /// Ensures this process is LISTENing on the given doorbell shard and returns a + /// wakeup receiver plus a drop guard that UNLISTENs once no receivers remain. + async fn ensure_shard_listen( + &self, + shard: usize, + ) -> (broadcast::Receiver<()>, tokio_util::sync::DropGuard) { + let channel = shard_channel(shard); - conn.execute( - "INSERT INTO ups_queue_messages (subject_hash, queue_hash, payload) \ - VALUES ($1, $2, $3)", - &[&subject_hash, &queue_hash, &payload], - ) - .await - .context("failed to insert queue message")?; + match self.shard_subscriptions.entry_async(channel.clone()).await { + scc::hash_map::Entry::Occupied(existing) => { + let rx = existing.subscribe(); + let drop_guard = + self.spawn_shard_cleanup_task(channel.clone(), existing.get().clone()); + (rx, drop_guard) + } + scc::hash_map::Entry::Vacant(e) => { + let (tx, rx) = broadcast::channel(1024); + e.insert_entry(tx.clone()); + metrics::POSTGRES_SUBSCRIPTION_COUNT.set(self.shard_subscriptions.len() as i64); - conn.execute(&format!("NOTIFY \"{}\"", channel), &[]) - .await - .context("failed to notify queue channel")?; - } + if let Some(client) = &*self.client.lock().await { + match client + .execute(&format!("LISTEN \"{channel}\""), &[]) + .instrument(tracing::trace_span!("pg_listen")) + .await + { + Result::Ok(_) => { + tracing::debug!(%channel, "successfully subscribed to shard"); + } + Result::Err(e) => { + tracing::warn!(?e, %channel, "failed to LISTEN, will retry on reconnection"); + } + } + } else { + tracing::debug!(%channel, "client not connected, will LISTEN on reconnection"); + } - Ok(()) + let drop_guard = self.spawn_shard_cleanup_task(channel.clone(), tx.clone()); + (rx, drop_guard) + } + } } - fn spawn_subscription_cleanup_task( + fn spawn_shard_cleanup_task( &self, - subject_hash: String, - tx: broadcast::Sender>, + channel: String, + tx: broadcast::Sender<()>, ) -> tokio_util::sync::DropGuard { let driver = self.clone(); let token = tokio_util::sync::CancellationToken::new(); @@ -434,47 +475,72 @@ impl PostgresDriver { token.cancelled().await; if tx.receiver_count() == 0 { if let Some(client) = &*driver.client.lock().await { - let sql = format!("UNLISTEN \"{}\"", subject_hash); + let sql = format!("UNLISTEN \"{}\"", channel); if let Err(err) = client.execute(sql.as_str(), &[]).await { - tracing::warn!(?err, %subject_hash, "failed to UNLISTEN channel"); + tracing::warn!(?err, %channel, "failed to UNLISTEN channel"); } else { - tracing::trace!(%subject_hash, "unlistened channel"); + tracing::trace!(%channel, "unlistened channel"); } } - driver.subscriptions.remove_async(&subject_hash).await; - metrics::POSTGRES_SUBSCRIPTION_COUNT.set(driver.subscriptions.len() as i64); + driver.shard_subscriptions.remove_async(&channel).await; + metrics::POSTGRES_SUBSCRIPTION_COUNT.set(driver.shard_subscriptions.len() as i64); } }); drop_guard } - fn spawn_queue_subscription_cleanup_task( + /// Inserts the broadcast row and any active queue-group rows in one transaction. + async fn try_publish_to_db( &self, - channel: String, - tx: broadcast::Sender>, - ) -> tokio_util::sync::DropGuard { - let driver = self.clone(); - let token = tokio_util::sync::CancellationToken::new(); - let drop_guard = token.clone().drop_guard(); + subject: &str, + subject_hash: &str, + payload: &[u8], + ) -> Result<()> { + let mut conn = self + .pool + .get() + .await + .context("failed to get connection for publish")?; + let tx = conn + .transaction() + .await + .context("failed to begin publish transaction")?; - tokio::spawn(async move { - token.cancelled().await; - if tx.receiver_count() == 0 { - if let Some(client) = &*driver.client.lock().await { - let sql = format!("UNLISTEN \"{}\"", channel); + // Broadcast row. + tx.execute( + "INSERT INTO ups_messages (subject_hash, subject, payload) VALUES ($1, $2, $3)", + &[&subject_hash, &subject, &payload], + ) + .await + .context("failed to insert broadcast message")?; - if let Err(err) = client.execute(sql.as_str(), &[]).await { - tracing::warn!(?err, %channel, "failed to UNLISTEN queue channel"); - } else { - tracing::trace!(%channel, "unlistened queue channel"); - } - } - driver.queue_subscriptions.remove_async(&channel).await; - } - }); + // Queue rows for every active queue group on this subject. Batched into the + // same transaction so a crash never strands a row mid-publish. + let rows = tx + .query( + "SELECT DISTINCT queue_hash FROM ups_queue_subs \ + WHERE subject_hash = $1 \ + AND heartbeat_at > NOW() - ($2::bigint * INTERVAL '1 second')", + &[&subject_hash, &QUEUE_SUB_TTL_SECS], + ) + .await + .context("failed to query active queue subs")?; - drop_guard + for row in rows { + let queue_hash: String = row.get(0); + tx.execute( + "INSERT INTO ups_queue_messages (subject_hash, queue_hash, payload) \ + VALUES ($1, $2, $3)", + &[&subject_hash, &queue_hash, &payload], + ) + .await + .context("failed to insert queue message")?; + } + + tx.commit().await.context("failed to commit publish")?; + + Ok(()) } } @@ -485,66 +551,23 @@ impl PubSubDriver for PostgresDriver { subject: &str, _reply_id: Option, ) -> Result { - // TODO: To match NATS implementation, LISTEN must be pipelined (i.e. wait for the command - // to reach the server, but not wait for it to respond). However, this has to ensure that - // NOTIFY & LISTEN are called on the same connection (not diff connections in a pool) or - // else there will be race conditions where messages might be published before - // subscriptions are registered. - // - // tokio-postgres currently does not expose the API for pipelining, so we are SOL. - // - // We might be able to use a background tokio task in combination with flush if we use the - // same Postgres connection, but unsure if that will create a bottleneck. - - let hashed = self.hash_subject(subject); - - // Check if we already have a subscription for this channel - let (rx, drop_guard) = match self.subscriptions.entry_async(hashed.clone()).await { - scc::hash_map::Entry::Occupied(existing_sub) => { - // Reuse the existing broadcast channel - let rx = existing_sub.tx.subscribe(); - let drop_guard = - self.spawn_subscription_cleanup_task(hashed.clone(), existing_sub.tx.clone()); - (rx, drop_guard) - } - scc::hash_map::Entry::Vacant(e) => { - // Create a new broadcast channel for this subject - let (tx, rx) = tokio::sync::broadcast::channel(1024); - let subscription = Subscription::new(tx.clone()); - - // Register subscription - e.insert_entry(subscription.clone()); - metrics::POSTGRES_SUBSCRIPTION_COUNT.set(self.subscriptions.len() as i64); - - // Execute LISTEN command on the async client (for receiving notifications) - // This only needs to be done once per channel - // Try to LISTEN if client is available, but don't fail if disconnected - // The reconnection logic will handle re-subscribing - if let Some(client) = &*self.client.lock().await { - match client - .execute(&format!("LISTEN \"{hashed}\""), &[]) - .instrument(tracing::trace_span!("pg_listen")) - .await - { - Result::Ok(_) => { - tracing::debug!(%hashed, "successfully subscribed to channel"); - } - Result::Err(e) => { - tracing::warn!(?e, %hashed, "failed to LISTEN, will retry on reconnection"); - } - } - } else { - tracing::debug!(%hashed, "client not connected, will LISTEN on reconnection"); - } + let subject_hash = self.hash_subject(subject); + let shard = shard_for(&subject_hash); - let drop_guard = self.spawn_subscription_cleanup_task(hashed.clone(), tx.clone()); - (rx, drop_guard) - } - }; + // Capture the cursor before LISTENing. Any message inserted after this point + // has a higher id and is delivered either by the doorbell wakeup or the poll + // backstop, so there is no subscribe/publish race. + let cursor = self.current_max_id().await?; + + let (rx, drop_guard) = self.ensure_shard_listen(shard).await; Ok(Box::new(PostgresSubscriber { subject: subject.to_string(), - rx: Some(rx), + subject_hash, + pool: self.pool.clone(), + cursor, + buffer: VecDeque::new(), + rx, _drop_guard: drop_guard, })) } @@ -552,7 +575,7 @@ impl PubSubDriver for PostgresDriver { async fn queue_subscribe(&self, subject: &str, queue: &str) -> Result { let subject_hash = self.hash_subject(subject); let queue_hash = self.hash_queue(queue); - let channel = self.queue_channel(&subject_hash, &queue_hash); + let shard = shard_for(&subject_hash); // Register this subscriber in the database so publishers know the queue exists let sub_id = Uuid::new_v4().to_string(); @@ -570,44 +593,7 @@ impl PubSubDriver for PostgresDriver { .context("failed to register queue subscriber")?; } - // Set up a shared LISTEN/broadcast channel for the wakeup signal - let (rx, drop_guard) = match self.queue_subscriptions.entry_async(channel.clone()).await { - scc::hash_map::Entry::Occupied(existing_sub) => { - let rx = existing_sub.tx.subscribe(); - let drop_guard = self.spawn_queue_subscription_cleanup_task( - channel.clone(), - existing_sub.tx.clone(), - ); - (rx, drop_guard) - } - scc::hash_map::Entry::Vacant(e) => { - let (tx, rx) = tokio::sync::broadcast::channel(1024); - let subscription = Subscription::new(tx.clone()); - - e.insert_entry(subscription.clone()); - - if let Some(client) = &*self.client.lock().await { - match client - .execute(&format!("LISTEN \"{}\"", channel), &[]) - .instrument(tracing::trace_span!("pg_listen_queue")) - .await - { - Result::Ok(_) => { - tracing::debug!(%channel, "successfully subscribed to queue channel"); - } - Result::Err(e) => { - tracing::warn!(?e, %channel, "failed to LISTEN queue channel, will retry on reconnection"); - } - } - } else { - tracing::debug!(%channel, "client not connected, will LISTEN queue channel on reconnection"); - } - - let drop_guard = - self.spawn_queue_subscription_cleanup_task(channel.clone(), tx.clone()); - (rx, drop_guard) - } - }; + let (rx, drop_guard) = self.ensure_shard_listen(shard).await; // Spawn heartbeat task to keep the registration alive let pool = self.pool.clone(); @@ -644,7 +630,7 @@ impl PubSubDriver for PostgresDriver { queue_hash, sub_id, pool: self.pool.clone(), - rx: Some(rx), + rx, _drop_guard: drop_guard, _heartbeat_token: heartbeat_token, })) @@ -656,78 +642,33 @@ impl PubSubDriver for PostgresDriver { payload: &[u8], _reply_subject: Option<&str>, ) -> Result<()> { - // TODO: See `subscribe` about pipelining - - // Encode payload to base64 and send NOTIFY - let encoded = BASE64.encode(payload); - let hashed = self.hash_subject(subject); - - tracing::trace!("attempting to get connection for publish"); - - // Wait for listen connection to be ready first if this channel has subscribers - // This ensures that if we're reconnecting, the LISTEN is re-registered before NOTIFY - if self.subscriptions.contains_async(&hashed).await { - self.wait_for_client().await?; - } + let subject_hash = self.hash_subject(subject); + let shard = shard_for(&subject_hash); - // Retry getting a connection from the pool with backoff in case the connection is - // currently disconnected + // Persist the message, retrying on transient connection errors. The row is + // committed before the doorbell rings so any wakeup observes it. let mut backoff = Backoff::default(); - let mut last_error; - loop { - match self.pool.get().await { - Result::Ok(conn) => { - // Test the connection with a simple query before using it - match conn.execute("SELECT 1", &[]).await { - Result::Ok(_) => { - // Connection is good; run NOTIFY and queue publish in parallel. - // publish_to_queues acquires its own pool connection so both - // can proceed concurrently. - let notify_sql = format!("NOTIFY \"{hashed}\", '{encoded}'"); - let (notify_result, queue_result) = tokio::join!( - conn.execute(notify_sql.as_str(), &[]) - .instrument(tracing::trace_span!("pg_notify")), - self.publish_to_queues(subject, payload), - ); - match notify_result { - Result::Ok(_) => { - if let Err(e) = queue_result { - tracing::warn!(?e, %subject, "failed to publish to queue subscribers"); - } - return Ok(()); - } - Result::Err(e) => { - tracing::debug!( - ?e, - "NOTIFY failed, retrying with new connection" - ); - last_error = Some(e.into()); - } - } - } - Result::Err(e) => { - tracing::debug!( - ?e, - "connection test failed, retrying with new connection" - ); - last_error = Some(e.into()); - } - } - } + match self + .try_publish_to_db(subject, &subject_hash, payload) + .await + { + Result::Ok(()) => break, Result::Err(e) => { - tracing::debug!(?e, "failed to get connection from pool, retrying"); - last_error = Some(e.into()); + if !backoff.tick().await { + tracing::warn!(?e, %subject, "failed to publish, cannot retry again"); + return Err(e); + } + tracing::debug!(?e, "publish failed, retrying"); } } - - // Check if we should continue retrying - if !backoff.tick().await { - return Err( - last_error.unwrap_or_else(|| anyhow!("failed to publish after retries")) - ); - } } + + // Ring the doorbell. Best-effort: the subscriber poll backstop covers a + // dropped or coalesced wakeup, so publish never blocks on NOTIFY. + self.doorbell.mark_dirty(shard); + + Ok(()) } async fn flush(&self) -> Result<()> { @@ -741,26 +682,80 @@ impl PubSubDriver for PostgresDriver { pub struct PostgresSubscriber { subject: String, - rx: Option>>, + subject_hash: String, + pool: Arc, + cursor: i64, + buffer: VecDeque>, + rx: broadcast::Receiver<()>, _drop_guard: tokio_util::sync::DropGuard, } +impl PostgresSubscriber { + /// Reads new rows past the cursor into the buffer, advancing the cursor. Rows + /// whose stored subject does not match are skipped (DefaultHasher collisions) but + /// still advance the cursor. + async fn fetch(&mut self) -> Result<()> { + let conn = self + .pool + .get() + .await + .context("failed to get connection for poll")?; + let rows = conn + .query( + "SELECT id, subject, payload FROM ups_messages \ + WHERE subject_hash = $1 AND id > $2 ORDER BY id", + &[&self.subject_hash, &self.cursor], + ) + .await + .context("failed to poll broadcast messages")?; + + for row in rows { + let id: i64 = row.get(0); + let subject: String = row.get(1); + let payload: Vec = row.get(2); + self.cursor = id; + if subject == self.subject { + self.buffer.push_back(payload); + } + } + + Ok(()) + } +} + #[async_trait] impl SubscriberDriver for PostgresSubscriber { async fn next(&mut self) -> Result { - let rx = match self.rx.as_mut() { - Some(rx) => rx, - None => return Ok(DriverOutput::Unsubscribed), - }; - match rx.recv().await { - std::result::Result::Ok(payload) => Ok(DriverOutput::Message { - subject: self.subject.clone(), - payload, - }), - Err(tokio::sync::broadcast::error::RecvError::Closed) => Ok(DriverOutput::Unsubscribed), - Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => { - // Try again - self.next().await + loop { + if let Some(payload) = self.buffer.pop_front() { + return Ok(DriverOutput::Message { + subject: self.subject.clone(), + payload, + }); + } + + if let Err(e) = self.fetch().await { + // Transient DB errors must not kill the subscriber; the next poll + // tick retries. + tracing::warn!(?e, subject = %self.subject, "failed to poll, will retry"); + } + + if !self.buffer.is_empty() { + continue; + } + + // Wait for a doorbell wakeup or the poll backstop, whichever is first. + tokio::select! { + res = self.rx.recv() => { + match res { + std::result::Result::Ok(()) => {} + Err(broadcast::error::RecvError::Lagged(_)) => {} + Err(broadcast::error::RecvError::Closed) => { + return Ok(DriverOutput::Unsubscribed); + } + } + } + _ = tokio::time::sleep(POLL_INTERVAL) => {} } } } @@ -772,7 +767,7 @@ pub struct PostgresQueueSubscriber { queue_hash: String, sub_id: String, pool: Arc, - rx: Option>>, + rx: broadcast::Receiver<()>, _drop_guard: tokio_util::sync::DropGuard, _heartbeat_token: tokio_util::sync::CancellationToken, } @@ -811,31 +806,32 @@ impl PostgresQueueSubscriber { impl SubscriberDriver for PostgresQueueSubscriber { async fn next(&mut self) -> Result { loop { - // Drain any messages that arrived before or between notifications. - // Do this before borrowing rx so claim_message can borrow self freely. - if let Some(payload) = self.claim_message().await? { - return Ok(DriverOutput::Message { - subject: self.subject.clone(), - payload, - }); - } - - // Wait for a wakeup notification, then loop back to claim. - let rx = match self.rx.as_mut() { - Some(rx) => rx, - None => return Ok(DriverOutput::Unsubscribed), - }; - match rx.recv().await { - std::result::Result::Ok(_) => { - // Wakeup received; loop back to claim + // Drain any messages that arrived before or between wakeups. + match self.claim_message().await { + Result::Ok(Some(payload)) => { + return Ok(DriverOutput::Message { + subject: self.subject.clone(), + payload, + }); } - Err(tokio::sync::broadcast::error::RecvError::Closed) => { - return Ok(DriverOutput::Unsubscribed); + Result::Ok(None) => {} + Result::Err(e) => { + tracing::warn!(?e, subject = %self.subject, "failed to claim, will retry"); } - Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => { - // Notifications were dropped while lagged; loop back to claim in case - // messages are waiting + } + + // Wait for a doorbell wakeup or the poll backstop, then loop back to claim. + tokio::select! { + res = self.rx.recv() => { + match res { + std::result::Result::Ok(()) => {} + Err(broadcast::error::RecvError::Lagged(_)) => {} + Err(broadcast::error::RecvError::Closed) => { + return Ok(DriverOutput::Unsubscribed); + } + } } + _ = tokio::time::sleep(POLL_INTERVAL) => {} } } } From 826a9de5b47efffd378d45d5187894a9a97fc334 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Thu, 25 Jun 2026 12:17:20 -0700 Subject: [PATCH 10/16] [slopfix] test(universaldb): postgres leader failover + abort resolver task on driver drop --- engine/artifacts/config-schema.json | 32 -- .../packages/config/src/config/api_public.rs | 25 -- engine/packages/config/src/config/mod.rs | 11 - engine/packages/engine/src/main.rs | 2 + engine/packages/metrics-server/src/server.rs | 2 +- .../pegboard-envoy/src/ws_to_tunnel_task.rs | 9 +- .../pegboard-gateway/src/shared_state.rs | 2 +- .../packages/pegboard-gateway2/src/metrics.rs | 6 - .../pegboard-gateway2/src/shared_state.rs | 20 +- engine/packages/perf/src/lib.rs | 2 +- engine/packages/service-manager/src/lib.rs | 2 - engine/packages/test-deps/src/datacenter.rs | 1 - engine/packages/universaldb/Cargo.toml | 1 + .../src/driver/postgres/database.rs | 7 +- .../src/driver/postgres/resolver/mod.rs | 7 +- engine/packages/universaldb/tests/failover.rs | 219 ++++++++++++ .../src/driver/postgres/mod.rs | 319 ++++++++++++++---- .../universalpubsub/tests/reconnect.rs | 35 +- scripts/run/engine-postgres.sh | 21 +- 19 files changed, 535 insertions(+), 188 deletions(-) delete mode 100644 engine/packages/config/src/config/api_public.rs create mode 100644 engine/packages/universaldb/tests/failover.rs diff --git a/engine/artifacts/config-schema.json b/engine/artifacts/config-schema.json index 998771acc1..af2af30077 100644 --- a/engine/artifacts/config-schema.json +++ b/engine/artifacts/config-schema.json @@ -40,17 +40,6 @@ } ] }, - "api_public": { - "default": null, - "anyOf": [ - { - "$ref": "#/definitions/ApiPublic" - }, - { - "type": "null" - } - ] - }, "auth": { "default": null, "anyOf": [ @@ -214,27 +203,6 @@ }, "additionalProperties": false }, - "ApiPublic": { - "description": "Configuration for the public API service.", - "type": "object", - "properties": { - "respect_forwarded_for": { - "description": "Flag to respect the X-Forwarded-For header for client IP addresses.\n\nWill be ignored in favor of CF-Connecting-IP if DNS provider is configured as Cloudflare.", - "type": [ - "boolean", - "null" - ] - }, - "verbose_errors": { - "description": "Flag to enable verbose error reporting.", - "type": [ - "boolean", - "null" - ] - } - }, - "additionalProperties": false - }, "Auth": { "type": "object", "required": [ diff --git a/engine/packages/config/src/config/api_public.rs b/engine/packages/config/src/config/api_public.rs deleted file mode 100644 index 53cb280a37..0000000000 --- a/engine/packages/config/src/config/api_public.rs +++ /dev/null @@ -1,25 +0,0 @@ -use schemars::JsonSchema; -use serde::{Deserialize, Serialize}; - -/// Configuration for the public API service. -#[derive(Debug, Serialize, Deserialize, Clone, Default, JsonSchema)] -#[serde(deny_unknown_fields)] -pub struct ApiPublic { - /// Flag to enable verbose error reporting. - pub verbose_errors: Option, - /// Flag to respect the X-Forwarded-For header for client IP addresses. - /// - /// Will be ignored in favor of CF-Connecting-IP if DNS provider is - /// configured as Cloudflare. - pub respect_forwarded_for: Option, -} - -impl ApiPublic { - pub fn verbose_errors(&self) -> bool { - self.verbose_errors.unwrap_or(true) - } - - pub fn respect_forwarded_for(&self) -> bool { - self.respect_forwarded_for.unwrap_or(false) - } -} diff --git a/engine/packages/config/src/config/mod.rs b/engine/packages/config/src/config/mod.rs index 3ee2aace1d..caa8bf196e 100644 --- a/engine/packages/config/src/config/mod.rs +++ b/engine/packages/config/src/config/mod.rs @@ -4,7 +4,6 @@ use serde::{Deserialize, Serialize}; use std::sync::LazyLock; pub mod api_peer; -pub mod api_public; pub mod auth; pub mod cache; pub mod clickhouse; @@ -21,7 +20,6 @@ pub mod telemetry; pub mod topology; pub use api_peer::*; -pub use api_public::*; pub use auth::*; pub use cache::*; pub use clickhouse::*; @@ -74,9 +72,6 @@ pub struct Root { #[serde(default)] pub guard: Option, - #[serde(default)] - pub api_public: Option, - #[serde(default)] pub api_peer: Option, @@ -122,7 +117,6 @@ impl Default for Root { Root { auth: None, guard: None, - api_public: None, api_peer: None, pegboard: None, logs: None, @@ -146,11 +140,6 @@ impl Root { self.guard.as_ref().unwrap_or(&DEFAULT) } - pub fn api_public(&self) -> &ApiPublic { - static DEFAULT: LazyLock = LazyLock::new(ApiPublic::default); - self.api_public.as_ref().unwrap_or(&DEFAULT) - } - pub fn api_peer(&self) -> &ApiPeer { static DEFAULT: LazyLock = LazyLock::new(ApiPeer::default); self.api_peer.as_ref().unwrap_or(&DEFAULT) diff --git a/engine/packages/engine/src/main.rs b/engine/packages/engine/src/main.rs index d0f38a6efb..88e74b0f1b 100644 --- a/engine/packages/engine/src/main.rs +++ b/engine/packages/engine/src/main.rs @@ -36,6 +36,8 @@ fn main() -> Result<()> { } async fn main_inner() -> Result<()> { + tracing::info!(version=%build_meta::VERSION, git_sha=%build_meta::GIT_SHA, built_at=%build_meta::BUILD_TIMESTAMP, "starting rivet"); + let cli = Cli::parse(); // Load config diff --git a/engine/packages/metrics-server/src/server.rs b/engine/packages/metrics-server/src/server.rs index 48b974c537..b2512b3e31 100644 --- a/engine/packages/metrics-server/src/server.rs +++ b/engine/packages/metrics-server/src/server.rs @@ -28,7 +28,7 @@ pub async fn run_standalone(config: rivet_config::Config) -> Result<()> { Ok::<_, hyper::Error>(service_fn(serve_req)) })); - tracing::info!(?host, ?port, "started metrics server"); + tracing::debug!(?host, ?port, "started metrics server"); server.await?; Ok(()) 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 fcbe313453..922225c282 100644 --- a/engine/packages/pegboard-envoy/src/ws_to_tunnel_task.rs +++ b/engine/packages/pegboard-envoy/src/ws_to_tunnel_task.rs @@ -474,9 +474,14 @@ pub async fn task_inner( // 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(), + requests: ctx + .config() + .pegboard() + .envoy_websocket_rate_limit_requests(), drip_rate: Duration::from_micros( - ctx.config().pegboard().envoy_websocket_rate_limit_drip_rate_us(), + ctx.config() + .pegboard() + .envoy_websocket_rate_limit_drip_rate_us(), ), }, ); diff --git a/engine/packages/pegboard-gateway/src/shared_state.rs b/engine/packages/pegboard-gateway/src/shared_state.rs index 580da8eabc..742b146921 100644 --- a/engine/packages/pegboard-gateway/src/shared_state.rs +++ b/engine/packages/pegboard-gateway/src/shared_state.rs @@ -142,7 +142,7 @@ pub struct SharedState(Arc); impl SharedState { pub fn new(config: &rivet_config::Config, ups: PubSub) -> Self { let gateway_id = protocol::util::generate_gateway_id(); - tracing::info!(gateway_id = %protocol::util::id_to_string(&gateway_id), "setting up shared state for gateway"); + tracing::debug!(gateway_id = %protocol::util::id_to_string(&gateway_id), "setting up shared state for gateway"); let receiver_subject = GatewayReceiverSubject::new(gateway_id); let pegboard_config = config.pegboard(); diff --git a/engine/packages/pegboard-gateway2/src/metrics.rs b/engine/packages/pegboard-gateway2/src/metrics.rs index 244df7f882..f087d2ac3b 100644 --- a/engine/packages/pegboard-gateway2/src/metrics.rs +++ b/engine/packages/pegboard-gateway2/src/metrics.rs @@ -53,12 +53,6 @@ lazy_static::lazy_static! { &["namespace_id", "pool_name", "protocol", "reason"], *REGISTRY ).unwrap(); - pub static ref SHUTDOWN_IN_FLIGHT_ABORTED_TOTAL: IntCounter = - register_int_counter_with_registry!( - "gateway2_shutdown_in_flight_aborted_total", - "In-flight gateway requests abandoned on pod shutdown without sending close.", - *REGISTRY - ).unwrap(); pub static ref MSG_SENT_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!( "gateway2_msg_sent_total", "Count of total of tunnel messages sent.", diff --git a/engine/packages/pegboard-gateway2/src/shared_state.rs b/engine/packages/pegboard-gateway2/src/shared_state.rs index b3de6e5887..1d1b7a95cd 100644 --- a/engine/packages/pegboard-gateway2/src/shared_state.rs +++ b/engine/packages/pegboard-gateway2/src/shared_state.rs @@ -163,7 +163,7 @@ impl SharedState { init_slow_ping_threshold_from_env(); let gateway_id = protocol::util::generate_gateway_id(); - tracing::info!(gateway_id = %display_id(&gateway_id), "setting up shared state for gateway"); + tracing::debug!(gateway_id = %display_id(&gateway_id), "setting up shared state for gateway"); let receiver_subject = GatewayReceiverSubject::new(gateway_id); let pegboard_config = config.pegboard(); @@ -194,27 +194,9 @@ impl SharedState { let self_clone = self.clone(); tokio::spawn(async move { self_clone.gc().await }); - let self_clone = self.clone(); - tokio::spawn(async move { self_clone.shutdown_watcher().await }); - Ok(()) } - #[tracing::instrument(skip_all)] - async fn shutdown_watcher(&self) { - let mut term_signal = __rivet_runtime::TermSignal::get(); - term_signal.recv().await; - - let in_flight_aborted = self.in_flight_requests.len(); - if in_flight_aborted > 0 { - metrics::SHUTDOWN_IN_FLIGHT_ABORTED_TOTAL.inc_by(in_flight_aborted as u64); - } - tracing::info!( - in_flight_aborted, - "gateway shutdown in-flight requests abandoned without close" - ); - } - #[tracing::instrument(skip_all)] async fn receiver(&self) { // Automatically resubscribe if unsubscribed diff --git a/engine/packages/perf/src/lib.rs b/engine/packages/perf/src/lib.rs index 7642c406ad..f7ebb0dc71 100644 --- a/engine/packages/perf/src/lib.rs +++ b/engine/packages/perf/src/lib.rs @@ -95,7 +95,7 @@ impl Drop for PerfMeasure { let elapsed = self.start.elapsed(); let _guard = self.span.enter(); - tracing::warn!( + tracing::debug!( name = self.name, elapsed_ms = PerfMeasure::__elapsed_ms(elapsed), "PerfMeasure dropped without finish() - measurement discarded", diff --git a/engine/packages/service-manager/src/lib.rs b/engine/packages/service-manager/src/lib.rs index 6c9f198d78..a60854f3e0 100644 --- a/engine/packages/service-manager/src/lib.rs +++ b/engine/packages/service-manager/src/lib.rs @@ -160,8 +160,6 @@ pub async fn start( let shutting_down = Arc::new(AtomicBool::new(false)); for service in services { - tracing::debug!(name=%service.name, kind=?service.kind, "server starting service"); - match service.kind.behavior() { ServiceBehavior::Service => { let config = config.clone(); diff --git a/engine/packages/test-deps/src/datacenter.rs b/engine/packages/test-deps/src/datacenter.rs index 33843062df..d7fdfcf040 100644 --- a/engine/packages/test-deps/src/datacenter.rs +++ b/engine/packages/test-deps/src/datacenter.rs @@ -73,7 +73,6 @@ pub async fn setup_single_datacenter( let mut root = rivet_config::config::Root::default(); root.database = Some(db_config); root.pubsub = Some(pubsub_config); - root.api_public = Some(Default::default()); root.api_peer = Some(rivet_config::config::ApiPeer { port: Some(api_peer_port), ..Default::default() diff --git a/engine/packages/universaldb/Cargo.toml b/engine/packages/universaldb/Cargo.toml index b2f6b5f045..57dba4ab33 100644 --- a/engine/packages/universaldb/Cargo.toml +++ b/engine/packages/universaldb/Cargo.toml @@ -39,4 +39,5 @@ rivet-config.workspace = true rivet-env.workspace = true rivet-pools.workspace = true rivet-test-deps-docker.workspace = true +tokio-postgres.workspace = true tracing-subscriber.workspace = true diff --git a/engine/packages/universaldb/src/driver/postgres/database.rs b/engine/packages/universaldb/src/driver/postgres/database.rs index 6e24cf647d..6cded0f320 100644 --- a/engine/packages/universaldb/src/driver/postgres/database.rs +++ b/engine/packages/universaldb/src/driver/postgres/database.rs @@ -59,6 +59,7 @@ impl PostgresConfig { pub struct PostgresDatabaseDriver { shared: Arc, max_retries: AtomicI32, + resolver_handle: JoinHandle<()>, gc_handle: JoinHandle<()>, } @@ -112,13 +113,14 @@ impl PostgresDatabaseDriver { let shared = PostgresShared::new(pool, node_id, listener); // Every node runs the resolver; only the elected leader drains the commit queue. - resolver::spawn(shared.clone()); + let resolver_handle = resolver::spawn(shared.clone()); let gc_handle = Self::spawn_gc(shared.clone()); Ok(PostgresDatabaseDriver { shared, max_retries: AtomicI32::new(100), + resolver_handle, gc_handle, }) } @@ -292,6 +294,9 @@ impl DatabaseDriver for PostgresDatabaseDriver { impl Drop for PostgresDatabaseDriver { fn drop(&mut self) { + // Abort the resolver so a dropped node stops renewing its lease; the lease then expires and + // another node can take over. Without this a dropped leader would renew its lease forever. + self.resolver_handle.abort(); self.gc_handle.abort(); } } diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs index 79f1f1f48b..3dfac50fce 100644 --- a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs +++ b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs @@ -33,9 +33,10 @@ enum DrainOutcome { } /// Spawn the per-process resolver task. Every node runs this; only the elected leader drains the -/// commit queue. -pub fn spawn(shared: Arc) { - tokio::spawn(run(shared)); +/// commit queue. The returned handle is aborted when the owning driver drops, which stops lease +/// renewal so the lease expires and another node can take over (node-death / failover path). +pub fn spawn(shared: Arc) -> tokio::task::JoinHandle<()> { + tokio::spawn(run(shared)) } async fn run(shared: Arc) { diff --git a/engine/packages/universaldb/tests/failover.rs b/engine/packages/universaldb/tests/failover.rs new file mode 100644 index 0000000000..7d122ca781 --- /dev/null +++ b/engine/packages/universaldb/tests/failover.rs @@ -0,0 +1,219 @@ +use std::{sync::Arc, time::Duration}; + +use rivet_test_deps_docker::TestDatabase; +use tokio_postgres::NoTls; +use universaldb::{Database, utils::IsolationLevel::*}; +use uuid::Uuid; + +const ALPHA_KEY: &[u8] = b"failover/alpha"; +const BETA_KEY: &[u8] = b"failover/beta"; + +/// Build a fresh Postgres-backed `Database`. Each call spins up an independent driver (its own pool, +/// node id, listener, and resolver), so two of them against one Postgres model two engine nodes. +async fn make_db(connection_string: &str) -> Database { + let driver = universaldb::driver::PostgresDatabaseDriver::new_with_config( + universaldb::driver::postgres::PostgresConfig::new(connection_string.to_string()), + ) + .await + .unwrap(); + Database::new(Arc::new(driver)) +} + +/// Raw verification connection used to inspect leader/lease/version state out of band. +async fn connect_raw(connection_string: &str) -> tokio_postgres::Client { + let (client, connection) = tokio_postgres::connect(connection_string, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client +} + +struct LeaseRow { + epoch: i64, + leader_addr: String, + durable_version: i64, +} + +async fn read_lease(client: &tokio_postgres::Client) -> Option { + let row = client + .query_opt( + "SELECT epoch, leader_addr, durable_version FROM udb_lease WHERE id = 1", + &[], + ) + .await + .unwrap()?; + Some(LeaseRow { + epoch: row.get(0), + leader_addr: row.get(1), + durable_version: row.get(2), + }) +} + +/// High-water of the LOGGED version sequence. A freshly elected leader must continue from at least +/// this value, never regress below it. +async fn read_seq_high(client: &tokio_postgres::Client) -> i64 { + client + .query_one("SELECT last_value FROM udb_version_seq", &[]) + .await + .unwrap() + .get(0) +} + +/// Poll `udb_lease` until `pred` holds or the deadline passes. +async fn wait_for_lease bool>( + client: &tokio_postgres::Client, + timeout: Duration, + pred: F, +) -> LeaseRow { + let deadline = tokio::time::Instant::now() + timeout; + loop { + if let Some(lease) = read_lease(client).await { + if pred(&lease) { + return lease; + } + } + if tokio::time::Instant::now() >= deadline { + panic!("timed out waiting for lease condition"); + } + tokio::time::sleep(Duration::from_millis(200)).await; + } +} + +async fn write_key(db: &Database, key: &'static [u8], value: &'static [u8]) { + db.txn("test_failover", move |tx| async move { + tx.set(key, value); + Ok(()) + }) + .await + .unwrap(); +} + +async fn read_key(db: &Database, key: &'static [u8]) -> Option> { + db.txn("test_failover", move |tx| async move { + let val = tx.get(key, Serializable).await?; + Ok(val) + }) + .await + .unwrap() + .map(|slice| slice.to_vec()) +} + +/// Exercises leader failover: two nodes share one Postgres, the elected leader is killed, the +/// survivor must take over the lease (new epoch), continue the crash-safe version sequence without +/// regression, preserve the dead leader's committed data, and resume accepting commits. +#[tokio::test] +async fn test_postgres_leader_failover() { + let _ = tracing_subscriber::fmt() + .with_env_filter("info") + .with_test_writer() + .try_init(); + + let (db_config, docker_config) = TestDatabase::Postgres + .config(Uuid::new_v4(), 1) + .await + .unwrap(); + let mut docker_config = docker_config.unwrap(); + docker_config.start().await.unwrap(); + + tokio::time::sleep(Duration::from_secs(4)).await; + + let rivet_config::config::Database::Postgres(postgres_config) = db_config else { + unreachable!(); + }; + let connection_string = postgres_config.url.read().clone(); + + let raw = connect_raw(&connection_string).await; + + // Node 1 comes up first and deterministically wins the first election (epoch 1). + let db1 = make_db(&connection_string).await; + let lease1 = wait_for_lease(&raw, Duration::from_secs(15), |l| l.epoch == 1).await; + let leader1_addr = lease1.leader_addr.clone(); + + // Node 2 joins while node 1 holds a valid lease, so it loses the election and runs as a + // follower. + let db2 = make_db(&connection_string).await; + + // Leader (node 1) commits data. The version sequence and watermark advance. + write_key(&db1, ALPHA_KEY, b"1").await; + + // The follower (node 2) reads through its own snapshot and sees the leader's committed write, + // proving cross-node reads work before any failover. + assert_eq!( + read_key(&db2, ALPHA_KEY).await, + Some(b"1".to_vec()), + "follower must see the leader's committed write" + ); + + let lease_before = read_lease(&raw).await.unwrap(); + let seq_before = read_seq_high(&raw).await; + assert!( + lease_before.durable_version >= 1, + "durable_version must have advanced after the first commit" + ); + + // Kill node 1. Dropping the driver aborts its resolver, so it stops renewing the lease. + drop(db1); + + // Node 2 must take over once node 1's lease expires (TTL is 10s). The epoch is bumped and the + // leader address changes to node 2. + let lease_after = wait_for_lease(&raw, Duration::from_secs(40), |l| { + l.epoch > lease_before.epoch + }) + .await; + assert!( + lease_after.epoch > lease_before.epoch, + "new leader must bump the epoch (was {}, now {})", + lease_before.epoch, + lease_after.epoch + ); + assert_ne!( + lease_after.leader_addr, leader1_addr, + "the surviving node must become the new leader" + ); + + // The crash-safe LOGGED sequence continues from the prior high-water; it never regresses. + let seq_after_takeover = read_seq_high(&raw).await; + assert!( + seq_after_takeover >= seq_before, + "version sequence regressed across failover ({} -> {})", + seq_before, + seq_after_takeover + ); + assert!( + lease_after.durable_version >= lease_before.durable_version, + "durable_version regressed across failover ({} -> {})", + lease_before.durable_version, + lease_after.durable_version + ); + + // The data the dead leader committed survives the failover. + assert_eq!( + read_key(&db2, ALPHA_KEY).await, + Some(b"1".to_vec()), + "committed data must survive leader failover" + ); + + // The new leader resumes accepting commits. + write_key(&db2, BETA_KEY, b"2").await; + assert_eq!( + read_key(&db2, BETA_KEY).await, + Some(b"2".to_vec()), + "new leader must accept and durably apply commits" + ); + + // The new commit advanced the version sequence and watermark past the pre-failover floor, + // confirming the new leader sequences from a strictly higher version. + let lease_final = read_lease(&raw).await.unwrap(); + assert!( + read_seq_high(&raw).await > seq_before, + "a post-failover commit must advance the version sequence" + ); + assert!( + lease_final.durable_version > lease_before.durable_version, + "a post-failover commit must advance the durable watermark" + ); + + drop(db2); +} diff --git a/engine/packages/universalpubsub/src/driver/postgres/mod.rs b/engine/packages/universalpubsub/src/driver/postgres/mod.rs index d361eb1149..2a492fbccf 100644 --- a/engine/packages/universalpubsub/src/driver/postgres/mod.rs +++ b/engine/packages/universalpubsub/src/driver/postgres/mod.rs @@ -40,9 +40,12 @@ const POLL_INTERVAL: Duration = Duration::from_secs(1); /// rather than a cluster outage. const LISTEN_IDLE_IN_TRANSACTION_TIMEOUT_MS: i64 = 30_000; -const QUEUE_SUB_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10); -/// How long a queue subscriber's heartbeat must be within to be considered active. -const QUEUE_SUB_TTL_SECS: i64 = 30; +/// How often this process refreshes its node liveness heartbeat. One heartbeat per +/// process keeps all of its subscriber registrations alive at once. +const NODE_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10); +/// How recent a node's heartbeat must be for its subscribers to count as live +/// responders. +const NODE_TTL_SECS: i64 = 30; /// How often to GC expired broadcast messages. const MESSAGE_GC_INTERVAL: Duration = Duration::from_secs(5); @@ -51,19 +54,34 @@ const MESSAGE_GC_INTERVAL: Duration = Duration::from_secs(5); /// messages, matching NATS-core at-most-once semantics for slow consumers. const MESSAGE_MAX_AGE_SECS: i64 = 10; +/// How often to GC dead nodes and the subscriber rows orphaned by them. +const REGISTRY_GC_INTERVAL: Duration = Duration::from_secs(30); + /// How often to GC orphaned queue messages. const QUEUE_MESSAGE_GC_INTERVAL: Duration = Duration::from_secs(300); /// Max age before an unconsumed queue message is garbage collected. const QUEUE_MESSAGE_MAX_AGE_SECS: i64 = 3600; +/// Per-shard signal carried over a subscriber's in-process wakeup channel. +#[derive(Clone)] +enum ShardSignal { + /// A doorbell NOTIFY landed for this shard. Poll the table. + Wakeup, + /// A local request found no responders for the given reply subject. The matching + /// reply subscriber surfaces a no-responders result. + NoResponders { subject: String }, +} + #[derive(Clone)] pub struct PostgresDriver { pool: Arc, client: Arc>>, + /// Identifies this process in the subscriber registry. A single heartbeat keeps + /// all of this node's registrations live. + node_id: String, /// Wakeup channels keyed by doorbell shard channel name. Shared by broadcast and - /// queue subscribers whose subjects map to the same shard. Carries empty wakeups - /// only; payload lives in the table. - shard_subscriptions: Arc>>, + /// queue subscribers whose subjects map to the same shard. + shard_subscriptions: Arc>>, doorbell: Arc, client_ready: tokio::sync::watch::Receiver, } @@ -105,9 +123,10 @@ impl PostgresDriver { tracing::debug!("postgres pool created successfully"); let pool = Arc::new(pool); - let shard_subscriptions: Arc>> = + let shard_subscriptions: Arc>> = Arc::new(HashMap::new()); let client: Arc>> = Arc::new(Mutex::new(None)); + let node_id = Uuid::new_v4().to_string(); // Create channel for client ready notifications let (ready_tx, client_ready) = tokio::sync::watch::channel(false); @@ -128,6 +147,7 @@ impl PostgresDriver { let driver = Self { pool, client, + node_id, shard_subscriptions, doorbell, client_ready, @@ -138,6 +158,7 @@ impl PostgresDriver { // Create tables eagerly so they exist before any publish or subscribe. { + tracing::debug!("configuring postgres udb tables"); let conn = driver .pool .get() @@ -157,11 +178,23 @@ impl PostgresDriver { ); \ CREATE INDEX IF NOT EXISTS ups_messages_subject_id \ ON ups_messages (subject_hash, id); \ + CREATE TABLE IF NOT EXISTS ups_nodes ( \ + node_id TEXT PRIMARY KEY, \ + heartbeat_at TIMESTAMPTZ NOT NULL DEFAULT NOW() \ + ); \ + CREATE TABLE IF NOT EXISTS ups_subs ( \ + id TEXT PRIMARY KEY, \ + node_id TEXT NOT NULL, \ + subject_hash TEXT NOT NULL, \ + subject TEXT NOT NULL \ + ); \ + CREATE INDEX IF NOT EXISTS ups_subs_subject \ + ON ups_subs (subject_hash); \ CREATE TABLE IF NOT EXISTS ups_queue_subs ( \ id TEXT PRIMARY KEY, \ + node_id TEXT NOT NULL, \ subject_hash TEXT NOT NULL, \ - queue_hash TEXT NOT NULL, \ - heartbeat_at TIMESTAMPTZ NOT NULL DEFAULT NOW() \ + queue_hash TEXT NOT NULL \ ); \ CREATE INDEX IF NOT EXISTS ups_queue_subs_subject_queue \ ON ups_queue_subs (subject_hash, queue_hash); \ @@ -177,9 +210,24 @@ impl PostgresDriver { ) .await .context("failed to create tables")?; - tracing::debug!("tables ready"); + tracing::debug!("postgres udb tables ready"); } + // Register this node and start its liveness heartbeat. + driver.heartbeat_node().await?; + let heartbeat_driver = driver.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(NODE_HEARTBEAT_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + interval.tick().await; + if let Err(e) = heartbeat_driver.heartbeat_node().await { + tracing::warn!(?e, "failed to heartbeat node"); + } + } + }); + // Spawn GC task for expired broadcast messages let message_gc_driver = driver.clone(); tokio::spawn(async move { @@ -203,6 +251,49 @@ impl PostgresDriver { } }); + // Spawn GC task for dead nodes and orphaned subscriber rows. + let registry_gc_driver = driver.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(REGISTRY_GC_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + interval.tick().await; + if let Ok(conn) = registry_gc_driver.pool.get().await { + if let Err(e) = conn + .execute( + "DELETE FROM ups_nodes \ + WHERE heartbeat_at < NOW() - ($1::bigint * INTERVAL '1 second')", + &[&NODE_TTL_SECS], + ) + .await + { + tracing::warn!(?e, "failed to gc dead nodes"); + } + if let Err(e) = conn + .execute( + "DELETE FROM ups_subs \ + WHERE node_id NOT IN (SELECT node_id FROM ups_nodes)", + &[], + ) + .await + { + tracing::warn!(?e, "failed to gc orphaned subs"); + } + if let Err(e) = conn + .execute( + "DELETE FROM ups_queue_subs \ + WHERE node_id NOT IN (SELECT node_id FROM ups_nodes)", + &[], + ) + .await + { + tracing::warn!(?e, "failed to gc orphaned queue subs"); + } + } + } + }); + // Spawn GC task for orphaned queue messages let gc_driver = driver.clone(); tokio::spawn(async move { @@ -232,7 +323,7 @@ impl PostgresDriver { /// Manages the connection lifecycle with automatic reconnection async fn spawn_connection_lifecycle( conn_str: String, - shard_subscriptions: Arc>>, + shard_subscriptions: Arc>>, client: Arc>>, ready_tx: tokio::sync::watch::Sender, ssl_root_cert_path: Option, @@ -337,7 +428,7 @@ impl PostgresDriver { /// Polls the connection for notifications until it closes or errors async fn poll_connection( mut conn: tokio_postgres::Connection, - shard_subscriptions: Arc>>, + shard_subscriptions: Arc>>, ) where T: tokio_postgres::tls::TlsStream + Unpin, { @@ -348,7 +439,7 @@ impl PostgresDriver { // Doorbell notifications are payload-free wakeup signals only. // Subscribers read their payload from the table. if let Some(sub) = shard_subscriptions.get_async(note.channel()).await { - let _ = sub.send(()); + let _ = sub.send(ShardSignal::Wakeup); } else { tracing::trace!(channel = %note.channel(), "wakeup for unknown shard"); } @@ -404,6 +495,24 @@ impl PostgresDriver { format!("{:x}", hasher.finish()) } + /// Upserts this node's liveness heartbeat. Re-inserts the row if a GC pass removed + /// it after a transient stall. + async fn heartbeat_node(&self) -> Result<()> { + let conn = self + .pool + .get() + .await + .context("failed to get connection for node heartbeat")?; + conn.execute( + "INSERT INTO ups_nodes (node_id, heartbeat_at) VALUES ($1, NOW()) \ + ON CONFLICT (node_id) DO UPDATE SET heartbeat_at = NOW()", + &[&self.node_id], + ) + .await + .context("failed to upsert node heartbeat")?; + Ok(()) + } + /// Returns the current max broadcast message id, used as a subscriber's starting /// cursor so it only sees future messages (NATS at-most-once, no replay). async fn current_max_id(&self) -> Result { @@ -424,7 +533,10 @@ impl PostgresDriver { async fn ensure_shard_listen( &self, shard: usize, - ) -> (broadcast::Receiver<()>, tokio_util::sync::DropGuard) { + ) -> ( + broadcast::Receiver, + tokio_util::sync::DropGuard, + ) { let channel = shard_channel(shard); match self.shard_subscriptions.entry_async(channel.clone()).await { @@ -465,7 +577,7 @@ impl PostgresDriver { fn spawn_shard_cleanup_task( &self, channel: String, - tx: broadcast::Sender<()>, + tx: broadcast::Sender, ) -> tokio_util::sync::DropGuard { let driver = self.clone(); let token = tokio_util::sync::CancellationToken::new(); @@ -515,14 +627,15 @@ impl PostgresDriver { .await .context("failed to insert broadcast message")?; - // Queue rows for every active queue group on this subject. Batched into the - // same transaction so a crash never strands a row mid-publish. + // Queue rows for every live queue group on this subject. Batched into the same + // transaction so a crash never strands a row mid-publish. let rows = tx .query( - "SELECT DISTINCT queue_hash FROM ups_queue_subs \ - WHERE subject_hash = $1 \ - AND heartbeat_at > NOW() - ($2::bigint * INTERVAL '1 second')", - &[&subject_hash, &QUEUE_SUB_TTL_SECS], + "SELECT DISTINCT s.queue_hash FROM ups_queue_subs s \ + JOIN ups_nodes n ON s.node_id = n.node_id \ + WHERE s.subject_hash = $1 \ + AND n.heartbeat_at > NOW() - ($2::bigint * INTERVAL '1 second')", + &[&subject_hash, &NODE_TTL_SECS], ) .await .context("failed to query active queue subs")?; @@ -542,6 +655,50 @@ impl PostgresDriver { Ok(()) } + + /// Returns whether any live subscriber (broadcast or queue) exists for the subject + /// anywhere in the fleet. Used to decide whether a request surfaces a no-responders + /// result instead of waiting out its timeout. + async fn has_responders(&self, subject_hash: &str, subject: &str) -> Result { + let conn = self + .pool + .get() + .await + .context("failed to get connection for responder check")?; + let row = conn + .query_one( + "SELECT \ + EXISTS( \ + SELECT 1 FROM ups_subs s \ + JOIN ups_nodes n ON s.node_id = n.node_id \ + WHERE s.subject_hash = $1 AND s.subject = $2 \ + AND n.heartbeat_at > NOW() - ($3::bigint * INTERVAL '1 second') \ + ) \ + OR EXISTS( \ + SELECT 1 FROM ups_queue_subs s \ + JOIN ups_nodes n ON s.node_id = n.node_id \ + WHERE s.subject_hash = $1 \ + AND n.heartbeat_at > NOW() - ($3::bigint * INTERVAL '1 second') \ + )", + &[&subject_hash, &subject, &NODE_TTL_SECS], + ) + .await + .context("failed to check responders")?; + Ok(row.get(0)) + } + + /// Delivers a no-responders result to the local reply subscriber. The requester is + /// always in this process, so the signal is routed in-memory over the reply + /// subject's shard channel rather than the table. + async fn signal_no_responders(&self, reply_subject: &str) { + let reply_hash = self.hash_subject(reply_subject); + let channel = shard_channel(shard_for(&reply_hash)); + if let Some(tx) = self.shard_subscriptions.get_async(&channel).await { + let _ = tx.send(ShardSignal::NoResponders { + subject: reply_subject.to_string(), + }); + } + } } #[async_trait] @@ -549,7 +706,7 @@ impl PubSubDriver for PostgresDriver { async fn subscribe( &self, subject: &str, - _reply_id: Option, + reply_id: Option, ) -> Result { let subject_hash = self.hash_subject(subject); let shard = shard_for(&subject_hash); @@ -561,6 +718,28 @@ impl PubSubDriver for PostgresDriver { let (rx, drop_guard) = self.ensure_shard_listen(shard).await; + // Register in the responder registry so requests to this subject can detect + // responders. Reply inboxes are never request targets, so they skip the + // registry to keep request latency off this path. + let sub_id = if reply_id.is_none() { + let sub_id = Uuid::new_v4().to_string(); + let conn = self + .pool + .get() + .await + .context("failed to get connection for subscribe")?; + conn.execute( + "INSERT INTO ups_subs (id, node_id, subject_hash, subject) \ + VALUES ($1, $2, $3, $4)", + &[&sub_id, &self.node_id, &subject_hash, &subject], + ) + .await + .context("failed to register subscriber")?; + Some(sub_id) + } else { + None + }; + Ok(Box::new(PostgresSubscriber { subject: subject.to_string(), subject_hash, @@ -568,6 +747,7 @@ impl PubSubDriver for PostgresDriver { cursor, buffer: VecDeque::new(), rx, + sub_id, _drop_guard: drop_guard, })) } @@ -586,8 +766,9 @@ impl PubSubDriver for PostgresDriver { .await .context("failed to get connection for queue subscribe")?; conn.execute( - "INSERT INTO ups_queue_subs (id, subject_hash, queue_hash) VALUES ($1, $2, $3)", - &[&sub_id, &subject_hash, &queue_hash], + "INSERT INTO ups_queue_subs (id, node_id, subject_hash, queue_hash) \ + VALUES ($1, $2, $3, $4)", + &[&sub_id, &self.node_id, &subject_hash, &queue_hash], ) .await .context("failed to register queue subscriber")?; @@ -595,35 +776,6 @@ impl PubSubDriver for PostgresDriver { let (rx, drop_guard) = self.ensure_shard_listen(shard).await; - // Spawn heartbeat task to keep the registration alive - let pool = self.pool.clone(); - let sub_id_for_heartbeat = sub_id.clone(); - let heartbeat_token = tokio_util::sync::CancellationToken::new(); - let heartbeat_token_child = heartbeat_token.clone(); - tokio::spawn(async move { - let mut interval = tokio::time::interval(QUEUE_SUB_HEARTBEAT_INTERVAL); - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - - loop { - tokio::select! { - _ = heartbeat_token_child.cancelled() => break, - _ = interval.tick() => { - if let Ok(conn) = pool.get().await { - if let Err(e) = conn - .execute( - "UPDATE ups_queue_subs SET heartbeat_at = NOW() WHERE id = $1", - &[&sub_id_for_heartbeat], - ) - .await - { - tracing::warn!(?e, id = %sub_id_for_heartbeat, "failed to heartbeat queue sub"); - } - } - } - } - } - }); - Ok(Box::new(PostgresQueueSubscriber { subject: subject.to_string(), subject_hash, @@ -632,7 +784,6 @@ impl PubSubDriver for PostgresDriver { pool: self.pool.clone(), rx, _drop_guard: drop_guard, - _heartbeat_token: heartbeat_token, })) } @@ -640,11 +791,29 @@ impl PubSubDriver for PostgresDriver { &self, subject: &str, payload: &[u8], - _reply_subject: Option<&str>, + reply_subject: Option<&str>, ) -> Result<()> { let subject_hash = self.hash_subject(subject); let shard = shard_for(&subject_hash); + // Request semantics: if a reply is expected and no responder exists anywhere, + // surface a no-responders result immediately instead of persisting a message + // nobody will read. + if let Some(reply_subject) = reply_subject { + match self.has_responders(&subject_hash, subject).await { + Result::Ok(false) => { + self.signal_no_responders(reply_subject).await; + return Ok(()); + } + Result::Ok(true) => {} + Result::Err(e) => { + // On a failed check, fall through to a normal publish rather than + // risk a false no-responders result. + tracing::warn!(?e, %subject, "responder check failed, publishing anyway"); + } + } + } + // Persist the message, retrying on transient connection errors. The row is // committed before the doorbell rings so any wakeup observes it. let mut backoff = Backoff::default(); @@ -686,7 +855,9 @@ pub struct PostgresSubscriber { pool: Arc, cursor: i64, buffer: VecDeque>, - rx: broadcast::Receiver<()>, + rx: broadcast::Receiver, + /// Responder-registry row id, present for non-inbox subscriptions. Deleted on drop. + sub_id: Option, _drop_guard: tokio_util::sync::DropGuard, } @@ -744,11 +915,17 @@ impl SubscriberDriver for PostgresSubscriber { continue; } - // Wait for a doorbell wakeup or the poll backstop, whichever is first. + // Wait for a doorbell wakeup, a no-responders signal, or the poll backstop. tokio::select! { res = self.rx.recv() => { match res { - std::result::Result::Ok(()) => {} + std::result::Result::Ok(ShardSignal::Wakeup) => {} + std::result::Result::Ok(ShardSignal::NoResponders { subject }) + if subject == self.subject => + { + return Ok(DriverOutput::NoResponders); + } + std::result::Result::Ok(ShardSignal::NoResponders { .. }) => {} Err(broadcast::error::RecvError::Lagged(_)) => {} Err(broadcast::error::RecvError::Closed) => { return Ok(DriverOutput::Unsubscribed); @@ -761,15 +938,33 @@ impl SubscriberDriver for PostgresSubscriber { } } +impl Drop for PostgresSubscriber { + fn drop(&mut self) { + let Some(sub_id) = self.sub_id.take() else { + return; + }; + let pool = self.pool.clone(); + tokio::spawn(async move { + if let Ok(conn) = pool.get().await { + if let Err(e) = conn + .execute("DELETE FROM ups_subs WHERE id = $1", &[&sub_id]) + .await + { + tracing::warn!(?e, %sub_id, "failed to deregister subscriber"); + } + } + }); + } +} + pub struct PostgresQueueSubscriber { subject: String, subject_hash: String, queue_hash: String, sub_id: String, pool: Arc, - rx: broadcast::Receiver<()>, + rx: broadcast::Receiver, _drop_guard: tokio_util::sync::DropGuard, - _heartbeat_token: tokio_util::sync::CancellationToken, } impl PostgresQueueSubscriber { @@ -820,11 +1015,11 @@ impl SubscriberDriver for PostgresQueueSubscriber { } } - // Wait for a doorbell wakeup or the poll backstop, then loop back to claim. + // Wait for any shard signal or the poll backstop, then loop back to claim. tokio::select! { res = self.rx.recv() => { match res { - std::result::Result::Ok(()) => {} + std::result::Result::Ok(_) => {} Err(broadcast::error::RecvError::Lagged(_)) => {} Err(broadcast::error::RecvError::Closed) => { return Ok(DriverOutput::Unsubscribed); diff --git a/engine/packages/universalpubsub/tests/reconnect.rs b/engine/packages/universalpubsub/tests/reconnect.rs index 35ace375d3..8c15f6ab54 100644 --- a/engine/packages/universalpubsub/tests/reconnect.rs +++ b/engine/packages/universalpubsub/tests/reconnect.rs @@ -43,7 +43,7 @@ async fn test_nats_driver_with_memory_reconnect() { .unwrap(); let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), true); - test_all_inner(&pubsub, &docker).await; + test_all_inner(&pubsub, &docker, true).await; } #[tokio::test] @@ -77,7 +77,7 @@ async fn test_nats_driver_without_memory_reconnect() { .unwrap(); let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false); - test_all_inner(&pubsub, &docker).await; + test_all_inner(&pubsub, &docker, true).await; } #[tokio::test] @@ -95,13 +95,12 @@ async fn test_postgres_driver_with_memory_reconnect() { }; let url = pg.url.read().clone(); - let driver = - universalpubsub::driver::postgres::PostgresDriver::connect(url, true, None, None, None) - .await - .unwrap(); + let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, None, None, None) + .await + .unwrap(); let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), true); - test_all_inner(&pubsub, &docker).await; + test_all_inner(&pubsub, &docker, false).await; } #[tokio::test] @@ -119,19 +118,27 @@ async fn test_postgres_driver_without_memory_reconnect() { }; let url = pg.url.read().clone(); - let driver = - universalpubsub::driver::postgres::PostgresDriver::connect(url, false, None, None, None) - .await - .unwrap(); + let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, None, None, None) + .await + .unwrap(); let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false); - test_all_inner(&pubsub, &docker).await; + test_all_inner(&pubsub, &docker, false).await; } -async fn test_all_inner(pubsub: &PubSub, docker: &rivet_test_deps_docker::DockerRunConfig) { +async fn test_all_inner( + pubsub: &PubSub, + docker: &rivet_test_deps_docker::DockerRunConfig, + supports_subscribe_while_stopped: bool, +) { test_reconnect_inner(&pubsub, &docker).await; test_publish_while_stopped(&pubsub, &docker).await; - test_subscribe_while_stopped(&pubsub, &docker).await; + // The table-backed Postgres driver must read its cursor and register in the + // responder table when subscribing, so it cannot subscribe while the backend is + // fully down. NATS buffers the subscribe and reconnects, so it can. + if supports_subscribe_while_stopped { + test_subscribe_while_stopped(&pubsub, &docker).await; + } } async fn test_reconnect_inner(pubsub: &PubSub, docker: &rivet_test_deps_docker::DockerRunConfig) { diff --git a/scripts/run/engine-postgres.sh b/scripts/run/engine-postgres.sh index 966c421bec..70689517bd 100755 --- a/scripts/run/engine-postgres.sh +++ b/scripts/run/engine-postgres.sh @@ -4,19 +4,26 @@ set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" -if ! command -v nc >/dev/null 2>&1; then - echo "error: required command 'nc' not found." - exit 1 -fi +POSTGRES_IMAGE="postgres:18" + +# pg_isready reports ready only once the server is actually accepting connections. +# The Postgres entrypoint binds the port during its bootstrap phase and then +# restarts, so a plain port check (nc -z) passes too early and the engine hits +# "connection reset" / "early eof" on first connect. Run pg_isready from a throwaway +# container on the host network so no client binary needs to be installed locally. +postgres_ready() { + docker run --rm --network host "${POSTGRES_IMAGE}" \ + pg_isready -h localhost -p 5432 -U postgres -d postgres >/dev/null 2>&1 +} -if ! nc -z localhost 5432 >/dev/null 2>&1; then - echo "Postgres is not reachable at localhost:5432." +if ! postgres_ready; then + echo "Postgres is not accepting connections." echo "Starting postgres container..." "${SCRIPT_DIR}/postgres.sh" echo "Waiting for postgres to be ready..." for i in {1..30}; do - if nc -z localhost 5432 >/dev/null 2>&1; then + if postgres_ready; then echo "Postgres is ready!" break fi From 833bfeb718b468e794b8083773393a20f19a957a Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Thu, 25 Jun 2026 13:07:04 -0700 Subject: [PATCH 11/16] [SLOP(claude-opus-4-8-high)] feat(universaldb): graceful postgres leader handoff on shutdown --- engine/packages/service-manager/src/lib.rs | 5 +- engine/packages/universaldb/src/database.rs | 5 ++ engine/packages/universaldb/src/driver/mod.rs | 7 ++ .../src/driver/postgres/database.rs | 11 +++ .../src/driver/postgres/resolver/lease.rs | 24 +++++++ .../src/driver/postgres/resolver/mod.rs | 65 ++++++++++++++++- .../universaldb/src/driver/postgres/shared.rs | 4 ++ engine/packages/universaldb/tests/failover.rs | 70 +++++++++++++++++++ 8 files changed, 187 insertions(+), 4 deletions(-) diff --git a/engine/packages/service-manager/src/lib.rs b/engine/packages/service-manager/src/lib.rs index a60854f3e0..0b84fbae82 100644 --- a/engine/packages/service-manager/src/lib.rs +++ b/engine/packages/service-manager/src/lib.rs @@ -393,7 +393,7 @@ pub async fn start( if abort { // Give time for services to handle final abort tokio::time::sleep(Duration::from_millis(50)).await; - rivet_runtime::shutdown().await; // TODO: Fix `JoinHandle polled after completion` error + rivet_runtime::shutdown().await; break; } @@ -401,6 +401,9 @@ pub async fn start( } } + // Shut down udb + pools.udb()?.shutdown().await; + // Stops term signal handler bg task rivet_runtime::TermSignal::stop(); diff --git a/engine/packages/universaldb/src/database.rs b/engine/packages/universaldb/src/database.rs index c2f95983a4..a5c65dd0f5 100644 --- a/engine/packages/universaldb/src/database.rs +++ b/engine/packages/universaldb/src/database.rs @@ -106,4 +106,9 @@ impl Database { pub fn checkpoint(&self, path: &Path) -> Result<()> { self.driver.checkpoint(path) } + + /// Gracefully release process-wide driver resources before shutdown. + pub async fn shutdown(&self) { + self.driver.shutdown().await; + } } diff --git a/engine/packages/universaldb/src/driver/mod.rs b/engine/packages/universaldb/src/driver/mod.rs index 01f7c22b2c..7b19e90fef 100644 --- a/engine/packages/universaldb/src/driver/mod.rs +++ b/engine/packages/universaldb/src/driver/mod.rs @@ -34,6 +34,13 @@ pub trait DatabaseDriver: Send + Sync { fn checkpoint(&self, _path: &Path) -> Result<()> { bail!("checkpoint not supported by this database driver") } + + /// Gracefully release any process-wide resources before shutdown. The Postgres driver hands off + /// its leader lease here so a standby node takes over immediately instead of waiting out the + /// lease TTL. Default is a no-op. + fn shutdown<'a>(&'a self) -> BoxFut<'a, ()> { + Box::pin(async {}) + } } pub trait TransactionDriver: Send + Sync { diff --git a/engine/packages/universaldb/src/driver/postgres/database.rs b/engine/packages/universaldb/src/driver/postgres/database.rs index 6cded0f320..3ff891011f 100644 --- a/engine/packages/universaldb/src/driver/postgres/database.rs +++ b/engine/packages/universaldb/src/driver/postgres/database.rs @@ -290,6 +290,17 @@ impl DatabaseDriver for PostgresDatabaseDriver { self.max_retries.store(limit, Ordering::SeqCst); Ok(()) } + + fn shutdown<'a>(&'a self) -> BoxFut<'a, ()> { + Box::pin(async move { + // Stop renewing the lease before releasing it so a racing renew cannot re-extend it. + self.resolver_handle.abort(); + self.gc_handle.abort(); + + // Hand off leadership immediately if we hold it, instead of waiting out the lease TTL. + resolver::handoff(&self.shared).await; + }) + } } impl Drop for PostgresDatabaseDriver { diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs b/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs index a17b6c7759..93f7f7884e 100644 --- a/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs +++ b/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs @@ -62,6 +62,30 @@ pub async fn renew(pool: &Pool, node_id: &str, epoch: i64) -> Result { Ok(updated == 1) } +/// Gracefully release the lease so a standby node can take over immediately instead of waiting out +/// the TTL. Expires the lease in place, fenced on this node's address so it never clobbers a +/// successor that already took over. Returns `true` if our lease was released (i.e. we were the +/// leader); `false` is the normal no-op when this node is a follower. Renewal must already be +/// stopped before calling this, otherwise a racing renew could re-extend the lease. +pub async fn release(pool: &Pool, node_id: &str) -> Result { + let conn = pool + .get() + .await + .context("failed to get connection for lease release")?; + + let updated = conn + .execute( + "UPDATE udb_lease + SET expires_at = now() + WHERE id = $1 AND leader_addr = $2", + &[&LEASE_ID, &node_id], + ) + .await + .context("failed to release lease")?; + + Ok(updated == 1) +} + /// Read the current durable version (`udb_lease.durable_version`). Used by a freshly elected leader /// to learn the watermark floor it must continue from. pub async fn current_durable_version(pool: &Pool) -> Result { diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs index 3dfac50fce..f9b6c45461 100644 --- a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs +++ b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs @@ -7,10 +7,13 @@ use std::{ }; use anyhow::{Context, Result}; +use tokio::sync::broadcast; use crate::{conflict_tracker::TransactionConflictTracker, transaction::TXN_TIMEOUT}; -use super::shared::{LEASE_ID, LeaseInfo, PostgresShared, WATERMARK_CHANNEL, commit_channel}; +use super::shared::{ + ELECTION_CHANNEL, LEASE_ID, LeaseInfo, PostgresShared, WATERMARK_CHANNEL, commit_channel, +}; /// Max commits resolved+applied per batch (group commit). Amortizes the resolver, Postgres /// round-trips, and fsync across the batch. @@ -40,6 +43,10 @@ pub fn spawn(shared: Arc) -> tokio::task::JoinHandle<()> { } async fn run(shared: Arc) { + // A departing leader NOTIFYs this channel after releasing its lease so we elect immediately + // rather than waiting out the full `ELECTION_RETRY` tick. + let mut election_rx = shared.listener.listen(ELECTION_CHANNEL).await; + loop { match lease::try_acquire(&shared.pool, &shared.node_id).await { Ok(Some(acquired)) => { @@ -50,16 +57,68 @@ async fn run(shared: Arc) { tracing::info!(epoch = acquired.epoch, "stepped down from udb leader"); } Ok(None) => { - tokio::time::sleep(ELECTION_RETRY).await; + wait_for_election_retry(&shared, &mut election_rx).await; } Err(err) => { tracing::warn!(?err, "failed udb lease acquire attempt"); - tokio::time::sleep(ELECTION_RETRY).await; + wait_for_election_retry(&shared, &mut election_rx).await; + } + } + } +} + +/// Wait before retrying the election: either the `ELECTION_RETRY` backstop elapses, or a departing +/// leader wakes us via `ELECTION_CHANNEL` so handoff is near-instant. +async fn wait_for_election_retry( + shared: &Arc, + election_rx: &mut broadcast::Receiver, +) { + tokio::select! { + _ = tokio::time::sleep(ELECTION_RETRY) => {} + res = election_rx.recv() => { + if matches!(res, Err(broadcast::error::RecvError::Closed)) { + // The listener recreates the channel on reconnect; re-subscribe. + *election_rx = shared.listener.listen(ELECTION_CHANNEL).await; } } } } +/// Best-effort graceful leadership handoff invoked on shutdown. If this node currently holds the +/// lease, expire it and wake a standby so it takes over immediately instead of waiting out the TTL. +/// Safe to call on a follower: the fenced release matches no row and nothing is notified. The +/// caller must already have stopped lease renewal before calling this. +pub async fn handoff(shared: &Arc) { + match lease::release(&shared.pool, &shared.node_id).await { + Ok(true) => { + tracing::info!(node_id = %shared.node_id, "released udb leader lease for graceful handoff"); + notify_election(shared).await; + } + Ok(false) => {} + Err(err) => { + tracing::warn!(?err, "failed to release udb lease on shutdown"); + } + } +} + +/// Wake standby candidates so the next election fires immediately after a graceful release. +async fn notify_election(shared: &Arc) { + let conn = match shared.pool.get().await { + Ok(conn) => conn, + Err(err) => { + tracing::debug!(?err, "failed to get connection for election notify"); + return; + } + }; + + if let Err(err) = conn + .execute("SELECT pg_notify($1, '')", &[&ELECTION_CHANNEL]) + .await + { + tracing::debug!(?err, "failed to notify election channel"); + } +} + /// Leader main loop: hold the lease, drain the commit queue on wake or poll, and renew the lease. async fn lead(shared: &Arc, epoch: i64) -> Result<()> { // Publish our own lease into the cache immediately so our local commits route to us. diff --git a/engine/packages/universaldb/src/driver/postgres/shared.rs b/engine/packages/universaldb/src/driver/postgres/shared.rs index 0ab2963158..5b14c4df08 100644 --- a/engine/packages/universaldb/src/driver/postgres/shared.rs +++ b/engine/packages/universaldb/src/driver/postgres/shared.rs @@ -32,6 +32,10 @@ pub fn reply_channel(node_id: &str) -> String { /// Channel the leader NOTIFYs on every watermark advance; all nodes LISTEN. pub const WATERMARK_CHANNEL: &str = "udb_watermark"; +/// Channel a departing leader NOTIFYs after releasing its lease so a standby candidate elects +/// immediately instead of waiting out `ELECTION_RETRY`. All non-leader candidates LISTEN. +pub const ELECTION_CHANNEL: &str = "udb_election"; + /// Cached view of the current leader lease, as seen by a follower. #[derive(Clone, Debug)] pub struct LeaseInfo { diff --git a/engine/packages/universaldb/tests/failover.rs b/engine/packages/universaldb/tests/failover.rs index 7d122ca781..5fb394453c 100644 --- a/engine/packages/universaldb/tests/failover.rs +++ b/engine/packages/universaldb/tests/failover.rs @@ -217,3 +217,73 @@ async fn test_postgres_leader_failover() { drop(db2); } + +/// Exercises graceful leader handoff: a leader that is shut down cleanly (SIGTERM path) releases its +/// lease immediately instead of letting it expire, so a standby takes over well within the lease TTL +/// rather than after it. This is what turns a rolling deploy from a ~TTL commit stall into a +/// near-instant handoff. +#[tokio::test] +async fn test_postgres_graceful_handoff() { + let _ = tracing_subscriber::fmt() + .with_env_filter("info") + .with_test_writer() + .try_init(); + + let (db_config, docker_config) = TestDatabase::Postgres + .config(Uuid::new_v4(), 1) + .await + .unwrap(); + let mut docker_config = docker_config.unwrap(); + docker_config.start().await.unwrap(); + + tokio::time::sleep(Duration::from_secs(4)).await; + + let rivet_config::config::Database::Postgres(postgres_config) = db_config else { + unreachable!(); + }; + let connection_string = postgres_config.url.read().clone(); + + let raw = connect_raw(&connection_string).await; + + // Node 1 wins the first election; node 2 joins as a follower. + let db1 = make_db(&connection_string).await; + let lease1 = wait_for_lease(&raw, Duration::from_secs(15), |l| l.epoch == 1).await; + let leader1_addr = lease1.leader_addr.clone(); + let db2 = make_db(&connection_string).await; + + write_key(&db1, ALPHA_KEY, b"1").await; + let lease_before = read_lease(&raw).await.unwrap(); + + // Gracefully shut down the leader. Unlike a hard drop, this releases the lease in place and + // wakes the standby, so takeover must complete in well under the 10s TTL. + let handoff_start = tokio::time::Instant::now(); + db1.shutdown().await; + + // The lease TTL is 10s; a graceful handoff must take over well under that. The 5s deadline here + // is itself below the TTL, so reaching this line already proves the lease was not waited out. + let lease_after = wait_for_lease(&raw, Duration::from_secs(5), |l| { + l.epoch > lease_before.epoch + }) + .await; + let handoff_elapsed = handoff_start.elapsed(); + assert!( + handoff_elapsed < Duration::from_secs(8), + "graceful handoff must beat the lease TTL (took {handoff_elapsed:?})" + ); + assert_ne!( + lease_after.leader_addr, leader1_addr, + "the standby must become the new leader after a graceful handoff" + ); + + // The new leader serves the old leader's data and accepts fresh commits. + assert_eq!( + read_key(&db2, ALPHA_KEY).await, + Some(b"1".to_vec()), + "committed data must survive graceful handoff" + ); + write_key(&db2, BETA_KEY, b"2").await; + assert_eq!(read_key(&db2, BETA_KEY).await, Some(b"2".to_vec())); + + drop(db1); + drop(db2); +} From dce315dde40e16ff745d891900046c74485be437 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Thu, 25 Jun 2026 13:36:10 -0700 Subject: [PATCH 12/16] [slopfix] fix(self-host): repair dev compose generation, kitchen-sink serverful runner, and postgres/udb bootstrap --- engine/packages/metrics/src/registry.rs | 4 +- examples/kitchen-sink/Dockerfile | 7 +- .../compose/prod-file-system/.gitattributes | 1 - self-host/compose/prod-file-system/README.md | 76 ------ .../prod-file-system/docker-compose.yml | 41 ---- .../rivet-engine/config.jsonc | 2 - .../compose/template/src/docker-compose.ts | 75 +++++- .../src/services/edge/rivet-engine.ts | 16 +- .../template/src/services/edge/runner.ts | 5 +- .../{compose => }/dev-host/.gitattributes | 0 self-host/{compose => }/dev-host/README.md | 0 .../dev-host/clickhouse/client-config.xml | 0 .../dev-host/clickhouse/config.xml | 0 .../clickhouse/init/01-create-otel-table.sql | 0 .../dev-host/clickhouse/users.xml | 0 .../{compose => }/dev-host/docker-compose.yml | 44 +++- .../dev-host/grafana/dashboards/api.json | 0 .../dev-host/grafana/dashboards/cache.json | 0 .../dev-host/grafana/dashboards/epoxy.json | 0 .../dev-host/grafana/dashboards/futures.json | 0 .../dev-host/grafana/dashboards/gasoline.json | 0 .../dev-host/grafana/dashboards/guard.json | 0 .../grafana/dashboards/operation.json | 0 .../dev-host/grafana/dashboards/pegboard.json | 0 .../dev-host/grafana/dashboards/tokio.json | 0 .../dev-host/grafana/dashboards/traces.json | 0 .../dev-host/grafana/grafana.ini | 0 .../provisioning/dashboards/dashboards.yaml | 0 .../provisioning/datasources/datasources.yaml | 0 .../dev-host/otel-collector/config.yaml | 0 .../dev-host/postgres/init-db.sh | 0 .../dev-host/prometheus/prometheus.yml | 0 .../dev-host/rivet-engine/config.jsonc | 12 +- .../dev-host/vector-client/vector.yaml | 0 .../dev-host/vector-server/vector.yaml | 0 .../dev-multidc-multinode/.gitattributes | 0 .../dev-multidc-multinode/README.md | 0 .../core/clickhouse/client-config.xml | 0 .../core/clickhouse/config.xml | 0 .../clickhouse/init/01-create-otel-table.sql | 0 .../core/clickhouse/users.xml | 0 .../core/grafana/dashboards/api.json | 0 .../core/grafana/dashboards/cache.json | 0 .../core/grafana/dashboards/epoxy.json | 0 .../core/grafana/dashboards/futures.json | 0 .../core/grafana/dashboards/gasoline.json | 0 .../core/grafana/dashboards/guard.json | 0 .../core/grafana/dashboards/operation.json | 0 .../core/grafana/dashboards/pegboard.json | 0 .../core/grafana/dashboards/tokio.json | 0 .../core/grafana/dashboards/traces.json | 0 .../core/grafana/grafana.ini | 0 .../provisioning/dashboards/dashboards.yaml | 0 .../provisioning/datasources/datasources.yaml | 0 .../core/prometheus/prometheus.yml | 0 .../dc-a/otel-collector/config.yaml | 0 .../datacenters/dc-a/postgres/init-db.sh | 0 .../dc-a/rivet-engine/0}/config.jsonc | 12 +- .../dc-a/rivet-engine/1}/config.jsonc | 12 +- .../dc-a/rivet-engine/2}/config.jsonc | 12 +- .../dc-a/vector-client/vector.yaml | 0 .../dc-a/vector-server/vector.yaml | 0 .../dc-b/otel-collector/config.yaml | 0 .../datacenters/dc-b/postgres/init-db.sh | 0 .../dc-b/rivet-engine/0}/config.jsonc | 12 +- .../dc-b/rivet-engine/1}/config.jsonc | 12 +- .../dc-b/rivet-engine/2/config.jsonc | 12 +- .../dc-b/vector-client/vector.yaml | 0 .../dc-b/vector-server/vector.yaml | 0 .../dc-c/otel-collector/config.yaml | 0 .../datacenters/dc-c/postgres/init-db.sh | 0 .../dc-c/rivet-engine/0/config.jsonc | 12 +- .../dc-c/rivet-engine/1/config.jsonc | 12 +- .../dc-c/rivet-engine/2/config.jsonc | 12 +- .../dc-c/vector-client/vector.yaml | 0 .../dc-c/vector-server/vector.yaml | 0 .../dev-multidc-multinode/docker-compose.yml | 216 ++++++++++++------ .../{compose => }/dev-multidc/.gitattributes | 0 self-host/{compose => }/dev-multidc/README.md | 0 .../core/clickhouse/client-config.xml | 0 .../dev-multidc/core/clickhouse/config.xml | 0 .../clickhouse/init/01-create-otel-table.sql | 0 .../dev-multidc/core/clickhouse/users.xml | 0 .../core/grafana/dashboards/api.json | 0 .../core/grafana/dashboards/cache.json | 0 .../core/grafana/dashboards/epoxy.json | 0 .../core/grafana/dashboards/futures.json | 0 .../core/grafana/dashboards/gasoline.json | 0 .../core/grafana/dashboards/guard.json | 0 .../core/grafana/dashboards/operation.json | 0 .../core/grafana/dashboards/pegboard.json | 0 .../core/grafana/dashboards/tokio.json | 0 .../core/grafana/dashboards/traces.json | 0 .../dev-multidc/core/grafana/grafana.ini | 0 .../provisioning/dashboards/dashboards.yaml | 0 .../provisioning/datasources/datasources.yaml | 0 .../core/prometheus/prometheus.yml | 0 .../dc-a/otel-collector/config.yaml | 0 .../datacenters/dc-a/postgres/init-db.sh | 0 .../dc-a/rivet-engine/config.jsonc | 12 +- .../dc-a/vector-client/vector.yaml | 0 .../dc-a/vector-server/vector.yaml | 0 .../dc-b/otel-collector/config.yaml | 0 .../datacenters/dc-b/postgres/init-db.sh | 0 .../dc-b/rivet-engine/config.jsonc | 12 +- .../dc-b/vector-client/vector.yaml | 0 .../dc-b/vector-server/vector.yaml | 0 .../dc-c/otel-collector/config.yaml | 0 .../datacenters/dc-c/postgres/init-db.sh | 0 .../dc-c/rivet-engine/config.jsonc | 12 +- .../dc-c/vector-client/vector.yaml | 0 .../dc-c/vector-server/vector.yaml | 0 .../dev-multidc/docker-compose.yml | 102 ++++++--- .../dev-multinode/.gitattributes | 0 .../{compose => }/dev-multinode/README.md | 0 .../clickhouse/client-config.xml | 0 .../dev-multinode/clickhouse/config.xml | 0 .../clickhouse/init/01-create-otel-table.sql | 0 .../dev-multinode/clickhouse/users.xml | 0 .../dev-multinode/docker-compose.yml | 85 +++++-- .../dev-multinode/grafana/dashboards/api.json | 0 .../grafana/dashboards/cache.json | 0 .../grafana/dashboards/epoxy.json | 0 .../grafana/dashboards/futures.json | 0 .../grafana/dashboards/gasoline.json | 0 .../grafana/dashboards/guard.json | 0 .../grafana/dashboards/operation.json | 0 .../grafana/dashboards/pegboard.json | 0 .../grafana/dashboards/tokio.json | 0 .../grafana/dashboards/traces.json | 0 .../dev-multinode/grafana/grafana.ini | 0 .../provisioning/dashboards/dashboards.yaml | 0 .../provisioning/datasources/datasources.yaml | 0 .../dev-multinode/otel-collector/config.yaml | 0 .../dev-multinode/postgres/init-db.sh | 0 .../dev-multinode/prometheus/prometheus.yml | 0 .../rivet-engine/0}/config.jsonc | 12 +- .../rivet-engine/1}/config.jsonc | 12 +- .../rivet-engine/2}/config.jsonc | 12 +- .../dev-multinode/vector-client/vector.yaml | 0 .../dev-multinode/vector-server/vector.yaml | 0 self-host/{compose => }/dev/.gitattributes | 0 self-host/{compose => }/dev/README.md | 0 .../dev/clickhouse/client-config.xml | 0 .../{compose => }/dev/clickhouse/config.xml | 0 .../clickhouse/init/01-create-otel-table.sql | 0 .../{compose => }/dev/clickhouse/users.xml | 0 .../{compose => }/dev/docker-compose.yml | 47 +++- .../dev/grafana/dashboards/api.json | 0 .../dev/grafana/dashboards/cache.json | 0 .../dev/grafana/dashboards/epoxy.json | 0 .../dev/grafana/dashboards/futures.json | 0 .../dev/grafana/dashboards/gasoline.json | 0 .../dev/grafana/dashboards/guard.json | 0 .../dev/grafana/dashboards/operation.json | 0 .../dev/grafana/dashboards/pegboard.json | 0 .../dev/grafana/dashboards/tokio.json | 0 .../dev/grafana/dashboards/traces.json | 0 .../{compose => }/dev/grafana/grafana.ini | 0 .../provisioning/dashboards/dashboards.yaml | 0 .../provisioning/datasources/datasources.yaml | 0 .../dev/otel-collector/config.yaml | 0 .../{compose => }/dev/postgres/init-db.sh | 0 .../dev/prometheus/prometheus.yml | 0 .../dev/rivet-engine/config.jsonc | 12 +- .../dev/vector-client/vector.yaml | 0 .../dev/vector-server/vector.yaml | 0 167 files changed, 468 insertions(+), 457 deletions(-) delete mode 100644 self-host/compose/prod-file-system/.gitattributes delete mode 100644 self-host/compose/prod-file-system/README.md delete mode 100644 self-host/compose/prod-file-system/docker-compose.yml delete mode 100644 self-host/compose/prod-file-system/rivet-engine/config.jsonc rename self-host/{compose => }/dev-host/.gitattributes (100%) rename self-host/{compose => }/dev-host/README.md (100%) rename self-host/{compose => }/dev-host/clickhouse/client-config.xml (100%) rename self-host/{compose => }/dev-host/clickhouse/config.xml (100%) rename self-host/{compose => }/dev-host/clickhouse/init/01-create-otel-table.sql (100%) rename self-host/{compose => }/dev-host/clickhouse/users.xml (100%) rename self-host/{compose => }/dev-host/docker-compose.yml (83%) rename self-host/{compose => }/dev-host/grafana/dashboards/api.json (100%) rename self-host/{compose => }/dev-host/grafana/dashboards/cache.json (100%) rename self-host/{compose => }/dev-host/grafana/dashboards/epoxy.json (100%) rename self-host/{compose => }/dev-host/grafana/dashboards/futures.json (100%) rename self-host/{compose => }/dev-host/grafana/dashboards/gasoline.json (100%) rename self-host/{compose => }/dev-host/grafana/dashboards/guard.json (100%) rename self-host/{compose => }/dev-host/grafana/dashboards/operation.json (100%) rename self-host/{compose => }/dev-host/grafana/dashboards/pegboard.json (100%) rename self-host/{compose => }/dev-host/grafana/dashboards/tokio.json (100%) rename self-host/{compose => }/dev-host/grafana/dashboards/traces.json (100%) rename self-host/{compose => }/dev-host/grafana/grafana.ini (100%) rename self-host/{compose => }/dev-host/grafana/provisioning/dashboards/dashboards.yaml (100%) rename self-host/{compose => }/dev-host/grafana/provisioning/datasources/datasources.yaml (100%) rename self-host/{compose => }/dev-host/otel-collector/config.yaml (100%) rename self-host/{compose => }/dev-host/postgres/init-db.sh (100%) rename self-host/{compose => }/dev-host/prometheus/prometheus.yml (100%) rename self-host/{compose => }/dev-host/rivet-engine/config.jsonc (79%) rename self-host/{compose => }/dev-host/vector-client/vector.yaml (100%) rename self-host/{compose => }/dev-host/vector-server/vector.yaml (100%) rename self-host/{compose => }/dev-multidc-multinode/.gitattributes (100%) rename self-host/{compose => }/dev-multidc-multinode/README.md (100%) rename self-host/{compose => }/dev-multidc-multinode/core/clickhouse/client-config.xml (100%) rename self-host/{compose => }/dev-multidc-multinode/core/clickhouse/config.xml (100%) rename self-host/{compose => }/dev-multidc-multinode/core/clickhouse/init/01-create-otel-table.sql (100%) rename self-host/{compose => }/dev-multidc-multinode/core/clickhouse/users.xml (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/dashboards/api.json (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/dashboards/cache.json (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/dashboards/epoxy.json (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/dashboards/futures.json (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/dashboards/gasoline.json (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/dashboards/guard.json (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/dashboards/operation.json (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/dashboards/pegboard.json (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/dashboards/tokio.json (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/dashboards/traces.json (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/grafana.ini (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/provisioning/dashboards/dashboards.yaml (100%) rename self-host/{compose => }/dev-multidc-multinode/core/grafana/provisioning/datasources/datasources.yaml (100%) rename self-host/{compose => }/dev-multidc-multinode/core/prometheus/prometheus.yml (100%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-a/otel-collector/config.yaml (100%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-a/postgres/init-db.sh (100%) rename self-host/{compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2 => dev-multidc-multinode/datacenters/dc-a/rivet-engine/0}/config.jsonc (88%) rename self-host/{compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0 => dev-multidc-multinode/datacenters/dc-a/rivet-engine/1}/config.jsonc (88%) rename self-host/{compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1 => dev-multidc-multinode/datacenters/dc-a/rivet-engine/2}/config.jsonc (88%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-a/vector-client/vector.yaml (100%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-a/vector-server/vector.yaml (100%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-b/otel-collector/config.yaml (100%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-b/postgres/init-db.sh (100%) rename self-host/{compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1 => dev-multidc-multinode/datacenters/dc-b/rivet-engine/0}/config.jsonc (88%) rename self-host/{compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0 => dev-multidc-multinode/datacenters/dc-b/rivet-engine/1}/config.jsonc (88%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc (88%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-b/vector-client/vector.yaml (100%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-b/vector-server/vector.yaml (100%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-c/otel-collector/config.yaml (100%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-c/postgres/init-db.sh (100%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc (88%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc (88%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc (88%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-c/vector-client/vector.yaml (100%) rename self-host/{compose => }/dev-multidc-multinode/datacenters/dc-c/vector-server/vector.yaml (100%) rename self-host/{compose => }/dev-multidc-multinode/docker-compose.yml (84%) rename self-host/{compose => }/dev-multidc/.gitattributes (100%) rename self-host/{compose => }/dev-multidc/README.md (100%) rename self-host/{compose => }/dev-multidc/core/clickhouse/client-config.xml (100%) rename self-host/{compose => }/dev-multidc/core/clickhouse/config.xml (100%) rename self-host/{compose => }/dev-multidc/core/clickhouse/init/01-create-otel-table.sql (100%) rename self-host/{compose => }/dev-multidc/core/clickhouse/users.xml (100%) rename self-host/{compose => }/dev-multidc/core/grafana/dashboards/api.json (100%) rename self-host/{compose => }/dev-multidc/core/grafana/dashboards/cache.json (100%) rename self-host/{compose => }/dev-multidc/core/grafana/dashboards/epoxy.json (100%) rename self-host/{compose => }/dev-multidc/core/grafana/dashboards/futures.json (100%) rename self-host/{compose => }/dev-multidc/core/grafana/dashboards/gasoline.json (100%) rename self-host/{compose => }/dev-multidc/core/grafana/dashboards/guard.json (100%) rename self-host/{compose => }/dev-multidc/core/grafana/dashboards/operation.json (100%) rename self-host/{compose => }/dev-multidc/core/grafana/dashboards/pegboard.json (100%) rename self-host/{compose => }/dev-multidc/core/grafana/dashboards/tokio.json (100%) rename self-host/{compose => }/dev-multidc/core/grafana/dashboards/traces.json (100%) rename self-host/{compose => }/dev-multidc/core/grafana/grafana.ini (100%) rename self-host/{compose => }/dev-multidc/core/grafana/provisioning/dashboards/dashboards.yaml (100%) rename self-host/{compose => }/dev-multidc/core/grafana/provisioning/datasources/datasources.yaml (100%) rename self-host/{compose => }/dev-multidc/core/prometheus/prometheus.yml (100%) rename self-host/{compose => }/dev-multidc/datacenters/dc-a/otel-collector/config.yaml (100%) rename self-host/{compose => }/dev-multidc/datacenters/dc-a/postgres/init-db.sh (100%) rename self-host/{compose => }/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc (88%) rename self-host/{compose => }/dev-multidc/datacenters/dc-a/vector-client/vector.yaml (100%) rename self-host/{compose => }/dev-multidc/datacenters/dc-a/vector-server/vector.yaml (100%) rename self-host/{compose => }/dev-multidc/datacenters/dc-b/otel-collector/config.yaml (100%) rename self-host/{compose => }/dev-multidc/datacenters/dc-b/postgres/init-db.sh (100%) rename self-host/{compose => }/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc (88%) rename self-host/{compose => }/dev-multidc/datacenters/dc-b/vector-client/vector.yaml (100%) rename self-host/{compose => }/dev-multidc/datacenters/dc-b/vector-server/vector.yaml (100%) rename self-host/{compose => }/dev-multidc/datacenters/dc-c/otel-collector/config.yaml (100%) rename self-host/{compose => }/dev-multidc/datacenters/dc-c/postgres/init-db.sh (100%) rename self-host/{compose => }/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc (88%) rename self-host/{compose => }/dev-multidc/datacenters/dc-c/vector-client/vector.yaml (100%) rename self-host/{compose => }/dev-multidc/datacenters/dc-c/vector-server/vector.yaml (100%) rename self-host/{compose => }/dev-multidc/docker-compose.yml (87%) rename self-host/{compose => }/dev-multinode/.gitattributes (100%) rename self-host/{compose => }/dev-multinode/README.md (100%) rename self-host/{compose => }/dev-multinode/clickhouse/client-config.xml (100%) rename self-host/{compose => }/dev-multinode/clickhouse/config.xml (100%) rename self-host/{compose => }/dev-multinode/clickhouse/init/01-create-otel-table.sql (100%) rename self-host/{compose => }/dev-multinode/clickhouse/users.xml (100%) rename self-host/{compose => }/dev-multinode/docker-compose.yml (82%) rename self-host/{compose => }/dev-multinode/grafana/dashboards/api.json (100%) rename self-host/{compose => }/dev-multinode/grafana/dashboards/cache.json (100%) rename self-host/{compose => }/dev-multinode/grafana/dashboards/epoxy.json (100%) rename self-host/{compose => }/dev-multinode/grafana/dashboards/futures.json (100%) rename self-host/{compose => }/dev-multinode/grafana/dashboards/gasoline.json (100%) rename self-host/{compose => }/dev-multinode/grafana/dashboards/guard.json (100%) rename self-host/{compose => }/dev-multinode/grafana/dashboards/operation.json (100%) rename self-host/{compose => }/dev-multinode/grafana/dashboards/pegboard.json (100%) rename self-host/{compose => }/dev-multinode/grafana/dashboards/tokio.json (100%) rename self-host/{compose => }/dev-multinode/grafana/dashboards/traces.json (100%) rename self-host/{compose => }/dev-multinode/grafana/grafana.ini (100%) rename self-host/{compose => }/dev-multinode/grafana/provisioning/dashboards/dashboards.yaml (100%) rename self-host/{compose => }/dev-multinode/grafana/provisioning/datasources/datasources.yaml (100%) rename self-host/{compose => }/dev-multinode/otel-collector/config.yaml (100%) rename self-host/{compose => }/dev-multinode/postgres/init-db.sh (100%) rename self-host/{compose => }/dev-multinode/prometheus/prometheus.yml (100%) rename self-host/{compose/dev-multinode/rivet-engine/2 => dev-multinode/rivet-engine/0}/config.jsonc (79%) rename self-host/{compose/dev-multinode/rivet-engine/0 => dev-multinode/rivet-engine/1}/config.jsonc (79%) rename self-host/{compose/dev-multinode/rivet-engine/1 => dev-multinode/rivet-engine/2}/config.jsonc (79%) rename self-host/{compose => }/dev-multinode/vector-client/vector.yaml (100%) rename self-host/{compose => }/dev-multinode/vector-server/vector.yaml (100%) rename self-host/{compose => }/dev/.gitattributes (100%) rename self-host/{compose => }/dev/README.md (100%) rename self-host/{compose => }/dev/clickhouse/client-config.xml (100%) rename self-host/{compose => }/dev/clickhouse/config.xml (100%) rename self-host/{compose => }/dev/clickhouse/init/01-create-otel-table.sql (100%) rename self-host/{compose => }/dev/clickhouse/users.xml (100%) rename self-host/{compose => }/dev/docker-compose.yml (84%) rename self-host/{compose => }/dev/grafana/dashboards/api.json (100%) rename self-host/{compose => }/dev/grafana/dashboards/cache.json (100%) rename self-host/{compose => }/dev/grafana/dashboards/epoxy.json (100%) rename self-host/{compose => }/dev/grafana/dashboards/futures.json (100%) rename self-host/{compose => }/dev/grafana/dashboards/gasoline.json (100%) rename self-host/{compose => }/dev/grafana/dashboards/guard.json (100%) rename self-host/{compose => }/dev/grafana/dashboards/operation.json (100%) rename self-host/{compose => }/dev/grafana/dashboards/pegboard.json (100%) rename self-host/{compose => }/dev/grafana/dashboards/tokio.json (100%) rename self-host/{compose => }/dev/grafana/dashboards/traces.json (100%) rename self-host/{compose => }/dev/grafana/grafana.ini (100%) rename self-host/{compose => }/dev/grafana/provisioning/dashboards/dashboards.yaml (100%) rename self-host/{compose => }/dev/grafana/provisioning/datasources/datasources.yaml (100%) rename self-host/{compose => }/dev/otel-collector/config.yaml (100%) rename self-host/{compose => }/dev/postgres/init-db.sh (100%) rename self-host/{compose => }/dev/prometheus/prometheus.yml (100%) rename self-host/{compose => }/dev/rivet-engine/config.jsonc (79%) rename self-host/{compose => }/dev/vector-client/vector.yaml (100%) rename self-host/{compose => }/dev/vector-server/vector.yaml (100%) diff --git a/engine/packages/metrics/src/registry.rs b/engine/packages/metrics/src/registry.rs index 880cc495ae..c6bd268591 100644 --- a/engine/packages/metrics/src/registry.rs +++ b/engine/packages/metrics/src/registry.rs @@ -1,5 +1,7 @@ use prometheus::*; lazy_static::lazy_static! { - pub static ref REGISTRY: Registry = Registry::new_custom(None, Some(labels! { })).unwrap(); + pub static ref REGISTRY: Registry = Registry::new_custom( + Some("rivet".to_string()), + Some(labels! { })).unwrap(); } diff --git a/examples/kitchen-sink/Dockerfile b/examples/kitchen-sink/Dockerfile index e1a6f2a3d8..21377a6c9d 100644 --- a/examples/kitchen-sink/Dockerfile +++ b/examples/kitchen-sink/Dockerfile @@ -1,4 +1,9 @@ -FROM node:22-slim +# Base image is overridable so local builds can match the host glibc. The cloud +# build keeps the node:22-slim default and supplies a glibc-compatible napi +# binary; a modern dev host (newer glibc) should override this with a newer base +# such as node:22-trixie-slim so the host-built napi binary loads. +ARG NODE_IMAGE=node:22-slim +FROM ${NODE_IMAGE} RUN corepack enable && corepack prepare pnpm@10.13.1 --activate WORKDIR /app ENV NODE_OPTIONS=--max-old-space-size=7168 diff --git a/self-host/compose/prod-file-system/.gitattributes b/self-host/compose/prod-file-system/.gitattributes deleted file mode 100644 index 447edeb5c2..0000000000 --- a/self-host/compose/prod-file-system/.gitattributes +++ /dev/null @@ -1 +0,0 @@ -. linguist-generated=true diff --git a/self-host/compose/prod-file-system/README.md b/self-host/compose/prod-file-system/README.md deleted file mode 100644 index 4b02525eeb..0000000000 --- a/self-host/compose/prod-file-system/README.md +++ /dev/null @@ -1,76 +0,0 @@ -# dev - Auto-generated Docker Compose Template - -> ! **Auto-generated**: This directory and its contents are automatically generated by `docker/template/`. Do not edit these files directly as your changes will be overwritten. - -## Overview - -This Docker Compose configuration provides a complete development environment for Rivet with the following services: - -- **Rivet Engine**: Main orchestration service -- **Rivet Shell**: Interactive shell for debugging -- **Runner**: Executes user code -- **ClickHouse**: Analytics and time-series database -- **NATS**: Message broker -- **PostgreSQL**: Relational database -- **Vector Server**: Log aggregation and processing -- **OpenTelemetry Collector**: Observability data collection - -## Port Configuration - -| Service | Port(s) | Description | -|---------|---------|-------------| -| Rivet Engine | 6420 | Public endpoint | -| Runner | 5050 | Code execution service | -| NATS | 4222 | Message broker | -| PostgreSQL | 5432 | Database | -| ClickHouse HTTP | 9300 | Database HTTP interface | -| ClickHouse Native | 9301 | Database native protocol | -| OpenTelemetry gRPC | 4317 | OTLP gRPC endpoint | -| OpenTelemetry HTTP | 4318 | OTLP HTTP endpoint | - -## Template Configuration - -**Template Name**: `dev` -**Base Port**: `6420` -**Network Mode**: `bridge` - -### Datacenters -- **1**: 1 engine(s), 1 runner(s) - -## Usage - -1. Start all services: - ```bash - docker-compose up -d - ``` - -2. Check service health: - ```bash - docker-compose ps - ``` - -3. View logs: - ```bash - docker-compose logs -f [service-name] - ``` - -4. Stop all services: - ```bash - docker-compose down - ``` - -## Generated Files - -This template generates the following files and directories: -- `docker-compose.yml` - Main Docker Compose configuration -- `core/` - Core services shared across datacenters: - - `clickhouse/` - ClickHouse configuration and initialization - - `vector-server/` - Vector aggregator configuration - - `otel-collector-server/` - OpenTelemetry Collector server configuration -- `datacenters/` - Datacenter-specific configurations: - - `1/` - Configuration for datacenter 1: - - `postgres/` - PostgreSQL setup scripts - - `rivet-engine/` - Rivet Engine configuration - - `vector-client/` - Vector client configuration - - `otel-collector-client/` - OpenTelemetry Collector client configuration -- `README.md` - This file diff --git a/self-host/compose/prod-file-system/docker-compose.yml b/self-host/compose/prod-file-system/docker-compose.yml deleted file mode 100644 index 2e4a315c33..0000000000 --- a/self-host/compose/prod-file-system/docker-compose.yml +++ /dev/null @@ -1,41 +0,0 @@ -services: - rivet-engine: - build: - context: ../../.. - dockerfile: docker/engine/Dockerfile - target: engine-full - restart: unless-stopped - command: /usr/bin/rivet-engine start - environment: - - RIVET__FILE_SYSTEM__PATH=/var/lib/rivet-engine - volumes: - - ./rivet-engine/config.jsonc:/etc/rivet/config.jsonc:ro - - rivet-engine-data:/var/lib/rivet-engine - ports: - - '6420:6420' - healthcheck: - test: - - CMD - - curl - - '-f' - - http://127.0.0.1:6421/health - interval: 2s - timeout: 10s - retries: 10 - start_period: 30s - runner: - build: - context: ../.. - dockerfile: docker/runner/Dockerfile - platform: linux/amd64 - restart: unless-stopped - environment: - - RIVET_ENDPOINT=http://rivet-engine:6420 - stop_grace_period: 4s - ports: - - '5050:5050' - depends_on: - rivet-engine: - condition: service_healthy -volumes: - rivet-engine-data: diff --git a/self-host/compose/prod-file-system/rivet-engine/config.jsonc b/self-host/compose/prod-file-system/rivet-engine/config.jsonc deleted file mode 100644 index 2c63c08510..0000000000 --- a/self-host/compose/prod-file-system/rivet-engine/config.jsonc +++ /dev/null @@ -1,2 +0,0 @@ -{ -} diff --git a/self-host/compose/template/src/docker-compose.ts b/self-host/compose/template/src/docker-compose.ts index 6db4a8ab9a..4e3781aab2 100644 --- a/self-host/compose/template/src/docker-compose.ts +++ b/self-host/compose/template/src/docker-compose.ts @@ -1,6 +1,8 @@ import * as yaml from "js-yaml"; import { CORE_NETWORK_NAME, type TemplateContext } from "./context"; +const RUNNER_CONFIG_INIT_SERVICE = "runner-config-init"; + export function generateDockerCompose(context: TemplateContext) { const config = context.config; @@ -172,6 +174,10 @@ export function generateDockerCompose(context: TemplateContext) { services[postgresServiceName] = { restart: "unless-stopped", image: "postgres:18-alpine", + // Each engine opens a UDB connection pool (up to 64 connections) plus a + // dedicated LISTEN connection and pubsub, so a multi-engine datacenter + // needs far more than the default max_connections of 100. + command: ["postgres", "-c", "max_connections=500"], environment: [ "POSTGRES_USER=postgres", "POSTGRES_PASSWORD=postgres", @@ -179,7 +185,7 @@ export function generateDockerCompose(context: TemplateContext) { ], volumes: [ `./${context.getDatacenterServicePath("postgres", datacenter.name)}/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh`, - `${postgresVolumeName}:/var/lib/postgresql/data`, + `${postgresVolumeName}:/var/lib/postgresql`, ], ports: isPrimary ? [`5432:5432`] : undefined, healthcheck: { @@ -194,7 +200,7 @@ export function generateDockerCompose(context: TemplateContext) { services[shellServiceName] = { build: { - context: "../../..", + context: "../..", dockerfile: "docker/engine/Dockerfile", target: "engine-full", args: { @@ -275,7 +281,7 @@ export function generateDockerCompose(context: TemplateContext) { services[serviceName] = { build: { - context: "../../..", + context: "../..", dockerfile: "docker/engine/Dockerfile", target: "engine-full", args: { @@ -335,30 +341,79 @@ export function generateDockerCompose(context: TemplateContext) { services[serviceName] = { build: { - context: "../../..", - dockerfile: "engine/sdks/rust/test-envoy/Dockerfile", + context: "../..", + dockerfile: "examples/kitchen-sink/Dockerfile", + // The runner copies the host-built napi binary, so the base + // image must have a glibc at least as new as the build host. + args: { + NODE_IMAGE: "node:22-trixie-slim", + }, }, platform: "linux/amd64", restart: "unless-stopped", environment: [ + `RIVET_KITCHEN_SINK_MODE=serverful`, `RIVET_ENDPOINT=http://${context.getServiceHost("rivet-engine", datacenter.name, 0)}:6420`, - `INTERNAL_SERVER_PORT=5050`, - `RIVET_POOL_NAME=test-envoy`, - `AUTOSTART_ENVOY=1`, - `AUTOCONFIGURE_SERVERLESS=0` + `RIVET_TOKEN=dev`, + `RIVET_NAMESPACE=default`, + `RIVET_POOL=default`, + `PORT=8080`, ], stop_grace_period: "4s", - ports: isPrimary && i === 0 ? [`5050:5050`] : undefined, + ports: isPrimary && i === 0 ? [`5050:8080`] : undefined, depends_on: { [engineServiceName]: { condition: "service_healthy", }, + [RUNNER_CONFIG_INIT_SERVICE]: { + condition: "service_completed_successfully", + }, }, networks: [dcNetworkName], }; } }); + // Serverful runners are rejected with `no_runner_config` until a runner + // config exists for their pool, so a one-shot init container upserts a + // `normal` runner config for the `default` pool across every datacenter once + // the leader engine is healthy. The runners wait for this to finish. + const primaryDc = config.datacenters[0]; + const primaryEngineHost = context.getServiceHost( + "rivet-engine", + primaryDc.name, + 0, + ); + const primaryEngineService = context.getServiceName( + "rivet-engine", + primaryDc.name, + 0, + ); + const runnerConfigBody = JSON.stringify({ + datacenters: Object.fromEntries( + config.datacenters.map((dc) => [dc.name, { normal: {} }]), + ), + }); + const runnerConfigScript = [ + `until curl -fsS -X PUT`, + `"http://${primaryEngineHost}:6420/runner-configs/default?namespace=default"`, + `-H "Authorization: Bearer dev"`, + `-H "Content-Type: application/json"`, + `-d '${runnerConfigBody}';`, + `do echo "waiting for engine to accept runner config"; sleep 2; done;`, + `echo "runner config upserted"`, + ].join(" "); + services[RUNNER_CONFIG_INIT_SERVICE] = { + image: "curlimages/curl:latest", + restart: "no", + depends_on: { + [primaryEngineService]: { condition: "service_healthy" }, + }, + entrypoint: ["sh", "-c"], + command: [runnerConfigScript], + networks: [context.getDatacenterNetworkName(primaryDc.name)], + }; + const dockerComposeConfig = { services, networks, diff --git a/self-host/compose/template/src/services/edge/rivet-engine.ts b/self-host/compose/template/src/services/edge/rivet-engine.ts index 25df2ce5ca..52d33dfc1b 100644 --- a/self-host/compose/template/src/services/edge/rivet-engine.ts +++ b/self-host/compose/template/src/services/edge/rivet-engine.ts @@ -35,32 +35,24 @@ export function generateDatacenterRivetEngine( datacenters, }; - // Config structure matching Rust schema in packages/common/config/src/config/mod.rs + // Config structure matching Rust schema in engine/packages/config/src/config/mod.rs. + // Values that match the engine's defaults are omitted. const config = { auth: { admin_token: "dev", }, - guard: { - port: GUARD_PORT, - // https is optional and not configured for local development - }, api_peer: { host: "0.0.0.0", - port: API_PEER_PORT, }, topology, postgres: { url: `postgresql://postgres:postgres@${context.getServiceHost("postgres", datacenter.name)}:5432/rivet_engine`, }, - cache: { - driver: "in_memory", - }, clickhouse: { - http_url: `http://${clickhouseHost}:9300`, // TODO: - native_url: `http://${clickhouseHost}:9301`, // TODO: + http_url: `http://${clickhouseHost}:9300`, + native_url: `http://${clickhouseHost}:9301`, username: "system", password: "default", - secure: false, }, }; diff --git a/self-host/compose/template/src/services/edge/runner.ts b/self-host/compose/template/src/services/edge/runner.ts index c04f936dd3..9a111e637c 100644 --- a/self-host/compose/template/src/services/edge/runner.ts +++ b/self-host/compose/template/src/services/edge/runner.ts @@ -1,6 +1,7 @@ import type { TemplateContext } from "../../context"; export function generateRunner(context: TemplateContext) { - // The test runner service now uses the Rust test-envoy binary. - // The docker-compose template points at the Rust Dockerfile directly. + // The runner service runs the kitchen-sink example in serverful mode, + // connecting to the engine as a long-lived runner. The docker-compose + // template builds examples/kitchen-sink/Dockerfile directly. } diff --git a/self-host/compose/dev-host/.gitattributes b/self-host/dev-host/.gitattributes similarity index 100% rename from self-host/compose/dev-host/.gitattributes rename to self-host/dev-host/.gitattributes diff --git a/self-host/compose/dev-host/README.md b/self-host/dev-host/README.md similarity index 100% rename from self-host/compose/dev-host/README.md rename to self-host/dev-host/README.md diff --git a/self-host/compose/dev-host/clickhouse/client-config.xml b/self-host/dev-host/clickhouse/client-config.xml similarity index 100% rename from self-host/compose/dev-host/clickhouse/client-config.xml rename to self-host/dev-host/clickhouse/client-config.xml diff --git a/self-host/compose/dev-host/clickhouse/config.xml b/self-host/dev-host/clickhouse/config.xml similarity index 100% rename from self-host/compose/dev-host/clickhouse/config.xml rename to self-host/dev-host/clickhouse/config.xml diff --git a/self-host/compose/dev-host/clickhouse/init/01-create-otel-table.sql b/self-host/dev-host/clickhouse/init/01-create-otel-table.sql similarity index 100% rename from self-host/compose/dev-host/clickhouse/init/01-create-otel-table.sql rename to self-host/dev-host/clickhouse/init/01-create-otel-table.sql diff --git a/self-host/compose/dev-host/clickhouse/users.xml b/self-host/dev-host/clickhouse/users.xml similarity index 100% rename from self-host/compose/dev-host/clickhouse/users.xml rename to self-host/dev-host/clickhouse/users.xml diff --git a/self-host/compose/dev-host/docker-compose.yml b/self-host/dev-host/docker-compose.yml similarity index 83% rename from self-host/compose/dev-host/docker-compose.yml rename to self-host/dev-host/docker-compose.yml index 96db5a57d0..94619e9324 100644 --- a/self-host/compose/dev-host/docker-compose.yml +++ b/self-host/dev-host/docker-compose.yml @@ -73,13 +73,17 @@ services: postgres: restart: unless-stopped image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres - POSTGRES_DB=postgres volumes: - ./postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data:/var/lib/postgresql/data + - postgres-data:/var/lib/postgresql healthcheck: test: - CMD-SHELL @@ -90,7 +94,7 @@ services: network_mode: host rivet-shell: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -144,7 +148,7 @@ services: network_mode: host rivet-engine: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -179,20 +183,42 @@ services: network_mode: host runner: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://127.0.0.1:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine: condition: service_healthy + runner-config-init: + condition: service_completed_successfully + network_mode: host + runner-config-init: + image: curlimages/curl:latest + restart: 'no' + depends_on: + rivet-engine: + condition: service_healthy + entrypoint: + - sh + - '-c' + command: + - >- + until curl -fsS -X PUT + "http://127.0.0.1:6420/runner-configs/default?namespace=default" -H + "Authorization: Bearer dev" -H "Content-Type: application/json" -d + '{"datacenters":{"default":{"normal":{}}}}'; do echo "waiting for engine + to accept runner config"; sleep 2; done; echo "runner config upserted" network_mode: host networks: rivet-core-network: diff --git a/self-host/compose/dev-host/grafana/dashboards/api.json b/self-host/dev-host/grafana/dashboards/api.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/api.json rename to self-host/dev-host/grafana/dashboards/api.json diff --git a/self-host/compose/dev-host/grafana/dashboards/cache.json b/self-host/dev-host/grafana/dashboards/cache.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/cache.json rename to self-host/dev-host/grafana/dashboards/cache.json diff --git a/self-host/compose/dev-host/grafana/dashboards/epoxy.json b/self-host/dev-host/grafana/dashboards/epoxy.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/epoxy.json rename to self-host/dev-host/grafana/dashboards/epoxy.json diff --git a/self-host/compose/dev-host/grafana/dashboards/futures.json b/self-host/dev-host/grafana/dashboards/futures.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/futures.json rename to self-host/dev-host/grafana/dashboards/futures.json diff --git a/self-host/compose/dev-host/grafana/dashboards/gasoline.json b/self-host/dev-host/grafana/dashboards/gasoline.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/gasoline.json rename to self-host/dev-host/grafana/dashboards/gasoline.json diff --git a/self-host/compose/dev-host/grafana/dashboards/guard.json b/self-host/dev-host/grafana/dashboards/guard.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/guard.json rename to self-host/dev-host/grafana/dashboards/guard.json diff --git a/self-host/compose/dev-host/grafana/dashboards/operation.json b/self-host/dev-host/grafana/dashboards/operation.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/operation.json rename to self-host/dev-host/grafana/dashboards/operation.json diff --git a/self-host/compose/dev-host/grafana/dashboards/pegboard.json b/self-host/dev-host/grafana/dashboards/pegboard.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/pegboard.json rename to self-host/dev-host/grafana/dashboards/pegboard.json diff --git a/self-host/compose/dev-host/grafana/dashboards/tokio.json b/self-host/dev-host/grafana/dashboards/tokio.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/tokio.json rename to self-host/dev-host/grafana/dashboards/tokio.json diff --git a/self-host/compose/dev-host/grafana/dashboards/traces.json b/self-host/dev-host/grafana/dashboards/traces.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/traces.json rename to self-host/dev-host/grafana/dashboards/traces.json diff --git a/self-host/compose/dev-host/grafana/grafana.ini b/self-host/dev-host/grafana/grafana.ini similarity index 100% rename from self-host/compose/dev-host/grafana/grafana.ini rename to self-host/dev-host/grafana/grafana.ini diff --git a/self-host/compose/dev-host/grafana/provisioning/dashboards/dashboards.yaml b/self-host/dev-host/grafana/provisioning/dashboards/dashboards.yaml similarity index 100% rename from self-host/compose/dev-host/grafana/provisioning/dashboards/dashboards.yaml rename to self-host/dev-host/grafana/provisioning/dashboards/dashboards.yaml diff --git a/self-host/compose/dev-host/grafana/provisioning/datasources/datasources.yaml b/self-host/dev-host/grafana/provisioning/datasources/datasources.yaml similarity index 100% rename from self-host/compose/dev-host/grafana/provisioning/datasources/datasources.yaml rename to self-host/dev-host/grafana/provisioning/datasources/datasources.yaml diff --git a/self-host/compose/dev-host/otel-collector/config.yaml b/self-host/dev-host/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-host/otel-collector/config.yaml rename to self-host/dev-host/otel-collector/config.yaml diff --git a/self-host/compose/dev-host/postgres/init-db.sh b/self-host/dev-host/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-host/postgres/init-db.sh rename to self-host/dev-host/postgres/init-db.sh diff --git a/self-host/compose/dev-host/prometheus/prometheus.yml b/self-host/dev-host/prometheus/prometheus.yml similarity index 100% rename from self-host/compose/dev-host/prometheus/prometheus.yml rename to self-host/dev-host/prometheus/prometheus.yml diff --git a/self-host/compose/dev-host/rivet-engine/config.jsonc b/self-host/dev-host/rivet-engine/config.jsonc similarity index 79% rename from self-host/compose/dev-host/rivet-engine/config.jsonc rename to self-host/dev-host/rivet-engine/config.jsonc index d5a3095269..ee79af406a 100644 --- a/self-host/compose/dev-host/rivet-engine/config.jsonc +++ b/self-host/dev-host/rivet-engine/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -28,14 +24,10 @@ "postgres": { "url": "postgresql://postgres:postgres@127.0.0.1:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://127.0.0.1:9300", "native_url": "http://127.0.0.1:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-host/vector-client/vector.yaml b/self-host/dev-host/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-host/vector-client/vector.yaml rename to self-host/dev-host/vector-client/vector.yaml diff --git a/self-host/compose/dev-host/vector-server/vector.yaml b/self-host/dev-host/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-host/vector-server/vector.yaml rename to self-host/dev-host/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/.gitattributes b/self-host/dev-multidc-multinode/.gitattributes similarity index 100% rename from self-host/compose/dev-multidc-multinode/.gitattributes rename to self-host/dev-multidc-multinode/.gitattributes diff --git a/self-host/compose/dev-multidc-multinode/README.md b/self-host/dev-multidc-multinode/README.md similarity index 100% rename from self-host/compose/dev-multidc-multinode/README.md rename to self-host/dev-multidc-multinode/README.md diff --git a/self-host/compose/dev-multidc-multinode/core/clickhouse/client-config.xml b/self-host/dev-multidc-multinode/core/clickhouse/client-config.xml similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/clickhouse/client-config.xml rename to self-host/dev-multidc-multinode/core/clickhouse/client-config.xml diff --git a/self-host/compose/dev-multidc-multinode/core/clickhouse/config.xml b/self-host/dev-multidc-multinode/core/clickhouse/config.xml similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/clickhouse/config.xml rename to self-host/dev-multidc-multinode/core/clickhouse/config.xml diff --git a/self-host/compose/dev-multidc-multinode/core/clickhouse/init/01-create-otel-table.sql b/self-host/dev-multidc-multinode/core/clickhouse/init/01-create-otel-table.sql similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/clickhouse/init/01-create-otel-table.sql rename to self-host/dev-multidc-multinode/core/clickhouse/init/01-create-otel-table.sql diff --git a/self-host/compose/dev-multidc-multinode/core/clickhouse/users.xml b/self-host/dev-multidc-multinode/core/clickhouse/users.xml similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/clickhouse/users.xml rename to self-host/dev-multidc-multinode/core/clickhouse/users.xml diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/api.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/api.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/api.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/api.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/cache.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/cache.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/cache.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/cache.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/epoxy.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/epoxy.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/epoxy.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/epoxy.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/futures.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/futures.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/futures.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/futures.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/gasoline.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/gasoline.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/gasoline.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/gasoline.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/guard.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/guard.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/guard.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/guard.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/operation.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/operation.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/operation.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/operation.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/pegboard.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/pegboard.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/pegboard.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/pegboard.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/tokio.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/tokio.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/tokio.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/tokio.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/traces.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/traces.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/traces.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/traces.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/grafana.ini b/self-host/dev-multidc-multinode/core/grafana/grafana.ini similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/grafana.ini rename to self-host/dev-multidc-multinode/core/grafana/grafana.ini diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/provisioning/dashboards/dashboards.yaml b/self-host/dev-multidc-multinode/core/grafana/provisioning/dashboards/dashboards.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/provisioning/dashboards/dashboards.yaml rename to self-host/dev-multidc-multinode/core/grafana/provisioning/dashboards/dashboards.yaml diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/provisioning/datasources/datasources.yaml b/self-host/dev-multidc-multinode/core/grafana/provisioning/datasources/datasources.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/provisioning/datasources/datasources.yaml rename to self-host/dev-multidc-multinode/core/grafana/provisioning/datasources/datasources.yaml diff --git a/self-host/compose/dev-multidc-multinode/core/prometheus/prometheus.yml b/self-host/dev-multidc-multinode/core/prometheus/prometheus.yml similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/prometheus/prometheus.yml rename to self-host/dev-multidc-multinode/core/prometheus/prometheus.yml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/otel-collector/config.yaml b/self-host/dev-multidc-multinode/datacenters/dc-a/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-a/otel-collector/config.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-a/otel-collector/config.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/postgres/init-db.sh b/self-host/dev-multidc-multinode/datacenters/dc-a/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-a/postgres/init-db.sh rename to self-host/dev-multidc-multinode/datacenters/dc-a/postgres/init-db.sh diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc index ab9cfa618a..0532bce5c3 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-a:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc index ab9cfa618a..0532bce5c3 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-a:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc index ab9cfa618a..0532bce5c3 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-a:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/vector-client/vector.yaml b/self-host/dev-multidc-multinode/datacenters/dc-a/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-a/vector-client/vector.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-a/vector-client/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/vector-server/vector.yaml b/self-host/dev-multidc-multinode/datacenters/dc-a/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-a/vector-server/vector.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-a/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/otel-collector/config.yaml b/self-host/dev-multidc-multinode/datacenters/dc-b/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/otel-collector/config.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-b/otel-collector/config.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/postgres/init-db.sh b/self-host/dev-multidc-multinode/datacenters/dc-b/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/postgres/init-db.sh rename to self-host/dev-multidc-multinode/datacenters/dc-b/postgres/init-db.sh diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc index 2a34bc83ff..e75a22f77c 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 2, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-b:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc index 2a34bc83ff..e75a22f77c 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 2, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-b:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc index 2a34bc83ff..e75a22f77c 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 2, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-b:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/vector-client/vector.yaml b/self-host/dev-multidc-multinode/datacenters/dc-b/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/vector-client/vector.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-b/vector-client/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/vector-server/vector.yaml b/self-host/dev-multidc-multinode/datacenters/dc-b/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/vector-server/vector.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-b/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/otel-collector/config.yaml b/self-host/dev-multidc-multinode/datacenters/dc-c/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/otel-collector/config.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-c/otel-collector/config.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/postgres/init-db.sh b/self-host/dev-multidc-multinode/datacenters/dc-c/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/postgres/init-db.sh rename to self-host/dev-multidc-multinode/datacenters/dc-c/postgres/init-db.sh diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc index 970454f7f6..cb4c587a19 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 3, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-c:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc index 970454f7f6..cb4c587a19 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 3, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-c:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc index 970454f7f6..cb4c587a19 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 3, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-c:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/vector-client/vector.yaml b/self-host/dev-multidc-multinode/datacenters/dc-c/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/vector-client/vector.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-c/vector-client/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/vector-server/vector.yaml b/self-host/dev-multidc-multinode/datacenters/dc-c/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/vector-server/vector.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-c/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/docker-compose.yml b/self-host/dev-multidc-multinode/docker-compose.yml similarity index 84% rename from self-host/compose/dev-multidc-multinode/docker-compose.yml rename to self-host/dev-multidc-multinode/docker-compose.yml index 7e421144c8..e4ccd4fc70 100644 --- a/self-host/compose/dev-multidc-multinode/docker-compose.yml +++ b/self-host/dev-multidc-multinode/docker-compose.yml @@ -86,6 +86,10 @@ services: postgres-dc-a: restart: unless-stopped image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -93,7 +97,7 @@ services: volumes: - >- ./datacenters/dc-a/postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data-dc-a:/var/lib/postgresql/data + - postgres-data-dc-a:/var/lib/postgresql ports: - '5432:5432' healthcheck: @@ -107,7 +111,7 @@ services: - rivet-network-dc-a rivet-shell-dc-a: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -174,7 +178,7 @@ services: - '4317:4317' rivet-engine-dc-a-0: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -217,7 +221,7 @@ services: start_period: 30s rivet-engine-dc-a-1: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -258,7 +262,7 @@ services: start_period: 30s rivet-engine-dc-a-2: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -299,63 +303,82 @@ services: start_period: 30s runner-dc-a-0: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-a-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s ports: - - '5050:5050' + - '5050:8080' depends_on: rivet-engine-dc-a-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-a runner-dc-a-1: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-a-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-a-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-a runner-dc-a-2: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-a-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-a-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-a postgres-dc-b: restart: unless-stopped image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -363,7 +386,7 @@ services: volumes: - >- ./datacenters/dc-b/postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data-dc-b:/var/lib/postgresql/data + - postgres-data-dc-b:/var/lib/postgresql healthcheck: test: - CMD-SHELL @@ -375,7 +398,7 @@ services: - rivet-network-dc-b rivet-shell-dc-b: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -440,7 +463,7 @@ services: - rivet-network-dc-b-to-core rivet-engine-dc-b-0: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -481,7 +504,7 @@ services: start_period: 30s rivet-engine-dc-b-1: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -522,7 +545,7 @@ services: start_period: 30s rivet-engine-dc-b-2: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -563,61 +586,80 @@ services: start_period: 30s runner-dc-b-0: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-b-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-b-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-b runner-dc-b-1: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-b-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-b-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-b runner-dc-b-2: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-b-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-b-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-b postgres-dc-c: restart: unless-stopped image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -625,7 +667,7 @@ services: volumes: - >- ./datacenters/dc-c/postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data-dc-c:/var/lib/postgresql/data + - postgres-data-dc-c:/var/lib/postgresql healthcheck: test: - CMD-SHELL @@ -637,7 +679,7 @@ services: - rivet-network-dc-c rivet-shell-dc-c: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -702,7 +744,7 @@ services: - rivet-network-dc-c-to-core rivet-engine-dc-c-0: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -743,7 +785,7 @@ services: start_period: 30s rivet-engine-dc-c-1: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -784,7 +826,7 @@ services: start_period: 30s rivet-engine-dc-c-2: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -825,58 +867,92 @@ services: start_period: 30s runner-dc-c-0: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-c-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-c-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-c runner-dc-c-1: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-c-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-c-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-c runner-dc-c-2: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-c-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-c-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-c + runner-config-init: + image: curlimages/curl:latest + restart: 'no' + depends_on: + rivet-engine-dc-a-0: + condition: service_healthy + entrypoint: + - sh + - '-c' + command: + - >- + until curl -fsS -X PUT + "http://rivet-engine-dc-a-0:6420/runner-configs/default?namespace=default" + -H "Authorization: Bearer dev" -H "Content-Type: application/json" -d + '{"datacenters":{"dc-a":{"normal":{}},"dc-b":{"normal":{}},"dc-c":{"normal":{}}}}'; + do echo "waiting for engine to accept runner config"; sleep 2; done; + echo "runner config upserted" + networks: + - rivet-network-dc-a networks: rivet-core-network: driver: bridge diff --git a/self-host/compose/dev-multidc/.gitattributes b/self-host/dev-multidc/.gitattributes similarity index 100% rename from self-host/compose/dev-multidc/.gitattributes rename to self-host/dev-multidc/.gitattributes diff --git a/self-host/compose/dev-multidc/README.md b/self-host/dev-multidc/README.md similarity index 100% rename from self-host/compose/dev-multidc/README.md rename to self-host/dev-multidc/README.md diff --git a/self-host/compose/dev-multidc/core/clickhouse/client-config.xml b/self-host/dev-multidc/core/clickhouse/client-config.xml similarity index 100% rename from self-host/compose/dev-multidc/core/clickhouse/client-config.xml rename to self-host/dev-multidc/core/clickhouse/client-config.xml diff --git a/self-host/compose/dev-multidc/core/clickhouse/config.xml b/self-host/dev-multidc/core/clickhouse/config.xml similarity index 100% rename from self-host/compose/dev-multidc/core/clickhouse/config.xml rename to self-host/dev-multidc/core/clickhouse/config.xml diff --git a/self-host/compose/dev-multidc/core/clickhouse/init/01-create-otel-table.sql b/self-host/dev-multidc/core/clickhouse/init/01-create-otel-table.sql similarity index 100% rename from self-host/compose/dev-multidc/core/clickhouse/init/01-create-otel-table.sql rename to self-host/dev-multidc/core/clickhouse/init/01-create-otel-table.sql diff --git a/self-host/compose/dev-multidc/core/clickhouse/users.xml b/self-host/dev-multidc/core/clickhouse/users.xml similarity index 100% rename from self-host/compose/dev-multidc/core/clickhouse/users.xml rename to self-host/dev-multidc/core/clickhouse/users.xml diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/api.json b/self-host/dev-multidc/core/grafana/dashboards/api.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/api.json rename to self-host/dev-multidc/core/grafana/dashboards/api.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/cache.json b/self-host/dev-multidc/core/grafana/dashboards/cache.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/cache.json rename to self-host/dev-multidc/core/grafana/dashboards/cache.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/epoxy.json b/self-host/dev-multidc/core/grafana/dashboards/epoxy.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/epoxy.json rename to self-host/dev-multidc/core/grafana/dashboards/epoxy.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/futures.json b/self-host/dev-multidc/core/grafana/dashboards/futures.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/futures.json rename to self-host/dev-multidc/core/grafana/dashboards/futures.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/gasoline.json b/self-host/dev-multidc/core/grafana/dashboards/gasoline.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/gasoline.json rename to self-host/dev-multidc/core/grafana/dashboards/gasoline.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/guard.json b/self-host/dev-multidc/core/grafana/dashboards/guard.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/guard.json rename to self-host/dev-multidc/core/grafana/dashboards/guard.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/operation.json b/self-host/dev-multidc/core/grafana/dashboards/operation.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/operation.json rename to self-host/dev-multidc/core/grafana/dashboards/operation.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/pegboard.json b/self-host/dev-multidc/core/grafana/dashboards/pegboard.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/pegboard.json rename to self-host/dev-multidc/core/grafana/dashboards/pegboard.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/tokio.json b/self-host/dev-multidc/core/grafana/dashboards/tokio.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/tokio.json rename to self-host/dev-multidc/core/grafana/dashboards/tokio.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/traces.json b/self-host/dev-multidc/core/grafana/dashboards/traces.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/traces.json rename to self-host/dev-multidc/core/grafana/dashboards/traces.json diff --git a/self-host/compose/dev-multidc/core/grafana/grafana.ini b/self-host/dev-multidc/core/grafana/grafana.ini similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/grafana.ini rename to self-host/dev-multidc/core/grafana/grafana.ini diff --git a/self-host/compose/dev-multidc/core/grafana/provisioning/dashboards/dashboards.yaml b/self-host/dev-multidc/core/grafana/provisioning/dashboards/dashboards.yaml similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/provisioning/dashboards/dashboards.yaml rename to self-host/dev-multidc/core/grafana/provisioning/dashboards/dashboards.yaml diff --git a/self-host/compose/dev-multidc/core/grafana/provisioning/datasources/datasources.yaml b/self-host/dev-multidc/core/grafana/provisioning/datasources/datasources.yaml similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/provisioning/datasources/datasources.yaml rename to self-host/dev-multidc/core/grafana/provisioning/datasources/datasources.yaml diff --git a/self-host/compose/dev-multidc/core/prometheus/prometheus.yml b/self-host/dev-multidc/core/prometheus/prometheus.yml similarity index 100% rename from self-host/compose/dev-multidc/core/prometheus/prometheus.yml rename to self-host/dev-multidc/core/prometheus/prometheus.yml diff --git a/self-host/compose/dev-multidc/datacenters/dc-a/otel-collector/config.yaml b/self-host/dev-multidc/datacenters/dc-a/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-a/otel-collector/config.yaml rename to self-host/dev-multidc/datacenters/dc-a/otel-collector/config.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-a/postgres/init-db.sh b/self-host/dev-multidc/datacenters/dc-a/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-a/postgres/init-db.sh rename to self-host/dev-multidc/datacenters/dc-a/postgres/init-db.sh diff --git a/self-host/compose/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc b/self-host/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc rename to self-host/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc index f989acaae9..ac3c4bc3c6 100644 --- a/self-host/compose/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc +++ b/self-host/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-a:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc/datacenters/dc-a/vector-client/vector.yaml b/self-host/dev-multidc/datacenters/dc-a/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-a/vector-client/vector.yaml rename to self-host/dev-multidc/datacenters/dc-a/vector-client/vector.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-a/vector-server/vector.yaml b/self-host/dev-multidc/datacenters/dc-a/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-a/vector-server/vector.yaml rename to self-host/dev-multidc/datacenters/dc-a/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-b/otel-collector/config.yaml b/self-host/dev-multidc/datacenters/dc-b/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-b/otel-collector/config.yaml rename to self-host/dev-multidc/datacenters/dc-b/otel-collector/config.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-b/postgres/init-db.sh b/self-host/dev-multidc/datacenters/dc-b/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-b/postgres/init-db.sh rename to self-host/dev-multidc/datacenters/dc-b/postgres/init-db.sh diff --git a/self-host/compose/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc b/self-host/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc rename to self-host/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc index fa082ce527..8e5c6aaa45 100644 --- a/self-host/compose/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc +++ b/self-host/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 2, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-b:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc/datacenters/dc-b/vector-client/vector.yaml b/self-host/dev-multidc/datacenters/dc-b/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-b/vector-client/vector.yaml rename to self-host/dev-multidc/datacenters/dc-b/vector-client/vector.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-b/vector-server/vector.yaml b/self-host/dev-multidc/datacenters/dc-b/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-b/vector-server/vector.yaml rename to self-host/dev-multidc/datacenters/dc-b/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-c/otel-collector/config.yaml b/self-host/dev-multidc/datacenters/dc-c/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-c/otel-collector/config.yaml rename to self-host/dev-multidc/datacenters/dc-c/otel-collector/config.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-c/postgres/init-db.sh b/self-host/dev-multidc/datacenters/dc-c/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-c/postgres/init-db.sh rename to self-host/dev-multidc/datacenters/dc-c/postgres/init-db.sh diff --git a/self-host/compose/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc b/self-host/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc rename to self-host/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc index f3c3c6ae38..6ba3be6d65 100644 --- a/self-host/compose/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc +++ b/self-host/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 3, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-c:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc/datacenters/dc-c/vector-client/vector.yaml b/self-host/dev-multidc/datacenters/dc-c/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-c/vector-client/vector.yaml rename to self-host/dev-multidc/datacenters/dc-c/vector-client/vector.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-c/vector-server/vector.yaml b/self-host/dev-multidc/datacenters/dc-c/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-c/vector-server/vector.yaml rename to self-host/dev-multidc/datacenters/dc-c/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc/docker-compose.yml b/self-host/dev-multidc/docker-compose.yml similarity index 87% rename from self-host/compose/dev-multidc/docker-compose.yml rename to self-host/dev-multidc/docker-compose.yml index 79cb619e0a..437ad5bdee 100644 --- a/self-host/compose/dev-multidc/docker-compose.yml +++ b/self-host/dev-multidc/docker-compose.yml @@ -86,6 +86,10 @@ services: postgres-dc-a: restart: unless-stopped image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -93,7 +97,7 @@ services: volumes: - >- ./datacenters/dc-a/postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data-dc-a:/var/lib/postgresql/data + - postgres-data-dc-a:/var/lib/postgresql ports: - '5432:5432' healthcheck: @@ -107,7 +111,7 @@ services: - rivet-network-dc-a rivet-shell-dc-a: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -173,7 +177,7 @@ services: - '4317:4317' rivet-engine-dc-a: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -215,27 +219,36 @@ services: start_period: 30s runner-dc-a: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-a:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s ports: - - '5050:5050' + - '5050:8080' depends_on: rivet-engine-dc-a: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-a postgres-dc-b: restart: unless-stopped image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -243,7 +256,7 @@ services: volumes: - >- ./datacenters/dc-b/postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data-dc-b:/var/lib/postgresql/data + - postgres-data-dc-b:/var/lib/postgresql healthcheck: test: - CMD-SHELL @@ -255,7 +268,7 @@ services: - rivet-network-dc-b rivet-shell-dc-b: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -319,7 +332,7 @@ services: - rivet-network-dc-b-to-core rivet-engine-dc-b: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -359,25 +372,34 @@ services: start_period: 30s runner-dc-b: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-b:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-b: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-b postgres-dc-c: restart: unless-stopped image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -385,7 +407,7 @@ services: volumes: - >- ./datacenters/dc-c/postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data-dc-c:/var/lib/postgresql/data + - postgres-data-dc-c:/var/lib/postgresql healthcheck: test: - CMD-SHELL @@ -397,7 +419,7 @@ services: - rivet-network-dc-c rivet-shell-dc-c: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -461,7 +483,7 @@ services: - rivet-network-dc-c-to-core rivet-engine-dc-c: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -501,22 +523,46 @@ services: start_period: 30s runner-dc-c: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-c:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-c: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-c + runner-config-init: + image: curlimages/curl:latest + restart: 'no' + depends_on: + rivet-engine-dc-a: + condition: service_healthy + entrypoint: + - sh + - '-c' + command: + - >- + until curl -fsS -X PUT + "http://rivet-engine-dc-a:6420/runner-configs/default?namespace=default" + -H "Authorization: Bearer dev" -H "Content-Type: application/json" -d + '{"datacenters":{"dc-a":{"normal":{}},"dc-b":{"normal":{}},"dc-c":{"normal":{}}}}'; + do echo "waiting for engine to accept runner config"; sleep 2; done; + echo "runner config upserted" + networks: + - rivet-network-dc-a networks: rivet-core-network: driver: bridge diff --git a/self-host/compose/dev-multinode/.gitattributes b/self-host/dev-multinode/.gitattributes similarity index 100% rename from self-host/compose/dev-multinode/.gitattributes rename to self-host/dev-multinode/.gitattributes diff --git a/self-host/compose/dev-multinode/README.md b/self-host/dev-multinode/README.md similarity index 100% rename from self-host/compose/dev-multinode/README.md rename to self-host/dev-multinode/README.md diff --git a/self-host/compose/dev-multinode/clickhouse/client-config.xml b/self-host/dev-multinode/clickhouse/client-config.xml similarity index 100% rename from self-host/compose/dev-multinode/clickhouse/client-config.xml rename to self-host/dev-multinode/clickhouse/client-config.xml diff --git a/self-host/compose/dev-multinode/clickhouse/config.xml b/self-host/dev-multinode/clickhouse/config.xml similarity index 100% rename from self-host/compose/dev-multinode/clickhouse/config.xml rename to self-host/dev-multinode/clickhouse/config.xml diff --git a/self-host/compose/dev-multinode/clickhouse/init/01-create-otel-table.sql b/self-host/dev-multinode/clickhouse/init/01-create-otel-table.sql similarity index 100% rename from self-host/compose/dev-multinode/clickhouse/init/01-create-otel-table.sql rename to self-host/dev-multinode/clickhouse/init/01-create-otel-table.sql diff --git a/self-host/compose/dev-multinode/clickhouse/users.xml b/self-host/dev-multinode/clickhouse/users.xml similarity index 100% rename from self-host/compose/dev-multinode/clickhouse/users.xml rename to self-host/dev-multinode/clickhouse/users.xml diff --git a/self-host/compose/dev-multinode/docker-compose.yml b/self-host/dev-multinode/docker-compose.yml similarity index 82% rename from self-host/compose/dev-multinode/docker-compose.yml rename to self-host/dev-multinode/docker-compose.yml index a432743057..389510f7b3 100644 --- a/self-host/compose/dev-multinode/docker-compose.yml +++ b/self-host/dev-multinode/docker-compose.yml @@ -82,13 +82,17 @@ services: postgres: restart: unless-stopped image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres - POSTGRES_DB=postgres volumes: - ./postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data:/var/lib/postgresql/data + - postgres-data:/var/lib/postgresql ports: - '5432:5432' healthcheck: @@ -102,7 +106,7 @@ services: - rivet-network rivet-shell: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -166,7 +170,7 @@ services: - '4317:4317' rivet-engine-0: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -206,7 +210,7 @@ services: start_period: 30s rivet-engine-1: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -244,7 +248,7 @@ services: start_period: 30s rivet-engine-2: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -282,58 +286,91 @@ services: start_period: 30s runner-0: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s ports: - - '5050:5050' + - '5050:8080' depends_on: rivet-engine-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network runner-1: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network runner-2: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully + networks: + - rivet-network + runner-config-init: + image: curlimages/curl:latest + restart: 'no' + depends_on: + rivet-engine-0: + condition: service_healthy + entrypoint: + - sh + - '-c' + command: + - >- + until curl -fsS -X PUT + "http://rivet-engine-0:6420/runner-configs/default?namespace=default" -H + "Authorization: Bearer dev" -H "Content-Type: application/json" -d + '{"datacenters":{"default":{"normal":{}}}}'; do echo "waiting for engine + to accept runner config"; sleep 2; done; echo "runner config upserted" networks: - rivet-network networks: diff --git a/self-host/compose/dev-multinode/grafana/dashboards/api.json b/self-host/dev-multinode/grafana/dashboards/api.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/api.json rename to self-host/dev-multinode/grafana/dashboards/api.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/cache.json b/self-host/dev-multinode/grafana/dashboards/cache.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/cache.json rename to self-host/dev-multinode/grafana/dashboards/cache.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/epoxy.json b/self-host/dev-multinode/grafana/dashboards/epoxy.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/epoxy.json rename to self-host/dev-multinode/grafana/dashboards/epoxy.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/futures.json b/self-host/dev-multinode/grafana/dashboards/futures.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/futures.json rename to self-host/dev-multinode/grafana/dashboards/futures.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/gasoline.json b/self-host/dev-multinode/grafana/dashboards/gasoline.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/gasoline.json rename to self-host/dev-multinode/grafana/dashboards/gasoline.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/guard.json b/self-host/dev-multinode/grafana/dashboards/guard.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/guard.json rename to self-host/dev-multinode/grafana/dashboards/guard.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/operation.json b/self-host/dev-multinode/grafana/dashboards/operation.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/operation.json rename to self-host/dev-multinode/grafana/dashboards/operation.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/pegboard.json b/self-host/dev-multinode/grafana/dashboards/pegboard.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/pegboard.json rename to self-host/dev-multinode/grafana/dashboards/pegboard.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/tokio.json b/self-host/dev-multinode/grafana/dashboards/tokio.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/tokio.json rename to self-host/dev-multinode/grafana/dashboards/tokio.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/traces.json b/self-host/dev-multinode/grafana/dashboards/traces.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/traces.json rename to self-host/dev-multinode/grafana/dashboards/traces.json diff --git a/self-host/compose/dev-multinode/grafana/grafana.ini b/self-host/dev-multinode/grafana/grafana.ini similarity index 100% rename from self-host/compose/dev-multinode/grafana/grafana.ini rename to self-host/dev-multinode/grafana/grafana.ini diff --git a/self-host/compose/dev-multinode/grafana/provisioning/dashboards/dashboards.yaml b/self-host/dev-multinode/grafana/provisioning/dashboards/dashboards.yaml similarity index 100% rename from self-host/compose/dev-multinode/grafana/provisioning/dashboards/dashboards.yaml rename to self-host/dev-multinode/grafana/provisioning/dashboards/dashboards.yaml diff --git a/self-host/compose/dev-multinode/grafana/provisioning/datasources/datasources.yaml b/self-host/dev-multinode/grafana/provisioning/datasources/datasources.yaml similarity index 100% rename from self-host/compose/dev-multinode/grafana/provisioning/datasources/datasources.yaml rename to self-host/dev-multinode/grafana/provisioning/datasources/datasources.yaml diff --git a/self-host/compose/dev-multinode/otel-collector/config.yaml b/self-host/dev-multinode/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multinode/otel-collector/config.yaml rename to self-host/dev-multinode/otel-collector/config.yaml diff --git a/self-host/compose/dev-multinode/postgres/init-db.sh b/self-host/dev-multinode/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multinode/postgres/init-db.sh rename to self-host/dev-multinode/postgres/init-db.sh diff --git a/self-host/compose/dev-multinode/prometheus/prometheus.yml b/self-host/dev-multinode/prometheus/prometheus.yml similarity index 100% rename from self-host/compose/dev-multinode/prometheus/prometheus.yml rename to self-host/dev-multinode/prometheus/prometheus.yml diff --git a/self-host/compose/dev-multinode/rivet-engine/2/config.jsonc b/self-host/dev-multinode/rivet-engine/0/config.jsonc similarity index 79% rename from self-host/compose/dev-multinode/rivet-engine/2/config.jsonc rename to self-host/dev-multinode/rivet-engine/0/config.jsonc index 655fc83ed7..b25680c399 100644 --- a/self-host/compose/dev-multinode/rivet-engine/2/config.jsonc +++ b/self-host/dev-multinode/rivet-engine/0/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -28,14 +24,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multinode/rivet-engine/0/config.jsonc b/self-host/dev-multinode/rivet-engine/1/config.jsonc similarity index 79% rename from self-host/compose/dev-multinode/rivet-engine/0/config.jsonc rename to self-host/dev-multinode/rivet-engine/1/config.jsonc index 655fc83ed7..b25680c399 100644 --- a/self-host/compose/dev-multinode/rivet-engine/0/config.jsonc +++ b/self-host/dev-multinode/rivet-engine/1/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -28,14 +24,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multinode/rivet-engine/1/config.jsonc b/self-host/dev-multinode/rivet-engine/2/config.jsonc similarity index 79% rename from self-host/compose/dev-multinode/rivet-engine/1/config.jsonc rename to self-host/dev-multinode/rivet-engine/2/config.jsonc index 655fc83ed7..b25680c399 100644 --- a/self-host/compose/dev-multinode/rivet-engine/1/config.jsonc +++ b/self-host/dev-multinode/rivet-engine/2/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -28,14 +24,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multinode/vector-client/vector.yaml b/self-host/dev-multinode/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multinode/vector-client/vector.yaml rename to self-host/dev-multinode/vector-client/vector.yaml diff --git a/self-host/compose/dev-multinode/vector-server/vector.yaml b/self-host/dev-multinode/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multinode/vector-server/vector.yaml rename to self-host/dev-multinode/vector-server/vector.yaml diff --git a/self-host/compose/dev/.gitattributes b/self-host/dev/.gitattributes similarity index 100% rename from self-host/compose/dev/.gitattributes rename to self-host/dev/.gitattributes diff --git a/self-host/compose/dev/README.md b/self-host/dev/README.md similarity index 100% rename from self-host/compose/dev/README.md rename to self-host/dev/README.md diff --git a/self-host/compose/dev/clickhouse/client-config.xml b/self-host/dev/clickhouse/client-config.xml similarity index 100% rename from self-host/compose/dev/clickhouse/client-config.xml rename to self-host/dev/clickhouse/client-config.xml diff --git a/self-host/compose/dev/clickhouse/config.xml b/self-host/dev/clickhouse/config.xml similarity index 100% rename from self-host/compose/dev/clickhouse/config.xml rename to self-host/dev/clickhouse/config.xml diff --git a/self-host/compose/dev/clickhouse/init/01-create-otel-table.sql b/self-host/dev/clickhouse/init/01-create-otel-table.sql similarity index 100% rename from self-host/compose/dev/clickhouse/init/01-create-otel-table.sql rename to self-host/dev/clickhouse/init/01-create-otel-table.sql diff --git a/self-host/compose/dev/clickhouse/users.xml b/self-host/dev/clickhouse/users.xml similarity index 100% rename from self-host/compose/dev/clickhouse/users.xml rename to self-host/dev/clickhouse/users.xml diff --git a/self-host/compose/dev/docker-compose.yml b/self-host/dev/docker-compose.yml similarity index 84% rename from self-host/compose/dev/docker-compose.yml rename to self-host/dev/docker-compose.yml index 91f82dce45..7511952a55 100644 --- a/self-host/compose/dev/docker-compose.yml +++ b/self-host/dev/docker-compose.yml @@ -82,13 +82,17 @@ services: postgres: restart: unless-stopped image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres - POSTGRES_DB=postgres volumes: - ./postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data:/var/lib/postgresql/data + - postgres-data:/var/lib/postgresql ports: - '5432:5432' healthcheck: @@ -102,7 +106,7 @@ services: - rivet-network rivet-shell: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -166,7 +170,7 @@ services: - '4317:4317' rivet-engine: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -206,22 +210,45 @@ services: start_period: 30s runner: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s ports: - - '5050:5050' + - '5050:8080' + depends_on: + rivet-engine: + condition: service_healthy + runner-config-init: + condition: service_completed_successfully + networks: + - rivet-network + runner-config-init: + image: curlimages/curl:latest + restart: 'no' depends_on: rivet-engine: condition: service_healthy + entrypoint: + - sh + - '-c' + command: + - >- + until curl -fsS -X PUT + "http://rivet-engine:6420/runner-configs/default?namespace=default" -H + "Authorization: Bearer dev" -H "Content-Type: application/json" -d + '{"datacenters":{"default":{"normal":{}}}}'; do echo "waiting for engine + to accept runner config"; sleep 2; done; echo "runner config upserted" networks: - rivet-network networks: diff --git a/self-host/compose/dev/grafana/dashboards/api.json b/self-host/dev/grafana/dashboards/api.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/api.json rename to self-host/dev/grafana/dashboards/api.json diff --git a/self-host/compose/dev/grafana/dashboards/cache.json b/self-host/dev/grafana/dashboards/cache.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/cache.json rename to self-host/dev/grafana/dashboards/cache.json diff --git a/self-host/compose/dev/grafana/dashboards/epoxy.json b/self-host/dev/grafana/dashboards/epoxy.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/epoxy.json rename to self-host/dev/grafana/dashboards/epoxy.json diff --git a/self-host/compose/dev/grafana/dashboards/futures.json b/self-host/dev/grafana/dashboards/futures.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/futures.json rename to self-host/dev/grafana/dashboards/futures.json diff --git a/self-host/compose/dev/grafana/dashboards/gasoline.json b/self-host/dev/grafana/dashboards/gasoline.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/gasoline.json rename to self-host/dev/grafana/dashboards/gasoline.json diff --git a/self-host/compose/dev/grafana/dashboards/guard.json b/self-host/dev/grafana/dashboards/guard.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/guard.json rename to self-host/dev/grafana/dashboards/guard.json diff --git a/self-host/compose/dev/grafana/dashboards/operation.json b/self-host/dev/grafana/dashboards/operation.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/operation.json rename to self-host/dev/grafana/dashboards/operation.json diff --git a/self-host/compose/dev/grafana/dashboards/pegboard.json b/self-host/dev/grafana/dashboards/pegboard.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/pegboard.json rename to self-host/dev/grafana/dashboards/pegboard.json diff --git a/self-host/compose/dev/grafana/dashboards/tokio.json b/self-host/dev/grafana/dashboards/tokio.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/tokio.json rename to self-host/dev/grafana/dashboards/tokio.json diff --git a/self-host/compose/dev/grafana/dashboards/traces.json b/self-host/dev/grafana/dashboards/traces.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/traces.json rename to self-host/dev/grafana/dashboards/traces.json diff --git a/self-host/compose/dev/grafana/grafana.ini b/self-host/dev/grafana/grafana.ini similarity index 100% rename from self-host/compose/dev/grafana/grafana.ini rename to self-host/dev/grafana/grafana.ini diff --git a/self-host/compose/dev/grafana/provisioning/dashboards/dashboards.yaml b/self-host/dev/grafana/provisioning/dashboards/dashboards.yaml similarity index 100% rename from self-host/compose/dev/grafana/provisioning/dashboards/dashboards.yaml rename to self-host/dev/grafana/provisioning/dashboards/dashboards.yaml diff --git a/self-host/compose/dev/grafana/provisioning/datasources/datasources.yaml b/self-host/dev/grafana/provisioning/datasources/datasources.yaml similarity index 100% rename from self-host/compose/dev/grafana/provisioning/datasources/datasources.yaml rename to self-host/dev/grafana/provisioning/datasources/datasources.yaml diff --git a/self-host/compose/dev/otel-collector/config.yaml b/self-host/dev/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev/otel-collector/config.yaml rename to self-host/dev/otel-collector/config.yaml diff --git a/self-host/compose/dev/postgres/init-db.sh b/self-host/dev/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev/postgres/init-db.sh rename to self-host/dev/postgres/init-db.sh diff --git a/self-host/compose/dev/prometheus/prometheus.yml b/self-host/dev/prometheus/prometheus.yml similarity index 100% rename from self-host/compose/dev/prometheus/prometheus.yml rename to self-host/dev/prometheus/prometheus.yml diff --git a/self-host/compose/dev/rivet-engine/config.jsonc b/self-host/dev/rivet-engine/config.jsonc similarity index 79% rename from self-host/compose/dev/rivet-engine/config.jsonc rename to self-host/dev/rivet-engine/config.jsonc index 74135c72fa..4c6312b6d3 100644 --- a/self-host/compose/dev/rivet-engine/config.jsonc +++ b/self-host/dev/rivet-engine/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -28,14 +24,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev/vector-client/vector.yaml b/self-host/dev/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev/vector-client/vector.yaml rename to self-host/dev/vector-client/vector.yaml diff --git a/self-host/compose/dev/vector-server/vector.yaml b/self-host/dev/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev/vector-server/vector.yaml rename to self-host/dev/vector-server/vector.yaml From 5dfeb48b3466da9a19a4c238b0606e645c5bdd69 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Thu, 25 Jun 2026 16:04:12 -0700 Subject: [PATCH 13/16] [slopfix] docs(self-hosting): promote postgres from experimental to recommended OSS multi-node backend --- .../src/content/cookbook/vpc-air-gapped.mdx | 4 +- .../docs/self-hosting/configuration.mdx | 4 +- .../docs/self-hosting/foundationdb.mdx | 2 +- .../content/docs/self-hosting/postgres.mdx | 57 +++++++++++++++++++ .../self-hosting/production-checklist.mdx | 12 ++-- 5 files changed, 70 insertions(+), 9 deletions(-) diff --git a/website/src/content/cookbook/vpc-air-gapped.mdx b/website/src/content/cookbook/vpc-air-gapped.mdx index 843aba522d..db26b1e0ed 100644 --- a/website/src/content/cookbook/vpc-air-gapped.mdx +++ b/website/src/content/cookbook/vpc-air-gapped.mdx @@ -107,10 +107,10 @@ If you ship software that runs inside your customers' VPCs, the same setup turns | Backend | Use when | Status | | --- | --- | --- | | [File System](/docs/self-hosting/filesystem) (RocksDB-based) | Single-node deployments, including air-gapped installs | Production-ready, single node only | -| [PostgreSQL](/docs/self-hosting/postgres) | Multi-node deployments | Recommended for multi-node today, but experimental | +| [PostgreSQL](/docs/self-hosting/postgres) | Multi-node and multi-region deployments | Production-ready for multi-node | | FoundationDB | Largest production deployments | [Enterprise](/sales) | -For multi-node deployments, run two or more engine nodes behind a load balancer and add NATS for pub/sub, which replaces the default PostgreSQL `LISTEN`/`NOTIFY` path at high throughput. Neither is needed for a single-node file system install. See the [Production Checklist](/docs/self-hosting/production-checklist). +For multi-node deployments, run two or more engine nodes behind a load balancer, all sharing one PostgreSQL instance. The built-in PostgreSQL pub/sub is sufficient for most deployments; very high-throughput deployments can add NATS as a dedicated pub/sub layer. Neither is needed for a single-node file system install. See the [Production Checklist](/docs/self-hosting/production-checklist). ## Perimeter Checklist diff --git a/website/src/content/docs/self-hosting/configuration.mdx b/website/src/content/docs/self-hosting/configuration.mdx index 781f8bba6a..8fb19f129c 100644 --- a/website/src/content/docs/self-hosting/configuration.mdx +++ b/website/src/content/docs/self-hosting/configuration.mdx @@ -76,5 +76,5 @@ Use `samples: 1` for a uniform random pick that skips slot reads. Use `samples > ## Related - RivetKit actor runtime persistence lives in SQLite. Existing actor KV data is imported into SQLite the first time an actor wakes on the migrated runtime, then the original KV data is left frozen for downgrade safety. -- [PostgreSQL](/docs/self-hosting/postgres): Configure the PostgreSQL backend for multi-node deployments -- [File System](/docs/self-hosting/filesystem): Configure file system storage for development +- [PostgreSQL](/docs/self-hosting/postgres): Configure the PostgreSQL backend for multi-node and multi-region deployments +- [File System](/docs/self-hosting/filesystem): Configure file system storage for single-node deployments diff --git a/website/src/content/docs/self-hosting/foundationdb.mdx b/website/src/content/docs/self-hosting/foundationdb.mdx index bce99b66dd..5de8408f51 100644 --- a/website/src/content/docs/self-hosting/foundationdb.mdx +++ b/website/src/content/docs/self-hosting/foundationdb.mdx @@ -25,7 +25,7 @@ Its strict serializability guarantees, fault tolerance, and ability to scale lin | | RocksDB (File System) | PostgreSQL | FoundationDB | |---|---|---|---| -| **Scalability** | Single node | Primary/replica failover | Linear horizontal scaling | +| **Scalability** | Single node | Multi-node and multi-region | Linear horizontal scaling | | **Fault tolerance** | None | Primary/replica failover | Automatic recovery with no data loss | | **Production readiness** | Development and small deployments | Production-ready for light-to-moderate multi-node workloads | Battle-tested at global scale | diff --git a/website/src/content/docs/self-hosting/postgres.mdx b/website/src/content/docs/self-hosting/postgres.mdx index 34b2fc5809..ef737738fc 100644 --- a/website/src/content/docs/self-hosting/postgres.mdx +++ b/website/src/content/docs/self-hosting/postgres.mdx @@ -8,6 +8,16 @@ skill: true PostgreSQL is the recommended backend for multi-node self-hosted deployments. It is production-ready for light-to-moderate workloads, up to roughly 1,000 concurrent actors, but is not built for enterprise scale beyond that. For a single-node deployment, use the file system backend (RocksDB-based). Teams running larger or high-throughput realtime workloads should contact [enterprise support](https://rivet.dev/sales) about FoundationDB. +## Overview + +PostgreSQL is the storage and coordination backend for self-hosted Rivet deployments that run more than one engine node. Multiple engine nodes can share a single PostgreSQL instance with no extra coordination service to deploy. Rivet handles leader election, failover, and version sequencing internally. + +Use PostgreSQL when you need: + +- **Multiple engine nodes** behind a load balancer for redundancy and horizontal scaling. +- **Multi-region deployments** (deploy one PostgreSQL instance per region, see [Multi-Region](/docs/self-hosting/multi-region)). +- **High availability** with a managed or self-managed primary/replica failover setup. + ## Choosing a Backend Pick your database backend based on how many engine nodes you run: @@ -59,6 +69,41 @@ Multi-node PostgreSQL deployments require NATS as the pub/sub backend so engine See the [production checklist](/docs/self-hosting/production-checklist#nats) and [Configuration](/docs/self-hosting/configuration) for details. +## Requirements and Recommendations + +### Version + +Use PostgreSQL 14 or newer. Rivet is tested against PostgreSQL 18, which is recommended for new deployments. + +### Connection Limits + +Each Rivet engine node opens a pool of direct connections to PostgreSQL and can use well over a hundred connections per node under load. PostgreSQL's default `max_connections` of `100` is too low for even a single busy engine node. + +- Set PostgreSQL `max_connections` to comfortably exceed `(number of engine nodes × 150)` plus headroom for backups, monitoring, and your own queries. +- If you use a managed PostgreSQL service, confirm its connection limit is high enough or pick a tier that allows raising it. Connection exhaustion shows up as engine startup failures or stalled requests under load. + + +Do not work around the connection limit with a connection pooler. See [Do Not Use Connection Poolers](#do-not-use-connection-poolers) below. + + +### Resources + +PostgreSQL is the system of record for the entire deployment, so size it accordingly: + +- Give PostgreSQL dedicated CPU, memory, and fast disk (SSD/NVMe with high IOPS). Avoid co-locating it with other heavy workloads. +- Rivet generates steady write and row-turnover on its internal tables. Keep autovacuum enabled and healthy so dead tuples do not accumulate. + +### High Availability and Backups + +A single PostgreSQL instance is a single point of failure for your whole deployment. + +- Configure a standby replica with automatic failover (managed services such as Amazon RDS, Cloud SQL, and Azure Database provide this). +- Enable automated backups and point-in-time recovery, and periodically test restoring from them. + +### Multi-Region + +Deploy one PostgreSQL instance per region or datacenter. Engine nodes connect to the PostgreSQL instance in their own region. See [Multi-Region](/docs/self-hosting/multi-region) for the full topology. + ## Managed Postgres Compatibility Some hosted PostgreSQL platforms require additional configuration due to platform-specific restrictions. @@ -209,3 +254,15 @@ Do not use: - PgBouncer - Supavisor - AWS RDS Proxy + + +## Troubleshooting + +### Too Many Connections + +Errors like `FATAL: sorry, too many clients already` or engine nodes failing to start under load mean PostgreSQL's `max_connections` is too low. Raise it to account for every engine node (see [Connection Limits](#connection-limits)). Do not add a connection pooler to work around this. + +### Connection Refused or TLS Errors + +- Confirm the engine connects directly to PostgreSQL and not through a pooler (PgBouncer, Supavisor, RDS Proxy). Rivet requires direct connections. +- For TLS errors, verify `sslmode` matches your server and, for custom certificate authorities, that `ssl.root_cert_path` points to the correct CA certificate. See [SSL/TLS Support](#ssltls-support). diff --git a/website/src/content/docs/self-hosting/production-checklist.mdx b/website/src/content/docs/self-hosting/production-checklist.mdx index c554dc2c3c..ad9c653e0d 100644 --- a/website/src/content/docs/self-hosting/production-checklist.mdx +++ b/website/src/content/docs/self-hosting/production-checklist.mdx @@ -34,10 +34,14 @@ Also review the [general production checklist](/docs/general/production-checklis ## PostgreSQL -- **PostgreSQL is recommended for multi-node deployments.** It is production-ready for light-to-moderate workloads (up to roughly 1,000 concurrent actors) but is not built for enterprise scale. Validate the deployment carefully before rollout. -- **Configure automated backups.** Set up regular backups for your PostgreSQL database to prevent data loss. -- **Configure failover.** Set up a standby replica with automatic failover to ensure high availability. -- **Use FoundationDB for the most scalable production-ready deployments.** FoundationDB provides the best performance, scalability, and uptime for Rivet. Contact [enterprise support](https://rivet.dev/sales) for FoundationDB guidance. +- **Use PostgreSQL for multi-node and multi-region deployments.** Multiple engine nodes can share one PostgreSQL instance; no extra coordination service is required. PostgreSQL is production-ready for light-to-moderate workloads (up to roughly 1,000 concurrent actors) but is not built for enterprise scale. See [PostgreSQL](/docs/self-hosting/postgres). +- **Raise `max_connections`.** Each engine node opens well over a hundred connections under load. Size `max_connections` to at least `(number of engine nodes × 150)` plus headroom. PostgreSQL's default of `100` is too low. See [Connection Limits](/docs/self-hosting/postgres#connection-limits). +- **Do not use a connection pooler.** Rivet requires direct connections. Do not put PgBouncer, Supavisor, or RDS Proxy in front of PostgreSQL. +- **Give PostgreSQL dedicated resources.** Provision dedicated CPU, memory, and fast disk, and keep autovacuum healthy. PostgreSQL is the system of record for the whole deployment. +- **Configure automated backups.** Set up regular backups and point-in-time recovery, and test restoring from them. +- **Configure failover.** Set up a standby replica with automatic failover to ensure high availability. A single instance is a single point of failure. +- **Use one PostgreSQL instance per region.** For multi-region deployments, deploy a separate PostgreSQL instance in each region. +- **Use FoundationDB for the largest deployments.** Enterprise teams running at very large scale can contact [enterprise support](https://rivet.dev/sales) for FoundationDB guidance. ## NATS From a783a31b8b2221a7925d467f27830faf486e0153 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Thu, 25 Jun 2026 16:04:12 -0700 Subject: [PATCH 14/16] [SLOP(claude-opus-4-8-high)] perf(universaldb): batch leader apply and fold follower commit round-trips --- engine/packages/universaldb/Cargo.toml | 2 +- .../universaldb/src/driver/postgres/commit.rs | 206 +++++++-- .../src/driver/postgres/resolver/apply.rs | 239 +++++++--- .../driver/postgres/resolver/apply_tests.rs | 299 +++++++++++++ .../src/driver/postgres/resolver/mod.rs | 411 +++++++++++++----- .../universaldb/src/driver/postgres/shared.rs | 15 + .../compose/template/src/docker-compose.ts | 1 + self-host/dev-multinode/docker-compose.yml | 3 + 8 files changed, 956 insertions(+), 220 deletions(-) create mode 100644 engine/packages/universaldb/src/driver/postgres/resolver/apply_tests.rs diff --git a/engine/packages/universaldb/Cargo.toml b/engine/packages/universaldb/Cargo.toml index 57dba4ab33..6781c7cde1 100644 --- a/engine/packages/universaldb/Cargo.toml +++ b/engine/packages/universaldb/Cargo.toml @@ -27,7 +27,7 @@ tempfile.workspace = true thiserror.workspace = true tokio-postgres-rustls.workspace = true tokio-postgres.workspace = true -tokio-util.workspace = true +tokio-util = { workspace = true, features = ["rt"] } tokio.workspace = true tracing.workspace = true url.workspace = true diff --git a/engine/packages/universaldb/src/driver/postgres/commit.rs b/engine/packages/universaldb/src/driver/postgres/commit.rs index e61714dd0e..d191c815a3 100644 --- a/engine/packages/universaldb/src/driver/postgres/commit.rs +++ b/engine/packages/universaldb/src/driver/postgres/commit.rs @@ -52,36 +52,43 @@ pub async fn submit( .await .context("failed to get connection for commit submit")?; + // Enqueue the request and wake the leader's drain loop in one round-trip. The autocommit + // statement durably inserts the row and fires the NOTIFY together, so there is no separate + // notify round-trip or second pool acquire. let id: i64 = conn .query_one( - "INSERT INTO udb_commit_requests (epoch, read_version, payload, reply_channel) - VALUES ($1, $2, $3, $4) - RETURNING id", - &[&lease.epoch, &read_version, &payload, &reply_channel], + "WITH ins AS ( + INSERT INTO udb_commit_requests (epoch, read_version, payload, reply_channel) + VALUES ($1, $2, $3, $4) + RETURNING id + ) + SELECT pg_notify($5, id::text), id FROM ins", + &[ + &lease.epoch, + &read_version, + &payload, + &reply_channel, + &commit_channel(&lease.leader_addr), + ], ) .await - .context("failed to enqueue commit request")? - .get(0); - - // Wake the leader's drain loop. - if let Err(err) = conn - .execute( - "SELECT pg_notify($1, $2)", - &[&commit_channel(&lease.leader_addr), &id.to_string()], - ) - .await - { - tracing::debug!( - ?err, - "failed to notify leader; relying on its poll backstop" - ); - } + .context("failed to enqueue and notify commit request")? + .get(1); // Release the connection before waiting so a long wait does not pin a pool slot. The request // row is durable, so await_result re-acquires a connection per poll. drop(conn); - await_result(shared, id, lease.epoch, &mut reply_rx).await + let submit_start = Instant::now(); + let result = await_result(shared, id, lease.epoch, &mut reply_rx).await; + tracing::debug!( + id, + epoch = lease.epoch, + wait_ms = submit_start.elapsed().as_millis() as u64, + ok = result.is_ok(), + "udb commit submit completed" + ); + result } /// Wait for a known leader, returning a retryable error if none is elected in time. @@ -98,53 +105,166 @@ async fn wait_for_leader(shared: &Arc) -> Result { } } -/// Poll the request row until it reaches a terminal status, woken by reply NOTIFYs with a polling -/// backstop. Bails as retryable if the leader epoch advances (our request is now orphaned and will -/// never be applied, so it is definitively not committed). +/// Wait for the commit result, resolved directly from the leader's reply NOTIFY payload on the happy +/// path. A polling `read_status` backstop covers a missed/lagged NOTIFY, and an epoch advance orphans +/// the request (it will never be applied, so it is definitively not committed). async fn await_result( shared: &Arc, id: i64, submit_epoch: i64, reply_rx: &mut tokio::sync::broadcast::Receiver, ) -> Result<()> { + let start = Instant::now(); + // Diagnostics: count how the waiter is driven so we can tell whether the reply NOTIFY is doing + // its job or whether commits are riding the slow poll backstop / getting orphaned by failover. + let mut status_reads = 0u32; + let mut notify_wakes = 0u32; + let mut poll_wakes = 0u32; + + // The backstop runs on a fixed-cadence interval, not a per-iteration sleep: under a flood of + // other commits' replies on this node's shared reply channel (which we skip past), a fresh + // per-iteration sleep would keep resetting and never fire, starving the backstop if our own + // NOTIFY was lost. An interval ticks on wall-clock cadence regardless of loop churn. + let mut poll_interval = tokio::time::interval(RESULT_POLL_INTERVAL); + poll_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + // Consume the immediate first tick so the first backstop is one interval out, after the reply + // has had a chance to arrive. + poll_interval.tick().await; + loop { - // Re-acquire a connection per poll: the request row is durable, so a transient pool/query - // error just means we retry the poll rather than failing a possibly-applied commit. + // Wait for our reply NOTIFY, falling back to a status read on a poll tick or a lagged + // broadcast. The happy path resolves straight from the payload with no status SELECT. + tokio::select! { + res = reply_rx.recv() => { + match res { + Ok(payload) => { + notify_wakes += 1; + match parse_reply(&payload, id) { + Some(ReplyOutcome::Committed) => { + tracing::debug!( + id, + wait_ms = start.elapsed().as_millis() as u64, + status_reads, + notify_wakes, + poll_wakes, + "udb commit resolved: committed" + ); + return Ok(()); + } + Some(ReplyOutcome::Conflict) => { + tracing::debug!( + id, + wait_ms = start.elapsed().as_millis() as u64, + status_reads, + notify_wakes, + poll_wakes, + "udb commit resolved: conflict" + ); + return Err(DatabaseError::NotCommitted.into()); + } + // Reply for another waiter on the shared channel; keep waiting for ours. + None => continue, + } + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => { + // We may have missed our own reply; fall through to the status backstop. + notify_wakes += 1; + tracing::debug!(id, lagged = n, "udb reply broadcast lagged"); + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + tracing::warn!(id, "udb reply broadcast closed; re-subscribing"); + *reply_rx = shared + .listener + .listen(&reply_channel(&shared.node_id)) + .await; + continue; + } + } + } + _ = poll_interval.tick() => { + poll_wakes += 1; + } + } + + // Backstop path (poll tick or lagged broadcast): re-acquire a connection and read the durable + // status. A transient pool/query error just means we retry rather than failing a + // possibly-applied commit. + status_reads += 1; match read_status(shared, id).await { - Ok(Some(Status::Committed)) => return Ok(()), - Ok(Some(Status::Conflict)) => return Err(DatabaseError::NotCommitted.into()), + Ok(Some(Status::Committed)) => { + tracing::debug!( + id, + wait_ms = start.elapsed().as_millis() as u64, + status_reads, + notify_wakes, + poll_wakes, + "udb commit resolved: committed (backstop)" + ); + return Ok(()); + } + Ok(Some(Status::Conflict)) => { + tracing::debug!( + id, + wait_ms = start.elapsed().as_millis() as u64, + status_reads, + notify_wakes, + poll_wakes, + "udb commit resolved: conflict (backstop)" + ); + return Err(DatabaseError::NotCommitted.into()); + } Ok(Some(Status::Pending)) => {} Ok(None) => { // The row was GC'd before we observed a terminal status. Treat as not committed // and let the retry loop resubmit. + tracing::warn!( + id, + wait_ms = start.elapsed().as_millis() as u64, + status_reads, + "udb commit row missing before terminal status (gc'd); treating as not committed" + ); return Err(DatabaseError::NotCommitted.into()); } Err(err) => { - tracing::debug!(?err, "transient error polling commit status, retrying"); + tracing::debug!(?err, id, "transient error polling commit status, retrying"); } } // If a new leader took over, our old-epoch request will never be claimed. if let Some(current) = shared.current_lease() { if current.epoch != submit_epoch { + tracing::warn!( + id, + submit_epoch, + current_epoch = current.epoch, + wait_ms = start.elapsed().as_millis() as u64, + notify_wakes, + poll_wakes, + "udb commit orphaned by leader failover; treating as not committed" + ); return Err(DatabaseError::NotCommitted.into()); } } + } +} - tokio::select! { - res = reply_rx.recv() => { - match res { - Ok(_) | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} - Err(tokio::sync::broadcast::error::RecvError::Closed) => { - *reply_rx = shared - .listener - .listen(&reply_channel(&shared.node_id)) - .await; - } - } - } - _ = tokio::time::sleep(RESULT_POLL_INTERVAL) => {} - } +enum ReplyOutcome { + Committed, + Conflict, +} + +/// Parse a leader reply payload (`":committed:"` or `":conflict"`). Returns +/// `None` when the payload is for a different waiter on the shared reply channel or is unparseable. +fn parse_reply(payload: &str, id: i64) -> Option { + let mut parts = payload.split(':'); + let reply_id: i64 = parts.next()?.parse().ok()?; + if reply_id != id { + return None; + } + match parts.next()? { + "committed" => Some(ReplyOutcome::Committed), + "conflict" => Some(ReplyOutcome::Conflict), + _ => None, } } diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs b/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs index d9729de884..278bb70fe7 100644 --- a/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs +++ b/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs @@ -1,66 +1,142 @@ +use std::collections::{BTreeMap, HashMap}; + use anyhow::{Context, Result}; -use deadpool_postgres::Transaction; use crate::{ atomic::apply_atomic_op, options::MutationType, tuple::Versionstamp, tx_ops::Operation, versionstamp::substitute_raw_versionstamp, }; -/// Apply a winning transaction's operations to `kv` inside the leader's batch txn. -/// -/// `commit_version` is the Postgres-resolved version assigned to this commit (`nextval`). It is -/// substituted into the 8-byte committed-version slot of any versionstamped key/value so -/// versionstamps are globally monotonic with commit order across all follower processes. -pub async fn apply( - txn: &Transaction<'_>, - operations: Vec, - commit_version: u64, -) -> Result<()> { - // Distinguishes multiple versionstamped operations within a single commit so their 10-byte - // stamps stay unique (8-byte version shared, 2-byte counter incremented). - let mut versionstamp_counter: u16 = 0; - - for op in operations { - match op { - Operation::SetValue { key, value } => { - upsert(txn, &key, &value).await?; - } - Operation::Clear { key } => { - txn.execute("DELETE FROM kv WHERE key = $1", &[&key]) - .await - .context("failed to clear key")?; - } - Operation::ClearRange { begin, end } => { - txn.execute( - "DELETE FROM kv WHERE key >= $1 AND key < $2", - &[&begin, &end], - ) - .await - .context("failed to clear range")?; +/// A winning commit's Postgres-resolved version and its decoded operations. Winners are folded in id +/// order. +pub struct Winner { + pub commit_version: u64, + pub operations: Vec, +} + +/// The materialized result of folding a batch of winners over the current `kv` state. Each distinct +/// key appears at most once across `upserts` and `point_deletes`, so the batch's point writes +/// collapse to a fixed number of statements regardless of batch size. +pub struct WriteSet { + pub upserts: Vec<(Vec, Vec)>, + pub point_deletes: Vec>, + pub range_deletes: Vec<(Vec, Vec)>, +} + +/// In-memory working state layered over the pre-batch `kv` snapshot. Folding the batch's winners +/// through this overlay in id order reproduces the exact serial semantics of applying each commit +/// one at a time: a later commit's atomic read observes an earlier commit's write because the +/// overlay is the live working state. +struct Overlay<'a> { + /// Pre-batch values for every key read by an atomic op in the batch. The only base read. + base: &'a HashMap, Vec>, + /// Point writes layered over the base. `Some` is a set, `None` is a point tombstone. + points: BTreeMap, Option>>, + /// Range tombstones in fold order. A key with no overlaying point write that falls in any range + /// reads as absent. + ranges: Vec<(Vec, Vec)>, +} + +impl<'a> Overlay<'a> { + fn new(base: &'a HashMap, Vec>) -> Self { + Overlay { + base, + points: BTreeMap::new(), + ranges: Vec::new(), + } + } + + /// Read-through lookup: an overlaying point write wins, then a range tombstone, then the base. + fn get(&self, key: &[u8]) -> Option> { + if let Some(value) = self.points.get(key) { + return value.clone(); + } + if self + .ranges + .iter() + .any(|(begin, end)| key >= begin.as_slice() && key < end.as_slice()) + { + return None; + } + self.base.get(key).cloned() + } + + fn set(&mut self, key: Vec, value: Vec) { + self.points.insert(key, Some(value)); + } + + fn clear(&mut self, key: Vec) { + self.points.insert(key, None); + } + + fn clear_range(&mut self, begin: Vec, end: Vec) { + // Drop any point writes inside the range; the range delete subsumes them, and a later set of + // a key in the range re-adds a point write that wins on read-through again. + let covered: Vec> = self + .points + .range(begin.clone()..end.clone()) + .map(|(key, _)| key.clone()) + .collect(); + for key in covered { + self.points.remove(&key); + } + self.ranges.push((begin, end)); + } + + fn into_write_set(self) -> WriteSet { + let mut upserts = Vec::new(); + let mut point_deletes = Vec::new(); + for (key, value) in self.points { + match value { + Some(value) => upserts.push((key, value)), + None => point_deletes.push(key), } - Operation::AtomicOp { - key, - param, - op_type, - } => { - apply_atomic( - txn, + } + WriteSet { + upserts, + point_deletes, + range_deletes: self.ranges, + } + } +} + +/// Fold every winner's operations, in id order, into a single materialized write-set over the +/// pre-batch `base` snapshot. `base` must contain the current value of every key returned by +/// [`atomic_read_keys`]. +pub fn fold_winners(winners: Vec, base: &HashMap, Vec>) -> Result { + let mut overlay = Overlay::new(base); + + for winner in winners { + // Distinguishes multiple versionstamped operations within a single commit so their 10-byte + // stamps stay unique (8-byte version shared, 2-byte counter incremented). Resets per winner. + let mut versionstamp_counter: u16 = 0; + + for op in winner.operations { + match op { + Operation::SetValue { key, value } => overlay.set(key, value), + Operation::Clear { key } => overlay.clear(key), + Operation::ClearRange { begin, end } => overlay.clear_range(begin, end), + Operation::AtomicOp { key, param, op_type, - commit_version, + } => fold_atomic( + &mut overlay, + key, + param, + op_type, + winner.commit_version, &mut versionstamp_counter, - ) - .await?; + )?, } } } - Ok(()) + Ok(overlay.into_write_set()) } -async fn apply_atomic( - txn: &Transaction<'_>, +fn fold_atomic( + overlay: &mut Overlay<'_>, key: Vec, param: Vec, op_type: MutationType, @@ -73,17 +149,17 @@ async fn apply_atomic( let key = substitute_raw_versionstamp(key, &versionstamp) .map_err(anyhow::Error::msg) .context("failed substituting versionstamped key")?; - upsert(txn, &key, ¶m).await?; + overlay.set(key, param); } MutationType::SetVersionstampedValue => { let versionstamp = build_versionstamp(commit_version, versionstamp_counter); let value = substitute_raw_versionstamp(param, &versionstamp) .map_err(anyhow::Error::msg) .context("failed substituting versionstamped value")?; - upsert(txn, &key, &value).await?; + overlay.set(key, value); } - // Read-modify-write atomics: the leader is the single writer, so reading the live value - // inside the apply txn and writing the result is serializable with no lost update. + // Read-modify-write atomics: the leader is the single writer, so reading the overlay's + // working value and writing the result is serializable with no lost update. MutationType::Add | MutationType::And | MutationType::BitAnd @@ -97,20 +173,10 @@ async fn apply_atomic( | MutationType::ByteMin | MutationType::ByteMax | MutationType::CompareAndClear => { - let current = txn - .query_opt("SELECT value FROM kv WHERE key = $1", &[&key]) - .await - .context("failed to read current value for atomic op")? - .map(|row| row.get::<_, Vec>(0)); - - let new_value = apply_atomic_op(current.as_deref(), ¶m, op_type); - - if let Some(new_value) = new_value { - upsert(txn, &key, &new_value).await?; - } else { - txn.execute("DELETE FROM kv WHERE key = $1", &[&key]) - .await - .context("failed to clear key after atomic op")?; + let current = overlay.get(&key); + match apply_atomic_op(current.as_deref(), ¶m, op_type) { + Some(new_value) => overlay.set(key, new_value), + None => overlay.clear(key), } } } @@ -118,14 +184,39 @@ async fn apply_atomic( Ok(()) } -async fn upsert(txn: &Transaction<'_>, key: &[u8], value: &[u8]) -> Result<()> { - txn.execute( - "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2", - &[&key, &value], - ) - .await - .context("failed to upsert kv")?; - Ok(()) +/// Every key a winner's atomic op reads, so the leader can fetch them all in one bulk query before +/// folding. Versionstamped ops do not read, so their keys are skipped. +pub fn atomic_read_keys(winners: &[Winner]) -> Vec> { + let mut keys = Vec::new(); + for winner in winners { + for op in &winner.operations { + if let Operation::AtomicOp { key, op_type, .. } = op { + if reads_current_value(*op_type) { + keys.push(key.clone()); + } + } + } + } + keys +} + +fn reads_current_value(op_type: MutationType) -> bool { + match op_type { + MutationType::SetVersionstampedKey | MutationType::SetVersionstampedValue => false, + MutationType::Add + | MutationType::And + | MutationType::BitAnd + | MutationType::Or + | MutationType::BitOr + | MutationType::Xor + | MutationType::BitXor + | MutationType::AppendIfFits + | MutationType::Max + | MutationType::Min + | MutationType::ByteMin + | MutationType::ByteMax + | MutationType::CompareAndClear => true, + } } /// Build a 10-byte versionstamp (plus the 2 user-version bytes the substitution helper ignores) @@ -137,3 +228,9 @@ fn build_versionstamp(commit_version: u64, counter: &mut u16) -> Versionstamp { *counter = counter.wrapping_add(1); Versionstamp::from(bytes) } + +// The fold operates on private types (`Overlay`, `Winner`, `WriteSet`) that integration tests under +// `tests/` cannot reach, so its unit tests live in a source-owned sibling file via this shim. +#[cfg(test)] +#[path = "apply_tests.rs"] +mod tests; diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/apply_tests.rs b/engine/packages/universaldb/src/driver/postgres/resolver/apply_tests.rs new file mode 100644 index 0000000000..95d991176d --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/resolver/apply_tests.rs @@ -0,0 +1,299 @@ +use std::collections::HashMap; + +use super::*; + +fn k(s: &str) -> Vec { + s.as_bytes().to_vec() +} + +fn add_param(n: i64) -> Vec { + n.to_le_bytes().to_vec() +} + +fn read_int(value: &[u8]) -> i64 { + let mut buf = [0u8; 8]; + let len = value.len().min(8); + buf[..len].copy_from_slice(&value[..len]); + i64::from_le_bytes(buf) +} + +fn set(key: &str, value: &str) -> Operation { + Operation::SetValue { + key: k(key), + value: k(value), + } +} + +fn clear(key: &str) -> Operation { + Operation::Clear { key: k(key) } +} + +fn clear_range(begin: &str, end: &str) -> Operation { + Operation::ClearRange { + begin: k(begin), + end: k(end), + } +} + +fn add(key: &str, n: i64) -> Operation { + Operation::AtomicOp { + key: k(key), + param: add_param(n), + op_type: MutationType::Add, + } +} + +fn winner(commit_version: u64, operations: Vec) -> Winner { + Winner { + commit_version, + operations, + } +} + +/// Materialize a write-set over a base map to inspect the resulting `kv` state. +fn materialize(base: &HashMap, Vec>, write_set: WriteSet) -> HashMap, Vec> { + let mut state = base.clone(); + for (begin, end) in &write_set.range_deletes { + state.retain(|key, _| { + !(key.as_slice() >= begin.as_slice() && key.as_slice() < end.as_slice()) + }); + } + for key in &write_set.point_deletes { + state.remove(key); + } + for (key, value) in write_set.upserts { + state.insert(key, value); + } + state +} + +/// Reference oracle: apply each winner's operations one at a time directly to a working state, the +/// exact serial semantics the fold must reproduce. +fn reference(base: &HashMap, Vec>, winners: &[Winner]) -> HashMap, Vec> { + let mut state = base.clone(); + for w in winners { + let mut counter: u16 = 0; + for op in &w.operations { + match op { + Operation::SetValue { key, value } => { + state.insert(key.clone(), value.clone()); + } + Operation::Clear { key } => { + state.remove(key); + } + Operation::ClearRange { begin, end } => { + state.retain(|key, _| { + !(key.as_slice() >= begin.as_slice() && key.as_slice() < end.as_slice()) + }); + } + Operation::AtomicOp { + key, + param, + op_type, + } => match op_type { + MutationType::SetVersionstampedKey => { + let vs = build_versionstamp(w.commit_version, &mut counter); + let new_key = substitute_raw_versionstamp(key.clone(), &vs).unwrap(); + state.insert(new_key, param.clone()); + } + MutationType::SetVersionstampedValue => { + let vs = build_versionstamp(w.commit_version, &mut counter); + let new_value = substitute_raw_versionstamp(param.clone(), &vs).unwrap(); + state.insert(key.clone(), new_value); + } + _ => { + let current = state.get(key).map(|v| v.as_slice()); + match apply_atomic_op(current, param, *op_type) { + Some(v) => { + state.insert(key.clone(), v); + } + None => { + state.remove(key); + } + } + } + }, + } + } + } + state +} + +/// Build a 4-byte-offset-trailed buffer that `substitute_raw_versionstamp` can stamp at `offset`. +fn stampable(prefix: &[u8], offset: u32) -> Vec { + let mut buf = Vec::new(); + buf.extend_from_slice(prefix); + buf.extend_from_slice(&[0u8; 10]); + buf.extend_from_slice(&offset.to_le_bytes()); + buf +} + +/// `fold_winners` consumes its input, so clone for tests that also run the oracle on the same data. +fn fold_winners_clone(winners: &[Winner], base: &HashMap, Vec>) -> WriteSet { + let cloned: Vec = winners + .iter() + .map(|w| Winner { + commit_version: w.commit_version, + operations: w.operations.clone(), + }) + .collect(); + fold_winners(cloned, base).unwrap() +} + +#[test] +fn same_key_set_later_id_wins() { + let winners = vec![ + winner(1, vec![set("a", "first")]), + winner(2, vec![set("a", "second")]), + ]; + let base = HashMap::new(); + let ws = fold_winners_clone(&winners, &base); + let state = materialize(&base, ws); + assert_eq!( + state.get(&k("a")).map(|v| v.as_slice()), + Some(b"second".as_slice()) + ); + assert_eq!(state, reference(&base, &winners)); +} + +#[test] +fn two_adds_same_key_fold_sequentially() { + // 5 -> 6 -> 7 across two commits in one batch. + let mut base = HashMap::new(); + base.insert(k("n"), add_param(5)); + let winners = vec![winner(1, vec![add("n", 1)]), winner(2, vec![add("n", 1)])]; + let ws = fold_winners_clone(&winners, &base); + let state = materialize(&base, ws); + assert_eq!(read_int(state.get(&k("n")).unwrap()), 7); + assert_eq!(state, reference(&base, &winners)); +} + +#[test] +fn set_then_add_atomic_sees_the_set() { + let base = HashMap::new(); + let winners = vec![ + winner(1, vec![set("n", "\x0a\0\0\0\0\0\0\0")]), + winner(2, vec![add("n", 1)]), + ]; + let ws = fold_winners_clone(&winners, &base); + let state = materialize(&base, ws); + assert_eq!(read_int(state.get(&k("n")).unwrap()), 11); + assert_eq!(state, reference(&base, &winners)); +} + +#[test] +fn clear_range_then_add_sees_none() { + let mut base = HashMap::new(); + base.insert(k("r/x"), add_param(100)); + let winners = vec![ + winner(1, vec![clear_range("r/", "r0")]), + winner(2, vec![add("r/x", 1)]), + ]; + let ws = fold_winners_clone(&winners, &base); + let state = materialize(&base, ws); + // The range clear wiped the base 100, so the add starts from absent/0 -> 1. + assert_eq!(read_int(state.get(&k("r/x")).unwrap()), 1); + assert_eq!(state, reference(&base, &winners)); +} + +#[test] +fn set_inside_cleared_range_is_reinserted() { + let mut base = HashMap::new(); + base.insert(k("r/x"), k("old")); + let winners = vec![winner(1, vec![clear_range("r/", "r0"), set("r/x", "new")])]; + let ws = fold_winners_clone(&winners, &base); + let state = materialize(&base, ws); + assert_eq!( + state.get(&k("r/x")).map(|v| v.as_slice()), + Some(b"new".as_slice()) + ); + assert_eq!(state, reference(&base, &winners)); +} + +#[test] +fn versionstamped_key_and_value_distinct_stamps() { + let base = HashMap::new(); + let winners = vec![winner( + 42, + vec![ + Operation::AtomicOp { + key: stampable(b"vk/", 3), + param: k("v1"), + op_type: MutationType::SetVersionstampedKey, + }, + Operation::AtomicOp { + key: k("vv/"), + param: stampable(b"", 0), + op_type: MutationType::SetVersionstampedValue, + }, + ], + )]; + let ws = fold_winners_clone(&winners, &base); + let state = materialize(&base, ws); + assert_eq!(state, reference(&base, &winners)); + + // The versionstamped key embeds commit_version 42 with per-commit counter 0. + let stamped_key: Vec = { + let mut key = b"vk/".to_vec(); + key.extend_from_slice(&42u64.to_be_bytes()); + key.extend_from_slice(&0u16.to_be_bytes()); + key + }; + assert_eq!( + state.get(&stamped_key).map(|v| v.as_slice()), + Some(b"v1".as_slice()) + ); + + // The versionstamped value uses the same commit_version but the next counter (1). + let stamped_value = state.get(&k("vv/")).unwrap(); + assert_eq!(&stamped_value[0..8], &42u64.to_be_bytes()); + assert_eq!(&stamped_value[8..10], &1u16.to_be_bytes()); +} + +#[test] +fn point_delete_and_upsert_disjoint() { + let mut base = HashMap::new(); + base.insert(k("keep"), k("base")); + base.insert(k("drop"), k("base")); + let winners = vec![winner(1, vec![clear("drop"), set("keep", "new")])]; + let ws = fold_winners_clone(&winners, &base); + assert_eq!(ws.point_deletes, vec![k("drop")]); + assert_eq!(ws.upserts, vec![(k("keep"), k("new"))]); + let state = materialize(&base, ws); + assert_eq!(state, reference(&base, &winners)); +} + +/// A multi-op, multi-winner batch mixing every operation kind must equal the serial oracle. +#[test] +fn mixed_batch_matches_oracle() { + let mut base = HashMap::new(); + base.insert(k("c1"), add_param(10)); + base.insert(k("c2"), add_param(20)); + base.insert(k("old"), k("x")); + base.insert(k("range/a"), k("ra")); + base.insert(k("range/b"), k("rb")); + + let winners = vec![ + winner(1, vec![set("c1", "\x01\0\0\0\0\0\0\0"), add("c1", 4)]), + winner(2, vec![clear("old"), add("c2", 5)]), + winner( + 3, + vec![clear_range("range/", "range0"), set("range/a", "back")], + ), + winner(4, vec![add("c2", 100), add("c1", 1)]), + ]; + + let ws = fold_winners_clone(&winners, &base); + let state = materialize(&base, ws); + assert_eq!(state, reference(&base, &winners)); + + // Spot checks of the folded result. + assert_eq!(read_int(state.get(&k("c1")).unwrap()), 6); // 1 +4 +1 + assert_eq!(read_int(state.get(&k("c2")).unwrap()), 125); // 20 +5 +100 + assert!(state.get(&k("old")).is_none()); + assert_eq!( + state.get(&k("range/a")).map(|v| v.as_slice()), + Some(b"back".as_slice()) + ); + assert!(state.get(&k("range/b")).is_none()); +} diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs index f9b6c45461..e07304180c 100644 --- a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs +++ b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs @@ -2,15 +2,19 @@ mod apply; mod lease; use std::{ + collections::HashMap, sync::Arc, time::{Duration, Instant}, }; use anyhow::{Context, Result}; use tokio::sync::broadcast; +use tokio_util::task::AbortOnDropHandle; use crate::{conflict_tracker::TransactionConflictTracker, transaction::TXN_TIMEOUT}; +use lease::LEASE_TTL_SECS; + use super::shared::{ ELECTION_CHANNEL, LEASE_ID, LeaseInfo, PostgresShared, WATERMARK_CHANNEL, commit_channel, }; @@ -28,13 +32,6 @@ const POLL_BACKSTOP: Duration = Duration::from_millis(50); /// How long a candidate waits before retrying election when another node holds the lease. const ELECTION_RETRY: Duration = Duration::from_secs(2); -enum DrainOutcome { - /// Processed zero or more requests; still leader. - Drained, - /// Lost the lease (epoch bumped by a new leader). Step down. - LostLease, -} - /// Spawn the per-process resolver task. Every node runs this; only the elected leader drains the /// commit queue. The returned handle is aborted when the owning driver drops, which stops lease /// renewal so the lease expires and another node can take over (node-death / failover path). @@ -119,7 +116,12 @@ async fn notify_election(shared: &Arc) { } } -/// Leader main loop: hold the lease, drain the commit queue on wake or poll, and renew the lease. +/// Leader entry point: publish our lease, compute the cold-window floor, then run renewal and +/// draining as two sibling tasks. They coordinate purely by completion: when either returns (lease +/// lost or error), the other is aborted and the leader steps down. Both operations are safe to +/// hard-abort, so no explicit cancellation signalling is needed. A renew is a single fenced +/// `UPDATE`; a drain batch runs in one Postgres transaction that rolls back cleanly when dropped, +/// leaving its claimed requests `pending` for the next leader. async fn lead(shared: &Arc, epoch: i64) -> Result<()> { // Publish our own lease into the cache immediately so our local commits route to us. shared.set_lease(LeaseInfo { @@ -133,97 +135,191 @@ async fn lead(shared: &Arc, epoch: i64) -> Result<()> { let recovery_version = recovery_floor(shared).await?; let recovery_deadline = Instant::now() + TXN_TIMEOUT; - let tracker = TransactionConflictTracker::new(); + tracing::info!( + epoch, + recovery_version, + cold_window_ms = TXN_TIMEOUT.as_millis() as u64, + "udb leader entering lead loop" + ); + + // Renewal runs in its own task so the drain loop can never starve it: if renewal shared the + // drain loop, a single long drain under sustained load would block renewal past the lease TTL + // and the lease would be lost mid-drain, thrashing leadership. Both are held in abort-on-drop + // handles so a hard abort of this `run` task (driver drop / node death) tears them down. A leaked + // renew task would keep this dead leader's lease alive and block failover. + let mut renew = AbortOnDropHandle::new(tokio::spawn(renew_loop(shared.clone(), epoch))); + let mut drain = AbortOnDropHandle::new(tokio::spawn(drain_loop( + shared.clone(), + epoch, + recovery_version, + recovery_deadline, + ))); - let mut wake_rx = shared - .listener - .listen(&commit_channel(&shared.node_id)) - .await; + // Whichever task returns first (lease lost or error), step down; the other is aborted when its + // handle drops at the end of this scope. A clean exit yields its inner result; a panic surfaces + // through `?` as a join error. + tokio::select! { + res = &mut renew => res?, + res = &mut drain => res?, + } +} - let mut renew_interval = tokio::time::interval(RENEW_INTERVAL); - renew_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - let mut poll_interval = tokio::time::interval(POLL_BACKSTOP); - poll_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); +/// Lease-renewal loop. Runs on its own task and pool connection so it cannot be starved by drain +/// work. Returns when the lease is definitively gone (epoch bumped by another node, or renewal +/// failing for the whole lease TTL), which causes [`lead`] to abort the drain task and step down. +async fn renew_loop(shared: Arc, epoch: i64) -> Result<()> { + let mut interval = tokio::time::interval(RENEW_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + // The lease was just acquired/renewed (expires_at = now + TTL), so consume the immediate first + // tick and renew after one interval. + interval.tick().await; - // Drain anything already queued before our first wake. - if matches!( - drain(shared, epoch, &tracker, recovery_version, recovery_deadline).await?, - DrainOutcome::LostLease - ) { - return Ok(()); - } + let mut last_renew = Instant::now(); loop { - tokio::select! { - _ = renew_interval.tick() => { - if !lease::renew(&shared.pool, &shared.node_id, epoch).await? { - tracing::warn!(epoch, "lost udb lease on renew"); - return Ok(()); + interval.tick().await; + + let gap_ms = last_renew.elapsed().as_millis() as u64; + let renew_start = Instant::now(); + match lease::renew(&shared.pool, &shared.node_id, epoch).await { + Ok(true) => { + last_renew = Instant::now(); + let renew_query_ms = renew_start.elapsed().as_millis() as u64; + // With renewal on its own task this gap should track RENEW_INTERVAL closely; a large + // gap now points at pool or Postgres contention. + if gap_ms > RENEW_INTERVAL.as_millis() as u64 * 2 { + tracing::warn!( + epoch, + gap_since_last_renew_ms = gap_ms, + renew_query_ms, + "udb leader renew was delayed (pool or postgres contention)" + ); + } else { + tracing::debug!( + epoch, + gap_since_last_renew_ms = gap_ms, + renew_query_ms, + "udb leader renewed lease" + ); } } - res = wake_rx.recv() => { - match res { - Ok(_) | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} - Err(tokio::sync::broadcast::error::RecvError::Closed) => { - wake_rx = shared.listener.listen(&commit_channel(&shared.node_id)).await; - } - } - if matches!( - drain(shared, epoch, &tracker, recovery_version, recovery_deadline).await?, - DrainOutcome::LostLease - ) { - return Ok(()); - } + Ok(false) => { + tracing::warn!( + epoch, + gap_since_last_renew_ms = gap_ms, + "udb leader lost lease on renew (epoch bumped by another node); stepping down" + ); + return Ok(()); } - _ = poll_interval.tick() => { - if matches!( - drain(shared, epoch, &tracker, recovery_version, recovery_deadline).await?, - DrainOutcome::LostLease - ) { + Err(err) => { + // A transient renew error is tolerable within the TTL; keep retrying. Only give up + // if we have been unable to renew for the whole lease TTL, at which point we can no + // longer assume we hold the lease. + if last_renew.elapsed() >= Duration::from_secs(LEASE_TTL_SECS as u64) { + tracing::warn!( + ?err, + epoch, + "udb leader renew failing past lease TTL; assuming lease lost and stepping down" + ); return Ok(()); } + tracing::warn!(?err, epoch, "udb leader lease renew errored; will retry"); } } } } -/// The version floor a freshly elected leader continues from: the higher of the durable watermark -/// and the sequence high-water. The LOGGED `udb_version_seq` is crash-safe, so this never regresses. -async fn recovery_floor(shared: &Arc) -> Result { - let durable = lease::current_durable_version(&shared.pool).await?; - - let conn = shared - .pool - .get() - .await - .context("failed to get connection for recovery floor")?; - let seq_high: i64 = conn - .query_one("SELECT last_value FROM udb_version_seq", &[]) - .await - .context("failed to read sequence high water")? - .get(0); - - Ok(durable.max(seq_high).max(0) as u64) -} - -/// Drain pending commit requests in id-ordered batches until none remain. Each batch resolves and -/// applies inside a single Postgres transaction (group commit), fenced on the leader's epoch. -async fn drain( - shared: &Arc, +/// Drain loop. Processes one batch per iteration. Draining until the queue emptied in a single call +/// could run for many seconds under sustained load; processing one batch at a time keeps the loop +/// at a clean await point between batches. A non-empty queue still drains back-to-back with no idle +/// wait, so throughput is unchanged. It blocks on `select!` (wake NOTIFY or poll backstop) only +/// when the queue is empty. Returns when this leader's epoch is fenced out; otherwise [`lead`] +/// aborts it when renewal reports the lease is lost. +async fn drain_loop( + shared: Arc, epoch: i64, - tracker: &TransactionConflictTracker, recovery_version: u64, recovery_deadline: Instant, -) -> Result { +) -> Result<()> { + let tracker = TransactionConflictTracker::new(); + + let mut wake_rx = shared + .listener + .listen(&commit_channel(&shared.node_id)) + .await; + + let mut poll_interval = tokio::time::interval(POLL_BACKSTOP); + poll_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { - match drain_batch(shared, epoch, tracker, recovery_version, recovery_deadline).await? { - BatchOutcome::Empty => return Ok(DrainOutcome::Drained), - BatchOutcome::Processed => {} - BatchOutcome::LostLease => return Ok(DrainOutcome::LostLease), + match drain_batch( + &shared, + epoch, + &tracker, + recovery_version, + recovery_deadline, + ) + .await? + { + BatchOutcome::LostLease => { + tracing::warn!( + epoch, + "udb leader stepping down: lost lease during drain (epoch fenced on watermark)" + ); + return Ok(()); + } + // More work may be pending; loop immediately to keep throughput up. + BatchOutcome::Processed => continue, + BatchOutcome::Empty => { + tokio::select! { + res = wake_rx.recv() => { + match res { + Ok(_) | Err(broadcast::error::RecvError::Lagged(_)) => {} + Err(broadcast::error::RecvError::Closed) => { + wake_rx = shared + .listener + .listen(&commit_channel(&shared.node_id)) + .await; + } + } + } + _ = poll_interval.tick() => {} + } + } } } } +/// The cold-window rejection floor for a freshly elected leader: the durable watermark +/// (`udb_lease.durable_version`) at election time. +/// +/// Reasoning (do NOT change this back to `max(durable, seq_high)`): +/// +/// A new leader starts with an empty conflict tracker, so it cannot detect a read-write conflict +/// against any committed write it does not already know about. The writes it is missing are exactly +/// the previous leader's winners, and every winner's write is applied to `kv` AND its +/// `commit_version` folded into `durable_version` in the SAME apply transaction. So every missing +/// write has `commit_version <= durable_version`. A committing transaction `T` is therefore safe +/// iff `T.read_version >= durable_version`: every write above its read_version was committed by THIS +/// leader and is in the tracker. Only `T.read_version < durable_version` can race a missing winner, +/// so that is the exact set the cold window must reject. +/// +/// `udb_version_seq.last_value` (the sequence high-water) is NOT a valid floor. Every drained +/// request consumes a `nextval` BEFORE the conflict check, including conflicts and cold rejects, so +/// the sequence races far ahead of `durable_version` with versions that never produced any write. +/// Using `max(durable, seq_high)` rejects essentially every commit for the whole cold window +/// (followers read at `durable_version`, which is always `< seq_high`), turning each failover into +/// a 5s mass-reject storm. The gap `(durable_version, seq_high]` holds only thrown-away loser +/// versions, so nothing in it is a missing write to guard against. +/// +/// Version ASSIGNMENT is unaffected: commit versions still come from `nextval('udb_version_seq')` +/// in `drain_batch`, which is always above the sequence high-water, so uniqueness and monotonicity +/// across failover are preserved independently of this floor. +async fn recovery_floor(shared: &Arc) -> Result { + let durable = lease::current_durable_version(&shared.pool).await?; + Ok(durable.max(0) as u64) +} + enum BatchOutcome { Empty, Processed, @@ -232,7 +328,9 @@ enum BatchOutcome { struct Reply { channel: String, - id: i64, + /// The follower's reply payload, encoding the outcome so the waiter resolves without a status + /// SELECT: `":committed:"` or `":conflict"`. + payload: String, } async fn drain_batch( @@ -273,30 +371,55 @@ async fn drain_batch( return Ok(BatchOutcome::Empty); } + let batch_start = Instant::now(); let cold_window = Instant::now() < recovery_deadline; + let batch_len = rows.len(); let mut max_winner_cv: i64 = 0; - let mut replies = Vec::with_capacity(rows.len()); - - for row in &rows { + let mut replies = Vec::with_capacity(batch_len); + let mut committed_count = 0u32; + let mut conflict_count = 0u32; + let mut cold_reject_count = 0u32; + + // Allocate all commit versions for the batch in one round-trip instead of a `nextval` per row. + // They are assigned to rows in id order (winners and losers alike; losers' versions are + // harmlessly skipped), so versionstamps stay monotonic with commit order. The defensive sort + // keeps assignment monotonic regardless of how Postgres orders the per-row `nextval` evaluation. + let mut versions: Vec = txn + .query( + "SELECT nextval('udb_version_seq') FROM generate_series(1, $1::bigint)", + &[&(batch_len as i64)], + ) + .await + .context("failed to allocate commit versions")? + .iter() + .map(|row| row.get::<_, i64>(0)) + .collect(); + versions.sort_unstable(); + + // Resolve every request in memory in id order. Winners are collected with their version and + // operations for the fold; the bulk status stamp is built for all rows at once. + let mut winners: Vec = Vec::new(); + let mut stamp_ids: Vec = Vec::with_capacity(batch_len); + let mut stamp_statuses: Vec<&str> = Vec::with_capacity(batch_len); + let mut stamp_versions: Vec> = Vec::with_capacity(batch_len); + + for (i, row) in rows.iter().enumerate() { let id: i64 = row.get(0); let read_version: i64 = row.get(1); let payload: Vec = row.get(2); let reply_channel: String = row.get(3); + let commit_version = versions[i]; let decoded = super::codec::decode_commit_request(&payload) .context("failed to decode commit payload")?; - let commit_version: i64 = txn - .query_one("SELECT nextval('udb_version_seq')", &[]) - .await - .context("failed to get next commit version")? - .get(0); - let start_version = read_version.max(0) as u64; // Cold-window guard: a commit whose read_version predates the recovery floor cannot be // safely resolved against this leader's empty window. Reject it as retryable. - let conflicted = if cold_window && start_version < recovery_version { + let cold_rejected = cold_window && start_version < recovery_version; + let conflicted = if cold_rejected { + cold_reject_count += 1; true } else { tracker @@ -308,32 +431,98 @@ async fn drain_batch( .await }; + stamp_ids.push(id); if conflicted { - txn.execute( - "UPDATE udb_commit_requests SET status = 'conflict' WHERE id = $1", - &[&id], - ) - .await - .context("failed to stamp conflict")?; + if !cold_rejected { + conflict_count += 1; + } + stamp_statuses.push("conflict"); + stamp_versions.push(None); } else { - apply::apply(&txn, decoded.operations, commit_version.max(0) as u64) - .await - .context("failed to apply commit")?; - txn.execute( - "UPDATE udb_commit_requests SET status = 'committed', commit_version = $1 WHERE id = $2", - &[&commit_version, &id], - ) - .await - .context("failed to stamp committed")?; + committed_count += 1; + stamp_statuses.push("committed"); + stamp_versions.push(Some(commit_version)); max_winner_cv = max_winner_cv.max(commit_version); + winners.push(apply::Winner { + commit_version: commit_version.max(0) as u64, + operations: decoded.operations, + }); } + let reply_payload = if conflicted { + format!("{id}:conflict") + } else { + format!("{id}:committed:{commit_version}") + }; replies.push(Reply { channel: reply_channel, - id, + payload: reply_payload, }); } + // Bulk-read the pre-batch value of every key a winner's atomic op reads in one query, then fold + // all winners into a single materialized write-set in memory. This collapses the per-row apply + // round-trips to a fixed count independent of batch size. + let atomic_keys = apply::atomic_read_keys(&winners); + let base = if atomic_keys.is_empty() { + HashMap::new() + } else { + txn.query( + "SELECT key, value FROM kv WHERE key = ANY($1::bytea[])", + &[&atomic_keys], + ) + .await + .context("failed to bulk-read atomic op keys")? + .into_iter() + .map(|row| (row.get::<_, Vec>(0), row.get::<_, Vec>(1))) + .collect() + }; + + let apply::WriteSet { + upserts, + point_deletes, + range_deletes, + } = apply::fold_winners(winners, &base).context("failed to fold batch winners")?; + + // Materialize the write-set in O(1) statements per kind. Range deletes run first so a key whose + // final state is a set but that fell inside an earlier range clear is re-inserted by the upsert, + // not removed. + for (begin, end) in &range_deletes { + txn.execute("DELETE FROM kv WHERE key >= $1 AND key < $2", &[begin, end]) + .await + .context("failed to clear range")?; + } + if !point_deletes.is_empty() { + txn.execute( + "DELETE FROM kv WHERE key = ANY($1::bytea[])", + &[&point_deletes], + ) + .await + .context("failed to bulk-delete cleared keys")?; + } + if !upserts.is_empty() { + let (keys, values): (Vec>, Vec>) = upserts.into_iter().unzip(); + txn.execute( + "INSERT INTO kv (key, value) + SELECT * FROM unnest($1::bytea[], $2::bytea[]) + ON CONFLICT (key) DO UPDATE SET value = excluded.value", + &[&keys, &values], + ) + .await + .context("failed to bulk-upsert kv")?; + } + + // Stamp every request's terminal status in one statement instead of a per-row UPDATE. + txn.execute( + "UPDATE udb_commit_requests AS r + SET status = b.status, commit_version = b.cv + FROM unnest($1::bigint[], $2::text[], $3::bigint[]) AS b(id, status, cv) + WHERE r.id = b.id", + &[&stamp_ids, &stamp_statuses, &stamp_versions], + ) + .await + .context("failed to stamp commit statuses")?; + // Advance the watermark, fenced on our epoch. A zombie old leader whose epoch was bumped sees // zero rows updated and must step down before any of its writes become visible. let new_durable: i64 = match txn @@ -362,6 +551,18 @@ async fn drain_batch( notify_after_commit(&conn, new_durable, &replies).await; + tracing::info!( + epoch, + batch_len, + committed = committed_count, + conflict = conflict_count, + cold_reject = cold_reject_count, + cold_window, + new_durable, + batch_ms = batch_start.elapsed().as_millis() as u64, + "udb leader processed commit batch" + ); + Ok(BatchOutcome::Processed) } @@ -383,11 +584,11 @@ async fn notify_after_commit( } let channels: Vec<&str> = replies.iter().map(|r| r.channel.as_str()).collect(); - let ids: Vec = replies.iter().map(|r| r.id.to_string()).collect(); + let payloads: Vec<&str> = replies.iter().map(|r| r.payload.as_str()).collect(); if let Err(err) = conn .execute( "SELECT pg_notify(c, p) FROM unnest($1::text[], $2::text[]) AS t(c, p)", - &[&channels, &ids], + &[&channels, &payloads], ) .await { diff --git a/engine/packages/universaldb/src/driver/postgres/shared.rs b/engine/packages/universaldb/src/driver/postgres/shared.rs index 5b14c4df08..b26f14a361 100644 --- a/engine/packages/universaldb/src/driver/postgres/shared.rs +++ b/engine/packages/universaldb/src/driver/postgres/shared.rs @@ -97,6 +97,21 @@ impl PostgresShared { /// Publish a freshly observed/elected lease into the cache. pub fn set_lease(&self, lease: LeaseInfo) { + let changed = self + .lease_rx + .borrow() + .as_ref() + .map(|prev| prev.epoch != lease.epoch || prev.leader_addr != lease.leader_addr) + .unwrap_or(true); + if changed { + tracing::info!( + epoch = lease.epoch, + leader_addr = %lease.leader_addr, + self_node = %self.node_id, + is_self = (lease.leader_addr == self.node_id), + "udb follower observed leader lease change" + ); + } let _ = self.lease_tx.send(Some(lease)); } diff --git a/self-host/compose/template/src/docker-compose.ts b/self-host/compose/template/src/docker-compose.ts index 4e3781aab2..2245a31c75 100644 --- a/self-host/compose/template/src/docker-compose.ts +++ b/self-host/compose/template/src/docker-compose.ts @@ -292,6 +292,7 @@ export function generateDockerCompose(context: TemplateContext) { restart: "unless-stopped", environment: [ "RUST_LOG_ANSI_COLOR=1", + "RUST_LOG=universaldb::driver::postgres=debug", "RIVET_OTEL_ENABLED=1", "RIVET_OTEL_SAMPLER_RATIO=1", `RIVET_OTEL_GRPC_ENDPOINT=http://${context.getServiceHost("otel-collector", datacenter.name)}:4317`, diff --git a/self-host/dev-multinode/docker-compose.yml b/self-host/dev-multinode/docker-compose.yml index 389510f7b3..28234658d2 100644 --- a/self-host/dev-multinode/docker-compose.yml +++ b/self-host/dev-multinode/docker-compose.yml @@ -179,6 +179,7 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector:4317 @@ -219,6 +220,7 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector:4317 @@ -257,6 +259,7 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector:4317 From 0c861be1cc95dae0ae6def2b97f6c0e809de6128 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Fri, 26 Jun 2026 10:29:03 -0700 Subject: [PATCH 15/16] [slopfix] perf(universaldb): fold postgres drain_batch into claim+nextval and single apply CTE --- .../universaldb/src/conflict_tracker.rs | 66 ++++--- .../src/driver/postgres/resolver/mod.rs | 160 +++++++++------ .../src/driver/rocksdb/transaction_task.rs | 2 +- .../universaldb/tests/kv_upsert_scaling.rs | 185 ++++++++++++++++++ 4 files changed, 325 insertions(+), 88 deletions(-) create mode 100644 engine/packages/universaldb/tests/kv_upsert_scaling.rs diff --git a/engine/packages/universaldb/src/conflict_tracker.rs b/engine/packages/universaldb/src/conflict_tracker.rs index 127e2bf615..3e543a99ee 100644 --- a/engine/packages/universaldb/src/conflict_tracker.rs +++ b/engine/packages/universaldb/src/conflict_tracker.rs @@ -1,4 +1,6 @@ use std::{ + collections::BTreeMap, + ops::Bound, sync::{ Arc, atomic::{AtomicU64, Ordering}, @@ -18,7 +20,6 @@ const TXN_CONFLICT_TTL: Duration = Duration::from_secs(10); struct PreviousTransaction { insert_instant: Instant, start_version: u64, - commit_version: u64, conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, } @@ -33,18 +34,18 @@ struct PreviousTransaction { /// `check_and_insert` takes the commit version from the caller instead of generating it. #[derive(Clone)] pub struct TransactionConflictTracker { - // NOTE: We use a mutex because we need to lock reads across all active txns. This could be optimized to - // only lock txns that have overlapping ranges with the currently checking one, but its a small - // optimization because most txns are going to be very recent and this only stores the last 10 seconds of - // txns. - txns: Arc>>, + // Keyed by commit version, which is unique per committed transaction (unlike start version, which + // concurrent transactions can share). The ordering lets the conflict scan skip transactions whose + // commit version cannot overlap the committing transaction, and lets pruning drop expired entries + // from the front since commit versions grow with commit time. + txns: Arc>>, global_version: Arc, } impl TransactionConflictTracker { pub fn new() -> Self { TransactionConflictTracker { - txns: Arc::new(Mutex::new(Vec::new())), + txns: Arc::new(Mutex::new(BTreeMap::new())), global_version: Arc::new(AtomicU64::new(0)), } } @@ -67,13 +68,23 @@ impl TransactionConflictTracker { ) -> bool { let mut txns = self.txns.lock().await; - // Prune old entries - txns.retain(|txn| txn.insert_instant.elapsed() < TXN_CONFLICT_TTL); + // Prune old entries. Commit versions grow with commit time, so expired entries are + // contiguous at the front of the map. + while let Some((_, txn)) = txns.first_key_value() { + if txn.insert_instant.elapsed() < TXN_CONFLICT_TTL { + break; + } + + txns.pop_first(); + } - for txn2 in &*txns { + // A retained transaction can only conflict if its commit version is greater than this + // transaction's start version, so skip everything at or below it. + for (txn2_commit_version, txn2) in + txns.range((Bound::Excluded(txn1_start_version), Bound::Unbounded)) + { // Check txn versions overlap (intersection or encapsulation) - if txn1_start_version < txn2.commit_version && txn2.start_version < txn1_commit_version - { + if txn2.start_version < txn1_commit_version { for (cr1_start, cr1_end, cr1_type) in &txn1_conflict_ranges { for (cr2_start, cr2_end, cr2_type) in &txn2.conflict_ranges { // Check conflict ranges overlap @@ -88,7 +99,7 @@ impl TransactionConflictTracker { txn1_start_version, txn1_commit_version, txn2_start_version = txn2.start_version, - txn2_commit_version = txn2.commit_version, + txn2_commit_version = %txn2_commit_version, "transaction conflict detected" ); return true; @@ -99,25 +110,26 @@ impl TransactionConflictTracker { } // If no conflicts were detected, save txn data - txns.push(PreviousTransaction { - insert_instant: Instant::now(), - start_version: txn1_start_version, - commit_version: txn1_commit_version, - conflict_ranges: txn1_conflict_ranges, - }); + txns.insert( + txn1_commit_version, + PreviousTransaction { + insert_instant: Instant::now(), + start_version: txn1_start_version, + conflict_ranges: txn1_conflict_ranges, + }, + ); false } - pub async fn remove(&self, txn_start_version: u64) { + pub async fn remove(&self, txn_commit_version: u64) { let mut txns = self.txns.lock().await; + txns.remove(&txn_commit_version); + } - if let Some(i) = txns - .iter() - .enumerate() - .find_map(|(i, txn)| (txn.start_version == txn_start_version).then_some(i)) - { - txns.remove(i); - } + /// Current retained transaction count. Diagnostic: the conflict scan in `check_and_insert` is + /// O(this) per commit, so a growing map directly inflates per-commit service time. + pub async fn len(&self) -> usize { + self.txns.lock().await.len() } } diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs index e07304180c..098f022663 100644 --- a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs +++ b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs @@ -351,11 +351,19 @@ async fn drain_batch( .await .context("failed to start drain batch txn")?; - // Claim a batch in id order. FOR UPDATE SKIP LOCKED holds the rows for this txn so they are - // stamped terminal on COMMIT with no intermediate 'claimed' state to clean up. + // Claim a batch in id order and allocate every commit version in the SAME round-trip via an + // inline `nextval`, instead of a separate version-allocation query. FOR UPDATE SKIP LOCKED holds + // the rows for this txn so they are stamped terminal on COMMIT with no intermediate 'claimed' + // state to clean up. + // + // Postgres does NOT guarantee `nextval` is evaluated in output (id) order, so the per-row `cv` + // here may not be monotonic in id. The versions are collected, sorted, and re-assigned to rows in + // id order below, exactly as the separate-query path did, because versionstamp monotonicity with + // commit order is load-bearing (epoxy changelog catch-up + depot PITR). let rows = txn .query( - "SELECT id, read_version, payload, reply_channel + "SELECT id, read_version, payload, reply_channel, + nextval('udb_version_seq') AS cv FROM udb_commit_requests WHERE status = 'pending' AND epoch = $1 ORDER BY id @@ -380,20 +388,10 @@ async fn drain_batch( let mut conflict_count = 0u32; let mut cold_reject_count = 0u32; - // Allocate all commit versions for the batch in one round-trip instead of a `nextval` per row. - // They are assigned to rows in id order (winners and losers alike; losers' versions are - // harmlessly skipped), so versionstamps stay monotonic with commit order. The defensive sort - // keeps assignment monotonic regardless of how Postgres orders the per-row `nextval` evaluation. - let mut versions: Vec = txn - .query( - "SELECT nextval('udb_version_seq') FROM generate_series(1, $1::bigint)", - &[&(batch_len as i64)], - ) - .await - .context("failed to allocate commit versions")? - .iter() - .map(|row| row.get::<_, i64>(0)) - .collect(); + // Re-assign the inline-allocated versions to rows in id order (winners and losers alike; losers' + // versions are harmlessly skipped) so versionstamps stay monotonic with commit order. The sort + // keeps assignment monotonic regardless of how Postgres ordered the per-row `nextval` evaluation. + let mut versions: Vec = rows.iter().map(|row| row.get::<_, i64>(4)).collect(); versions.sort_unstable(); // Resolve every request in memory in id order. Winners are collected with their version and @@ -403,6 +401,11 @@ async fn drain_batch( let mut stamp_statuses: Vec<&str> = Vec::with_capacity(batch_len); let mut stamp_versions: Vec> = Vec::with_capacity(batch_len); + // Sub-phase timers to decompose batch_ms and confirm whether the service time is constant (pure + // M/M/1 queueing) or itself inflates under load. + let mut decode_dur = Duration::ZERO; + let mut conflict_dur = Duration::ZERO; + for (i, row) in rows.iter().enumerate() { let id: i64 = row.get(0); let read_version: i64 = row.get(1); @@ -410,8 +413,10 @@ async fn drain_batch( let reply_channel: String = row.get(3); let commit_version = versions[i]; + let decode_start = Instant::now(); let decoded = super::codec::decode_commit_request(&payload) .context("failed to decode commit payload")?; + decode_dur += decode_start.elapsed(); let start_version = read_version.max(0) as u64; @@ -422,13 +427,16 @@ async fn drain_batch( cold_reject_count += 1; true } else { - tracker + let conflict_start = Instant::now(); + let res = tracker .check_and_insert( start_version, commit_version.max(0) as u64, decoded.conflict_ranges, ) - .await + .await; + conflict_dur += conflict_start.elapsed(); + res }; stamp_ids.push(id); @@ -460,6 +468,8 @@ async fn drain_batch( }); } + let t_resolved = Instant::now(); + // Bulk-read the pre-batch value of every key a winner's atomic op reads in one query, then fold // all winners into a single materialized write-set in memory. This collapses the per-row apply // round-trips to a fixed count independent of batch size. @@ -478,63 +488,74 @@ async fn drain_batch( .collect() }; + let t_read = Instant::now(); + let apply::WriteSet { upserts, point_deletes, range_deletes, } = apply::fold_winners(winners, &base).context("failed to fold batch winners")?; - // Materialize the write-set in O(1) statements per kind. Range deletes run first so a key whose - // final state is a set but that fell inside an earlier range clear is re-inserted by the upsert, - // not removed. - for (begin, end) in &range_deletes { - txn.execute("DELETE FROM kv WHERE key >= $1 AND key < $2", &[begin, end]) - .await - .context("failed to clear range")?; - } - if !point_deletes.is_empty() { - txn.execute( - "DELETE FROM kv WHERE key = ANY($1::bytea[])", - &[&point_deletes], - ) - .await - .context("failed to bulk-delete cleared keys")?; - } - if !upserts.is_empty() { - let (keys, values): (Vec>, Vec>) = upserts.into_iter().unzip(); + let t_fold = Instant::now(); + + let (upsert_keys, upsert_values): (Vec>, Vec>) = upserts.into_iter().unzip(); + let (range_begins, range_ends): (Vec>, Vec>) = + range_deletes.into_iter().unzip(); + + // Range deletes run in their OWN statement BEFORE the apply CTE. Postgres data-modifying CTE + // sub-statements all observe the same snapshot and never see each other's effects, so a range + // delete and an in-range upsert in one CTE would have unspecified results. `range_deletes` are + // ranges (not materialized per-key) and can overlap an upsert key (a key inside a cleared range + // that is also re-set), so the range clear must commit its effect first; the upsert then + // re-inserts the key. Collapsed from a per-range loop into one statement over a pair of arrays. + if !range_begins.is_empty() { txn.execute( - "INSERT INTO kv (key, value) - SELECT * FROM unnest($1::bytea[], $2::bytea[]) - ON CONFLICT (key) DO UPDATE SET value = excluded.value", - &[&keys, &values], + "DELETE FROM kv USING unnest($1::bytea[], $2::bytea[]) AS r(b, e) + WHERE key >= r.b AND key < r.e", + &[&range_begins, &range_ends], ) .await - .context("failed to bulk-upsert kv")?; + .context("failed to clear ranges")?; } - // Stamp every request's terminal status in one statement instead of a per-row UPDATE. - txn.execute( - "UPDATE udb_commit_requests AS r - SET status = b.status, commit_version = b.cv - FROM unnest($1::bigint[], $2::text[], $3::bigint[]) AS b(id, status, cv) - WHERE r.id = b.id", - &[&stamp_ids, &stamp_statuses, &stamp_versions], - ) - .await - .context("failed to stamp commit statuses")?; - - // Advance the watermark, fenced on our epoch. A zombie old leader whose epoch was bumped sees - // zero rows updated and must step down before any of its writes become visible. + // Apply the rest of the batch in one CTE: point deletes, the kv upsert, the terminal status + // stamp, and the epoch-fenced watermark advance. This is safe to fold because the write sets are + // disjoint: `apply::WriteSet` guarantees each key appears at most once across `upserts` and + // `point_deletes` (see apply.rs), and the three tables (`kv`, `udb_commit_requests`, `udb_lease`) + // are independent. The watermark UPDATE is fenced on our epoch: a zombie old leader whose epoch + // was bumped sees zero rows returned and must step down before any of its writes become visible. let new_durable: i64 = match txn .query_opt( - "UPDATE udb_lease - SET durable_version = GREATEST(durable_version, $1) - WHERE id = $2 AND epoch = $3 + "WITH pdel AS ( + DELETE FROM kv WHERE key = ANY($1::bytea[]) + ), up AS ( + INSERT INTO kv (key, value) + SELECT * FROM unnest($2::bytea[], $3::bytea[]) + ON CONFLICT (key) DO UPDATE SET value = excluded.value + ), stamp AS ( + UPDATE udb_commit_requests AS r + SET status = b.status, commit_version = b.cv + FROM unnest($4::bigint[], $5::text[], $6::bigint[]) AS b(id, status, cv) + WHERE r.id = b.id + ) + UPDATE udb_lease + SET durable_version = GREATEST(durable_version, $7) + WHERE id = $8 AND epoch = $9 RETURNING durable_version", - &[&max_winner_cv, &LEASE_ID, &epoch], + &[ + &point_deletes, + &upsert_keys, + &upsert_values, + &stamp_ids, + &stamp_statuses, + &stamp_versions, + &max_winner_cv, + &LEASE_ID, + &epoch, + ], ) .await - .context("failed to advance watermark")? + .context("failed to apply batch and advance watermark")? { Some(row) => row.get(0), None => { @@ -543,14 +564,22 @@ async fn drain_batch( } }; + let t_applied = Instant::now(); + txn.commit().await.context("failed to commit drain batch")?; + let t_committed = Instant::now(); + // Watermark advances strictly after the apply txn is durably committed and visible, so a // reader handed this read_version can never miss a write with commit_version <= read_version. shared.advance_durable_version(new_durable); notify_after_commit(&conn, new_durable, &replies).await; + let t_notified = Instant::now(); + + let tracker_len = tracker.len().await; + tracing::info!( epoch, batch_len, @@ -560,6 +589,17 @@ async fn drain_batch( cold_window, new_durable, batch_ms = batch_start.elapsed().as_millis() as u64, + // Sub-phase decomposition of batch_ms (micros) to separate constant service time from + // load-dependent service inflation. + tracker_len, + decode_us = decode_dur.as_micros() as u64, + conflict_us = conflict_dur.as_micros() as u64, + resolve_us = (t_resolved - batch_start).as_micros() as u64, + read_us = (t_read - t_resolved).as_micros() as u64, + fold_us = (t_fold - t_read).as_micros() as u64, + apply_us = (t_applied - t_fold).as_micros() as u64, + commit_us = (t_committed - t_applied).as_micros() as u64, + notify_us = (t_notified - t_committed).as_micros() as u64, "udb leader processed commit batch" ); diff --git a/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs b/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs index 986e6ba2bd..662a1a1f39 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs @@ -427,7 +427,7 @@ impl TransactionTask { Ok(_) => Ok(()), Err(e) => { // If the txn failed due to a rocksdb error, remove it from the conflict tracker - self.txn_conflict_tracker.remove(start_version).await; + self.txn_conflict_tracker.remove(commit_version).await; let err_str = e.to_string(); diff --git a/engine/packages/universaldb/tests/kv_upsert_scaling.rs b/engine/packages/universaldb/tests/kv_upsert_scaling.rs new file mode 100644 index 0000000000..2039258602 --- /dev/null +++ b/engine/packages/universaldb/tests/kv_upsert_scaling.rs @@ -0,0 +1,185 @@ +//! Throwaway diagnostic: isolates the pure-Postgres cost of the resolver's kv write from all of the +//! coordination code around it (conflict tracker, fold, claim SELECT, doorbell/reply NOTIFY, queue +//! round-trips). It does NOTHING but the `kv` upsert the leader's apply CTE performs, so the delta +//! between this floor and the real `batch_ms` is "our code", not Postgres. +//! +//! Run: +//! cargo test -p universaldb --test kv_upsert_scaling -- --ignored --nocapture +//! +//! Two sweeps: +//! 1. batch-size sweep, single serial writer (mirrors the single leader drain): one txn per batch, +//! distinct keys, `unnest` upsert. Shows the per-batch fixed floor + marginal per-row cost of a +//! pure PG write as batch length grows. +//! 2. concurrency sweep: N independent writers doing single-row upsert txns. Shows the raw PG write +//! concurrency the single-leader design forgoes (PG itself scales writers; the leader serializes +//! them through one drain task). + +use std::time::{Duration, Instant}; + +use rivet_test_deps_docker::TestDatabase; +use tokio_postgres::NoTls; +use uuid::Uuid; + +async fn connect(conn_str: &str) -> tokio_postgres::Client { + let (client, connection) = tokio_postgres::connect(conn_str, NoTls) + .await + .expect("connect"); + tokio::spawn(async move { + let _ = connection.await; + }); + client +} + +/// Percentile from a pre-sorted slice of microsecond samples, returned as milliseconds. +fn pct_ms(sorted_us: &[u128], p: f64) -> f64 { + if sorted_us.is_empty() { + return 0.0; + } + let idx = ((sorted_us.len() as f64 * p) as usize).min(sorted_us.len() - 1); + sorted_us[idx] as f64 / 1000.0 +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore = "diagnostic benchmark; run explicitly with --ignored --nocapture"] +async fn kv_upsert_scaling() { + let _ = tracing_subscriber::fmt() + .with_env_filter("warn") + .with_test_writer() + .try_init(); + + let (db_config, docker_config) = TestDatabase::Postgres + .config(Uuid::new_v4(), 1) + .await + .unwrap(); + let mut docker_config = docker_config.unwrap(); + docker_config.start().await.unwrap(); + tokio::time::sleep(Duration::from_secs(4)).await; + + let rivet_config::config::Database::Postgres(postgres_config) = db_config else { + unreachable!(); + }; + let conn_str = postgres_config.url.read().clone(); + + let setup = connect(&conn_str).await; + setup + .batch_execute( + "CREATE TABLE IF NOT EXISTS kv (key BYTEA PRIMARY KEY, value BYTEA NOT NULL)", + ) + .await + .unwrap(); + + // Value size roughly matching a small UDB kv write. + const VALUE_LEN: usize = 64; + + // ===== Sweep 1: batch-size, single serial writer (mirrors the leader drain) ===== + println!("\n=== batch-size sweep (single serial writer, 1 txn/batch + COMMIT) ==="); + println!( + "{:>10} {:>10} {:>16} {:>16} {:>16}", + "batch_len", "batches", "batch_ms p50", "batch_ms p95", "per_row_ms p50" + ); + { + let mut client = connect(&conn_str).await; + let mut key_ctr: u64 = 0; + for &bs in &[1usize, 2, 4, 8, 16, 32, 64, 128, 256] { + // Aim for a similar number of total rows per size so timings are comparable. + let iters = (4096 / bs).max(64); + // Warm up so the first-statement plan/parse cost is not counted. + for _ in 0..3 { + let (keys, vals) = next_batch(&mut key_ctr, bs, VALUE_LEN); + upsert_batch(&mut client, &keys, &vals).await; + } + let mut batch_us = Vec::with_capacity(iters); + for _ in 0..iters { + let (keys, vals) = next_batch(&mut key_ctr, bs, VALUE_LEN); + let t = Instant::now(); + upsert_batch(&mut client, &keys, &vals).await; + batch_us.push(t.elapsed().as_micros()); + } + batch_us.sort_unstable(); + let p50 = pct_ms(&batch_us, 0.5); + println!( + "{:>10} {:>10} {:>16.3} {:>16.3} {:>16.4}", + bs, + iters, + p50, + pct_ms(&batch_us, 0.95), + p50 / bs as f64 + ); + } + } + + // ===== Sweep 2: concurrency, N independent single-row writers ===== + println!("\n=== concurrency sweep (N parallel writers, single-row txn each, 3s) ==="); + println!( + "{:>6} {:>14} {:>14} {:>14}", + "N", "ops/s", "op_ms p50", "op_ms p95" + ); + for &n in &[1usize, 2, 4, 8, 16, 32, 64] { + let run = Duration::from_secs(3); + let mut handles = Vec::with_capacity(n); + for w in 0..n { + let cs = conn_str.clone(); + handles.push(tokio::spawn(async move { + let mut client = connect(&cs).await; + // Disjoint key space per worker so writers do not contend on the same row lock. + let mut key: u64 = (w as u64) << 40; + let mut lat_us = Vec::new(); + let deadline = Instant::now() + run; + while Instant::now() < deadline { + key += 1; + let k = key.to_be_bytes().to_vec(); + let v = vec![0u8; VALUE_LEN]; + let t = Instant::now(); + let txn = client.transaction().await.unwrap(); + txn.execute( + "INSERT INTO kv (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE SET value = excluded.value", + &[&k, &v], + ) + .await + .unwrap(); + txn.commit().await.unwrap(); + lat_us.push(t.elapsed().as_micros()); + } + lat_us + })); + } + let mut all = Vec::new(); + for h in handles { + all.extend(h.await.unwrap()); + } + all.sort_unstable(); + println!( + "{:>6} {:>14.0} {:>14.3} {:>14.3}", + n, + all.len() as f64 / 3.0, + pct_ms(&all, 0.5), + pct_ms(&all, 0.95) + ); + } +} + +/// Build a batch of `bs` distinct keys and equal-length values, advancing the shared counter. +fn next_batch(key_ctr: &mut u64, bs: usize, value_len: usize) -> (Vec>, Vec>) { + let mut keys = Vec::with_capacity(bs); + let mut vals = Vec::with_capacity(bs); + for _ in 0..bs { + *key_ctr += 1; + keys.push(key_ctr.to_be_bytes().to_vec()); + vals.push(vec![0u8; value_len]); + } + (keys, vals) +} + +/// One batch upsert in its own transaction with a real COMMIT, matching the leader's apply shape. +async fn upsert_batch(client: &mut tokio_postgres::Client, keys: &[Vec], vals: &[Vec]) { + let txn = client.transaction().await.unwrap(); + txn.execute( + "INSERT INTO kv (key, value) SELECT * FROM unnest($1::bytea[], $2::bytea[]) + ON CONFLICT (key) DO UPDATE SET value = excluded.value", + &[&keys, &vals], + ) + .await + .unwrap(); + txn.commit().await.unwrap(); +} From c2f534e2e1eee36a1b82bb93470bca9924d4b7a3 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Mon, 29 Jun 2026 14:16:28 -0700 Subject: [PATCH 16/16] [slopfix] refactor(universaldb,ups): single/multi-node udb over nats, remove ups postgres driver --- Cargo.lock | 7 +- docs-internal/engine/TEST_DEPENDENCIES.md | 1 - engine/artifacts/config-schema.json | 37 +- .../packages/config/src/config/db/postgres.rs | 12 + engine/packages/config/src/config/mod.rs | 17 +- engine/packages/config/src/config/pubsub.rs | 38 - engine/packages/pools/src/db/udb.rs | 14 + engine/packages/pools/src/db/ups.rs | 25 - .../packages/test-deps-docker/src/database.rs | 1 + engine/packages/universaldb/Cargo.toml | 1 + .../universaldb/src/conflict_tracker.rs | 8 +- .../universaldb/src/driver/postgres/codec.rs | 59 +- .../universaldb/src/driver/postgres/commit.rs | 366 ++---- .../src/driver/postgres/database.rs | 114 +- .../src/driver/postgres/listener.rs | 228 ---- .../universaldb/src/driver/postgres/mod.rs | 4 +- .../universaldb/src/driver/postgres/nats.rs | 149 +++ .../src/driver/postgres/resolver/mod.rs | 734 ++++++------ .../universaldb/src/driver/postgres/shared.rs | 115 +- .../src/driver/postgres/transport.rs | 87 ++ .../src/driver/rocksdb/database.rs | 2 +- engine/packages/universaldb/tests/failover.rs | 56 +- engine/packages/universalpubsub/Cargo.toml | 6 - .../universalpubsub/benches/simple.rs | 126 +- .../universalpubsub/src/driver/mod.rs | 1 - .../src/driver/postgres/doorbell.rs | 133 --- .../src/driver/postgres/mod.rs | 1050 ----------------- .../packages/universalpubsub/src/metrics.rs | 7 - .../universalpubsub/tests/integration.rs | 71 +- .../universalpubsub/tests/reconnect.rs | 48 +- .../rust/universaldb-commit/src/versioned.rs | 62 + .../sdks/schemas/universaldb-commit/v1.bare | 38 +- .../compose/template/src/docker-compose.ts | 52 +- .../src/services/edge/rivet-engine.ts | 12 +- self-host/dev-host/docker-compose.yml | 1 + .../dc-a/rivet-engine/0/config.jsonc | 5 + .../dc-a/rivet-engine/1/config.jsonc | 5 + .../dc-a/rivet-engine/2/config.jsonc | 5 + .../dc-b/rivet-engine/0/config.jsonc | 5 + .../dc-b/rivet-engine/1/config.jsonc | 5 + .../dc-b/rivet-engine/2/config.jsonc | 5 + .../dc-c/rivet-engine/0/config.jsonc | 5 + .../dc-c/rivet-engine/1/config.jsonc | 5 + .../dc-c/rivet-engine/2/config.jsonc | 5 + .../dev-multidc-multinode/docker-compose.yml | 80 ++ self-host/dev-multidc/docker-compose.yml | 3 + self-host/dev-multinode/docker-compose.yml | 25 + .../dev-multinode/rivet-engine/0/config.jsonc | 5 + .../dev-multinode/rivet-engine/1/config.jsonc | 5 + .../dev-multinode/rivet-engine/2/config.jsonc | 5 + self-host/dev/docker-compose.yml | 1 + 51 files changed, 1325 insertions(+), 2526 deletions(-) delete mode 100644 engine/packages/universaldb/src/driver/postgres/listener.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/nats.rs create mode 100644 engine/packages/universaldb/src/driver/postgres/transport.rs delete mode 100644 engine/packages/universalpubsub/src/driver/postgres/doorbell.rs delete mode 100644 engine/packages/universalpubsub/src/driver/postgres/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 7b0053c042..868daf76fb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8273,6 +8273,7 @@ dependencies = [ "hex", "lazy_static", "rand 0.8.5", + "rivet-async-nats", "rivet-config", "rivet-env", "rivet-metrics", @@ -8303,8 +8304,6 @@ version = "2.3.7" dependencies = [ "anyhow", "async-trait", - "base64 0.22.1", - "deadpool-postgres", "futures-util", "lazy_static", "rand 0.8.5", @@ -8314,19 +8313,15 @@ dependencies = [ "rivet-error", "rivet-metrics", "rivet-perf", - "rivet-postgres-util", "rivet-test-deps-docker", "rivet-ups-protocol", "rivet-util", "scc", "serde", "serde_json", - "sha2", "tabled", "tempfile", "tokio", - "tokio-postgres", - "tokio-postgres-rustls", "tokio-util", "tracing", "tracing-subscriber", diff --git a/docs-internal/engine/TEST_DEPENDENCIES.md b/docs-internal/engine/TEST_DEPENDENCIES.md index 7eb03f0f53..e4a7e1807a 100644 --- a/docs-internal/engine/TEST_DEPENDENCIES.md +++ b/docs-internal/engine/TEST_DEPENDENCIES.md @@ -13,7 +13,6 @@ Configure backends via environment variables: - `RIVET_TEST_PUBSUB`: Choose pub/sub backend - `nats` - Runs NATS in Docker - - `postgres_notify` - PostgreSQL in Docker - `memory` - In-memory channels (default) - `RUST_LOG`: Enable debug logs to see container lifecycle details diff --git a/engine/artifacts/config-schema.json b/engine/artifacts/config-schema.json index af2af30077..9147b7a7d1 100644 --- a/engine/artifacts/config-schema.json +++ b/engine/artifacts/config-schema.json @@ -1244,47 +1244,18 @@ "url" ], "properties": { - "ssl": { - "description": "SSL configuration options", + "nats": { + "description": "NATS configuration for UniversalDB multi-node mode.\n\nWhen set, UniversalDB runs in multi-node mode and uses NATS for follower-to-leader commit transport instead of an in-process resolver. When absent, UniversalDB runs single-node. If unset but the UPS pubsub is configured for NATS, this is inherited from that config at startup (see `Root::validate_and_set_defaults`).", "default": null, "anyOf": [ { - "$ref": "#/definitions/PostgresSsl" + "$ref": "#/definitions/Nats" }, { "type": "null" } ] }, - "url": { - "description": "URL to connect to Postgres with\n\nSupports standard PostgreSQL connection parameters including `sslmode`. Supported sslmode values: `disable`, `prefer` (default), `require`. To verify server certificates, use `sslmode=require` with `ssl.root_cert_path`.\n\nExample with sslmode: `postgresql://user:pass@host:5432/db?sslmode=require`\n\nSee: https://docs.rs/postgres/0.19.10/postgres/config/struct.Config.html#url", - "allOf": [ - { - "$ref": "#/definitions/Secret" - } - ] - } - }, - "additionalProperties": false - }, - "Postgres2": { - "type": "object", - "required": [ - "url" - ], - "properties": { - "disable_memory_optimization": { - "description": "When true, force every UPS publish to round-trip through the postgres driver instead of taking the in-process fast path for subjects that have a local subscriber on the same engine pod. Opt-in diagnostic; default false.", - "default": false, - "type": "boolean" - }, - "memory_optimization": { - "deprecated": true, - "type": [ - "boolean", - "null" - ] - }, "ssl": { "description": "SSL configuration options", "default": null, @@ -1298,7 +1269,7 @@ ] }, "url": { - "description": "URL to connect to Postgres with\n\nSupports standard PostgreSQL connection parameters including `sslmode`. Supported sslmode values: `disable`, `prefer` (default), `require`. To verify server certificates, use `sslmode=require` with `ssl.root_cert_path`.\n\nSee: https://docs.rs/postgres/0.19.10/postgres/config/struct.Config.html#url", + "description": "URL to connect to Postgres with\n\nSupports standard PostgreSQL connection parameters including `sslmode`. Supported sslmode values: `disable`, `prefer` (default), `require`. To verify server certificates, use `sslmode=require` with `ssl.root_cert_path`.\n\nExample with sslmode: `postgresql://user:pass@host:5432/db?sslmode=require`\n\nSee: https://docs.rs/postgres/0.19.10/postgres/config/struct.Config.html#url", "allOf": [ { "$ref": "#/definitions/Secret" diff --git a/engine/packages/config/src/config/db/postgres.rs b/engine/packages/config/src/config/db/postgres.rs index f80af52f37..4cf8ea2684 100644 --- a/engine/packages/config/src/config/db/postgres.rs +++ b/engine/packages/config/src/config/db/postgres.rs @@ -5,6 +5,8 @@ use serde::{Deserialize, Serialize}; use crate::secret::Secret; +use super::super::pubsub::Nats; + #[derive(Debug, Serialize, Deserialize, Clone, JsonSchema)] #[serde(deny_unknown_fields)] pub struct PostgresSsl { @@ -47,6 +49,15 @@ pub struct Postgres { /// SSL configuration options #[serde(default)] pub ssl: Option, + + /// NATS configuration for UniversalDB multi-node mode. + /// + /// When set, UniversalDB runs in multi-node mode and uses NATS for follower-to-leader commit + /// transport instead of an in-process resolver. When absent, UniversalDB runs single-node. + /// If unset but the UPS pubsub is configured for NATS, this is inherited from that config at + /// startup (see `Root::validate_and_set_defaults`). + #[serde(default)] + pub nats: Option, } impl Default for Postgres { @@ -54,6 +65,7 @@ impl Default for Postgres { Self { url: Secret::new("postgresql://postgres:postgres@127.0.0.1:5432/postgres".into()), ssl: None, + nats: None, } } } diff --git a/engine/packages/config/src/config/mod.rs b/engine/packages/config/src/config/mod.rs index caa8bf196e..5ee03c6d2a 100644 --- a/engine/packages/config/src/config/mod.rs +++ b/engine/packages/config/src/config/mod.rs @@ -185,17 +185,14 @@ impl Root { } pub fn validate_and_set_defaults(&mut self) -> Result<()> { - // Set default pubsub to Postgres if configured for database - if self.pubsub.is_none() - && let Some(Database::Postgres(pg)) = &self.database + // When UDB runs on Postgres without an explicit NATS config, inherit the UPS NATS config if + // one is set. The presence of NATS is what selects UDB multi-node mode, so this lets a + // single `pubsub: nats` config drive both UPS and UDB across nodes. + if let Some(PubSub::Nats(nats)) = self.pubsub.clone() + && let Some(Database::Postgres(pg)) = &mut self.database + && pg.nats.is_none() { - self.pubsub = Some(PubSub::PostgresNotify(pubsub::Postgres { - url: pg.url.clone(), - #[allow(deprecated)] - memory_optimization: None, - disable_memory_optimization: false, - ssl: pg.ssl.clone(), - })); + pg.nats = Some(nats); } self.pegboard().validate()?; diff --git a/engine/packages/config/src/config/pubsub.rs b/engine/packages/config/src/config/pubsub.rs index cae3526141..478275c675 100644 --- a/engine/packages/config/src/config/pubsub.rs +++ b/engine/packages/config/src/config/pubsub.rs @@ -3,13 +3,10 @@ use serde::{Deserialize, Serialize}; use crate::secret::Secret; -use super::db::PostgresSsl; - #[derive(Debug, Serialize, Deserialize, Clone, JsonSchema)] #[serde(rename_all = "snake_case", deny_unknown_fields)] pub enum PubSub { Nats(Nats), - PostgresNotify(Postgres), Memory(Memory), } @@ -19,41 +16,6 @@ impl Default for PubSub { } } -#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema)] -#[serde(deny_unknown_fields)] -pub struct Postgres { - /// URL to connect to Postgres with - /// - /// Supports standard PostgreSQL connection parameters including `sslmode`. - /// Supported sslmode values: `disable`, `prefer` (default), `require`. - /// To verify server certificates, use `sslmode=require` with `ssl.root_cert_path`. - /// - /// See: https://docs.rs/postgres/0.19.10/postgres/config/struct.Config.html#url - pub url: Secret, - #[deprecated] - pub memory_optimization: Option, - /// When true, force every UPS publish to round-trip through the postgres driver instead of - /// taking the in-process fast path for subjects that have a local subscriber on the same - /// engine pod. Opt-in diagnostic; default false. - #[serde(default)] - pub disable_memory_optimization: bool, - /// SSL configuration options - #[serde(default)] - pub ssl: Option, -} - -impl Default for Postgres { - fn default() -> Self { - Self { - url: Secret::new("postgresql://postgres:postgres@127.0.0.1:5432/postgres".into()), - #[allow(deprecated)] - memory_optimization: None, - disable_memory_optimization: false, - ssl: None, - } - } -} - #[derive(Debug, Serialize, Deserialize, Clone, JsonSchema)] #[serde(deny_unknown_fields)] pub struct Nats { diff --git a/engine/packages/pools/src/db/udb.rs b/engine/packages/pools/src/db/udb.rs index cf3f40edc7..8049ae7019 100644 --- a/engine/packages/pools/src/db/udb.rs +++ b/engine/packages/pools/src/db/udb.rs @@ -20,6 +20,19 @@ impl Deref for UdbPool { pub async fn setup(config: &Config) -> Result> { let db_driver = match config.database() { config::Database::Postgres(pg) => { + // A NATS config (set directly or inherited from the UPS config in + // `validate_and_set_defaults`) selects UniversalDB multi-node mode. + let nats = pg + .nats + .as_ref() + .map(|nats| universaldb::driver::postgres::NatsConfig { + addresses: nats.addresses.clone(), + username: nats.username.clone(), + password: nats.password.as_ref().map(|p| p.read().clone()), + client_capacity: nats.client_capacity, + subscription_capacity: nats.subscription_capacity, + }); + let postgres_config = universaldb::driver::postgres::PostgresConfig { connection_string: pg.url.read().clone(), ssl_config: pg.ssl.as_ref().map(|ssl| { @@ -29,6 +42,7 @@ pub async fn setup(config: &Config) -> Result> { ssl_client_key_path: ssl.client_key_path.clone(), } }), + nats, }; Arc::new( diff --git a/engine/packages/pools/src/db/ups.rs b/engine/packages/pools/src/db/ups.rs index 9bf1423c18..6b63e5dc00 100644 --- a/engine/packages/pools/src/db/ups.rs +++ b/engine/packages/pools/src/db/ups.rs @@ -94,30 +94,6 @@ pub async fn setup(config: &Config, client_name: &str) -> Result { Arc::new(driver) as ups::PubSubDriverHandle } - config::PubSub::PostgresNotify(pg) => { - tracing::debug!("creating postgres pubsub driver"); - - let (ssl_root_cert_path, ssl_client_cert_path, ssl_client_key_path) = - if let Some(ssl) = &pg.ssl { - ( - ssl.root_cert_path.clone(), - ssl.client_cert_path.clone(), - ssl.client_key_path.clone(), - ) - } else { - (None, None, None) - }; - - Arc::new( - ups::driver::postgres::PostgresDriver::connect( - pg.url.read().clone(), - ssl_root_cert_path, - ssl_client_cert_path, - ssl_client_key_path, - ) - .await?, - ) as ups::PubSubDriverHandle - } config::PubSub::Memory(memory) => { tracing::debug!(channel=%memory.channel, "creating memory pubsub driver"); Arc::new(ups::driver::memory::MemoryDriver::new( @@ -128,7 +104,6 @@ pub async fn setup(config: &Config, client_name: &str) -> Result { let disable_memory_optimization = match config.pubsub() { config::PubSub::Nats(nats) => nats.disable_memory_optimization, - config::PubSub::PostgresNotify(pg) => pg.disable_memory_optimization, config::PubSub::Memory(memory) => memory.disable_memory_optimization, }; Ok(ups::PubSub::new_with_memory_optimization( diff --git a/engine/packages/test-deps-docker/src/database.rs b/engine/packages/test-deps-docker/src/database.rs index 955e4d7897..2ddd0ee916 100644 --- a/engine/packages/test-deps-docker/src/database.rs +++ b/engine/packages/test-deps-docker/src/database.rs @@ -58,6 +58,7 @@ impl TestDatabase { rivet_config::config::Database::Postgres(rivet_config::config::db::Postgres { url: rivet_config::secret::Secret::new(connection_string.clone()), ssl: None, + nats: None, }); let docker_config = DockerRunConfig { diff --git a/engine/packages/universaldb/Cargo.toml b/engine/packages/universaldb/Cargo.toml index 6781c7cde1..7f6b323e62 100644 --- a/engine/packages/universaldb/Cargo.toml +++ b/engine/packages/universaldb/Cargo.toml @@ -8,6 +8,7 @@ edition.workspace = true [dependencies] anyhow.workspace = true +async-nats.workspace = true async-trait.workspace = true base64.workspace = true deadpool-postgres.workspace = true diff --git a/engine/packages/universaldb/src/conflict_tracker.rs b/engine/packages/universaldb/src/conflict_tracker.rs index 3e543a99ee..a257f14ecc 100644 --- a/engine/packages/universaldb/src/conflict_tracker.rs +++ b/engine/packages/universaldb/src/conflict_tracker.rs @@ -57,7 +57,7 @@ impl TransactionConflictTracker { self.global_version.fetch_add(1, Ordering::SeqCst) } - /// Returns `true` on conflict (same polarity as the original rocksdb tracker). The caller + /// Returns `true` on conflicts. The caller /// supplies `commit_version` (e.g. `nextval('udb_version_seq')` on the postgres leader, or /// `next_global_version()` on rocksdb) so version assignment stays the caller's responsibility. pub async fn check_and_insert( @@ -126,10 +126,4 @@ impl TransactionConflictTracker { let mut txns = self.txns.lock().await; txns.remove(&txn_commit_version); } - - /// Current retained transaction count. Diagnostic: the conflict scan in `check_and_insert` is - /// O(this) per commit, so a growing map directly inflates per-commit service time. - pub async fn len(&self) -> usize { - self.txns.lock().await.len() - } } diff --git a/engine/packages/universaldb/src/driver/postgres/codec.rs b/engine/packages/universaldb/src/driver/postgres/codec.rs index 432ede2f34..2e7b975ef0 100644 --- a/engine/packages/universaldb/src/driver/postgres/codec.rs +++ b/engine/packages/universaldb/src/driver/postgres/codec.rs @@ -7,13 +7,17 @@ use crate::{ tx_ops::Operation, }; -/// Decoded form of a `udb_commit_requests.payload` blob. -/// -/// `read_version` is intentionally omitted: it is also denormalized into the `read_version` column, -/// which is what the leader's drain reads, so decoding it here would be dead. +use super::transport::CommitOutcome; + +/// Decoded form of a commit request payload sent from a follower to the leader over NATS. pub struct DecodedCommit { + pub read_version: u64, pub conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, pub operations: Vec, + /// Submitting follower's node id (part of the failover dedup key). + pub client_node_id: Vec, + /// Per-process monotonic counter (part of the failover dedup key). + pub client_seq: u64, } /// Encode a follower's commit request to the versioned BARE wire format with an embedded version @@ -22,6 +26,8 @@ pub fn encode_commit_request( read_version: u64, conflict_ranges: &[(Vec, Vec, ConflictRangeType)], operations: &[Operation], + client_node_id: &[u8], + client_seq: u64, ) -> Result> { let request = proto::CommitRequest { read_version, @@ -34,13 +40,15 @@ pub fn encode_commit_request( }) .collect(), operations: operations.iter().map(operation_to_proto).collect(), + client_node_id: client_node_id.to_vec(), + client_seq, }; versioned::CommitRequest::wrap_latest(request) .serialize_with_embedded_version(proto::PROTOCOL_VERSION) } -/// Decode a `udb_commit_requests.payload` blob produced by [`encode_commit_request`]. +/// Decode a commit request payload produced by [`encode_commit_request`]. pub fn decode_commit_request(payload: &[u8]) -> Result { let request = versioned::CommitRequest::deserialize_with_embedded_version(payload)?; @@ -63,11 +71,52 @@ pub fn decode_commit_request(payload: &[u8]) -> Result { .collect(); Ok(DecodedCommit { + read_version: request.read_version, conflict_ranges, operations, + client_node_id: request.client_node_id, + client_seq: request.client_seq, + }) +} + +/// Encode a leader's commit reply to the versioned BARE wire format with an embedded version header +/// so a follower running older or newer code can still decode it during a rolling deploy. +pub fn encode_commit_reply(outcome: CommitOutcome) -> Result> { + let reply = match outcome { + CommitOutcome::Committed { commit_version } => { + proto::CommitReply::CommitCommitted(proto::CommitCommitted { commit_version }) + } + CommitOutcome::Conflict => proto::CommitReply::CommitConflict, + }; + + versioned::CommitReply::wrap_latest(reply) + .serialize_with_embedded_version(proto::PROTOCOL_VERSION) +} + +/// Decode a commit reply payload produced by [`encode_commit_reply`]. +pub fn decode_commit_reply(payload: &[u8]) -> Result { + let reply = versioned::CommitReply::deserialize_with_embedded_version(payload)?; + Ok(match reply { + proto::CommitReply::CommitCommitted(proto::CommitCommitted { commit_version }) => { + CommitOutcome::Committed { commit_version } + } + proto::CommitReply::CommitConflict => CommitOutcome::Conflict, }) } +/// Encode a durable-version watermark broadcast to the versioned BARE wire format with an embedded +/// version header. +pub fn encode_watermark(durable_version: i64) -> Result> { + versioned::Watermark::wrap_latest(proto::Watermark { durable_version }) + .serialize_with_embedded_version(proto::PROTOCOL_VERSION) +} + +/// Decode a watermark payload produced by [`encode_watermark`], returning the durable version. +pub fn decode_watermark(payload: &[u8]) -> Result { + let watermark = versioned::Watermark::deserialize_with_embedded_version(payload)?; + Ok(watermark.durable_version) +} + fn conflict_range_type_to_proto(kind: ConflictRangeType) -> proto::ConflictRangeType { match kind { ConflictRangeType::Read => proto::ConflictRangeType::Read, diff --git a/engine/packages/universaldb/src/driver/postgres/commit.rs b/engine/packages/universaldb/src/driver/postgres/commit.rs index d191c815a3..21d2335913 100644 --- a/engine/packages/universaldb/src/driver/postgres/commit.rs +++ b/engine/packages/universaldb/src/driver/postgres/commit.rs @@ -4,18 +4,28 @@ use std::{ }; use anyhow::{Context, Result}; +use tokio::sync::oneshot; use crate::{error::DatabaseError, options::ConflictRangeType, tx_ops::Operation}; use super::{ codec, - shared::{LeaseInfo, PostgresShared, commit_channel, reply_channel}, + shared::{LeaseInfo, PostgresShared}, + transport::{CommitJob, CommitOutcome, Responder, Transport}, }; /// How long to wait for a leader to be elected before giving up a submit as retryable. const LEADER_WAIT_TIMEOUT: Duration = Duration::from_secs(5); -/// Backstop poll cadence while waiting for a commit result, in case a reply NOTIFY is missed. -const RESULT_POLL_INTERVAL: Duration = Duration::from_millis(250); +/// Poll cadence while waiting for a leader to appear in the cache. +const LEADER_POLL_INTERVAL: Duration = Duration::from_millis(50); +/// Per-attempt timeout for a NATS commit request. +const REQUEST_TIMEOUT: Duration = Duration::from_secs(5); +/// How many times a multi-node commit resends the same request (same dedup key) across leader +/// failover / indeterminate failures before giving up as retryable. The dedup table makes the +/// resends exactly-once. +const MAX_SUBMIT_ATTEMPTS: usize = 8; +/// Backoff between multi-node resends. +const RESEND_BACKOFF: Duration = Duration::from_millis(100); /// Submit a follower transaction's commit to the leader and await the result. /// @@ -27,8 +37,8 @@ pub async fn submit( operations: Vec, conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, ) -> Result<()> { - // A transaction with no writes and no serializable read ranges has nothing to order or - // validate; it never needs the leader. + // A transaction with no writes and no serializable read ranges has nothing to order or validate; + // it never needs the leader. if operations.is_empty() && conflict_ranges .iter() @@ -37,269 +47,133 @@ pub async fn submit( return Ok(()); } - let lease = wait_for_leader(shared).await?; - let payload = - codec::encode_commit_request(read_version.max(0) as u64, &conflict_ranges, &operations) - .context("failed to encode commit request")?; - let reply_channel = reply_channel(&shared.node_id); - - // Subscribe to our reply channel before inserting so we cannot miss the leader's NOTIFY. - let mut reply_rx = shared.listener.listen(&reply_channel).await; - - let conn = shared - .pool - .get() - .await - .context("failed to get connection for commit submit")?; - - // Enqueue the request and wake the leader's drain loop in one round-trip. The autocommit - // statement durably inserts the row and fires the NOTIFY together, so there is no separate - // notify round-trip or second pool acquire. - let id: i64 = conn - .query_one( - "WITH ins AS ( - INSERT INTO udb_commit_requests (epoch, read_version, payload, reply_channel) - VALUES ($1, $2, $3, $4) - RETURNING id - ) - SELECT pg_notify($5, id::text), id FROM ins", - &[ - &lease.epoch, - &read_version, - &payload, - &reply_channel, - &commit_channel(&lease.leader_addr), - ], - ) - .await - .context("failed to enqueue and notify commit request")? - .get(1); + match &shared.transport { + Transport::SingleNode { commit_tx } => { + submit_local(commit_tx, read_version, operations, conflict_ranges).await + } + Transport::MultiNode(_) => { + submit_nats(shared, read_version, operations, conflict_ranges).await + } + } +} - // Release the connection before waiting so a long wait does not pin a pool slot. The request - // row is durable, so await_result re-acquires a connection per poll. - drop(conn); +/// Single-node: hand the job straight to the in-process leader drain loop and await its result. +async fn submit_local( + commit_tx: &tokio::sync::mpsc::Sender, + read_version: i64, + operations: Vec, + conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, +) -> Result<()> { + let (response_tx, response_rx) = oneshot::channel(); + let job = CommitJob { + read_version: read_version.max(0) as u64, + conflict_ranges, + operations, + dedup_key: None, + responder: Responder::Local(response_tx), + }; - let submit_start = Instant::now(); - let result = await_result(shared, id, lease.epoch, &mut reply_rx).await; - tracing::debug!( - id, - epoch = lease.epoch, - wait_ms = submit_start.elapsed().as_millis() as u64, - ok = result.is_ok(), - "udb commit submit completed" - ); - result -} + if commit_tx.send(job).await.is_err() { + // The leader drain loop is gone (driver shutting down). Retryable. + return Err(DatabaseError::NotCommitted.into()); + } -/// Wait for a known leader, returning a retryable error if none is elected in time. -async fn wait_for_leader(shared: &Arc) -> Result { - let deadline = Instant::now() + LEADER_WAIT_TIMEOUT; - loop { - if let Some(lease) = shared.current_lease() { - return Ok(lease); - } - if Instant::now() >= deadline { - return Err(DatabaseError::NotCommitted.into()); - } - tokio::time::sleep(RESULT_POLL_INTERVAL).await; + match response_rx.await { + Ok(CommitOutcome::Committed { .. }) => Ok(()), + Ok(CommitOutcome::Conflict) => Err(DatabaseError::NotCommitted.into()), + // The leader dropped the job without responding; it was not applied. + Err(_) => Err(DatabaseError::NotCommitted.into()), } } -/// Wait for the commit result, resolved directly from the leader's reply NOTIFY payload on the happy -/// path. A polling `read_status` backstop covers a missed/lagged NOTIFY, and an epoch advance orphans -/// the request (it will never be applied, so it is definitively not committed). -async fn await_result( +/// Multi-node: send the commit to the elected leader over NATS request/reply, resending the same +/// request (same dedup key) across leader failover. The reply carries the commit result directly. +async fn submit_nats( shared: &Arc, - id: i64, - submit_epoch: i64, - reply_rx: &mut tokio::sync::broadcast::Receiver, + read_version: i64, + operations: Vec, + conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, ) -> Result<()> { - let start = Instant::now(); - // Diagnostics: count how the waiter is driven so we can tell whether the reply NOTIFY is doing - // its job or whether commits are riding the slow poll backstop / getting orphaned by failover. - let mut status_reads = 0u32; - let mut notify_wakes = 0u32; - let mut poll_wakes = 0u32; + let Transport::MultiNode(nats) = &shared.transport else { + unreachable!("submit_nats requires the multi-node transport"); + }; - // The backstop runs on a fixed-cadence interval, not a per-iteration sleep: under a flood of - // other commits' replies on this node's shared reply channel (which we skip past), a fresh - // per-iteration sleep would keep resetting and never fire, starving the backstop if our own - // NOTIFY was lost. An interval ticks on wall-clock cadence regardless of loop churn. - let mut poll_interval = tokio::time::interval(RESULT_POLL_INTERVAL); - poll_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - // Consume the immediate first tick so the first backstop is one interval out, after the reply - // has had a chance to arrive. - poll_interval.tick().await; + // One dedup key for this logical commit, reused across every resend so the leader applies it at + // most once even if an earlier attempt was applied but its reply was lost to a failover. + let client_seq = shared.next_commit_seq(); + let payload = codec::encode_commit_request( + read_version.max(0) as u64, + &conflict_ranges, + &operations, + shared.node_id.as_bytes(), + client_seq as u64, + ) + .context("failed to encode commit request")?; - loop { - // Wait for our reply NOTIFY, falling back to a status read on a poll tick or a lagged - // broadcast. The happy path resolves straight from the payload with no status SELECT. - tokio::select! { - res = reply_rx.recv() => { - match res { - Ok(payload) => { - notify_wakes += 1; - match parse_reply(&payload, id) { - Some(ReplyOutcome::Committed) => { - tracing::debug!( - id, - wait_ms = start.elapsed().as_millis() as u64, - status_reads, - notify_wakes, - poll_wakes, - "udb commit resolved: committed" - ); - return Ok(()); - } - Some(ReplyOutcome::Conflict) => { - tracing::debug!( - id, - wait_ms = start.elapsed().as_millis() as u64, - status_reads, - notify_wakes, - poll_wakes, - "udb commit resolved: conflict" - ); - return Err(DatabaseError::NotCommitted.into()); - } - // Reply for another waiter on the shared channel; keep waiting for ours. - None => continue, - } - } - Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => { - // We may have missed our own reply; fall through to the status backstop. - notify_wakes += 1; - tracing::debug!(id, lagged = n, "udb reply broadcast lagged"); - } - Err(tokio::sync::broadcast::error::RecvError::Closed) => { - tracing::warn!(id, "udb reply broadcast closed; re-subscribing"); - *reply_rx = shared - .listener - .listen(&reply_channel(&shared.node_id)) - .await; - continue; - } + let submit_start = Instant::now(); + for attempt in 0..MAX_SUBMIT_ATTEMPTS { + let lease = wait_for_leader(shared).await?; + let subject = nats.subjects.commit(&lease.leader_addr); + + let request = nats.client.request(subject, payload.clone().into()); + match tokio::time::timeout(REQUEST_TIMEOUT, request).await { + Ok(Ok(msg)) => match codec::decode_commit_reply(&msg.payload) { + Ok(CommitOutcome::Committed { .. }) => { + tracing::debug!( + client_seq, + attempt, + wait_ms = submit_start.elapsed().as_millis() as u64, + "udb commit resolved: committed" + ); + return Ok(()); } - } - _ = poll_interval.tick() => { - poll_wakes += 1; - } - } - - // Backstop path (poll tick or lagged broadcast): re-acquire a connection and read the durable - // status. A transient pool/query error just means we retry rather than failing a - // possibly-applied commit. - status_reads += 1; - match read_status(shared, id).await { - Ok(Some(Status::Committed)) => { + Ok(CommitOutcome::Conflict) => { + return Err(DatabaseError::NotCommitted.into()); + } + Err(err) => { + tracing::warn!(?err, client_seq, "malformed udb commit reply; resending"); + } + }, + // Indeterminate (no responder / transport error / timeout): the leader may have died + // before or after applying. Resend the same dedup key; the leader dedups any double apply. + Ok(Err(err)) => { tracing::debug!( - id, - wait_ms = start.elapsed().as_millis() as u64, - status_reads, - notify_wakes, - poll_wakes, - "udb commit resolved: committed (backstop)" + ?err, + client_seq, + attempt, + "udb commit request errored; resending" ); - return Ok(()); } - Ok(Some(Status::Conflict)) => { + Err(_) => { tracing::debug!( - id, - wait_ms = start.elapsed().as_millis() as u64, - status_reads, - notify_wakes, - poll_wakes, - "udb commit resolved: conflict (backstop)" + client_seq, + attempt, + "udb commit request timed out; resending" ); - return Err(DatabaseError::NotCommitted.into()); - } - Ok(Some(Status::Pending)) => {} - Ok(None) => { - // The row was GC'd before we observed a terminal status. Treat as not committed - // and let the retry loop resubmit. - tracing::warn!( - id, - wait_ms = start.elapsed().as_millis() as u64, - status_reads, - "udb commit row missing before terminal status (gc'd); treating as not committed" - ); - return Err(DatabaseError::NotCommitted.into()); - } - Err(err) => { - tracing::debug!(?err, id, "transient error polling commit status, retrying"); } } - // If a new leader took over, our old-epoch request will never be claimed. - if let Some(current) = shared.current_lease() { - if current.epoch != submit_epoch { - tracing::warn!( - id, - submit_epoch, - current_epoch = current.epoch, - wait_ms = start.elapsed().as_millis() as u64, - notify_wakes, - poll_wakes, - "udb commit orphaned by leader failover; treating as not committed" - ); - return Err(DatabaseError::NotCommitted.into()); - } - } + tokio::time::sleep(RESEND_BACKOFF).await; } -} -enum ReplyOutcome { - Committed, - Conflict, + tracing::warn!( + client_seq, + wait_ms = submit_start.elapsed().as_millis() as u64, + "udb commit exhausted resend attempts; treating as not committed" + ); + Err(DatabaseError::NotCommitted.into()) } -/// Parse a leader reply payload (`":committed:"` or `":conflict"`). Returns -/// `None` when the payload is for a different waiter on the shared reply channel or is unparseable. -fn parse_reply(payload: &str, id: i64) -> Option { - let mut parts = payload.split(':'); - let reply_id: i64 = parts.next()?.parse().ok()?; - if reply_id != id { - return None; - } - match parts.next()? { - "committed" => Some(ReplyOutcome::Committed), - "conflict" => Some(ReplyOutcome::Conflict), - _ => None, +/// Wait for a known leader, returning a retryable error if none is elected in time. +async fn wait_for_leader(shared: &Arc) -> Result { + let deadline = Instant::now() + LEADER_WAIT_TIMEOUT; + loop { + if let Some(lease) = shared.current_lease() { + return Ok(lease); + } + if Instant::now() >= deadline { + return Err(DatabaseError::NotCommitted.into()); + } + tokio::time::sleep(LEADER_POLL_INTERVAL).await; } } - -enum Status { - Pending, - Committed, - Conflict, -} - -/// Read the current status of a commit request. `Ok(None)` means the row no longer exists. -async fn read_status(shared: &Arc, id: i64) -> Result> { - let conn = shared - .pool - .get() - .await - .context("failed to get connection for commit status poll")?; - - let row = conn - .query_opt( - "SELECT status FROM udb_commit_requests WHERE id = $1", - &[&id], - ) - .await - .context("failed to read commit request status")?; - - let Some(row) = row else { - return Ok(None); - }; - - let status: String = row.get(0); - let status = match status.as_str() { - "committed" => Status::Committed, - "conflict" => Status::Conflict, - // 'pending' or any in-flight state. - _ => Status::Pending, - }; - Ok(Some(status)) -} diff --git a/engine/packages/universaldb/src/driver/postgres/database.rs b/engine/packages/universaldb/src/driver/postgres/database.rs index 3ff891011f..dc1c42d598 100644 --- a/engine/packages/universaldb/src/driver/postgres/database.rs +++ b/engine/packages/universaldb/src/driver/postgres/database.rs @@ -10,7 +10,7 @@ use std::{ use anyhow::{Context, Result}; use deadpool_postgres::{Config, ManagerConfig, Pool, PoolConfig, RecyclingMethod, Runtime}; use rivet_postgres_util::build_tls_config; -use tokio::task::JoinHandle; +use tokio::{sync::mpsc, task::JoinHandle}; use tokio_postgres_rustls::MakeRustlsConnect; use url::Url; use uuid::Uuid; @@ -24,19 +24,26 @@ use crate::{ }; use super::{ - listener::PgListener, resolver, shared::PostgresShared, transaction::PostgresTransactionDriver, + nats::{self, NatsConfig, NatsTransport, Subjects}, + resolver::{self, ResolverInput}, + shared::PostgresShared, + transaction::PostgresTransactionDriver, + transport::{COMMIT_QUEUE_BOUND, Transport}, }; const GC_INTERVAL: Duration = Duration::from_secs(30); -/// Terminal and orphaned commit-request rows older than this are garbage collected. Must be well -/// beyond the longest a follower could spend awaiting a result, so a result is never deleted before -/// it is observed. -const COMMIT_ROW_MAX_AGE_SECS: i64 = 60; +/// Failover dedup rows older than this are garbage collected. Must be well beyond the longest a +/// follower could spend resending a commit across a leader failover, so a dedup record is never +/// deleted while a resend that needs it could still arrive. +const DEDUP_ROW_MAX_AGE_SECS: i64 = 60; #[derive(Clone, Debug)] pub struct PostgresConfig { pub connection_string: String, pub ssl_config: Option, + /// When set, UniversalDB runs in multi-node mode and uses NATS for follower-to-leader commit + /// transport. When `None`, it runs single-node with an in-process resolver. + pub nats: Option, } #[derive(Clone, Debug)] @@ -47,11 +54,12 @@ pub struct PostgresSslConfig { } impl PostgresConfig { - /// Create a new PostgreSQL configuration with sane defaults + /// Create a new PostgreSQL configuration with sane defaults (single-node). pub fn new(connection_string: String) -> Self { Self { connection_string, ssl_config: None, + nats: None, } } } @@ -89,37 +97,50 @@ impl PostgresDatabaseDriver { Self::init_schema(&conn).await?; } - // Unique per-process node id (no hyphens) used to name this node's NOTIFY channels. Kept - // short so `udb_commit_` stays within Postgres's 63-byte identifier limit. + // Unique per-process node id (no hyphens). Names this node's NATS commit subject and is the + // dedup `client_node_id`. let node_id = Uuid::new_v4().simple().to_string(); - let listener = PgListener::new( - config.connection_string.clone(), - ssl_disabled, - config - .ssl_config - .as_ref() - .and_then(|c| c.ssl_root_cert_path.clone()), - config - .ssl_config - .as_ref() - .and_then(|c| c.ssl_client_cert_path.clone()), - config - .ssl_config - .as_ref() - .and_then(|c| c.ssl_client_key_path.clone()), - ); - - let shared = PostgresShared::new(pool, node_id, listener); - - // Every node runs the resolver; only the elected leader drains the commit queue. - let resolver_handle = resolver::spawn(shared.clone()); + let (shared, resolver_input) = match &config.nats { + None => { + // Single-node: the follower commit path hands jobs straight to the in-process leader + // drain loop. + let (commit_tx, commit_rx) = mpsc::channel(COMMIT_QUEUE_BOUND); + let shared = + PostgresShared::new(pool, node_id, Transport::SingleNode { commit_tx }); + + // Acquire the leader lease behind the correctness gate BEFORE spawning the resolver so + // a failure to acquire fails startup loudly. + let initial_epoch = resolver::acquire_single_node_gate(&shared).await?; + + ( + shared, + ResolverInput::SingleNode { + rx: commit_rx, + initial_epoch, + }, + ) + } + Some(nats_config) => { + // Multi-node: commits travel to the elected leader over NATS request/reply. + let client = nats::connect(nats_config).await?; + let subjects = Subjects::new(&config.connection_string); + let shared = PostgresShared::new( + pool, + node_id, + Transport::MultiNode(NatsTransport { client, subjects }), + ); + + (shared, ResolverInput::MultiNode) + } + }; + let resolver_handle = resolver::spawn(shared.clone(), resolver_input); let gc_handle = Self::spawn_gc(shared.clone()); Ok(PostgresDatabaseDriver { shared, - max_retries: AtomicI32::new(100), + max_retries: AtomicI32::new(10), resolver_handle, gc_handle, }) @@ -180,19 +201,13 @@ impl PostgresDatabaseDriver { CREATE SEQUENCE IF NOT EXISTS udb_version_seq AS BIGINT START WITH 1 INCREMENT BY 1 MINVALUE 1; - CREATE TABLE IF NOT EXISTS udb_commit_requests ( - id BIGSERIAL PRIMARY KEY, - epoch BIGINT NOT NULL, - read_version BIGINT NOT NULL, - payload BYTEA NOT NULL, - reply_channel TEXT NOT NULL, - status TEXT NOT NULL DEFAULT 'pending', - commit_version BIGINT, - created_at TIMESTAMPTZ NOT NULL DEFAULT now() - ); - - CREATE INDEX IF NOT EXISTS udb_commit_requests_pending - ON udb_commit_requests (id) WHERE status = 'pending';", + CREATE TABLE IF NOT EXISTS udb_applied ( + client_node_id BYTEA NOT NULL, + client_seq BIGINT NOT NULL, + commit_version BIGINT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (client_node_id, client_seq) + );", ) .await .context("failed to initialize postgres schema")?; @@ -200,6 +215,9 @@ impl PostgresDatabaseDriver { Ok(()) } + /// Garbage-collect old failover dedup records. A dedup row only needs to outlive the longest a + /// follower could spend resending a single commit across a leader failover, so terminal rows past + /// [`DEDUP_ROW_MAX_AGE_SECS`] are safe to drop. (Single-node never writes this table.) fn spawn_gc(shared: Arc) -> JoinHandle<()> { tokio::spawn(async move { let mut interval = tokio::time::interval(GC_INTERVAL); @@ -211,20 +229,20 @@ impl PostgresDatabaseDriver { let conn = match shared.pool.get().await { Ok(conn) => conn, Err(err) => { - tracing::debug!(?err, "failed to get connection for commit gc"); + tracing::debug!(?err, "failed to get connection for dedup gc"); continue; } }; if let Err(err) = conn .execute( - "DELETE FROM udb_commit_requests + "DELETE FROM udb_applied WHERE created_at < now() - ($1::bigint * interval '1 second')", - &[&COMMIT_ROW_MAX_AGE_SECS], + &[&DEDUP_ROW_MAX_AGE_SECS], ) .await { - tracing::error!(?err, "failed postgres commit-queue gc"); + tracing::error!(?err, "failed postgres dedup gc"); } } }) diff --git a/engine/packages/universaldb/src/driver/postgres/listener.rs b/engine/packages/universaldb/src/driver/postgres/listener.rs deleted file mode 100644 index e5a7919dfe..0000000000 --- a/engine/packages/universaldb/src/driver/postgres/listener.rs +++ /dev/null @@ -1,228 +0,0 @@ -use std::{path::PathBuf, sync::Arc, time::Duration}; - -use futures_util::future::poll_fn; -use rivet_postgres_util::build_tls_config; -use scc::HashMap; -use tokio::{ - io::{AsyncRead, AsyncWrite}, - sync::{Mutex, broadcast}, -}; -use tokio_postgres::AsyncMessage; -use tokio_postgres_rustls::MakeRustlsConnect; - -/// How long to wait between reconnect attempts for the dedicated LISTEN connection. -const RECONNECT_BACKOFF: Duration = Duration::from_secs(1); -/// Capacity of each channel's broadcast buffer. Notifications are wakeup signals with a polling -/// backstop, so a lagged receiver only delays a wake, never drops a durable commit. -const BROADCAST_CAPACITY: usize = 1024; - -struct Subscription { - tx: broadcast::Sender, -} - -/// Owns a single dedicated Postgres connection used exclusively for `LISTEN`. Demultiplexes -/// incoming `NOTIFY` payloads to per-channel broadcast senders and re-`LISTEN`s every registered -/// channel after a reconnect. -/// -/// This is separate from the deadpool pool because deadpool recycles connections and drops the -/// async notification stream; LISTEN requires owning the connection's message stream directly. -pub struct PgListener { - conn_str: String, - ssl_disabled: bool, - ssl_root_cert_path: Option, - ssl_client_cert_path: Option, - ssl_client_key_path: Option, - channels: Arc>, - client: Arc>>, -} - -impl PgListener { - pub fn new( - conn_str: String, - ssl_disabled: bool, - ssl_root_cert_path: Option, - ssl_client_cert_path: Option, - ssl_client_key_path: Option, - ) -> Self { - let channels: Arc> = Arc::new(HashMap::new()); - let client: Arc>> = Arc::new(Mutex::new(None)); - - tokio::spawn(Self::connection_lifecycle( - conn_str.clone(), - ssl_disabled, - ssl_root_cert_path.clone(), - ssl_client_cert_path.clone(), - ssl_client_key_path.clone(), - channels.clone(), - client.clone(), - )); - - Self { - conn_str, - ssl_disabled, - ssl_root_cert_path, - ssl_client_cert_path, - ssl_client_key_path, - channels, - client, - } - } - - /// Subscribe to a channel, registering a `LISTEN` if this is the first subscriber. Returns a - /// broadcast receiver of notification payloads. Idempotent per channel. - pub async fn listen(&self, channel: &str) -> broadcast::Receiver { - match self.channels.entry_async(channel.to_string()).await { - scc::hash_map::Entry::Occupied(entry) => entry.get().tx.subscribe(), - scc::hash_map::Entry::Vacant(entry) => { - let (tx, rx) = broadcast::channel(BROADCAST_CAPACITY); - entry.insert_entry(Subscription { tx }); - - // Best-effort immediate LISTEN; the lifecycle task re-LISTENs on reconnect. - if let Some(client) = &*self.client.lock().await { - if let Err(err) = client.execute(&format!("LISTEN \"{channel}\""), &[]).await { - tracing::warn!(?err, %channel, "failed to LISTEN, will retry on reconnect"); - } - } - - rx - } - } - } - - async fn connection_lifecycle( - conn_str: String, - ssl_disabled: bool, - ssl_root_cert_path: Option, - ssl_client_cert_path: Option, - ssl_client_key_path: Option, - channels: Arc>, - client: Arc>>, - ) { - loop { - let connected = if ssl_disabled { - Self::connect_and_run(&conn_str, tokio_postgres::NoTls, &channels, &client).await - } else { - match build_tls_config( - ssl_root_cert_path.as_ref(), - ssl_client_cert_path.as_ref(), - ssl_client_key_path.as_ref(), - ) { - Ok(tls_config) => { - Self::connect_and_run( - &conn_str, - MakeRustlsConnect::new(tls_config), - &channels, - &client, - ) - .await - } - Err(err) => { - tracing::error!(?err, "failed to build listener TLS config"); - false - } - } - }; - - if !connected { - tokio::time::sleep(RECONNECT_BACKOFF).await; - } - } - } - - /// Connects, re-LISTENs all channels, then drives the notification poll loop until the - /// connection closes. Returns `true` if a connection was successfully established (so the caller - /// can skip the reconnect backoff). - async fn connect_and_run( - conn_str: &str, - tls: T, - channels: &Arc>, - client: &Arc>>, - ) -> bool - where - T: tokio_postgres::tls::MakeTlsConnect, - T::Stream: AsyncRead + AsyncWrite + Unpin + Send + 'static, - T::TlsConnect: Send, - >::Future: Send, - { - let (new_client, connection) = match tokio_postgres::connect(conn_str, tls).await { - Ok(pair) => pair, - Err(err) => { - tracing::error!(?err, "failed to connect postgres listener"); - return false; - } - }; - - let channels_poll = channels.clone(); - let poll_handle = - tokio::spawn(async move { Self::poll_connection(connection, channels_poll).await }); - - // Re-LISTEN all registered channels on the fresh connection. - let mut registered = Vec::new(); - channels - .iter_async(|k, _| { - registered.push(k.clone()); - true - }) - .await; - for channel in ®istered { - if let Err(err) = new_client - .execute(&format!("LISTEN \"{channel}\""), &[]) - .await - { - tracing::error!(?err, %channel, "failed to re-LISTEN channel after reconnect"); - } - } - - *client.lock().await = Some(new_client); - - // Block until the poll loop ends (connection closed or errored). - let _ = poll_handle.await; - - *client.lock().await = None; - - true - } - - async fn poll_connection( - mut connection: tokio_postgres::Connection, - channels: Arc>, - ) where - S: AsyncRead + AsyncWrite + Unpin, - T: AsyncRead + AsyncWrite + Unpin, - { - loop { - match poll_fn(|cx| connection.poll_message(cx)).await { - Some(Ok(AsyncMessage::Notification(note))) => { - if let Some(sub) = channels.get_async(note.channel()).await { - // Ignore send errors: no active receiver just means no one is waiting - // right now; the polling backstop covers them. - let _ = sub.tx.send(note.payload().to_string()); - } - } - Some(Ok(_)) => {} - Some(Err(err)) => { - tracing::warn!(?err, "postgres listener connection error"); - break; - } - None => { - tracing::warn!("postgres listener connection closed"); - break; - } - } - } - } -} - -impl Clone for PgListener { - fn clone(&self) -> Self { - Self { - conn_str: self.conn_str.clone(), - ssl_disabled: self.ssl_disabled, - ssl_root_cert_path: self.ssl_root_cert_path.clone(), - ssl_client_cert_path: self.ssl_client_cert_path.clone(), - ssl_client_key_path: self.ssl_client_key_path.clone(), - channels: self.channels.clone(), - client: self.client.clone(), - } - } -} diff --git a/engine/packages/universaldb/src/driver/postgres/mod.rs b/engine/packages/universaldb/src/driver/postgres/mod.rs index 64f4bbd1bf..20b80765c0 100644 --- a/engine/packages/universaldb/src/driver/postgres/mod.rs +++ b/engine/packages/universaldb/src/driver/postgres/mod.rs @@ -1,10 +1,12 @@ mod codec; mod commit; mod database; -mod listener; +mod nats; mod resolver; mod shared; mod transaction; mod transaction_task; +mod transport; pub use database::{PostgresConfig, PostgresDatabaseDriver, PostgresSslConfig}; +pub use nats::NatsConfig; diff --git a/engine/packages/universaldb/src/driver/postgres/nats.rs b/engine/packages/universaldb/src/driver/postgres/nats.rs new file mode 100644 index 0000000000..8fab78811a --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/nats.rs @@ -0,0 +1,149 @@ +use std::str::FromStr; + +use anyhow::{Context, Result}; +use futures_util::StreamExt; +use tokio::sync::mpsc; + +use super::{ + codec, + transport::{CommitJob, DedupKey, Responder}, +}; + +/// NATS connection settings for UniversalDB multi-node mode. Built from the resolved +/// `database.postgres.nats` config (which may be inherited from the UPS NATS config). +#[derive(Clone, Debug)] +pub struct NatsConfig { + /// `host:port` server addresses. + pub addresses: Vec, + pub username: Option, + pub password: Option, + pub client_capacity: usize, + pub subscription_capacity: usize, +} + +/// The multi-node transport handle: the NATS client plus the cluster-scoped subject names. +pub struct NatsTransport { + pub client: async_nats::Client, + pub subjects: Subjects, +} + +/// Cluster-scoped UniversalDB NATS subjects. The cluster prefix is derived from the Postgres +/// connection string so two separate clusters that happen to share one NATS deployment do not +/// cross-deliver watermark/election broadcasts or commits. +#[derive(Clone)] +pub struct Subjects { + prefix: String, +} + +impl Subjects { + pub fn new(connection_string: &str) -> Self { + Subjects { + prefix: format!("udb.{:016x}", fnv1a_64(connection_string.as_bytes())), + } + } + + /// Subject a follower sends a commit request to, and the elected leader subscribes to. Namespaced + /// by the leader's node id so only the current leader receives commits. + pub fn commit(&self, leader_id: &str) -> String { + format!("{}.commit.{leader_id}", self.prefix) + } + + /// Subject the leader publishes each watermark advance to; every node subscribes. + pub fn watermark(&self) -> String { + format!("{}.watermark", self.prefix) + } + + /// Subject a departing leader publishes to so standby candidates elect immediately. + pub fn election(&self) -> String { + format!("{}.election", self.prefix) + } +} + +/// Connect a NATS client for the UniversalDB multi-node transport. +pub async fn connect(config: &NatsConfig) -> Result { + let server_addrs = config + .addresses + .iter() + .map(|addr| format!("nats://{addr}")) + .map(|url| async_nats::ServerAddr::from_str(&url)) + .collect::, _>>() + .context("failed to parse udb nats addresses")?; + + let mut options = match (&config.username, &config.password) { + (Some(username), Some(password)) => { + async_nats::ConnectOptions::with_user_and_password(username.clone(), password.clone()) + } + _ => async_nats::ConnectOptions::new(), + }; + options = options + .client_capacity(config.client_capacity) + .subscription_capacity(config.subscription_capacity); + + options + .connect(&server_addrs[..]) + .await + .context("failed to connect udb nats client") +} + +/// Leader-side task: subscribe to this leader's commit subject, decode each request into a +/// [`CommitJob`], and forward it into the drain loop's job queue. Returns when the subscription ends +/// (client closed) or the drain loop's receiver is dropped (step-down). +pub async fn run_commit_subscriber( + client: async_nats::Client, + subject: String, + jobs_tx: mpsc::Sender, +) -> Result<()> { + let mut sub = client + .subscribe(subject.clone()) + .await + .with_context(|| format!("failed to subscribe to udb commit subject {subject}"))?; + + while let Some(msg) = sub.next().await { + let Some(reply) = msg.reply.clone() else { + tracing::warn!("udb commit request missing reply subject; dropping"); + continue; + }; + + let decoded = match codec::decode_commit_request(&msg.payload) { + Ok(decoded) => decoded, + Err(err) => { + tracing::warn!(?err, "failed to decode udb commit request; dropping"); + continue; + } + }; + + let job = CommitJob { + read_version: decoded.read_version, + conflict_ranges: decoded.conflict_ranges, + operations: decoded.operations, + dedup_key: Some(DedupKey { + client_node_id: decoded.client_node_id, + client_seq: decoded.client_seq as i64, + }), + responder: Responder::Nats { + client: client.clone(), + reply, + }, + }; + + // A full queue applies backpressure; a closed queue means the drain loop stepped down. + if jobs_tx.send(job).await.is_err() { + break; + } + } + + Ok(()) +} + +/// FNV-1a 64-bit hash. Deterministic across processes (unlike `DefaultHasher`), used only to derive a +/// stable cluster subject prefix. +fn fnv1a_64(bytes: &[u8]) -> u64 { + const OFFSET: u64 = 0xcbf2_9ce4_8422_2325; + const PRIME: u64 = 0x0000_0100_0000_01b3; + let mut hash = OFFSET; + for &b in bytes { + hash ^= b as u64; + hash = hash.wrapping_mul(PRIME); + } + hash +} diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs index 098f022663..3a779f838f 100644 --- a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs +++ b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs @@ -7,171 +7,292 @@ use std::{ time::{Duration, Instant}, }; -use anyhow::{Context, Result}; -use tokio::sync::broadcast; +use anyhow::{Context, Result, bail}; +use futures_util::StreamExt; +use tokio::sync::mpsc; use tokio_util::task::AbortOnDropHandle; use crate::{conflict_tracker::TransactionConflictTracker, transaction::TXN_TIMEOUT}; use lease::LEASE_TTL_SECS; -use super::shared::{ - ELECTION_CHANNEL, LEASE_ID, LeaseInfo, PostgresShared, WATERMARK_CHANNEL, commit_channel, +use super::{ + shared::{LEASE_ID, LeaseInfo, PostgresShared}, + transport::{CommitJob, CommitOutcome, Transport}, }; /// Max commits resolved+applied per batch (group commit). Amortizes the resolver, Postgres /// round-trips, and fsync across the batch. -const DRAIN_BATCH_SIZE: i64 = 256; +const DRAIN_BATCH_SIZE: usize = 256; /// How often a leader renews its lease. Must be comfortably under `LEASE_TTL_SECS`. const RENEW_INTERVAL: Duration = Duration::from_secs(3); -/// Backstop poll cadence so a missed `udb_commit` NOTIFY cannot stall the drain indefinitely. -const POLL_BACKSTOP: Duration = Duration::from_millis(50); - /// How long a candidate waits before retrying election when another node holds the lease. const ELECTION_RETRY: Duration = Duration::from_secs(2); -/// Spawn the per-process resolver task. Every node runs this; only the elected leader drains the -/// commit queue. The returned handle is aborted when the owning driver drops, which stops lease -/// renewal so the lease expires and another node can take over (node-death / failover path). -pub fn spawn(shared: Arc) -> tokio::task::JoinHandle<()> { - tokio::spawn(run(shared)) +/// Single-node leader-acquire gate: total time to keep retrying before giving up and failing +/// startup. Must exceed the lease TTL so a crashed predecessor's lease has time to expire. +const GATE_TOTAL: Duration = Duration::from_secs((LEASE_TTL_SECS as u64) * 2 + 5); +/// Backoff between single-node gate attempts. +const GATE_RETRY: Duration = Duration::from_secs(1); + +/// What feeds the leader drain loop. Single-node owns the process-wide commit receiver and an +/// already-acquired lease epoch (the startup gate ran before this task spawned). Multi-node creates a +/// fresh NATS-fed receiver each time it wins an election. +pub enum ResolverInput { + SingleNode { + rx: mpsc::Receiver, + initial_epoch: i64, + }, + MultiNode, } -async fn run(shared: Arc) { - // A departing leader NOTIFYs this channel after releasing its lease so we elect immediately - // rather than waiting out the full `ELECTION_RETRY` tick. - let mut election_rx = shared.listener.listen(ELECTION_CHANNEL).await; +/// Single-node startup gate: acquire the leader lease, retrying with backoff. Fails if it cannot be +/// acquired within [`GATE_TOTAL`], which means either another engine instance is running against this +/// Postgres (a real misconfiguration in single-node mode) or a previous instance crashed without +/// releasing its lease and it has not yet expired. Each failed attempt warns. +pub async fn acquire_single_node_gate(shared: &Arc) -> Result { + let deadline = Instant::now() + GATE_TOTAL; + let mut attempt = 0u32; + loop { + attempt += 1; + match lease::try_acquire(&shared.pool, &shared.node_id).await { + Ok(Some(acquired)) => { + tracing::debug!( + epoch = acquired.epoch, + attempt, + node_id = %shared.node_id, + "acquired udb postgres single-node leader lease" + ); + return Ok(acquired.epoch); + } + Ok(None) => { + tracing::warn!( + attempt, + "udb postgres single-node could not acquire leader lease. another instance may hold it \ + or a previous instance did not release it; backing off" + ); + } + Err(err) => { + tracing::warn!( + ?err, + attempt, + "udb postgres single-node lease acquire errored, retrying" + ); + } + } + + if Instant::now() >= deadline { + bail!( + "udb postgres single-node failed to acquire leader lease after {attempt} attempts; refusing \ + to start. another engine instance may be running against this postgres, or a \ + previous instance crashed without releasing its lease (wait for it to expire). if you intend \ + to run a multi-node setup you must configure NATS." + ); + } + + tokio::time::sleep(GATE_RETRY).await; + } +} + +/// Spawn the per-process resolver task. The returned handle is aborted when the owning driver drops, +/// which stops lease renewal so the lease expires and another node can take over. +pub fn spawn(shared: Arc, input: ResolverInput) -> tokio::task::JoinHandle<()> { + tokio::spawn(run(shared, input)) +} + +async fn run(shared: Arc, input: ResolverInput) { + match input { + ResolverInput::SingleNode { rx, initial_epoch } => { + run_single_node(shared, rx, initial_epoch).await + } + ResolverInput::MultiNode => run_multi_node(shared).await, + } +} + +/// Single-node: this node is the only node and is always the leader. The startup gate already +/// acquired the lease, so lead immediately. If the lease is ever lost (another node appeared, a real +/// misconfiguration), log loudly and re-acquire through the gate. +async fn run_single_node( + shared: Arc, + mut rx: mpsc::Receiver, + initial_epoch: i64, +) { + let mut epoch = initial_epoch; + loop { + if let Err(err) = lead(&shared, epoch, &mut rx).await { + tracing::error!(?err, "udb postgres single-node leader loop errored"); + } + tracing::error!( + epoch, + "udb postgres single-node lost the leader lease; re-acquiring (another engine instance may be \ + running against this postgres)" + ); + epoch = loop { + match acquire_single_node_gate(&shared).await { + Ok(epoch) => break epoch, + Err(err) => { + tracing::error!(?err, "udb postgres single-node re-acquire failed; retrying"); + } + } + }; + } +} +/// Multi-node: race the lease against other nodes; whoever wins leads until it loses the lease. +async fn run_multi_node(shared: Arc) { loop { match lease::try_acquire(&shared.pool, &shared.node_id).await { Ok(Some(acquired)) => { - tracing::info!(epoch = acquired.epoch, node_id = %shared.node_id, "acquired udb leader lease"); - if let Err(err) = lead(&shared, acquired.epoch).await { + tracing::info!(epoch = acquired.epoch, node_id = %shared.node_id, "acquired udb postgres leader lease"); + + // Each leadership term gets its own NATS-fed commit queue. The subscriber forwards + // decoded commit requests into `rx`; aborting it on step-down stops accepting commits. + let (tx, mut rx) = mpsc::channel(super::transport::COMMIT_QUEUE_BOUND); + let subscriber = spawn_commit_subscriber(&shared, tx); + + if let Err(err) = lead(&shared, acquired.epoch, &mut rx).await { tracing::error!(?err, "udb leader loop errored, stepping down"); } - tracing::info!(epoch = acquired.epoch, "stepped down from udb leader"); - } - Ok(None) => { - wait_for_election_retry(&shared, &mut election_rx).await; + if let Some(handle) = subscriber { + handle.abort(); + } + tracing::info!( + epoch = acquired.epoch, + "stepped down from udb postgres leader" + ); } + Ok(None) => wait_for_election_retry(&shared).await, Err(err) => { tracing::warn!(?err, "failed udb lease acquire attempt"); - wait_for_election_retry(&shared, &mut election_rx).await; + wait_for_election_retry(&shared).await; } } } } -/// Wait before retrying the election: either the `ELECTION_RETRY` backstop elapses, or a departing -/// leader wakes us via `ELECTION_CHANNEL` so handoff is near-instant. -async fn wait_for_election_retry( +/// Spawn the leader's NATS commit subscriber for multi-node. Returns `None` in single-node (no NATS). +fn spawn_commit_subscriber( shared: &Arc, - election_rx: &mut broadcast::Receiver, -) { - tokio::select! { - _ = tokio::time::sleep(ELECTION_RETRY) => {} - res = election_rx.recv() => { - if matches!(res, Err(broadcast::error::RecvError::Closed)) { - // The listener recreates the channel on reconnect; re-subscribe. - *election_rx = shared.listener.listen(ELECTION_CHANNEL).await; + tx: mpsc::Sender, +) -> Option> { + let Transport::MultiNode(nats) = &shared.transport else { + return None; + }; + let client = nats.client.clone(); + let subject = nats.subjects.commit(&shared.node_id); + Some(AbortOnDropHandle::new(tokio::spawn(async move { + if let Err(err) = super::nats::run_commit_subscriber(client, subject, tx).await { + tracing::warn!(?err, "udb commit subscriber ended"); + } + }))) +} + +/// Wait before retrying the election: either the `ELECTION_RETRY` backstop elapses, or a departing +/// leader wakes us via the election broadcast so handoff is near-instant. +async fn wait_for_election_retry(shared: &Arc) { + let Transport::MultiNode(nats) = &shared.transport else { + tokio::time::sleep(ELECTION_RETRY).await; + return; + }; + + let election = nats.client.subscribe(nats.subjects.election()).await; + match election { + Ok(mut sub) => { + tokio::select! { + _ = tokio::time::sleep(ELECTION_RETRY) => {} + _ = sub.next() => {} } } + Err(_) => tokio::time::sleep(ELECTION_RETRY).await, } } -/// Best-effort graceful leadership handoff invoked on shutdown. If this node currently holds the -/// lease, expire it and wake a standby so it takes over immediately instead of waiting out the TTL. -/// Safe to call on a follower: the fenced release matches no row and nothing is notified. The -/// caller must already have stopped lease renewal before calling this. +/// Best-effort graceful leadership handoff invoked on shutdown. If this node holds the lease, expire +/// it and wake standbys so they take over immediately instead of waiting out the TTL. Safe to call on +/// a follower. The caller must already have stopped lease renewal before calling this. pub async fn handoff(shared: &Arc) { match lease::release(&shared.pool, &shared.node_id).await { Ok(true) => { - tracing::info!(node_id = %shared.node_id, "released udb leader lease for graceful handoff"); - notify_election(shared).await; + tracing::info!(node_id = %shared.node_id, "released udb postgres leader lease for graceful handoff"); + if let Transport::MultiNode(nats) = &shared.transport { + if let Err(err) = nats + .client + .publish(nats.subjects.election(), Vec::new().into()) + .await + { + tracing::debug!(?err, "failed to publish udb election wake"); + } + } } Ok(false) => {} - Err(err) => { - tracing::warn!(?err, "failed to release udb lease on shutdown"); - } - } -} - -/// Wake standby candidates so the next election fires immediately after a graceful release. -async fn notify_election(shared: &Arc) { - let conn = match shared.pool.get().await { - Ok(conn) => conn, - Err(err) => { - tracing::debug!(?err, "failed to get connection for election notify"); - return; - } - }; - - if let Err(err) = conn - .execute("SELECT pg_notify($1, '')", &[&ELECTION_CHANNEL]) - .await - { - tracing::debug!(?err, "failed to notify election channel"); + Err(err) => tracing::warn!(?err, "failed to release udb lease on shutdown"), } } /// Leader entry point: publish our lease, compute the cold-window floor, then run renewal and -/// draining as two sibling tasks. They coordinate purely by completion: when either returns (lease -/// lost or error), the other is aborted and the leader steps down. Both operations are safe to -/// hard-abort, so no explicit cancellation signalling is needed. A renew is a single fenced -/// `UPDATE`; a drain batch runs in one Postgres transaction that rolls back cleanly when dropped, -/// leaving its claimed requests `pending` for the next leader. -async fn lead(shared: &Arc, epoch: i64) -> Result<()> { +/// draining. Renewal runs on its own task so a long drain cannot starve it past the lease TTL; the +/// drain loop runs inline so it can borrow the commit receiver. Returns when the lease is lost (renew +/// reports it gone, or an apply is epoch-fenced). +async fn lead( + shared: &Arc, + epoch: i64, + rx: &mut mpsc::Receiver, +) -> Result<()> { // Publish our own lease into the cache immediately so our local commits route to us. shared.set_lease(LeaseInfo { epoch, leader_addr: shared.node_id.clone(), }); - // The recovery floor: a freshly elected leader has a cold conflict window, so reject commits - // whose read_version predates the floor until the window warms (one TXN_TIMEOUT), forcing - // those followers to take a fresh read_version. + // The recovery floor: a freshly elected leader has a cold conflict window, so reject commits whose + // read_version predates the floor until the window warms (one TXN_TIMEOUT), forcing those + // followers to take a fresh read_version. let recovery_version = recovery_floor(shared).await?; let recovery_deadline = Instant::now() + TXN_TIMEOUT; - tracing::info!( + // Seed our read-version cache to the durable floor so our own follower reads are not cold-window + // rejected (a read at `read_version < recovery_version` is rejected during the cold window, and the + // cache otherwise starts at 0). Followers learn the floor from the watermark broadcast / lease poll. + shared.advance_durable_version(recovery_version as i64); + + tracing::debug!( epoch, recovery_version, cold_window_ms = TXN_TIMEOUT.as_millis() as u64, + multi_node = shared.is_multi_node(), "udb leader entering lead loop" ); - // Renewal runs in its own task so the drain loop can never starve it: if renewal shared the - // drain loop, a single long drain under sustained load would block renewal past the lease TTL - // and the lease would be lost mid-drain, thrashing leadership. Both are held in abort-on-drop - // handles so a hard abort of this `run` task (driver drop / node death) tears them down. A leaked - // renew task would keep this dead leader's lease alive and block failover. + let tracker = TransactionConflictTracker::new(); let mut renew = AbortOnDropHandle::new(tokio::spawn(renew_loop(shared.clone(), epoch))); - let mut drain = AbortOnDropHandle::new(tokio::spawn(drain_loop( - shared.clone(), + + let drain = drain_loop( + shared, epoch, + &tracker, recovery_version, recovery_deadline, - ))); + rx, + ); + tokio::pin!(drain); - // Whichever task returns first (lease lost or error), step down; the other is aborted when its - // handle drops at the end of this scope. A clean exit yields its inner result; a panic surfaces - // through `?` as a join error. + // Whichever finishes first (lease lost via renew, or epoch fenced during a drain apply) ends the + // leadership term. The renew handle yields a join result; the inline drain yields directly. tokio::select! { res = &mut renew => res?, - res = &mut drain => res?, + res = &mut drain => res, } } /// Lease-renewal loop. Runs on its own task and pool connection so it cannot be starved by drain -/// work. Returns when the lease is definitively gone (epoch bumped by another node, or renewal -/// failing for the whole lease TTL), which causes [`lead`] to abort the drain task and step down. +/// work. Returns when the lease is definitively gone. async fn renew_loop(shared: Arc, epoch: i64) -> Result<()> { let mut interval = tokio::time::interval(RENEW_INTERVAL); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - // The lease was just acquired/renewed (expires_at = now + TTL), so consume the immediate first - // tick and renew after one interval. + // The lease was just acquired (expires_at = now + TTL), so consume the immediate first tick and + // renew after one interval. interval.tick().await; let mut last_renew = Instant::now(); @@ -179,42 +300,16 @@ async fn renew_loop(shared: Arc, epoch: i64) -> Result<()> { loop { interval.tick().await; - let gap_ms = last_renew.elapsed().as_millis() as u64; - let renew_start = Instant::now(); match lease::renew(&shared.pool, &shared.node_id, epoch).await { - Ok(true) => { - last_renew = Instant::now(); - let renew_query_ms = renew_start.elapsed().as_millis() as u64; - // With renewal on its own task this gap should track RENEW_INTERVAL closely; a large - // gap now points at pool or Postgres contention. - if gap_ms > RENEW_INTERVAL.as_millis() as u64 * 2 { - tracing::warn!( - epoch, - gap_since_last_renew_ms = gap_ms, - renew_query_ms, - "udb leader renew was delayed (pool or postgres contention)" - ); - } else { - tracing::debug!( - epoch, - gap_since_last_renew_ms = gap_ms, - renew_query_ms, - "udb leader renewed lease" - ); - } - } + Ok(true) => last_renew = Instant::now(), Ok(false) => { tracing::warn!( epoch, - gap_since_last_renew_ms = gap_ms, - "udb leader lost lease on renew (epoch bumped by another node); stepping down" + "udb leader lost lease on renew (epoch bumped); stepping down" ); return Ok(()); } Err(err) => { - // A transient renew error is tolerable within the TTL; keep retrying. Only give up - // if we have been unable to renew for the whole lease TTL, at which point we can no - // longer assume we hold the lease. if last_renew.elapsed() >= Duration::from_secs(LEASE_TTL_SECS as u64) { tracing::warn!( ?err, @@ -229,117 +324,76 @@ async fn renew_loop(shared: Arc, epoch: i64) -> Result<()> { } } -/// Drain loop. Processes one batch per iteration. Draining until the queue emptied in a single call -/// could run for many seconds under sustained load; processing one batch at a time keeps the loop -/// at a clean await point between batches. A non-empty queue still drains back-to-back with no idle -/// wait, so throughput is unchanged. It blocks on `select!` (wake NOTIFY or poll backstop) only -/// when the queue is empty. Returns when this leader's epoch is fenced out; otherwise [`lead`] -/// aborts it when renewal reports the lease is lost. +/// Drain loop. Collects one batch of commit jobs from the queue and processes it per iteration. +/// Returns when this leader's epoch is fenced out during an apply, or the queue closes. async fn drain_loop( - shared: Arc, + shared: &Arc, epoch: i64, + tracker: &TransactionConflictTracker, recovery_version: u64, recovery_deadline: Instant, + rx: &mut mpsc::Receiver, ) -> Result<()> { - let tracker = TransactionConflictTracker::new(); - - let mut wake_rx = shared - .listener - .listen(&commit_channel(&shared.node_id)) - .await; - - let mut poll_interval = tokio::time::interval(POLL_BACKSTOP); - poll_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - loop { + let batch = collect_batch(rx).await; + if batch.is_empty() { + // The queue closed (all senders dropped). Step down. + return Ok(()); + } + match drain_batch( - &shared, + shared, epoch, - &tracker, + tracker, recovery_version, recovery_deadline, + batch, ) .await? { BatchOutcome::LostLease => { - tracing::warn!( - epoch, - "udb leader stepping down: lost lease during drain (epoch fenced on watermark)" - ); + tracing::warn!(epoch, "udb leader stepping down: epoch fenced during apply"); return Ok(()); } - // More work may be pending; loop immediately to keep throughput up. BatchOutcome::Processed => continue, - BatchOutcome::Empty => { - tokio::select! { - res = wake_rx.recv() => { - match res { - Ok(_) | Err(broadcast::error::RecvError::Lagged(_)) => {} - Err(broadcast::error::RecvError::Closed) => { - wake_rx = shared - .listener - .listen(&commit_channel(&shared.node_id)) - .await; - } - } - } - _ = poll_interval.tick() => {} - } - } } } } -/// The cold-window rejection floor for a freshly elected leader: the durable watermark -/// (`udb_lease.durable_version`) at election time. -/// -/// Reasoning (do NOT change this back to `max(durable, seq_high)`): -/// -/// A new leader starts with an empty conflict tracker, so it cannot detect a read-write conflict -/// against any committed write it does not already know about. The writes it is missing are exactly -/// the previous leader's winners, and every winner's write is applied to `kv` AND its -/// `commit_version` folded into `durable_version` in the SAME apply transaction. So every missing -/// write has `commit_version <= durable_version`. A committing transaction `T` is therefore safe -/// iff `T.read_version >= durable_version`: every write above its read_version was committed by THIS -/// leader and is in the tracker. Only `T.read_version < durable_version` can race a missing winner, -/// so that is the exact set the cold window must reject. -/// -/// `udb_version_seq.last_value` (the sequence high-water) is NOT a valid floor. Every drained -/// request consumes a `nextval` BEFORE the conflict check, including conflicts and cold rejects, so -/// the sequence races far ahead of `durable_version` with versions that never produced any write. -/// Using `max(durable, seq_high)` rejects essentially every commit for the whole cold window -/// (followers read at `durable_version`, which is always `< seq_high`), turning each failover into -/// a 5s mass-reject storm. The gap `(durable_version, seq_high]` holds only thrown-away loser -/// versions, so nothing in it is a missing write to guard against. -/// -/// Version ASSIGNMENT is unaffected: commit versions still come from `nextval('udb_version_seq')` -/// in `drain_batch`, which is always above the sequence high-water, so uniqueness and monotonicity -/// across failover are preserved independently of this floor. +/// Collect up to [`DRAIN_BATCH_SIZE`] jobs from the queue, blocking for the first one and then +/// draining any immediately-available followers without waiting. Returns an empty batch only when the +/// queue is closed and drained. +async fn collect_batch(rx: &mut mpsc::Receiver) -> Vec { + let mut batch = Vec::with_capacity(DRAIN_BATCH_SIZE); + rx.recv_many(&mut batch, DRAIN_BATCH_SIZE).await; + batch +} + +/// The cold-window rejection floor for a freshly elected leader: the durable watermark at election +/// time. A new leader's tracker is empty, so it cannot detect a conflict against a previous leader's +/// winner; every such winner has `commit_version <= durable_version` (applied and folded into +/// `durable_version` in one txn), so a commit is safe if `read_version >= durable_version`. async fn recovery_floor(shared: &Arc) -> Result { let durable = lease::current_durable_version(&shared.pool).await?; Ok(durable.max(0) as u64) } enum BatchOutcome { - Empty, Processed, LostLease, } -struct Reply { - channel: String, - /// The follower's reply payload, encoding the outcome so the waiter resolves without a status - /// SELECT: `":committed:"` or `":conflict"`. - payload: String, -} - async fn drain_batch( shared: &Arc, epoch: i64, tracker: &TransactionConflictTracker, recovery_version: u64, recovery_deadline: Instant, + mut jobs: Vec, ) -> Result { + let batch_start = Instant::now(); + let batch_len = jobs.len(); + let mut conn = shared .pool .get() @@ -351,128 +405,132 @@ async fn drain_batch( .await .context("failed to start drain batch txn")?; - // Claim a batch in id order and allocate every commit version in the SAME round-trip via an - // inline `nextval`, instead of a separate version-allocation query. FOR UPDATE SKIP LOCKED holds - // the rows for this txn so they are stamped terminal on COMMIT with no intermediate 'claimed' - // state to clean up. - // - // Postgres does NOT guarantee `nextval` is evaluated in output (id) order, so the per-row `cv` - // here may not be monotonic in id. The versions are collected, sorted, and re-assigned to rows in - // id order below, exactly as the separate-query path did, because versionstamp monotonicity with - // commit order is load-bearing (epoxy changelog catch-up + depot PITR). - let rows = txn - .query( - "SELECT id, read_version, payload, reply_channel, - nextval('udb_version_seq') AS cv - FROM udb_commit_requests - WHERE status = 'pending' AND epoch = $1 - ORDER BY id - LIMIT $2 - FOR UPDATE SKIP LOCKED", - &[&epoch, &DRAIN_BATCH_SIZE], - ) - .await - .context("failed to claim commit batch")?; + // Build the failover dedup keys: a job whose (client_node_id, client_seq) is already recorded in + // udb_applied was committed by a prior leader; respond with the recorded version and do not + // re-apply. Single-node jobs carry no dedup key and never hit this path. + let mut dedup_nids: Vec> = Vec::new(); + let mut dedup_seqs: Vec = Vec::new(); + for job in &jobs { + if let Some(key) = &job.dedup_key { + dedup_nids.push(key.client_node_id.clone()); + dedup_seqs.push(key.client_seq); + } + } - if rows.is_empty() { - txn.rollback().await.ok(); - return Ok(BatchOutcome::Empty); + // Pipeline the dedup pre-check and commit-version allocation on the same connection. Versions are + // allocated for every job rather than only to-resolve jobs, so the version count no longer depends + // on the dedup result and the two queries have no data dependency; tokio-postgres pipelines them + // into a single round-trip. A dedup hit wastes its allocated version, but sequence gaps are already + // expected (conflict losers and rolled-back batches burn versions too) and do not affect + // versionstamp monotonicity. + let job_count = jobs.len() as i64; + let dedup_fut = async { + if dedup_nids.is_empty() { + return anyhow::Ok(HashMap::<(Vec, i64), i64>::new()); + } + let applied = txn + .query( + "SELECT a.client_node_id, a.client_seq, a.commit_version + FROM udb_applied a + JOIN unnest($1::bytea[], $2::bigint[]) AS q(nid, seq) + ON a.client_node_id = q.nid AND a.client_seq = q.seq", + &[&dedup_nids, &dedup_seqs], + ) + .await + .context("failed dedup pre-check")? + .into_iter() + .map(|row| { + ( + (row.get::<_, Vec>(0), row.get::<_, i64>(1)), + row.get::<_, i64>(2), + ) + }) + .collect(); + anyhow::Ok(applied) + }; + let versions_fut = async { + let versions = txn + .query( + "SELECT nextval('udb_version_seq') FROM generate_series(1, $1::bigint)", + &[&job_count], + ) + .await + .context("failed to allocate commit versions")? + .into_iter() + .map(|row| row.get::<_, i64>(0)) + .collect::>(); + anyhow::Ok(versions) + }; + let (applied, mut versions) = tokio::try_join!(dedup_fut, versions_fut)?; + + // Postgres does not guarantee nextval is evaluated in row order, so the versions are sorted and + // assigned to to-resolve jobs in arrival order to keep versionstamps monotonic with commit order + // (load-bearing for epoxy changelog catch-up + depot PITR). + versions.sort_unstable(); + + // Classify each job: a dedup hit resolves immediately; everything else needs version assignment + // and conflict resolution. + let mut outcomes: Vec> = vec![None; jobs.len()]; + let mut resolve_indices: Vec = Vec::with_capacity(jobs.len()); + for (i, job) in jobs.iter().enumerate() { + if let Some(key) = &job.dedup_key { + if let Some(&cv) = applied.get(&(key.client_node_id.clone(), key.client_seq)) { + outcomes[i] = Some(CommitOutcome::Committed { commit_version: cv }); + continue; + } + } + resolve_indices.push(i); } - let batch_start = Instant::now(); let cold_window = Instant::now() < recovery_deadline; - let batch_len = rows.len(); + let mut winners: Vec = Vec::new(); + let mut winner_dedup_nids: Vec> = Vec::new(); + let mut winner_dedup_seqs: Vec = Vec::new(); + let mut winner_dedup_cvs: Vec = Vec::new(); let mut max_winner_cv: i64 = 0; - let mut replies = Vec::with_capacity(batch_len); let mut committed_count = 0u32; let mut conflict_count = 0u32; let mut cold_reject_count = 0u32; - // Re-assign the inline-allocated versions to rows in id order (winners and losers alike; losers' - // versions are harmlessly skipped) so versionstamps stay monotonic with commit order. The sort - // keeps assignment monotonic regardless of how Postgres ordered the per-row `nextval` evaluation. - let mut versions: Vec = rows.iter().map(|row| row.get::<_, i64>(4)).collect(); - versions.sort_unstable(); + for (slot, &i) in resolve_indices.iter().enumerate() { + let commit_version = versions[slot]; + let job = &mut jobs[i]; + let start_version = job.read_version; + let conflict_ranges = std::mem::take(&mut job.conflict_ranges); - // Resolve every request in memory in id order. Winners are collected with their version and - // operations for the fold; the bulk status stamp is built for all rows at once. - let mut winners: Vec = Vec::new(); - let mut stamp_ids: Vec = Vec::with_capacity(batch_len); - let mut stamp_statuses: Vec<&str> = Vec::with_capacity(batch_len); - let mut stamp_versions: Vec> = Vec::with_capacity(batch_len); - - // Sub-phase timers to decompose batch_ms and confirm whether the service time is constant (pure - // M/M/1 queueing) or itself inflates under load. - let mut decode_dur = Duration::ZERO; - let mut conflict_dur = Duration::ZERO; - - for (i, row) in rows.iter().enumerate() { - let id: i64 = row.get(0); - let read_version: i64 = row.get(1); - let payload: Vec = row.get(2); - let reply_channel: String = row.get(3); - let commit_version = versions[i]; - - let decode_start = Instant::now(); - let decoded = super::codec::decode_commit_request(&payload) - .context("failed to decode commit payload")?; - decode_dur += decode_start.elapsed(); - - let start_version = read_version.max(0) as u64; - - // Cold-window guard: a commit whose read_version predates the recovery floor cannot be - // safely resolved against this leader's empty window. Reject it as retryable. let cold_rejected = cold_window && start_version < recovery_version; let conflicted = if cold_rejected { cold_reject_count += 1; true } else { - let conflict_start = Instant::now(); - let res = tracker - .check_and_insert( - start_version, - commit_version.max(0) as u64, - decoded.conflict_ranges, - ) - .await; - conflict_dur += conflict_start.elapsed(); - res + tracker + .check_and_insert(start_version, commit_version.max(0) as u64, conflict_ranges) + .await }; - stamp_ids.push(id); if conflicted { if !cold_rejected { conflict_count += 1; } - stamp_statuses.push("conflict"); - stamp_versions.push(None); + outcomes[i] = Some(CommitOutcome::Conflict); } else { committed_count += 1; - stamp_statuses.push("committed"); - stamp_versions.push(Some(commit_version)); + outcomes[i] = Some(CommitOutcome::Committed { commit_version }); max_winner_cv = max_winner_cv.max(commit_version); + if let Some(key) = &job.dedup_key { + winner_dedup_nids.push(key.client_node_id.clone()); + winner_dedup_seqs.push(key.client_seq); + winner_dedup_cvs.push(commit_version); + } winners.push(apply::Winner { commit_version: commit_version.max(0) as u64, - operations: decoded.operations, + operations: std::mem::take(&mut job.operations), }); } - - let reply_payload = if conflicted { - format!("{id}:conflict") - } else { - format!("{id}:committed:{commit_version}") - }; - replies.push(Reply { - channel: reply_channel, - payload: reply_payload, - }); } - let t_resolved = Instant::now(); - - // Bulk-read the pre-batch value of every key a winner's atomic op reads in one query, then fold - // all winners into a single materialized write-set in memory. This collapses the per-row apply - // round-trips to a fixed count independent of batch size. + // Bulk-read the pre-batch value of every key a winner's atomic op reads, then fold all winners + // into one materialized write-set in memory. let atomic_keys = apply::atomic_read_keys(&winners); let base = if atomic_keys.is_empty() { HashMap::new() @@ -488,26 +546,19 @@ async fn drain_batch( .collect() }; - let t_read = Instant::now(); - let apply::WriteSet { upserts, point_deletes, range_deletes, } = apply::fold_winners(winners, &base).context("failed to fold batch winners")?; - let t_fold = Instant::now(); - let (upsert_keys, upsert_values): (Vec>, Vec>) = upserts.into_iter().unzip(); let (range_begins, range_ends): (Vec>, Vec>) = range_deletes.into_iter().unzip(); - // Range deletes run in their OWN statement BEFORE the apply CTE. Postgres data-modifying CTE - // sub-statements all observe the same snapshot and never see each other's effects, so a range - // delete and an in-range upsert in one CTE would have unspecified results. `range_deletes` are - // ranges (not materialized per-key) and can overlap an upsert key (a key inside a cleared range - // that is also re-set), so the range clear must commit its effect first; the upsert then - // re-inserts the key. Collapsed from a per-range loop into one statement over a pair of arrays. + // Range deletes run in their own statement before the apply CTE: a range delete and an in-range + // upsert in one CTE would have unspecified ordering, so the clear must commit its effect first and + // the upsert then re-inserts the key. if !range_begins.is_empty() { txn.execute( "DELETE FROM kv USING unnest($1::bytea[], $2::bytea[]) AS r(b, e) @@ -518,12 +569,9 @@ async fn drain_batch( .context("failed to clear ranges")?; } - // Apply the rest of the batch in one CTE: point deletes, the kv upsert, the terminal status - // stamp, and the epoch-fenced watermark advance. This is safe to fold because the write sets are - // disjoint: `apply::WriteSet` guarantees each key appears at most once across `upserts` and - // `point_deletes` (see apply.rs), and the three tables (`kv`, `udb_commit_requests`, `udb_lease`) - // are independent. The watermark UPDATE is fenced on our epoch: a zombie old leader whose epoch - // was bumped sees zero rows returned and must step down before any of its writes become visible. + // Apply the rest of the batch in one CTE: point deletes, the kv upsert, the dedup records for + // multi-node winners, and the epoch-fenced watermark advance. A zombie old leader whose epoch was + // bumped sees zero rows from the lease UPDATE and steps down before any write becomes visible. let new_durable: i64 = match txn .query_opt( "WITH pdel AS ( @@ -532,11 +580,10 @@ async fn drain_batch( INSERT INTO kv (key, value) SELECT * FROM unnest($2::bytea[], $3::bytea[]) ON CONFLICT (key) DO UPDATE SET value = excluded.value - ), stamp AS ( - UPDATE udb_commit_requests AS r - SET status = b.status, commit_version = b.cv - FROM unnest($4::bigint[], $5::text[], $6::bigint[]) AS b(id, status, cv) - WHERE r.id = b.id + ), applied AS ( + INSERT INTO udb_applied (client_node_id, client_seq, commit_version) + SELECT * FROM unnest($4::bytea[], $5::bigint[], $6::bigint[]) + ON CONFLICT (client_node_id, client_seq) DO NOTHING ) UPDATE udb_lease SET durable_version = GREATEST(durable_version, $7) @@ -546,9 +593,9 @@ async fn drain_batch( &point_deletes, &upsert_keys, &upsert_values, - &stamp_ids, - &stamp_statuses, - &stamp_versions, + &winner_dedup_nids, + &winner_dedup_seqs, + &winner_dedup_cvs, &max_winner_cv, &LEASE_ID, &epoch, @@ -564,23 +611,39 @@ async fn drain_batch( } }; - let t_applied = Instant::now(); - txn.commit().await.context("failed to commit drain batch")?; - let t_committed = Instant::now(); - - // Watermark advances strictly after the apply txn is durably committed and visible, so a + // The watermark advances strictly after the apply txn is durably committed and visible, so a // reader handed this read_version can never miss a write with commit_version <= read_version. shared.advance_durable_version(new_durable); - notify_after_commit(&conn, new_durable, &replies).await; - - let t_notified = Instant::now(); + if let Transport::MultiNode(nats) = &shared.transport { + match super::codec::encode_watermark(new_durable) { + Ok(payload) => { + if let Err(err) = nats + .client + .publish(nats.subjects.watermark(), payload.into()) + .await + { + tracing::debug!(?err, "failed to publish udb watermark"); + } + } + Err(err) => tracing::error!(?err, "failed to encode udb watermark"), + } + } - let tracker_len = tracker.len().await; + // Respond to every job (dedup hits, winners, losers). Responses are independent per job, so fan the + // replies out concurrently instead of awaiting each publish in series. + futures_util::stream::iter(jobs.into_iter().enumerate()) + .for_each_concurrent(None, |(i, job)| { + let outcome = outcomes[i].expect("every job must be resolved"); + async move { + job.responder.respond(outcome).await; + } + }) + .await; - tracing::info!( + tracing::debug!( epoch, batch_len, committed = committed_count, @@ -589,49 +652,8 @@ async fn drain_batch( cold_window, new_durable, batch_ms = batch_start.elapsed().as_millis() as u64, - // Sub-phase decomposition of batch_ms (micros) to separate constant service time from - // load-dependent service inflation. - tracker_len, - decode_us = decode_dur.as_micros() as u64, - conflict_us = conflict_dur.as_micros() as u64, - resolve_us = (t_resolved - batch_start).as_micros() as u64, - read_us = (t_read - t_resolved).as_micros() as u64, - fold_us = (t_fold - t_read).as_micros() as u64, - apply_us = (t_applied - t_fold).as_micros() as u64, - commit_us = (t_committed - t_applied).as_micros() as u64, - notify_us = (t_notified - t_committed).as_micros() as u64, "udb leader processed commit batch" ); Ok(BatchOutcome::Processed) } - -/// Wake watermark listeners and the followers waiting on each processed request. Best-effort: a -/// missed NOTIFY is covered by the follower's polling backstop and the watermark refresh timer. -async fn notify_after_commit( - conn: &deadpool_postgres::Client, - new_durable: i64, - replies: &[Reply], -) { - if let Err(err) = conn - .execute( - "SELECT pg_notify($1, $2)", - &[&WATERMARK_CHANNEL, &new_durable.to_string()], - ) - .await - { - tracing::debug!(?err, "failed to notify watermark"); - } - - let channels: Vec<&str> = replies.iter().map(|r| r.channel.as_str()).collect(); - let payloads: Vec<&str> = replies.iter().map(|r| r.payload.as_str()).collect(); - if let Err(err) = conn - .execute( - "SELECT pg_notify(c, p) FROM unnest($1::text[], $2::text[]) AS t(c, p)", - &[&channels, &payloads], - ) - .await - { - tracing::debug!(?err, "failed to notify commit replies"); - } -} diff --git a/engine/packages/universaldb/src/driver/postgres/shared.rs b/engine/packages/universaldb/src/driver/postgres/shared.rs index b26f14a361..5bb98fcbf2 100644 --- a/engine/packages/universaldb/src/driver/postgres/shared.rs +++ b/engine/packages/universaldb/src/driver/postgres/shared.rs @@ -1,87 +1,90 @@ use std::{ sync::{ Arc, - atomic::{AtomicI64, Ordering}, + atomic::{AtomicI64, AtomicU64, Ordering}, }, time::Duration, }; use deadpool_postgres::Pool; +use futures_util::StreamExt; use tokio::sync::{Notify, watch}; -use super::listener::PgListener; +use super::transport::Transport; /// The singleton row id of `udb_lease`. pub const LEASE_ID: i32 = 1; -/// How often the follower refreshes its cached lease row (epoch, leader channel, watermark) as a -/// backstop to the `udb_watermark` NOTIFY. A stale-but-older watermark only widens the conflict -/// window, so this can be loose. +/// How often a node refreshes its cached lease row (epoch, leader id, watermark) as a backstop to the +/// watermark broadcast. A stale-but-older watermark only widens the conflict window, so this can be +/// loose. const LEASE_REFRESH_INTERVAL: Duration = Duration::from_millis(500); -/// Channel a follower NOTIFYs (and the leader LISTENs) to wake the leader's drain loop. -pub fn commit_channel(node_id: &str) -> String { - format!("udb_commit_{node_id}") -} - -/// Channel the leader NOTIFYs (and a follower LISTENs) to deliver a commit result. -pub fn reply_channel(node_id: &str) -> String { - format!("udb_reply_{node_id}") -} - -/// Channel the leader NOTIFYs on every watermark advance; all nodes LISTEN. -pub const WATERMARK_CHANNEL: &str = "udb_watermark"; - -/// Channel a departing leader NOTIFYs after releasing its lease so a standby candidate elects -/// immediately instead of waiting out `ELECTION_RETRY`. All non-leader candidates LISTEN. -pub const ELECTION_CHANNEL: &str = "udb_election"; - /// Cached view of the current leader lease, as seen by a follower. #[derive(Clone, Debug)] pub struct LeaseInfo { pub epoch: i64, - /// Node id of the current leader, used to build its commit channel. + /// Node id of the current leader, used to build its commit subject. pub leader_addr: String, } -/// Process-wide state shared by the follower transaction tasks and the leader resolver. Every node -/// is both a follower (it submits its own commits) and a candidate leader. +/// Process-wide state shared by the follower transaction tasks and the leader resolver. Every node is +/// both a follower (it submits its own commits) and, in multi-node mode, a candidate leader. pub struct PostgresShared { pub pool: Pool, - /// Unique per-process id used to name this node's NOTIFY channels. + /// Unique per-process id. Names this node's commit subject and is the dedup `client_node_id`. pub node_id: String, - pub listener: PgListener, + /// How follower commits reach the leader (in-process channel or NATS). + pub transport: Transport, /// Highest durable commit version (`udb_lease.durable_version`); the follower read version. durable_version: AtomicI64, /// Pinged whenever `durable_version` advances. watermark_notify: Notify, + /// Per-process monotonic commit sequence, the dedup `client_seq`. + commit_seq: AtomicU64, lease_tx: watch::Sender>, lease_rx: watch::Receiver>, } impl PostgresShared { - pub fn new(pool: Pool, node_id: String, listener: PgListener) -> Arc { + pub fn new(pool: Pool, node_id: String, transport: Transport) -> Arc { let (lease_tx, lease_rx) = watch::channel(None); let shared = Arc::new(Self { pool, node_id, - listener, + transport, durable_version: AtomicI64::new(0), watermark_notify: Notify::new(), + commit_seq: AtomicU64::new(0), lease_tx, lease_rx, }); - tokio::spawn(Self::cache_refresh_task(shared.clone())); + // Single-node advances `durable_version` and the lease cache in-process, so it needs no + // cross-process refresh. Multi-node refreshes from the NATS watermark broadcast and a lease + // row poll. + if matches!(shared.transport, Transport::MultiNode(_)) { + tokio::spawn(Self::cache_refresh_task(shared.clone())); + } shared } + /// Whether this driver is running in multi-node mode. + pub fn is_multi_node(&self) -> bool { + matches!(self.transport, Transport::MultiNode(_)) + } + /// The cached follower read version (`durable_version`). pub fn read_version(&self) -> i64 { self.durable_version.load(Ordering::SeqCst) } + /// Allocate the next per-process commit sequence for the failover dedup key. + pub fn next_commit_seq(&self) -> i64 { + self.commit_seq.fetch_add(1, Ordering::Relaxed) as i64 + } + /// Advance the cached watermark monotonically and wake any waiters. pub fn advance_durable_version(&self, version: i64) { let prev = self.durable_version.fetch_max(version, Ordering::SeqCst); @@ -104,7 +107,7 @@ impl PostgresShared { .map(|prev| prev.epoch != lease.epoch || prev.leader_addr != lease.leader_addr) .unwrap_or(true); if changed { - tracing::info!( + tracing::debug!( epoch = lease.epoch, leader_addr = %lease.leader_addr, self_node = %self.node_id, @@ -115,27 +118,41 @@ impl PostgresShared { let _ = self.lease_tx.send(Some(lease)); } - /// Background task: keep `durable_version` and the cached lease fresh via the `udb_watermark` - /// NOTIFY plus a periodic poll of `udb_lease`. + /// Background task (multi-node only): keep `durable_version` and the cached lease fresh via the + /// NATS watermark broadcast plus a periodic poll of `udb_lease`. async fn cache_refresh_task(shared: Arc) { - let mut watermark_rx = shared.listener.listen(WATERMARK_CHANNEL).await; + let Transport::MultiNode(nats) = &shared.transport else { + return; + }; + + let mut watermark_sub = match nats.client.subscribe(nats.subjects.watermark()).await { + Ok(sub) => sub, + Err(err) => { + tracing::error!( + ?err, + "failed to subscribe to udb watermark; relying on lease poll" + ); + return shared.lease_poll_only().await; + } + }; + let mut interval = tokio::time::interval(LEASE_REFRESH_INTERVAL); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { tokio::select! { - notify = watermark_rx.recv() => { - match notify { - Ok(payload) => { - if let Ok(version) = payload.parse::() { - shared.advance_durable_version(version); + msg = watermark_sub.next() => { + match msg { + Some(msg) => { + match super::codec::decode_watermark(&msg.payload) { + Ok(version) => shared.advance_durable_version(version), + Err(err) => { + tracing::debug!(?err, "failed to decode udb watermark") + } } } - Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} - Err(tokio::sync::broadcast::error::RecvError::Closed) => { - // Re-subscribe; the listener recreates the channel on reconnect. - watermark_rx = shared.listener.listen(WATERMARK_CHANNEL).await; - } + // The subscription ended (client closed). Fall back to lease polling only. + None => return shared.lease_poll_only().await, } } _ = interval.tick() => { @@ -145,6 +162,16 @@ impl PostgresShared { } } + /// Degraded refresh path: poll the lease row when the watermark subscription is unavailable. + async fn lease_poll_only(self: Arc) { + let mut interval = tokio::time::interval(LEASE_REFRESH_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + interval.tick().await; + self.refresh_lease_row().await; + } + } + async fn refresh_lease_row(&self) { let conn = match self.pool.get().await { Ok(conn) => conn, diff --git a/engine/packages/universaldb/src/driver/postgres/transport.rs b/engine/packages/universaldb/src/driver/postgres/transport.rs new file mode 100644 index 0000000000..8e1b90ecda --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/transport.rs @@ -0,0 +1,87 @@ +use tokio::sync::{mpsc, oneshot}; + +use crate::{options::ConflictRangeType, tx_ops::Operation}; + +use super::{codec, nats::NatsTransport}; + +/// Bound on the in-process commit queue feeding the leader drain loop. The leader drains in large +/// batches, so this only needs to absorb a brief burst between batches. +pub const COMMIT_QUEUE_BOUND: usize = 4096; + +/// How the follower commit path reaches the leader resolver. +/// +/// Single-node keeps an in-process channel directly into the leader drain loop (this node is always +/// the leader). Multi-node sends the commit to the elected leader over NATS request/reply; the reply +/// carries the commit result. +pub enum Transport { + SingleNode { + /// Sender into the leader drain loop's job queue. The matching receiver is owned by the + /// resolver task spawned at startup. + commit_tx: mpsc::Sender, + }, + MultiNode(NatsTransport), +} + +/// The outcome of resolving a single commit. +#[derive(Clone, Copy, Debug)] +pub enum CommitOutcome { + /// The commit won resolution and its writes were durably applied at `commit_version`. + Committed { commit_version: i64 }, + /// The commit lost resolution (read-write conflict or cold-window reject). Maps to the + /// retryable `DatabaseError::NotCommitted` on the follower. + Conflict, +} + +/// How the leader delivers a commit result back to the waiting follower. +pub enum Responder { + /// Single-node: resolve the follower's `oneshot` directly. + Local(oneshot::Sender), + /// Multi-node: publish the encoded outcome to the NATS request's reply inbox. + Nats { + client: async_nats::Client, + reply: async_nats::Subject, + }, +} + +impl Responder { + /// Deliver the outcome to the follower. Best-effort: a follower that already gave up (oneshot + /// dropped, or no NATS responder) is covered by its own retry path. + pub async fn respond(self, outcome: CommitOutcome) { + match self { + Responder::Local(tx) => { + let _ = tx.send(outcome); + } + Responder::Nats { client, reply } => { + let payload = match codec::encode_commit_reply(outcome) { + Ok(payload) => payload, + Err(err) => { + tracing::error!(?err, "failed to encode udb commit reply"); + return; + } + }; + if let Err(err) = client.publish(reply, payload.into()).await { + tracing::debug!(?err, "failed to publish udb commit reply"); + } + } + } + } +} + +/// A single commit handed to the leader drain loop, transport-agnostic. Single-node jobs are built by +/// the follower commit path; multi-node jobs are built by the NATS commit subscriber from a decoded +/// request. +pub struct CommitJob { + pub read_version: u64, + pub conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, + pub operations: Vec, + /// Failover dedup key, present only in multi-node (single-node has no lost-reply window). + pub dedup_key: Option, + pub responder: Responder, +} + +/// Identifies a single logical commit across follower resends so the leader applies it exactly once. +#[derive(Clone, Debug)] +pub struct DedupKey { + pub client_node_id: Vec, + pub client_seq: i64, +} diff --git a/engine/packages/universaldb/src/driver/rocksdb/database.rs b/engine/packages/universaldb/src/driver/rocksdb/database.rs index ba186d872e..718ced8836 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/database.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/database.rs @@ -48,7 +48,7 @@ impl RocksDbDatabaseDriver { Ok(RocksDbDatabaseDriver { db: Arc::new(db), - max_retries: AtomicI32::new(100), + max_retries: AtomicI32::new(10), txn_conflict_tracker: TransactionConflictTracker::new(), }) } diff --git a/engine/packages/universaldb/tests/failover.rs b/engine/packages/universaldb/tests/failover.rs index 5fb394453c..6293161760 100644 --- a/engine/packages/universaldb/tests/failover.rs +++ b/engine/packages/universaldb/tests/failover.rs @@ -1,21 +1,45 @@ use std::{sync::Arc, time::Duration}; -use rivet_test_deps_docker::TestDatabase; +use rivet_test_deps_docker::{TestDatabase, TestPubSub}; use tokio_postgres::NoTls; -use universaldb::{Database, utils::IsolationLevel::*}; +use universaldb::{Database, driver::postgres::NatsConfig, utils::IsolationLevel::*}; use uuid::Uuid; const ALPHA_KEY: &[u8] = b"failover/alpha"; const BETA_KEY: &[u8] = b"failover/beta"; -/// Build a fresh Postgres-backed `Database`. Each call spins up an independent driver (its own pool, -/// node id, listener, and resolver), so two of them against one Postgres model two engine nodes. -async fn make_db(connection_string: &str) -> Database { - let driver = universaldb::driver::PostgresDatabaseDriver::new_with_config( - universaldb::driver::postgres::PostgresConfig::new(connection_string.to_string()), - ) - .await - .unwrap(); +/// Boot a NATS container and return the multi-node UniversalDB NATS config plus the docker handle +/// (kept alive for the test duration). Leader failover is a multi-node scenario, so the drivers must +/// share a NATS deployment for follower-to-leader commit transport and watermark/election broadcast. +async fn setup_nats() -> (NatsConfig, rivet_test_deps_docker::DockerRunConfig) { + let (pubsub_config, docker_config) = TestPubSub::Nats.config(Uuid::new_v4(), 1).await.unwrap(); + let mut docker_config = docker_config.unwrap(); + docker_config.start().await.unwrap(); + tokio::time::sleep(Duration::from_secs(1)).await; + + let rivet_config::config::PubSub::Nats(nats) = pubsub_config else { + unreachable!(); + }; + let config = NatsConfig { + addresses: nats.addresses.clone(), + username: nats.username.clone(), + password: nats.password.as_ref().map(|p| p.read().clone()), + client_capacity: nats.client_capacity, + subscription_capacity: nats.subscription_capacity, + }; + (config, docker_config) +} + +/// Build a fresh multi-node Postgres-backed `Database`. Each call spins up an independent driver (its +/// own pool, node id, NATS client, and resolver), so two of them against one Postgres + one NATS model +/// two engine nodes. +async fn make_db(connection_string: &str, nats: &NatsConfig) -> Database { + let mut config = + universaldb::driver::postgres::PostgresConfig::new(connection_string.to_string()); + config.nats = Some(nats.clone()); + let driver = universaldb::driver::PostgresDatabaseDriver::new_with_config(config) + .await + .unwrap(); Database::new(Arc::new(driver)) } @@ -124,16 +148,18 @@ async fn test_postgres_leader_failover() { }; let connection_string = postgres_config.url.read().clone(); + let (nats_config, _nats_docker) = setup_nats().await; + let raw = connect_raw(&connection_string).await; // Node 1 comes up first and deterministically wins the first election (epoch 1). - let db1 = make_db(&connection_string).await; + let db1 = make_db(&connection_string, &nats_config).await; let lease1 = wait_for_lease(&raw, Duration::from_secs(15), |l| l.epoch == 1).await; let leader1_addr = lease1.leader_addr.clone(); // Node 2 joins while node 1 holds a valid lease, so it loses the election and runs as a // follower. - let db2 = make_db(&connection_string).await; + let db2 = make_db(&connection_string, &nats_config).await; // Leader (node 1) commits data. The version sequence and watermark advance. write_key(&db1, ALPHA_KEY, b"1").await; @@ -243,13 +269,15 @@ async fn test_postgres_graceful_handoff() { }; let connection_string = postgres_config.url.read().clone(); + let (nats_config, _nats_docker) = setup_nats().await; + let raw = connect_raw(&connection_string).await; // Node 1 wins the first election; node 2 joins as a follower. - let db1 = make_db(&connection_string).await; + let db1 = make_db(&connection_string, &nats_config).await; let lease1 = wait_for_lease(&raw, Duration::from_secs(15), |l| l.epoch == 1).await; let leader1_addr = lease1.leader_addr.clone(); - let db2 = make_db(&connection_string).await; + let db2 = make_db(&connection_string, &nats_config).await; write_key(&db1, ALPHA_KEY, b"1").await; let lease_before = read_lease(&raw).await.unwrap(); diff --git a/engine/packages/universalpubsub/Cargo.toml b/engine/packages/universalpubsub/Cargo.toml index 7043f44fac..c9c9dfafe1 100644 --- a/engine/packages/universalpubsub/Cargo.toml +++ b/engine/packages/universalpubsub/Cargo.toml @@ -10,23 +10,17 @@ edition.workspace = true anyhow.workspace = true async-nats.workspace = true async-trait.workspace = true -base64.workspace = true -deadpool-postgres.workspace = true futures-util.workspace = true lazy_static.workspace = true rand.workspace = true rivet-error.workspace = true rivet-metrics.workspace = true rivet-perf.workspace = true -rivet-postgres-util.workspace = true rivet-ups-protocol.workspace = true rivet-util.workspace = true scc.workspace = true serde_json.workspace = true serde.workspace = true -sha2.workspace = true -tokio-postgres-rustls.workspace = true -tokio-postgres.workspace = true tokio-util.workspace = true tokio.workspace = true tracing.workspace = true diff --git a/engine/packages/universalpubsub/benches/simple.rs b/engine/packages/universalpubsub/benches/simple.rs index 7edc5ff926..efb5624d2e 100644 --- a/engine/packages/universalpubsub/benches/simple.rs +++ b/engine/packages/universalpubsub/benches/simple.rs @@ -6,7 +6,7 @@ use std::time::{Duration, Instant}; use anyhow::*; use futures_util::future::join_all; -use rivet_test_deps_docker::{TestDatabase, TestPubSub}; +use rivet_test_deps_docker::TestPubSub; use std::future::Future; use tabled::{builder::Builder, settings::Style}; use universalpubsub::{NextOutput, PubSub, PublishOpts}; @@ -654,98 +654,6 @@ async fn setup_nats_single_mem() -> Result<(PubSub, PubSub)> { Ok((pubsub.clone(), pubsub)) } -async fn setup_pg_pair() -> Result<(PubSub, PubSub)> { - let test_id = Uuid::new_v4(); - let (db_config, mut docker) = TestDatabase::Postgres.config(test_id, 1).await?; - if let Some(ref mut d) = docker { - d.start().await?; - } - tokio::time::sleep(Duration::from_secs(5)).await; - let rivet_config::config::Database::Postgres(pg) = db_config else { - unreachable!() - }; - let url = pg.url.read().clone(); - let driver_pub = universalpubsub::driver::postgres::PostgresDriver::connect( - url.clone(), - false, - None, - None, - None, - ) - .await?; - let driver_sub = - universalpubsub::driver::postgres::PostgresDriver::connect(url, false, None, None, None) - .await?; - Ok(( - PubSub::new_with_memory_optimization(Arc::new(driver_pub), false), - PubSub::new_with_memory_optimization(Arc::new(driver_sub), false), - )) -} - -async fn setup_pg_single() -> Result<(PubSub, PubSub)> { - let test_id = Uuid::new_v4(); - let (db_config, mut docker) = TestDatabase::Postgres.config(test_id, 1).await?; - if let Some(ref mut d) = docker { - d.start().await?; - } - tokio::time::sleep(Duration::from_secs(5)).await; - let rivet_config::config::Database::Postgres(pg) = db_config else { - unreachable!() - }; - let url = pg.url.read().clone(); - let driver = - universalpubsub::driver::postgres::PostgresDriver::connect(url, false, None, None, None) - .await?; - let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false); - Ok((pubsub.clone(), pubsub)) -} - -async fn setup_pg_pair_mem() -> Result<(PubSub, PubSub)> { - let test_id = Uuid::new_v4(); - let (db_config, mut docker) = TestDatabase::Postgres.config(test_id, 1).await?; - if let Some(ref mut d) = docker { - d.start().await?; - } - tokio::time::sleep(Duration::from_secs(5)).await; - let rivet_config::config::Database::Postgres(pg) = db_config else { - unreachable!() - }; - let url = pg.url.read().clone(); - let driver_pub = universalpubsub::driver::postgres::PostgresDriver::connect( - url.clone(), - true, - None, - None, - None, - ) - .await?; - let driver_sub = - universalpubsub::driver::postgres::PostgresDriver::connect(url, true, None, None, None) - .await?; - Ok(( - PubSub::new_with_memory_optimization(Arc::new(driver_pub), true), - PubSub::new_with_memory_optimization(Arc::new(driver_sub), true), - )) -} - -async fn setup_pg_single_mem() -> Result<(PubSub, PubSub)> { - let test_id = Uuid::new_v4(); - let (db_config, mut docker) = TestDatabase::Postgres.config(test_id, 1).await?; - if let Some(ref mut d) = docker { - d.start().await?; - } - tokio::time::sleep(Duration::from_secs(5)).await; - let rivet_config::config::Database::Postgres(pg) = db_config else { - unreachable!() - }; - let url = pg.url.read().clone(); - let driver = - universalpubsub::driver::postgres::PostgresDriver::connect(url, true, None, None, None) - .await?; - let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), true); - Ok((pubsub.clone(), pubsub)) -} - async fn setup_mem_pair() -> Result<(PubSub, PubSub)> { let test_id = Uuid::new_v4(); let (pubsub_config, _docker) = TestPubSub::Memory.config(test_id, 1).await?; @@ -891,16 +799,6 @@ async fn main() -> Result<()> { let results = run_benches("nats-mem", publisher.clone(), subscriber.clone(), iters).await?; all_results.insert("nats-mem".to_string(), results); - // Postgres (no memory optimization) - let (publisher, subscriber) = setup_pg_pair().await?; - let results = run_benches("pg-nomem", publisher.clone(), subscriber.clone(), iters).await?; - all_results.insert("pg-nomem".to_string(), results); - - // Postgres (memory optimization) - let (publisher, subscriber) = setup_pg_pair_mem().await?; - let results = run_benches("pg-mem", publisher.clone(), subscriber.clone(), iters).await?; - all_results.insert("pg-mem".to_string(), results); - // NATS single connection (no memory optimization) let (publisher, subscriber) = setup_nats_single().await?; let results = run_benches( @@ -923,28 +821,6 @@ async fn main() -> Result<()> { .await?; all_results.insert("nats-single-mem".to_string(), results); - // Postgres single connection (no memory optimization) - let (publisher, subscriber) = setup_pg_single().await?; - let results = run_benches( - "pg-single-nomem", - publisher.clone(), - subscriber.clone(), - iters, - ) - .await?; - all_results.insert("pg-single-nomem".to_string(), results); - - // Postgres single connection (memory optimization) - let (publisher, subscriber) = setup_pg_single_mem().await?; - let results = run_benches( - "pg-single-mem", - publisher.clone(), - subscriber.clone(), - iters, - ) - .await?; - all_results.insert("pg-single-mem".to_string(), results); - // Print results table print_results_table(&all_results); diff --git a/engine/packages/universalpubsub/src/driver/mod.rs b/engine/packages/universalpubsub/src/driver/mod.rs index 1e79e4383b..b19a3aa85f 100644 --- a/engine/packages/universalpubsub/src/driver/mod.rs +++ b/engine/packages/universalpubsub/src/driver/mod.rs @@ -8,7 +8,6 @@ use crate::InboxSubject; pub mod memory; pub mod nats; -pub mod postgres; pub type PubSubDriverHandle = Arc; diff --git a/engine/packages/universalpubsub/src/driver/postgres/doorbell.rs b/engine/packages/universalpubsub/src/driver/postgres/doorbell.rs deleted file mode 100644 index f08ff7bbf3..0000000000 --- a/engine/packages/universalpubsub/src/driver/postgres/doorbell.rs +++ /dev/null @@ -1,133 +0,0 @@ -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::time::Duration; - -use deadpool_postgres::Pool; -use tokio::sync::Notify; -use tokio::time::Instant; - -/// Number of doorbell shards. A subject maps to a shard via `hash(subject_hash) % K`. -/// Subscribers LISTEN their subject's shard channel; publishers wake the local -/// doorbell task which NOTIFYs the shard. -pub const DOORBELL_SHARD_COUNT: usize = 32; - -/// Debounce window. Caps each (process, shard) NOTIFY rate at one per window, which -/// bounds how many backends are woken per shard over time. -const DOORBELL_WINDOW: Duration = Duration::from_millis(5); - -/// Returns the NOTIFY channel name for a doorbell shard. -pub fn shard_channel(shard: usize) -> String { - format!("ups_db_{shard}") -} - -/// Returns the doorbell shard for a subject hash. -pub fn shard_for(subject_hash: &str) -> usize { - use std::hash::{DefaultHasher, Hash, Hasher}; - let mut hasher = DefaultHasher::new(); - subject_hash.hash(&mut hasher); - (hasher.finish() as usize) % DOORBELL_SHARD_COUNT -} - -/// Coalesced, payload-free NOTIFY doorbell. -/// -/// Publishers call [`Doorbell::mark_dirty`] after committing a row. A single -/// per-process task drains dirty shards and emits at most one NOTIFY per shard per -/// debounce window using leading-edge fire plus a trailing-edge flush. The doorbell -/// is a latency optimization only. Correctness comes from the table plus the -/// subscriber poll backstop, so a dropped or failed NOTIFY only adds latency. -pub struct Doorbell { - dirty: [AtomicBool; DOORBELL_SHARD_COUNT], - notify: Notify, - pool: Arc, -} - -impl Doorbell { - pub fn new(pool: Arc) -> Arc { - let doorbell = Arc::new(Self { - dirty: std::array::from_fn(|_| AtomicBool::new(false)), - notify: Notify::new(), - pool, - }); - - let task_doorbell = doorbell.clone(); - tokio::spawn(async move { task_doorbell.run().await }); - - doorbell - } - - /// Marks a shard dirty and wakes the doorbell task. Never blocks. - pub fn mark_dirty(&self, shard: usize) { - self.dirty[shard].store(true, Ordering::Release); - self.notify.notify_one(); - } - - async fn run(self: Arc) { - // Per-shard timestamp of the last NOTIFY emitted by this process. - let mut last_notify: [Option; DOORBELL_SHARD_COUNT] = [None; DOORBELL_SHARD_COUNT]; - // Per-shard deadline for a pending trailing-edge NOTIFY, if any. - let mut trailing: [Option; DOORBELL_SHARD_COUNT] = [None; DOORBELL_SHARD_COUNT]; - - loop { - // Arm on the next pending trailing deadline so the trailing edge fires - // even with no further publishes. Wait on the notify permit otherwise. - let next_deadline = trailing.iter().filter_map(|x| *x).min(); - match next_deadline { - Some(deadline) => { - tokio::select! { - _ = self.notify.notified() => {} - _ = tokio::time::sleep_until(deadline) => {} - } - } - None => { - self.notify.notified().await; - } - } - - let now = Instant::now(); - for shard in 0..DOORBELL_SHARD_COUNT { - let is_dirty = self.dirty[shard].swap(false, Ordering::AcqRel); - if is_dirty { - match last_notify[shard] { - Some(last) if now.duration_since(last) < DOORBELL_WINDOW => { - // Within the window. Defer to a trailing-edge NOTIFY at - // window end so at most one NOTIFY fires per shard per W. - if trailing[shard].is_none() { - trailing[shard] = Some(last + DOORBELL_WINDOW); - } - } - _ => { - // Leading edge. Fire immediately for low idle latency. - self.notify_shard(shard).await; - last_notify[shard] = Some(now); - trailing[shard] = None; - } - } - } - - // Flush a trailing-edge NOTIFY whose window has elapsed. - if let Some(deadline) = trailing[shard] { - if now >= deadline { - self.notify_shard(shard).await; - last_notify[shard] = Some(now); - trailing[shard] = None; - } - } - } - } - } - - async fn notify_shard(&self, shard: usize) { - let channel = shard_channel(shard); - match self.pool.get().await { - Ok(conn) => { - // Payload-free doorbell. The payload lives in the table. - if let Err(err) = conn.execute("SELECT pg_notify($1, '')", &[&channel]).await { - tracing::warn!(?err, %channel, "failed to emit doorbell notify"); - } - } - Err(err) => { - tracing::warn!(?err, %channel, "failed to get connection for doorbell notify"); - } - } - } -} diff --git a/engine/packages/universalpubsub/src/driver/postgres/mod.rs b/engine/packages/universalpubsub/src/driver/postgres/mod.rs deleted file mode 100644 index 2a492fbccf..0000000000 --- a/engine/packages/universalpubsub/src/driver/postgres/mod.rs +++ /dev/null @@ -1,1050 +0,0 @@ -use anyhow::{Context, Result}; -use async_trait::async_trait; -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::throttle::Backoff; -use scc::HashMap; -use std::collections::VecDeque; -use std::hash::{DefaultHasher, Hash, Hasher}; -use std::path::PathBuf; -use std::sync::Arc; -use std::time::Duration; -use tokio::sync::{Mutex, broadcast}; -use tokio_postgres::AsyncMessage; -use tokio_postgres_rustls::MakeRustlsConnect; -use tracing::Instrument; -use uuid::Uuid; - -use crate::driver::{PubSubDriver, SubscriberDriver, SubscriberDriverHandle}; -use crate::metrics; -use crate::pubsub::DriverOutput; - -mod doorbell; - -use doorbell::{Doorbell, shard_channel, shard_for}; - -/// The transport is the table, not the NOTIFY payload, so there is no per-message -/// size cap from the 8000-byte NOTIFY limit. Match the NATS ceiling so chunking -/// behaves identically across drivers. -pub const POSTGRES_MAX_MESSAGE_SIZE: usize = 1024 * 1024; - -/// Poll backstop interval. Every subscriber reads its table on this interval -/// regardless of doorbell wakeups. This is the correctness floor that makes delivery -/// independent of any NOTIFY arriving. -const POLL_INTERVAL: Duration = Duration::from_secs(1); - -/// Idle-in-transaction timeout applied to the LISTEN connection. A wedged listener -/// holding a transaction open would otherwise fill the shared notify queue and fail -/// NOTIFY cluster-wide. Bounding it keeps a stuck listener degrading to added latency -/// rather than a cluster outage. -const LISTEN_IDLE_IN_TRANSACTION_TIMEOUT_MS: i64 = 30_000; - -/// How often this process refreshes its node liveness heartbeat. One heartbeat per -/// process keeps all of its subscriber registrations alive at once. -const NODE_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10); -/// How recent a node's heartbeat must be for its subscribers to count as live -/// responders. -const NODE_TTL_SECS: i64 = 30; - -/// How often to GC expired broadcast messages. -const MESSAGE_GC_INTERVAL: Duration = Duration::from_secs(5); -/// Max age before a broadcast message row is garbage collected. Must exceed the poll -/// interval plus the reconnect gap. A subscriber that falls behind this misses -/// messages, matching NATS-core at-most-once semantics for slow consumers. -const MESSAGE_MAX_AGE_SECS: i64 = 10; - -/// How often to GC dead nodes and the subscriber rows orphaned by them. -const REGISTRY_GC_INTERVAL: Duration = Duration::from_secs(30); - -/// How often to GC orphaned queue messages. -const QUEUE_MESSAGE_GC_INTERVAL: Duration = Duration::from_secs(300); -/// Max age before an unconsumed queue message is garbage collected. -const QUEUE_MESSAGE_MAX_AGE_SECS: i64 = 3600; - -/// Per-shard signal carried over a subscriber's in-process wakeup channel. -#[derive(Clone)] -enum ShardSignal { - /// A doorbell NOTIFY landed for this shard. Poll the table. - Wakeup, - /// A local request found no responders for the given reply subject. The matching - /// reply subscriber surfaces a no-responders result. - NoResponders { subject: String }, -} - -#[derive(Clone)] -pub struct PostgresDriver { - pool: Arc, - client: Arc>>, - /// Identifies this process in the subscriber registry. A single heartbeat keeps - /// all of this node's registrations live. - node_id: String, - /// Wakeup channels keyed by doorbell shard channel name. Shared by broadcast and - /// queue subscribers whose subjects map to the same shard. - shard_subscriptions: Arc>>, - doorbell: Arc, - client_ready: tokio::sync::watch::Receiver, -} - -impl PostgresDriver { - #[tracing::instrument(skip(conn_str))] - pub async fn connect( - conn_str: String, - ssl_root_cert_path: Option, - ssl_client_cert_path: Option, - ssl_client_key_path: Option, - ) -> Result { - // Create deadpool config from connection string - let mut config = Config::new(); - config.url = Some(conn_str.clone()); - config.pool = Some(PoolConfig { - max_size: 64, - ..Default::default() - }); - config.manager = Some(ManagerConfig { - recycling_method: RecyclingMethod::Fast, - }); - - // Create the pool - tracing::debug!("creating postgres pool"); - - // Build TLS configuration with optional custom certificates - let tls_config = build_tls_config( - ssl_root_cert_path.as_ref(), - ssl_client_cert_path.as_ref(), - ssl_client_key_path.as_ref(), - )?; - - let tls = MakeRustlsConnect::new(tls_config); - - let pool = config - .create_pool(Some(Runtime::Tokio1), tls) - .context("failed to create postgres pool")?; - tracing::debug!("postgres pool created successfully"); - - let pool = Arc::new(pool); - let shard_subscriptions: Arc>> = - Arc::new(HashMap::new()); - let client: Arc>> = Arc::new(Mutex::new(None)); - let node_id = Uuid::new_v4().to_string(); - - // Create channel for client ready notifications - let (ready_tx, client_ready) = tokio::sync::watch::channel(false); - - // Spawn connection lifecycle task - tokio::spawn(Self::spawn_connection_lifecycle( - conn_str.clone(), - shard_subscriptions.clone(), - client.clone(), - ready_tx, - ssl_root_cert_path.clone(), - ssl_client_cert_path.clone(), - ssl_client_key_path.clone(), - )); - - let doorbell = Doorbell::new(pool.clone()); - - let driver = Self { - pool, - client, - node_id, - shard_subscriptions, - doorbell, - client_ready, - }; - - // Wait for initial connection to be established - driver.wait_for_client().await?; - - // Create tables eagerly so they exist before any publish or subscribe. - { - tracing::debug!("configuring postgres udb tables"); - let conn = driver - .pool - .get() - .await - .context("failed to get connection for table creation")?; - conn.batch_execute( - // Broadcast transport table. UNLOGGED gives at-most-once across a - // crash, matching NATS-core semantics, and avoids WAL fsync on every - // publish. The real subject is stored so receivers can verify it and - // reject DefaultHasher subject-hash collisions. - "CREATE UNLOGGED TABLE IF NOT EXISTS ups_messages ( \ - id BIGSERIAL PRIMARY KEY, \ - subject_hash TEXT NOT NULL, \ - subject TEXT NOT NULL, \ - payload BYTEA NOT NULL, \ - created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() \ - ); \ - CREATE INDEX IF NOT EXISTS ups_messages_subject_id \ - ON ups_messages (subject_hash, id); \ - CREATE TABLE IF NOT EXISTS ups_nodes ( \ - node_id TEXT PRIMARY KEY, \ - heartbeat_at TIMESTAMPTZ NOT NULL DEFAULT NOW() \ - ); \ - CREATE TABLE IF NOT EXISTS ups_subs ( \ - id TEXT PRIMARY KEY, \ - node_id TEXT NOT NULL, \ - subject_hash TEXT NOT NULL, \ - subject TEXT NOT NULL \ - ); \ - CREATE INDEX IF NOT EXISTS ups_subs_subject \ - ON ups_subs (subject_hash); \ - CREATE TABLE IF NOT EXISTS ups_queue_subs ( \ - id TEXT PRIMARY KEY, \ - node_id TEXT NOT NULL, \ - subject_hash TEXT NOT NULL, \ - queue_hash TEXT NOT NULL \ - ); \ - CREATE INDEX IF NOT EXISTS ups_queue_subs_subject_queue \ - ON ups_queue_subs (subject_hash, queue_hash); \ - CREATE UNLOGGED TABLE IF NOT EXISTS ups_queue_messages ( \ - id BIGSERIAL PRIMARY KEY, \ - subject_hash TEXT NOT NULL, \ - queue_hash TEXT NOT NULL, \ - payload BYTEA NOT NULL, \ - created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() \ - ); \ - CREATE INDEX IF NOT EXISTS ups_queue_messages_idx \ - ON ups_queue_messages (subject_hash, queue_hash, id);", - ) - .await - .context("failed to create tables")?; - tracing::debug!("postgres udb tables ready"); - } - - // Register this node and start its liveness heartbeat. - driver.heartbeat_node().await?; - let heartbeat_driver = driver.clone(); - tokio::spawn(async move { - let mut interval = tokio::time::interval(NODE_HEARTBEAT_INTERVAL); - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - - loop { - interval.tick().await; - if let Err(e) = heartbeat_driver.heartbeat_node().await { - tracing::warn!(?e, "failed to heartbeat node"); - } - } - }); - - // Spawn GC task for expired broadcast messages - let message_gc_driver = driver.clone(); - tokio::spawn(async move { - let mut interval = tokio::time::interval(MESSAGE_GC_INTERVAL); - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - - loop { - interval.tick().await; - if let Ok(conn) = message_gc_driver.pool.get().await { - let result = conn - .execute( - "DELETE FROM ups_messages \ - WHERE created_at < NOW() - ($1::bigint * INTERVAL '1 second')", - &[&MESSAGE_MAX_AGE_SECS], - ) - .await; - if let Err(e) = result { - tracing::warn!(?e, "failed to gc broadcast messages"); - } - } - } - }); - - // Spawn GC task for dead nodes and orphaned subscriber rows. - let registry_gc_driver = driver.clone(); - tokio::spawn(async move { - let mut interval = tokio::time::interval(REGISTRY_GC_INTERVAL); - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - - loop { - interval.tick().await; - if let Ok(conn) = registry_gc_driver.pool.get().await { - if let Err(e) = conn - .execute( - "DELETE FROM ups_nodes \ - WHERE heartbeat_at < NOW() - ($1::bigint * INTERVAL '1 second')", - &[&NODE_TTL_SECS], - ) - .await - { - tracing::warn!(?e, "failed to gc dead nodes"); - } - if let Err(e) = conn - .execute( - "DELETE FROM ups_subs \ - WHERE node_id NOT IN (SELECT node_id FROM ups_nodes)", - &[], - ) - .await - { - tracing::warn!(?e, "failed to gc orphaned subs"); - } - if let Err(e) = conn - .execute( - "DELETE FROM ups_queue_subs \ - WHERE node_id NOT IN (SELECT node_id FROM ups_nodes)", - &[], - ) - .await - { - tracing::warn!(?e, "failed to gc orphaned queue subs"); - } - } - } - }); - - // Spawn GC task for orphaned queue messages - let gc_driver = driver.clone(); - tokio::spawn(async move { - let mut interval = tokio::time::interval(QUEUE_MESSAGE_GC_INTERVAL); - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - - loop { - interval.tick().await; - if let Ok(conn) = gc_driver.pool.get().await { - let result = conn - .execute( - "DELETE FROM ups_queue_messages \ - WHERE created_at < NOW() - ($1::bigint * INTERVAL '1 second')", - &[&QUEUE_MESSAGE_MAX_AGE_SECS], - ) - .await; - if let Err(e) = result { - tracing::warn!(?e, "failed to gc queue messages"); - } - } - } - }); - - Ok(driver) - } - - /// Manages the connection lifecycle with automatic reconnection - async fn spawn_connection_lifecycle( - conn_str: String, - shard_subscriptions: Arc>>, - client: Arc>>, - ready_tx: tokio::sync::watch::Sender, - ssl_root_cert_path: Option, - ssl_client_cert_path: Option, - ssl_client_key_path: Option, - ) { - let mut backoff = Backoff::default(); - - // Build TLS configuration with optional custom certificates - let tls_config = match build_tls_config( - ssl_root_cert_path.as_ref(), - ssl_client_cert_path.as_ref(), - ssl_client_key_path.as_ref(), - ) { - std::result::Result::Ok(config) => config, - std::result::Result::Err(e) => { - tracing::error!(?e, "failed to build TLS config"); - return; - } - }; - - let tls = MakeRustlsConnect::new(tls_config); - - loop { - match tokio_postgres::connect(&conn_str, tls.clone()).await { - Result::Ok((new_client, conn)) => { - tracing::debug!("postgres listen connection established"); - // Reset backoff on successful connection - backoff = Backoff::default(); - - // Spawn the polling task immediately - // This must be done before any operations on the client - let shard_subscriptions_clone = shard_subscriptions.clone(); - let poll_handle = tokio::spawn(async move { - Self::poll_connection(conn, shard_subscriptions_clone).await; - }); - - // Bound a stuck listener so it cannot wedge the shared notify queue. - if let Result::Err(e) = new_client - .execute( - &format!( - "SET idle_in_transaction_session_timeout = '{}'", - LISTEN_IDLE_IN_TRANSACTION_TIMEOUT_MS - ), - &[], - ) - .await - { - tracing::warn!(?e, "failed to set idle_in_transaction_session_timeout"); - } - - // Get shard channels to re-subscribe to - let mut channels = Vec::new(); - shard_subscriptions - .iter_async(|k, _| { - channels.push(k.clone()); - true - }) - .await; - - if !channels.is_empty() { - tracing::debug!( - channels = channels.len(), - "re-subscribing to doorbell shards after reconnection" - ); - } - - for channel in channels.iter() { - tracing::debug!(?channel, "re-subscribing to channel"); - if let Result::Err(e) = new_client - .execute(&format!("LISTEN \"{}\"", channel), &[]) - .await - { - tracing::error!(?e, %channel, "failed to re-subscribe to channel"); - } else { - tracing::debug!(%channel, "successfully re-subscribed to channel"); - } - } - - // Update the client reference and signal ready - // Do this AFTER re-subscribing to ensure LISTEN is complete - *client.lock().await = Some(new_client); - let _ = ready_tx.send(true); - - // Wait for the polling task to complete (when the connection closes) - let _ = poll_handle.await; - - // Clear the client reference on disconnect - *client.lock().await = None; - - // Notify that client is disconnected - let _ = ready_tx.send(false); - } - Result::Err(e) => { - tracing::error!(?e, "failed to connect to postgres, retrying"); - backoff.tick().await; - } - } - } - } - - /// Polls the connection for notifications until it closes or errors - async fn poll_connection( - mut conn: tokio_postgres::Connection, - shard_subscriptions: Arc>>, - ) where - T: tokio_postgres::tls::TlsStream + Unpin, - { - loop { - match poll_fn(|cx| conn.poll_message(cx)).await { - Some(std::result::Result::Ok(AsyncMessage::Notification(note))) => { - tracing::trace!(channel = %note.channel(), "received doorbell wakeup"); - // Doorbell notifications are payload-free wakeup signals only. - // Subscribers read their payload from the table. - if let Some(sub) = shard_subscriptions.get_async(note.channel()).await { - let _ = sub.send(ShardSignal::Wakeup); - } else { - tracing::trace!(channel = %note.channel(), "wakeup for unknown shard"); - } - } - Some(std::result::Result::Ok(_)) => { - // Ignore other async messages - } - Some(std::result::Result::Err(err)) => { - tracing::error!(?err, "postgres connection error"); - break; - } - None => { - tracing::warn!("postgres connection closed"); - break; - } - } - } - } - - /// Wait for the client to be connected - async fn wait_for_client(&self) -> Result<()> { - let mut ready_rx = self.client_ready.clone(); - tokio::time::timeout(tokio::time::Duration::from_secs(5), async { - loop { - // Check if client is already available - if self.client.lock().await.is_some() { - return Ok(()); - } - - // Wait for the ready signal to change - ready_rx - .changed() - .await - .context("connection lifecycle task ended")?; - } - }) - .await - .context("timeout waiting for postgres client connection")? - } - - fn hash_subject(&self, subject: &str) -> String { - // Postgres channel names have a 64 character limit, but this hash is also the - // table index key. Collisions are possible and resolved by verifying the real - // subject stored alongside each row. - let mut hasher = DefaultHasher::new(); - subject.hash(&mut hasher); - format!("ups_{:x}", hasher.finish()) - } - - fn hash_queue(&self, queue: &str) -> String { - let mut hasher = DefaultHasher::new(); - queue.hash(&mut hasher); - format!("{:x}", hasher.finish()) - } - - /// Upserts this node's liveness heartbeat. Re-inserts the row if a GC pass removed - /// it after a transient stall. - async fn heartbeat_node(&self) -> Result<()> { - let conn = self - .pool - .get() - .await - .context("failed to get connection for node heartbeat")?; - conn.execute( - "INSERT INTO ups_nodes (node_id, heartbeat_at) VALUES ($1, NOW()) \ - ON CONFLICT (node_id) DO UPDATE SET heartbeat_at = NOW()", - &[&self.node_id], - ) - .await - .context("failed to upsert node heartbeat")?; - Ok(()) - } - - /// Returns the current max broadcast message id, used as a subscriber's starting - /// cursor so it only sees future messages (NATS at-most-once, no replay). - async fn current_max_id(&self) -> Result { - let conn = self - .pool - .get() - .await - .context("failed to get connection for cursor init")?; - let row = conn - .query_one("SELECT COALESCE(MAX(id), 0) FROM ups_messages", &[]) - .await - .context("failed to read current max id")?; - Ok(row.get(0)) - } - - /// Ensures this process is LISTENing on the given doorbell shard and returns a - /// wakeup receiver plus a drop guard that UNLISTENs once no receivers remain. - async fn ensure_shard_listen( - &self, - shard: usize, - ) -> ( - broadcast::Receiver, - tokio_util::sync::DropGuard, - ) { - let channel = shard_channel(shard); - - match self.shard_subscriptions.entry_async(channel.clone()).await { - scc::hash_map::Entry::Occupied(existing) => { - let rx = existing.subscribe(); - let drop_guard = - self.spawn_shard_cleanup_task(channel.clone(), existing.get().clone()); - (rx, drop_guard) - } - scc::hash_map::Entry::Vacant(e) => { - let (tx, rx) = broadcast::channel(1024); - e.insert_entry(tx.clone()); - metrics::POSTGRES_SUBSCRIPTION_COUNT.set(self.shard_subscriptions.len() as i64); - - if let Some(client) = &*self.client.lock().await { - match client - .execute(&format!("LISTEN \"{channel}\""), &[]) - .instrument(tracing::trace_span!("pg_listen")) - .await - { - Result::Ok(_) => { - tracing::debug!(%channel, "successfully subscribed to shard"); - } - Result::Err(e) => { - tracing::warn!(?e, %channel, "failed to LISTEN, will retry on reconnection"); - } - } - } else { - tracing::debug!(%channel, "client not connected, will LISTEN on reconnection"); - } - - let drop_guard = self.spawn_shard_cleanup_task(channel.clone(), tx.clone()); - (rx, drop_guard) - } - } - } - - fn spawn_shard_cleanup_task( - &self, - channel: String, - tx: broadcast::Sender, - ) -> tokio_util::sync::DropGuard { - let driver = self.clone(); - let token = tokio_util::sync::CancellationToken::new(); - let drop_guard = token.clone().drop_guard(); - - tokio::spawn(async move { - token.cancelled().await; - if tx.receiver_count() == 0 { - if let Some(client) = &*driver.client.lock().await { - let sql = format!("UNLISTEN \"{}\"", channel); - if let Err(err) = client.execute(sql.as_str(), &[]).await { - tracing::warn!(?err, %channel, "failed to UNLISTEN channel"); - } else { - tracing::trace!(%channel, "unlistened channel"); - } - } - driver.shard_subscriptions.remove_async(&channel).await; - metrics::POSTGRES_SUBSCRIPTION_COUNT.set(driver.shard_subscriptions.len() as i64); - } - }); - - drop_guard - } - - /// Inserts the broadcast row and any active queue-group rows in one transaction. - async fn try_publish_to_db( - &self, - subject: &str, - subject_hash: &str, - payload: &[u8], - ) -> Result<()> { - let mut conn = self - .pool - .get() - .await - .context("failed to get connection for publish")?; - let tx = conn - .transaction() - .await - .context("failed to begin publish transaction")?; - - // Broadcast row. - tx.execute( - "INSERT INTO ups_messages (subject_hash, subject, payload) VALUES ($1, $2, $3)", - &[&subject_hash, &subject, &payload], - ) - .await - .context("failed to insert broadcast message")?; - - // Queue rows for every live queue group on this subject. Batched into the same - // transaction so a crash never strands a row mid-publish. - let rows = tx - .query( - "SELECT DISTINCT s.queue_hash FROM ups_queue_subs s \ - JOIN ups_nodes n ON s.node_id = n.node_id \ - WHERE s.subject_hash = $1 \ - AND n.heartbeat_at > NOW() - ($2::bigint * INTERVAL '1 second')", - &[&subject_hash, &NODE_TTL_SECS], - ) - .await - .context("failed to query active queue subs")?; - - for row in rows { - let queue_hash: String = row.get(0); - tx.execute( - "INSERT INTO ups_queue_messages (subject_hash, queue_hash, payload) \ - VALUES ($1, $2, $3)", - &[&subject_hash, &queue_hash, &payload], - ) - .await - .context("failed to insert queue message")?; - } - - tx.commit().await.context("failed to commit publish")?; - - Ok(()) - } - - /// Returns whether any live subscriber (broadcast or queue) exists for the subject - /// anywhere in the fleet. Used to decide whether a request surfaces a no-responders - /// result instead of waiting out its timeout. - async fn has_responders(&self, subject_hash: &str, subject: &str) -> Result { - let conn = self - .pool - .get() - .await - .context("failed to get connection for responder check")?; - let row = conn - .query_one( - "SELECT \ - EXISTS( \ - SELECT 1 FROM ups_subs s \ - JOIN ups_nodes n ON s.node_id = n.node_id \ - WHERE s.subject_hash = $1 AND s.subject = $2 \ - AND n.heartbeat_at > NOW() - ($3::bigint * INTERVAL '1 second') \ - ) \ - OR EXISTS( \ - SELECT 1 FROM ups_queue_subs s \ - JOIN ups_nodes n ON s.node_id = n.node_id \ - WHERE s.subject_hash = $1 \ - AND n.heartbeat_at > NOW() - ($3::bigint * INTERVAL '1 second') \ - )", - &[&subject_hash, &subject, &NODE_TTL_SECS], - ) - .await - .context("failed to check responders")?; - Ok(row.get(0)) - } - - /// Delivers a no-responders result to the local reply subscriber. The requester is - /// always in this process, so the signal is routed in-memory over the reply - /// subject's shard channel rather than the table. - async fn signal_no_responders(&self, reply_subject: &str) { - let reply_hash = self.hash_subject(reply_subject); - let channel = shard_channel(shard_for(&reply_hash)); - if let Some(tx) = self.shard_subscriptions.get_async(&channel).await { - let _ = tx.send(ShardSignal::NoResponders { - subject: reply_subject.to_string(), - }); - } - } -} - -#[async_trait] -impl PubSubDriver for PostgresDriver { - async fn subscribe( - &self, - subject: &str, - reply_id: Option, - ) -> Result { - let subject_hash = self.hash_subject(subject); - let shard = shard_for(&subject_hash); - - // Capture the cursor before LISTENing. Any message inserted after this point - // has a higher id and is delivered either by the doorbell wakeup or the poll - // backstop, so there is no subscribe/publish race. - let cursor = self.current_max_id().await?; - - let (rx, drop_guard) = self.ensure_shard_listen(shard).await; - - // Register in the responder registry so requests to this subject can detect - // responders. Reply inboxes are never request targets, so they skip the - // registry to keep request latency off this path. - let sub_id = if reply_id.is_none() { - let sub_id = Uuid::new_v4().to_string(); - let conn = self - .pool - .get() - .await - .context("failed to get connection for subscribe")?; - conn.execute( - "INSERT INTO ups_subs (id, node_id, subject_hash, subject) \ - VALUES ($1, $2, $3, $4)", - &[&sub_id, &self.node_id, &subject_hash, &subject], - ) - .await - .context("failed to register subscriber")?; - Some(sub_id) - } else { - None - }; - - Ok(Box::new(PostgresSubscriber { - subject: subject.to_string(), - subject_hash, - pool: self.pool.clone(), - cursor, - buffer: VecDeque::new(), - rx, - sub_id, - _drop_guard: drop_guard, - })) - } - - async fn queue_subscribe(&self, subject: &str, queue: &str) -> Result { - let subject_hash = self.hash_subject(subject); - let queue_hash = self.hash_queue(queue); - let shard = shard_for(&subject_hash); - - // Register this subscriber in the database so publishers know the queue exists - let sub_id = Uuid::new_v4().to_string(); - { - let conn = self - .pool - .get() - .await - .context("failed to get connection for queue subscribe")?; - conn.execute( - "INSERT INTO ups_queue_subs (id, node_id, subject_hash, queue_hash) \ - VALUES ($1, $2, $3, $4)", - &[&sub_id, &self.node_id, &subject_hash, &queue_hash], - ) - .await - .context("failed to register queue subscriber")?; - } - - let (rx, drop_guard) = self.ensure_shard_listen(shard).await; - - Ok(Box::new(PostgresQueueSubscriber { - subject: subject.to_string(), - subject_hash, - queue_hash, - sub_id, - pool: self.pool.clone(), - rx, - _drop_guard: drop_guard, - })) - } - - async fn publish( - &self, - subject: &str, - payload: &[u8], - reply_subject: Option<&str>, - ) -> Result<()> { - let subject_hash = self.hash_subject(subject); - let shard = shard_for(&subject_hash); - - // Request semantics: if a reply is expected and no responder exists anywhere, - // surface a no-responders result immediately instead of persisting a message - // nobody will read. - if let Some(reply_subject) = reply_subject { - match self.has_responders(&subject_hash, subject).await { - Result::Ok(false) => { - self.signal_no_responders(reply_subject).await; - return Ok(()); - } - Result::Ok(true) => {} - Result::Err(e) => { - // On a failed check, fall through to a normal publish rather than - // risk a false no-responders result. - tracing::warn!(?e, %subject, "responder check failed, publishing anyway"); - } - } - } - - // Persist the message, retrying on transient connection errors. The row is - // committed before the doorbell rings so any wakeup observes it. - let mut backoff = Backoff::default(); - loop { - match self - .try_publish_to_db(subject, &subject_hash, payload) - .await - { - Result::Ok(()) => break, - Result::Err(e) => { - if !backoff.tick().await { - tracing::warn!(?e, %subject, "failed to publish, cannot retry again"); - return Err(e); - } - tracing::debug!(?e, "publish failed, retrying"); - } - } - } - - // Ring the doorbell. Best-effort: the subscriber poll backstop covers a - // dropped or coalesced wakeup, so publish never blocks on NOTIFY. - self.doorbell.mark_dirty(shard); - - Ok(()) - } - - async fn flush(&self) -> Result<()> { - Ok(()) - } - - fn max_message_size(&self) -> usize { - POSTGRES_MAX_MESSAGE_SIZE - } -} - -pub struct PostgresSubscriber { - subject: String, - subject_hash: String, - pool: Arc, - cursor: i64, - buffer: VecDeque>, - rx: broadcast::Receiver, - /// Responder-registry row id, present for non-inbox subscriptions. Deleted on drop. - sub_id: Option, - _drop_guard: tokio_util::sync::DropGuard, -} - -impl PostgresSubscriber { - /// Reads new rows past the cursor into the buffer, advancing the cursor. Rows - /// whose stored subject does not match are skipped (DefaultHasher collisions) but - /// still advance the cursor. - async fn fetch(&mut self) -> Result<()> { - let conn = self - .pool - .get() - .await - .context("failed to get connection for poll")?; - let rows = conn - .query( - "SELECT id, subject, payload FROM ups_messages \ - WHERE subject_hash = $1 AND id > $2 ORDER BY id", - &[&self.subject_hash, &self.cursor], - ) - .await - .context("failed to poll broadcast messages")?; - - for row in rows { - let id: i64 = row.get(0); - let subject: String = row.get(1); - let payload: Vec = row.get(2); - self.cursor = id; - if subject == self.subject { - self.buffer.push_back(payload); - } - } - - Ok(()) - } -} - -#[async_trait] -impl SubscriberDriver for PostgresSubscriber { - async fn next(&mut self) -> Result { - loop { - if let Some(payload) = self.buffer.pop_front() { - return Ok(DriverOutput::Message { - subject: self.subject.clone(), - payload, - }); - } - - if let Err(e) = self.fetch().await { - // Transient DB errors must not kill the subscriber; the next poll - // tick retries. - tracing::warn!(?e, subject = %self.subject, "failed to poll, will retry"); - } - - if !self.buffer.is_empty() { - continue; - } - - // Wait for a doorbell wakeup, a no-responders signal, or the poll backstop. - tokio::select! { - res = self.rx.recv() => { - match res { - std::result::Result::Ok(ShardSignal::Wakeup) => {} - std::result::Result::Ok(ShardSignal::NoResponders { subject }) - if subject == self.subject => - { - return Ok(DriverOutput::NoResponders); - } - std::result::Result::Ok(ShardSignal::NoResponders { .. }) => {} - Err(broadcast::error::RecvError::Lagged(_)) => {} - Err(broadcast::error::RecvError::Closed) => { - return Ok(DriverOutput::Unsubscribed); - } - } - } - _ = tokio::time::sleep(POLL_INTERVAL) => {} - } - } - } -} - -impl Drop for PostgresSubscriber { - fn drop(&mut self) { - let Some(sub_id) = self.sub_id.take() else { - return; - }; - let pool = self.pool.clone(); - tokio::spawn(async move { - if let Ok(conn) = pool.get().await { - if let Err(e) = conn - .execute("DELETE FROM ups_subs WHERE id = $1", &[&sub_id]) - .await - { - tracing::warn!(?e, %sub_id, "failed to deregister subscriber"); - } - } - }); - } -} - -pub struct PostgresQueueSubscriber { - subject: String, - subject_hash: String, - queue_hash: String, - sub_id: String, - pool: Arc, - rx: broadcast::Receiver, - _drop_guard: tokio_util::sync::DropGuard, -} - -impl PostgresQueueSubscriber { - /// Attempts to atomically claim and delete one pending message for this (subject, queue). - async fn claim_message(&self) -> Result>> { - let conn = self - .pool - .get() - .await - .context("failed to get connection for queue claim")?; - - let rows = conn - .query( - "WITH claimed AS ( \ - SELECT id, payload FROM ups_queue_messages \ - WHERE subject_hash = $1 AND queue_hash = $2 \ - ORDER BY id \ - LIMIT 1 \ - FOR UPDATE SKIP LOCKED \ - ) \ - DELETE FROM ups_queue_messages \ - WHERE id IN (SELECT id FROM claimed) \ - RETURNING payload", - &[&self.subject_hash, &self.queue_hash], - ) - .await - .context("failed to claim queue message")?; - - Ok(rows.into_iter().next().map(|row| row.get::<_, Vec>(0))) - } -} - -#[async_trait] -impl SubscriberDriver for PostgresQueueSubscriber { - async fn next(&mut self) -> Result { - loop { - // Drain any messages that arrived before or between wakeups. - match self.claim_message().await { - Result::Ok(Some(payload)) => { - return Ok(DriverOutput::Message { - subject: self.subject.clone(), - payload, - }); - } - Result::Ok(None) => {} - Result::Err(e) => { - tracing::warn!(?e, subject = %self.subject, "failed to claim, will retry"); - } - } - - // Wait for any shard signal or the poll backstop, then loop back to claim. - tokio::select! { - res = self.rx.recv() => { - match res { - std::result::Result::Ok(_) => {} - Err(broadcast::error::RecvError::Lagged(_)) => {} - Err(broadcast::error::RecvError::Closed) => { - return Ok(DriverOutput::Unsubscribed); - } - } - } - _ = tokio::time::sleep(POLL_INTERVAL) => {} - } - } - } -} - -impl Drop for PostgresQueueSubscriber { - fn drop(&mut self) { - let pool = self.pool.clone(); - let sub_id = self.sub_id.clone(); - tokio::spawn(async move { - if let Ok(conn) = pool.get().await { - if let Err(e) = conn - .execute("DELETE FROM ups_queue_subs WHERE id = $1", &[&sub_id]) - .await - { - tracing::warn!(?e, %sub_id, "failed to deregister queue subscriber"); - } - } - }); - } -} diff --git a/engine/packages/universalpubsub/src/metrics.rs b/engine/packages/universalpubsub/src/metrics.rs index 9c508e963a..93edf9eb1e 100644 --- a/engine/packages/universalpubsub/src/metrics.rs +++ b/engine/packages/universalpubsub/src/metrics.rs @@ -25,13 +25,6 @@ lazy_static::lazy_static! { "Number of subject entries in the memory driver.", *REGISTRY ).unwrap(); - // Postgres driver metrics - pub static ref POSTGRES_SUBSCRIPTION_COUNT: IntGauge = register_int_gauge_with_registry!( - "ups_postgres_subscription_count", - "Number of subscription entries in the postgres driver.", - *REGISTRY - ).unwrap(); - // Message metrics pub static ref MESSAGE_RECV_COUNT: IntCounterVec = register_int_counter_vec_with_registry!( "ups_message_recv_count", diff --git a/engine/packages/universalpubsub/tests/integration.rs b/engine/packages/universalpubsub/tests/integration.rs index 368fd4a537..f849935e45 100644 --- a/engine/packages/universalpubsub/tests/integration.rs +++ b/engine/packages/universalpubsub/tests/integration.rs @@ -1,7 +1,7 @@ use anyhow::Result; use futures_util::StreamExt; use rivet_error::RivetError; -use rivet_test_deps_docker::{TestDatabase, TestPubSub}; +use rivet_test_deps_docker::TestPubSub; use std::{ sync::Arc, time::{Duration, Instant}, @@ -119,29 +119,6 @@ async fn test_nats_no_responders() { test_no_responders(&pubsub).await.unwrap(); } -#[tokio::test] -async fn test_postgres_no_responders() { - setup_logging(); - - let test_id = Uuid::new_v4(); - let (db_config, docker_config) = TestDatabase::Postgres.config(test_id, 1).await.unwrap(); - let mut docker = docker_config.unwrap(); - docker.start().await.unwrap(); - tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; - - let rivet_config::config::Database::Postgres(pg) = db_config else { - unreachable!(); - }; - let url = pg.url.read().clone(); - - let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, None, None, None) - .await - .unwrap(); - let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false); - - test_no_responders(&pubsub).await.unwrap(); -} - #[tokio::test] async fn test_memory_no_responders() { setup_logging(); @@ -158,52 +135,6 @@ async fn test_memory_no_responders() { test_no_responders(&pubsub).await.unwrap(); } -#[tokio::test] -async fn test_postgres_driver_with_memory() { - setup_logging(); - - let test_id = Uuid::new_v4(); - let (db_config, docker_config) = TestDatabase::Postgres.config(test_id, 1).await.unwrap(); - let mut docker = docker_config.unwrap(); - docker.start().await.unwrap(); - tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; - - let rivet_config::config::Database::Postgres(pg) = db_config else { - unreachable!(); - }; - let url = pg.url.read().clone(); - - let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, None, None, None) - .await - .unwrap(); - let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), true); - - test_inner(&pubsub).await; -} - -#[tokio::test] -async fn test_postgres_driver_without_memory() { - setup_logging(); - - let test_id = Uuid::new_v4(); - let (db_config, docker_config) = TestDatabase::Postgres.config(test_id, 1).await.unwrap(); - let mut docker = docker_config.unwrap(); - docker.start().await.unwrap(); - tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; - - let rivet_config::config::Database::Postgres(pg) = db_config else { - unreachable!(); - }; - let url = pg.url.read().clone(); - - let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, None, None, None) - .await - .unwrap(); - let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false); - - test_inner(&pubsub).await; -} - #[tokio::test] async fn test_memory_driver() { setup_logging(); diff --git a/engine/packages/universalpubsub/tests/reconnect.rs b/engine/packages/universalpubsub/tests/reconnect.rs index 8c15f6ab54..c08892d8d3 100644 --- a/engine/packages/universalpubsub/tests/reconnect.rs +++ b/engine/packages/universalpubsub/tests/reconnect.rs @@ -1,5 +1,5 @@ use anyhow::Result; -use rivet_test_deps_docker::{TestDatabase, TestPubSub}; +use rivet_test_deps_docker::TestPubSub; use std::{sync::Arc, time::Duration}; use universalpubsub::{NextOutput, PubSub, PublishOpts}; use uuid::Uuid; @@ -80,52 +80,6 @@ async fn test_nats_driver_without_memory_reconnect() { test_all_inner(&pubsub, &docker, true).await; } -#[tokio::test] -async fn test_postgres_driver_with_memory_reconnect() { - setup_logging(); - - let test_id = Uuid::new_v4(); - let (db_config, docker_config) = TestDatabase::Postgres.config(test_id, 1).await.unwrap(); - let mut docker = docker_config.unwrap(); - docker.start().await.unwrap(); - tokio::time::sleep(Duration::from_secs(5)).await; - - let rivet_config::config::Database::Postgres(pg) = db_config else { - unreachable!(); - }; - let url = pg.url.read().clone(); - - let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, None, None, None) - .await - .unwrap(); - let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), true); - - test_all_inner(&pubsub, &docker, false).await; -} - -#[tokio::test] -async fn test_postgres_driver_without_memory_reconnect() { - setup_logging(); - - let test_id = Uuid::new_v4(); - let (db_config, docker_config) = TestDatabase::Postgres.config(test_id, 1).await.unwrap(); - let mut docker = docker_config.unwrap(); - docker.start().await.unwrap(); - tokio::time::sleep(Duration::from_secs(5)).await; - - let rivet_config::config::Database::Postgres(pg) = db_config else { - unreachable!(); - }; - let url = pg.url.read().clone(); - - let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, None, None, None) - .await - .unwrap(); - let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false); - - test_all_inner(&pubsub, &docker, false).await; -} - async fn test_all_inner( pubsub: &PubSub, docker: &rivet_test_deps_docker::DockerRunConfig, diff --git a/engine/sdks/rust/universaldb-commit/src/versioned.rs b/engine/sdks/rust/universaldb-commit/src/versioned.rs index a9d637aa81..804f533a64 100644 --- a/engine/sdks/rust/universaldb-commit/src/versioned.rs +++ b/engine/sdks/rust/universaldb-commit/src/versioned.rs @@ -36,3 +36,65 @@ impl OwnedVersionedData for CommitRequest { } } } + +pub enum CommitReply { + V1(v1::CommitReply), +} + +impl OwnedVersionedData for CommitReply { + type Latest = v1::CommitReply; + + fn wrap_latest(latest: v1::CommitReply) -> Self { + CommitReply::V1(latest) + } + + fn unwrap_latest(self) -> Result { + match self { + CommitReply::V1(data) => Ok(data), + } + } + + fn deserialize_version(payload: &[u8], version: u16) -> Result { + match version { + 1 => Ok(CommitReply::V1(serde_bare::from_slice(payload)?)), + _ => bail!("invalid version: {version}"), + } + } + + fn serialize_version(self, _version: u16) -> Result> { + match self { + CommitReply::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + } + } +} + +pub enum Watermark { + V1(v1::Watermark), +} + +impl OwnedVersionedData for Watermark { + type Latest = v1::Watermark; + + fn wrap_latest(latest: v1::Watermark) -> Self { + Watermark::V1(latest) + } + + fn unwrap_latest(self) -> Result { + match self { + Watermark::V1(data) => Ok(data), + } + } + + fn deserialize_version(payload: &[u8], version: u16) -> Result { + match version { + 1 => Ok(Watermark::V1(serde_bare::from_slice(payload)?)), + _ => bail!("invalid version: {version}"), + } + } + + fn serialize_version(self, _version: u16) -> Result> { + match self { + Watermark::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + } + } +} diff --git a/engine/sdks/schemas/universaldb-commit/v1.bare b/engine/sdks/schemas/universaldb-commit/v1.bare index 2377312f84..c85d43f447 100644 --- a/engine/sdks/schemas/universaldb-commit/v1.bare +++ b/engine/sdks/schemas/universaldb-commit/v1.bare @@ -1,9 +1,9 @@ -# Commit-queue wire format for the Postgres leader-resolver UDB driver. +# Commit wire format for the multi-node Postgres leader-resolver UDB driver. # -# Followers encode a CommitRequest into the `payload` column of -# `udb_commit_requests`; the leader decodes it to resolve and apply. -# Rust-only (never leaves the engine), but versioned so rolling deploys can -# skew follower vs leader code. +# In multi-node mode a follower encodes a CommitRequest and sends it to the +# elected leader over NATS request/reply; the leader decodes it to resolve and +# apply. Rust-only (never leaves the engine), but versioned so rolling deploys +# can skew follower vs leader code. type ConflictRangeType enum { READ @@ -67,4 +67,32 @@ type CommitRequest struct { readVersion: u64 conflictRanges: list operations: list + # Failover dedup key. `clientNodeId` is the submitting follower's node id and + # `clientSeq` a per-process monotonic counter, unique per logical commit. The + # leader records committed (clientNodeId, clientSeq) so a follower that resends + # the same commit after an indeterminate leader failure is applied exactly once. + clientNodeId: data + clientSeq: u64 +} + +# Reply the leader sends back to the follower over the NATS request/reply inbox +# after resolving a commit. `CommitCommitted` carries the assigned commit +# version; `CommitConflict` (read-write conflict or cold-window reject) carries +# nothing and maps to the retryable NotCommitted on the follower. +type CommitCommitted struct { + commitVersion: i64 +} + +type CommitConflict void + +type CommitReply union { + CommitCommitted | + CommitConflict +} + +# Durable-version watermark the leader broadcasts on every applied batch so +# followers advance their read-snapshot floor. Best-effort; a lease-row poll is +# the backstop. +type Watermark struct { + durableVersion: i64 } diff --git a/self-host/compose/template/src/docker-compose.ts b/self-host/compose/template/src/docker-compose.ts index 2245a31c75..703559ced7 100644 --- a/self-host/compose/template/src/docker-compose.ts +++ b/self-host/compose/template/src/docker-compose.ts @@ -120,7 +120,12 @@ export function generateDockerCompose(context: TemplateContext) { const dcEnginePeerNetworkName = `${dcNetworkName}-engine-peer`; const dcToCoreNetworkName = `${dcNetworkName}-to-core`; - //const natsServiceName = context.getServiceName("nats", datacenter.name); + // A datacenter with more than one engine needs NATS: the engines coordinate the UniversalDB + // leader/follower commit transport and the UPS pubsub across nodes through it. A single-engine + // datacenter runs UniversalDB single-node (in-process resolver) and UPS in-process Memory, so it + // needs no NATS. + const useNats = datacenter.engines > 1; + const natsServiceName = context.getServiceName("nats", datacenter.name); const vectorServerServiceName = context.getServiceName( "vector-server", datacenter.name, @@ -154,18 +159,25 @@ export function generateDockerCompose(context: TemplateContext) { driver: "bridge", }; - //services[natsServiceName] = { - // restart: "unless-stopped", - // image: "nats:2.10.22-scratch", - // networks: [dcNetworkName], - // ports: isPrimary ? [`4222:4222`] : undefined, - // healthcheck: { - // test: ["CMD", "nats-server", "--health"], - // interval: "2s", - // timeout: "10s", - // retries: 10, - // }, - //}; + if (useNats) { + services[natsServiceName] = { + restart: "unless-stopped", + image: "nats:2.10.22-alpine", + // Enable the HTTP monitoring port so the healthcheck can hit /healthz. + command: ["-m", "8222"], + networks: [dcNetworkName], + ports: isPrimary ? [`4222:4222`] : undefined, + healthcheck: { + test: [ + "CMD-SHELL", + "wget -q -O /dev/null http://127.0.0.1:8222/healthz || exit 1", + ], + interval: "2s", + timeout: "10s", + retries: 10, + }, + }; + } const postgresVolumeName = context.getVolumeName( "postgres", @@ -174,9 +186,9 @@ export function generateDockerCompose(context: TemplateContext) { services[postgresServiceName] = { restart: "unless-stopped", image: "postgres:18-alpine", - // Each engine opens a UDB connection pool (up to 64 connections) plus a - // dedicated LISTEN connection and pubsub, so a multi-engine datacenter - // needs far more than the default max_connections of 100. + // Each engine opens a UDB connection pool (up to 64 connections), so a + // multi-engine datacenter needs far more than the default max_connections + // of 100. command: ["postgres", "-c", "max_connections=500"], environment: [ "POSTGRES_USER=postgres", @@ -213,7 +225,9 @@ export function generateDockerCompose(context: TemplateContext) { command: "infinity", stop_grace_period: "0s", depends_on: { - //[natsServiceName]: { condition: "service_healthy" }, + ...(useNats + ? { [natsServiceName]: { condition: "service_healthy" } } + : {}), [postgresServiceName]: { condition: "service_healthy" }, }, volumes: [ @@ -300,7 +314,9 @@ export function generateDockerCompose(context: TemplateContext) { ], stop_grace_period: "0s", depends_on: { - //[natsServiceName]: { condition: "service_healthy" }, + ...(useNats + ? { [natsServiceName]: { condition: "service_healthy" } } + : {}), [vectorClientServiceName]: { condition: "service_started", }, diff --git a/self-host/compose/template/src/services/edge/rivet-engine.ts b/self-host/compose/template/src/services/edge/rivet-engine.ts index 52d33dfc1b..1c5b8247da 100644 --- a/self-host/compose/template/src/services/edge/rivet-engine.ts +++ b/self-host/compose/template/src/services/edge/rivet-engine.ts @@ -37,7 +37,7 @@ export function generateDatacenterRivetEngine( // Config structure matching Rust schema in engine/packages/config/src/config/mod.rs. // Values that match the engine's defaults are omitted. - const config = { + const config: Record = { auth: { admin_token: "dev", }, @@ -56,6 +56,16 @@ export function generateDatacenterRivetEngine( }, }; + // A multi-engine datacenter coordinates through NATS: UPS uses it for cross-node pubsub, and + // UniversalDB inherits this config (see Root::validate_and_set_defaults) to run multi-node, + // routing follower commits to the elected leader over NATS. A single-engine datacenter omits + // this and falls back to in-process Memory pubsub + single-node UniversalDB. + if (datacenter.engines > 1) { + config.nats = { + addresses: [`${context.getServiceHost("nats", datacenter.name)}:4222`], + }; + } + context.writeDatacenterServiceFile( "rivet-engine", datacenter.name, diff --git a/self-host/dev-host/docker-compose.yml b/self-host/dev-host/docker-compose.yml index 94619e9324..fc44c89a26 100644 --- a/self-host/dev-host/docker-compose.yml +++ b/self-host/dev-host/docker-compose.yml @@ -157,6 +157,7 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://127.0.0.1:4317 diff --git a/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc index 0532bce5c3..29ca3ee404 100644 --- a/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc @@ -51,5 +51,10 @@ "native_url": "http://clickhouse:9301", "username": "system", "password": "default" + }, + "nats": { + "addresses": [ + "nats-dc-a:4222" + ] } } \ No newline at end of file diff --git a/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc index 0532bce5c3..29ca3ee404 100644 --- a/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc @@ -51,5 +51,10 @@ "native_url": "http://clickhouse:9301", "username": "system", "password": "default" + }, + "nats": { + "addresses": [ + "nats-dc-a:4222" + ] } } \ No newline at end of file diff --git a/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc index 0532bce5c3..29ca3ee404 100644 --- a/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc @@ -51,5 +51,10 @@ "native_url": "http://clickhouse:9301", "username": "system", "password": "default" + }, + "nats": { + "addresses": [ + "nats-dc-a:4222" + ] } } \ No newline at end of file diff --git a/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc index e75a22f77c..0fe743310a 100644 --- a/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc @@ -51,5 +51,10 @@ "native_url": "http://clickhouse:9301", "username": "system", "password": "default" + }, + "nats": { + "addresses": [ + "nats-dc-b:4222" + ] } } \ No newline at end of file diff --git a/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc index e75a22f77c..0fe743310a 100644 --- a/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc @@ -51,5 +51,10 @@ "native_url": "http://clickhouse:9301", "username": "system", "password": "default" + }, + "nats": { + "addresses": [ + "nats-dc-b:4222" + ] } } \ No newline at end of file diff --git a/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc index e75a22f77c..0fe743310a 100644 --- a/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc @@ -51,5 +51,10 @@ "native_url": "http://clickhouse:9301", "username": "system", "password": "default" + }, + "nats": { + "addresses": [ + "nats-dc-b:4222" + ] } } \ No newline at end of file diff --git a/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc index cb4c587a19..5a639b2fd4 100644 --- a/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc @@ -51,5 +51,10 @@ "native_url": "http://clickhouse:9301", "username": "system", "password": "default" + }, + "nats": { + "addresses": [ + "nats-dc-c:4222" + ] } } \ No newline at end of file diff --git a/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc index cb4c587a19..5a639b2fd4 100644 --- a/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc @@ -51,5 +51,10 @@ "native_url": "http://clickhouse:9301", "username": "system", "password": "default" + }, + "nats": { + "addresses": [ + "nats-dc-c:4222" + ] } } \ No newline at end of file diff --git a/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc index cb4c587a19..5a639b2fd4 100644 --- a/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc @@ -51,5 +51,10 @@ "native_url": "http://clickhouse:9301", "username": "system", "password": "default" + }, + "nats": { + "addresses": [ + "nats-dc-c:4222" + ] } } \ No newline at end of file diff --git a/self-host/dev-multidc-multinode/docker-compose.yml b/self-host/dev-multidc-multinode/docker-compose.yml index e4ccd4fc70..81b39d39eb 100644 --- a/self-host/dev-multidc-multinode/docker-compose.yml +++ b/self-host/dev-multidc-multinode/docker-compose.yml @@ -83,6 +83,23 @@ services: condition: service_healthy prometheus: condition: service_healthy + nats-dc-a: + restart: unless-stopped + image: nats:2.10.22-alpine + command: + - '-m' + - '8222' + networks: + - rivet-network-dc-a + ports: + - '4222:4222' + healthcheck: + test: + - CMD-SHELL + - wget -q -O /dev/null http://127.0.0.1:8222/healthz || exit 1 + interval: 2s + timeout: 10s + retries: 10 postgres-dc-a: restart: unless-stopped image: postgres:18-alpine @@ -122,6 +139,8 @@ services: command: infinity stop_grace_period: 0s depends_on: + nats-dc-a: + condition: service_healthy postgres-dc-a: condition: service_healthy volumes: @@ -187,11 +206,14 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector-dc-a:4317 stop_grace_period: 0s depends_on: + nats-dc-a: + condition: service_healthy vector-client-dc-a: condition: service_started otel-collector-dc-a: @@ -230,11 +252,14 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector-dc-a:4317 stop_grace_period: 0s depends_on: + nats-dc-a: + condition: service_healthy vector-client-dc-a: condition: service_started otel-collector-dc-a: @@ -271,11 +296,14 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector-dc-a:4317 stop_grace_period: 0s depends_on: + nats-dc-a: + condition: service_healthy vector-client-dc-a: condition: service_started otel-collector-dc-a: @@ -372,6 +400,21 @@ services: condition: service_completed_successfully networks: - rivet-network-dc-a + nats-dc-b: + restart: unless-stopped + image: nats:2.10.22-alpine + command: + - '-m' + - '8222' + networks: + - rivet-network-dc-b + healthcheck: + test: + - CMD-SHELL + - wget -q -O /dev/null http://127.0.0.1:8222/healthz || exit 1 + interval: 2s + timeout: 10s + retries: 10 postgres-dc-b: restart: unless-stopped image: postgres:18-alpine @@ -409,6 +452,8 @@ services: command: infinity stop_grace_period: 0s depends_on: + nats-dc-b: + condition: service_healthy postgres-dc-b: condition: service_healthy volumes: @@ -472,11 +517,14 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector-dc-b:4317 stop_grace_period: 0s depends_on: + nats-dc-b: + condition: service_healthy vector-client-dc-b: condition: service_started otel-collector-dc-b: @@ -513,11 +561,14 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector-dc-b:4317 stop_grace_period: 0s depends_on: + nats-dc-b: + condition: service_healthy vector-client-dc-b: condition: service_started otel-collector-dc-b: @@ -554,11 +605,14 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector-dc-b:4317 stop_grace_period: 0s depends_on: + nats-dc-b: + condition: service_healthy vector-client-dc-b: condition: service_started otel-collector-dc-b: @@ -653,6 +707,21 @@ services: condition: service_completed_successfully networks: - rivet-network-dc-b + nats-dc-c: + restart: unless-stopped + image: nats:2.10.22-alpine + command: + - '-m' + - '8222' + networks: + - rivet-network-dc-c + healthcheck: + test: + - CMD-SHELL + - wget -q -O /dev/null http://127.0.0.1:8222/healthz || exit 1 + interval: 2s + timeout: 10s + retries: 10 postgres-dc-c: restart: unless-stopped image: postgres:18-alpine @@ -690,6 +759,8 @@ services: command: infinity stop_grace_period: 0s depends_on: + nats-dc-c: + condition: service_healthy postgres-dc-c: condition: service_healthy volumes: @@ -753,11 +824,14 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector-dc-c:4317 stop_grace_period: 0s depends_on: + nats-dc-c: + condition: service_healthy vector-client-dc-c: condition: service_started otel-collector-dc-c: @@ -794,11 +868,14 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector-dc-c:4317 stop_grace_period: 0s depends_on: + nats-dc-c: + condition: service_healthy vector-client-dc-c: condition: service_started otel-collector-dc-c: @@ -835,11 +912,14 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector-dc-c:4317 stop_grace_period: 0s depends_on: + nats-dc-c: + condition: service_healthy vector-client-dc-c: condition: service_started otel-collector-dc-c: diff --git a/self-host/dev-multidc/docker-compose.yml b/self-host/dev-multidc/docker-compose.yml index 437ad5bdee..77938647f1 100644 --- a/self-host/dev-multidc/docker-compose.yml +++ b/self-host/dev-multidc/docker-compose.yml @@ -186,6 +186,7 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector-dc-a:4317 @@ -341,6 +342,7 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector-dc-b:4317 @@ -492,6 +494,7 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector-dc-c:4317 diff --git a/self-host/dev-multinode/docker-compose.yml b/self-host/dev-multinode/docker-compose.yml index 28234658d2..349b19edbc 100644 --- a/self-host/dev-multinode/docker-compose.yml +++ b/self-host/dev-multinode/docker-compose.yml @@ -79,6 +79,23 @@ services: condition: service_healthy prometheus: condition: service_healthy + nats: + restart: unless-stopped + image: nats:2.10.22-alpine + command: + - '-m' + - '8222' + networks: + - rivet-network + ports: + - '4222:4222' + healthcheck: + test: + - CMD-SHELL + - wget -q -O /dev/null http://127.0.0.1:8222/healthz || exit 1 + interval: 2s + timeout: 10s + retries: 10 postgres: restart: unless-stopped image: postgres:18-alpine @@ -117,6 +134,8 @@ services: command: infinity stop_grace_period: 0s depends_on: + nats: + condition: service_healthy postgres: condition: service_healthy volumes: @@ -185,6 +204,8 @@ services: - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector:4317 stop_grace_period: 0s depends_on: + nats: + condition: service_healthy vector-client: condition: service_started otel-collector: @@ -226,6 +247,8 @@ services: - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector:4317 stop_grace_period: 0s depends_on: + nats: + condition: service_healthy vector-client: condition: service_started otel-collector: @@ -265,6 +288,8 @@ services: - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector:4317 stop_grace_period: 0s depends_on: + nats: + condition: service_healthy vector-client: condition: service_started otel-collector: diff --git a/self-host/dev-multinode/rivet-engine/0/config.jsonc b/self-host/dev-multinode/rivet-engine/0/config.jsonc index b25680c399..13e9996d6e 100644 --- a/self-host/dev-multinode/rivet-engine/0/config.jsonc +++ b/self-host/dev-multinode/rivet-engine/0/config.jsonc @@ -29,5 +29,10 @@ "native_url": "http://clickhouse:9301", "username": "system", "password": "default" + }, + "nats": { + "addresses": [ + "nats:4222" + ] } } \ No newline at end of file diff --git a/self-host/dev-multinode/rivet-engine/1/config.jsonc b/self-host/dev-multinode/rivet-engine/1/config.jsonc index b25680c399..13e9996d6e 100644 --- a/self-host/dev-multinode/rivet-engine/1/config.jsonc +++ b/self-host/dev-multinode/rivet-engine/1/config.jsonc @@ -29,5 +29,10 @@ "native_url": "http://clickhouse:9301", "username": "system", "password": "default" + }, + "nats": { + "addresses": [ + "nats:4222" + ] } } \ No newline at end of file diff --git a/self-host/dev-multinode/rivet-engine/2/config.jsonc b/self-host/dev-multinode/rivet-engine/2/config.jsonc index b25680c399..13e9996d6e 100644 --- a/self-host/dev-multinode/rivet-engine/2/config.jsonc +++ b/self-host/dev-multinode/rivet-engine/2/config.jsonc @@ -29,5 +29,10 @@ "native_url": "http://clickhouse:9301", "username": "system", "password": "default" + }, + "nats": { + "addresses": [ + "nats:4222" + ] } } \ No newline at end of file diff --git a/self-host/dev/docker-compose.yml b/self-host/dev/docker-compose.yml index 7511952a55..2234f54130 100644 --- a/self-host/dev/docker-compose.yml +++ b/self-host/dev/docker-compose.yml @@ -179,6 +179,7 @@ services: restart: unless-stopped environment: - RUST_LOG_ANSI_COLOR=1 + - RUST_LOG=universaldb::driver::postgres=debug - RIVET_OTEL_ENABLED=1 - RIVET_OTEL_SAMPLER_RATIO=1 - RIVET_OTEL_GRPC_ENDPOINT=http://otel-collector:4317