diff --git a/Cargo.lock b/Cargo.lock index 12ad23aba3..868daf76fb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3938,6 +3938,7 @@ dependencies = [ "futures-util", "gasoline", "lazy_static", + "moka", "namespace", "nix 0.30.1", "portpicker", @@ -4010,6 +4011,7 @@ dependencies = [ "rivet-pools", "rivet-runtime", "rivet-types", + "rivet-util", "rusqlite", "scc", "serde", @@ -4147,6 +4149,7 @@ dependencies = [ "rivet-runner-protocol", "rivet-runtime", "rivet-types", + "rivet-util", "scc", "serde", "serde_bare", @@ -5382,6 +5385,7 @@ name = "rivet-depot-protocol" version = "2.3.7" dependencies = [ "anyhow", + "rivet-util", "rivet-vbare-compiler", "serde", "serde_bare", @@ -5516,7 +5520,7 @@ dependencies = [ "anyhow", "hex", "rand 0.8.5", - "rivet-util-serde", + "rivet-util", "rivet-vbare-compiler", "serde", "serde_bare", @@ -5943,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" @@ -8251,12 +8266,14 @@ version = "2.3.7" dependencies = [ "anyhow", "async-trait", + "base64 0.22.1", "deadpool-postgres", "foundationdb-tuple", "futures-util", "hex", "lazy_static", "rand 0.8.5", + "rivet-async-nats", "rivet-config", "rivet-env", "rivet-metrics", @@ -8264,17 +8281,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]] @@ -8283,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", @@ -8294,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/Cargo.toml b/Cargo.toml index 3f544252f3..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 @@ -672,6 +676,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/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 ead08e9e6f..9147b7a7d1 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": [ @@ -855,6 +823,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 +972,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 +1061,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": [ @@ -1222,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, @@ -1276,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/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/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/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/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 3ee2aace1d..5ee03c6d2a 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) @@ -196,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/pegboard.rs b/engine/packages/config/src/config/pegboard.rs index 616af06a36..f0b660401e 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. @@ -143,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. @@ -162,6 +176,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 +391,30 @@ 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 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) + } + + 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/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/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/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/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/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-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..4c7fcf5e50 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, @@ -138,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/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/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/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..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 = serde_json::from_str::(guard.0.get())?; + let value = rivet_util::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/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/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/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/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/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..922225c282 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,30 @@ 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 +1489,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-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/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/metrics.rs b/engine/packages/pegboard-gateway2/src/metrics.rs index e9c7a4acfe..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.", @@ -66,42 +60,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 ff870d330d..1d1b7a95cd 100644 --- a/engine/packages/pegboard-gateway2/src/shared_state.rs +++ b/engine/packages/pegboard-gateway2/src/shared_state.rs @@ -160,11 +160,10 @@ 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(); - 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(); @@ -195,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 @@ -661,7 +642,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..27dec6914c 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,9 +24,27 @@ pub async fn task( ) -> Result { let mut ws_rx = ws_rx.lock().await; + // Leaky bucket rate limit on consuming ws messages + let mut rate_limit = rivet_util::throttle::RateLimiter::new( + rivet_util::throttle::RateLimitMethod::LeakyBucket { + requests: ctx + .config() + .pegboard() + .gateway_websocket_rate_limit_requests(), + drip_rate: Duration::from_millis( + ctx.config() + .pegboard() + .gateway_websocket_rate_limit_drip_rate_ms(), + ), + }, + ); + loop { 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-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/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/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/pegboard/src/ops/actor/create.rs b/engine/packages/pegboard/src/ops/actor/create.rs index c21e878e80..b2695c16d9 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_limit = 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_limit.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/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/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/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/service-manager/src/lib.rs b/engine/packages/service-manager/src/lib.rs index 6c9f198d78..0b84fbae82 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(); @@ -395,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; } @@ -403,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/test-deps-docker/src/database.rs b/engine/packages/test-deps-docker/src/database.rs index 11532ebf25..2ddd0ee916 100644 --- a/engine/packages/test-deps-docker/src/database.rs +++ b/engine/packages/test-deps-docker/src/database.rs @@ -58,10 +58,11 @@ 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 { - 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/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 1fdc9b1c99..7f6b323e62 100644 --- a/engine/packages/universaldb/Cargo.toml +++ b/engine/packages/universaldb/Cargo.toml @@ -8,7 +8,9 @@ edition.workspace = true [dependencies] anyhow.workspace = true +async-nats.workspace = true async-trait.workspace = true +base64.workspace = true deadpool-postgres.workspace = true foundationdb-tuple.workspace = true futures-util.workspace = true @@ -18,20 +20,25 @@ 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, features = ["rt"] } tokio.workspace = true tracing.workspace = true url.workspace = true uuid.workspace = true +vbare.workspace = true [dev-dependencies] 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/conflict_tracker.rs b/engine/packages/universaldb/src/conflict_tracker.rs new file mode 100644 index 0000000000..a257f14ecc --- /dev/null +++ b/engine/packages/universaldb/src/conflict_tracker.rs @@ -0,0 +1,129 @@ +use std::{ + collections::BTreeMap, + ops::Bound, + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::{Duration, Instant}, +}; + +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. +const TXN_CONFLICT_TTL: Duration = Duration::from_secs(10); + +#[derive(Debug)] +struct PreviousTransaction { + insert_instant: Instant, + start_version: u64, + 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 { + // 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(BTreeMap::new())), + global_version: Arc::new(AtomicU64::new(0)), + } + } + + /// 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 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( + &self, + txn1_start_version: u64, + txn1_commit_version: u64, + txn1_conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, + ) -> bool { + let mut txns = self.txns.lock().await; + + // 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(); + } + + // 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 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 + if cr1_start < cr2_end && cr2_start < cr1_end && cr1_type != cr2_type { + tracing::debug!( + cr1_start=%hex::encode(cr1_start), + cr1_end=%hex::encode(cr1_end), + ?cr1_type, + cr2_start=%hex::encode(cr2_start), + cr2_end=%hex::encode(cr2_end), + ?cr2_type, + txn1_start_version, + txn1_commit_version, + txn2_start_version = txn2.start_version, + txn2_commit_version = %txn2_commit_version, + "transaction conflict detected" + ); + return true; + } + } + } + } + } + + // If no conflicts were detected, save txn data + 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_commit_version: u64) { + let mut txns = self.txns.lock().await; + txns.remove(&txn_commit_version); + } +} 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/codec.rs b/engine/packages/universaldb/src/driver/postgres/codec.rs new file mode 100644 index 0000000000..2e7b975ef0 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/codec.rs @@ -0,0 +1,216 @@ +use anyhow::Result; +use rivet_universaldb_commit::{self as proto, versioned}; +use vbare::OwnedVersionedData; + +use crate::{ + options::{ConflictRangeType, MutationType}, + tx_ops::Operation, +}; + +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 +/// 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], + client_node_id: &[u8], + client_seq: u64, +) -> 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(), + client_node_id: client_node_id.to_vec(), + client_seq, + }; + + versioned::CommitRequest::wrap_latest(request) + .serialize_with_embedded_version(proto::PROTOCOL_VERSION) +} + +/// 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)?; + + 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 { + 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, + 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..21d2335913 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/commit.rs @@ -0,0 +1,179 @@ +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result}; +use tokio::sync::oneshot; + +use crate::{error::DatabaseError, options::ConflictRangeType, tx_ops::Operation}; + +use super::{ + codec, + 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); +/// 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. +/// +/// `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(()); + } + + 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 + } + } +} + +/// 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), + }; + + if commit_tx.send(job).await.is_err() { + // The leader drain loop is gone (driver shutting down). Retryable. + return Err(DatabaseError::NotCommitted.into()); + } + + 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()), + } +} + +/// 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, + read_version: i64, + operations: Vec, + conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, +) -> Result<()> { + let Transport::MultiNode(nats) = &shared.transport else { + unreachable!("submit_nats requires the multi-node transport"); + }; + + // 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")?; + + 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(()); + } + 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!( + ?err, + client_seq, + attempt, + "udb commit request errored; resending" + ); + } + Err(_) => { + tracing::debug!( + client_seq, + attempt, + "udb commit request timed out; resending" + ); + } + } + + tokio::time::sleep(RESEND_BACKOFF).await; + } + + 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()) +} + +/// 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; + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/database.rs b/engine/packages/universaldb/src/driver/postgres/database.rs index 4d50d08ab3..dc1c42d598 100644 --- a/engine/packages/universaldb/src/driver/postgres/database.rs +++ b/engine/packages/universaldb/src/driver/postgres/database.rs @@ -10,9 +10,10 @@ 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; use crate::{ RetryableTransaction, Transaction, @@ -22,14 +23,27 @@ use crate::{ utils::{MaybeCommitted, calculate_tx_retry_backoff}, }; -use super::transaction::PostgresTransactionDriver; +use super::{ + nats::{self, NatsConfig, NatsTransport, Subjects}, + resolver::{self, ResolverInput}, + shared::PostgresShared, + transaction::PostgresTransactionDriver, + transport::{COMMIT_QUEUE_BOUND, Transport}, +}; -const GC_INTERVAL: Duration = Duration::from_secs(5); +const GC_INTERVAL: Duration = Duration::from_secs(30); +/// 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)] @@ -40,18 +54,20 @@ 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, } } } pub struct PostgresDatabaseDriver { - pool: Pool, + shared: Arc, max_retries: AtomicI32, + resolver_handle: JoinHandle<()>, gc_handle: JoinHandle<()>, } @@ -63,7 +79,74 @@ 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). Names this node's NATS commit subject and is the + // dedup `client_node_id`. + let node_id = Uuid::new_v4().simple().to_string(); + + 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(10), + resolver_handle, + 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 +157,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 +176,84 @@ 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 $$", - &[], - ) - .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, + ); + + 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_applied ( + client_node_id BYTEA NOT NULL, + client_seq 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 <> - ) - )", - &[], + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (client_node_id, client_seq) + );", ) .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 { + /// 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); 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 dedup 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_applied + WHERE created_at < now() - ($1::bigint * interval '1 second')", + &[&DEDUP_ROW_MAX_AGE_SECS], ) .await { - tracing::error!(?err, "failed postgres gc task"); + tracing::error!(?err, "failed postgres dedup 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>( @@ -291,10 +308,24 @@ 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 { 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/mod.rs b/engine/packages/universaldb/src/driver/postgres/mod.rs index 1c24f9cb94..20b80765c0 100644 --- a/engine/packages/universaldb/src/driver/postgres/mod.rs +++ b/engine/packages/universaldb/src/driver/postgres/mod.rs @@ -1,5 +1,12 @@ +mod codec; +mod commit; mod database; +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/apply.rs b/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs new file mode 100644 index 0000000000..278bb70fe7 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs @@ -0,0 +1,236 @@ +use std::collections::{BTreeMap, HashMap}; + +use anyhow::{Context, Result}; + +use crate::{ + atomic::apply_atomic_op, options::MutationType, tuple::Versionstamp, tx_ops::Operation, + versionstamp::substitute_raw_versionstamp, +}; + +/// 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), + } + } + 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, + } => fold_atomic( + &mut overlay, + key, + param, + op_type, + winner.commit_version, + &mut versionstamp_counter, + )?, + } + } + } + + Ok(overlay.into_write_set()) +} + +fn fold_atomic( + overlay: &mut Overlay<'_>, + 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")?; + 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")?; + overlay.set(key, value); + } + // 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 + | MutationType::Or + | MutationType::BitOr + | MutationType::Xor + | MutationType::BitXor + | MutationType::AppendIfFits + | MutationType::Max + | MutationType::Min + | MutationType::ByteMin + | MutationType::ByteMax + | MutationType::CompareAndClear => { + 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), + } + } + } + + 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) +/// 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) +} + +// 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/lease.rs b/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs new file mode 100644 index 0000000000..93f7f7884e --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs @@ -0,0 +1,106 @@ +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) +} + +/// 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 { + 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..3a779f838f --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs @@ -0,0 +1,659 @@ +mod apply; +mod lease; + +use std::{ + collections::HashMap, + sync::Arc, + time::{Duration, Instant}, +}; + +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::{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: usize = 256; + +/// How often a leader renews its lease. Must be comfortably under `LEASE_TTL_SECS`. +const RENEW_INTERVAL: Duration = Duration::from_secs(3); + +/// How long a candidate waits before retrying election when another node holds the lease. +const ELECTION_RETRY: Duration = Duration::from_secs(2); + +/// 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, +} + +/// 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 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"); + } + 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).await; + } + } + } +} + +/// Spawn the leader's NATS commit subscriber for multi-node. Returns `None` in single-node (no NATS). +fn spawn_commit_subscriber( + shared: &Arc, + 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 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 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"), + } +} + +/// Leader entry point: publish our lease, compute the cold-window floor, then run renewal and +/// 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. + let recovery_version = recovery_floor(shared).await?; + let recovery_deadline = Instant::now() + TXN_TIMEOUT; + + // 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" + ); + + let tracker = TransactionConflictTracker::new(); + let mut renew = AbortOnDropHandle::new(tokio::spawn(renew_loop(shared.clone(), epoch))); + + let drain = drain_loop( + shared, + epoch, + &tracker, + recovery_version, + recovery_deadline, + rx, + ); + tokio::pin!(drain); + + // 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, + } +} + +/// 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. +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 (expires_at = now + TTL), so consume the immediate first tick and + // renew after one interval. + interval.tick().await; + + let mut last_renew = Instant::now(); + + loop { + interval.tick().await; + + match lease::renew(&shared.pool, &shared.node_id, epoch).await { + Ok(true) => last_renew = Instant::now(), + Ok(false) => { + tracing::warn!( + epoch, + "udb leader lost lease on renew (epoch bumped); stepping down" + ); + return Ok(()); + } + Err(err) => { + 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"); + } + } + } +} + +/// 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, + epoch: i64, + tracker: &TransactionConflictTracker, + recovery_version: u64, + recovery_deadline: Instant, + rx: &mut mpsc::Receiver, +) -> Result<()> { + 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, + epoch, + tracker, + recovery_version, + recovery_deadline, + batch, + ) + .await? + { + BatchOutcome::LostLease => { + tracing::warn!(epoch, "udb leader stepping down: epoch fenced during apply"); + return Ok(()); + } + BatchOutcome::Processed => continue, + } + } +} + +/// 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 { + Processed, + LostLease, +} + +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() + .await + .context("failed to get connection for drain batch")?; + let txn = conn + .build_transaction() + .start() + .await + .context("failed to start drain batch txn")?; + + // 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); + } + } + + // 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 cold_window = Instant::now() < recovery_deadline; + 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 committed_count = 0u32; + let mut conflict_count = 0u32; + let mut cold_reject_count = 0u32; + + 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); + + let cold_rejected = cold_window && start_version < recovery_version; + let conflicted = if cold_rejected { + cold_reject_count += 1; + true + } else { + tracker + .check_and_insert(start_version, commit_version.max(0) as u64, conflict_ranges) + .await + }; + + if conflicted { + if !cold_rejected { + conflict_count += 1; + } + outcomes[i] = Some(CommitOutcome::Conflict); + } else { + committed_count += 1; + 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: std::mem::take(&mut job.operations), + }); + } + } + + // 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() + } 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")?; + + 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: 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) + WHERE key >= r.b AND key < r.e", + &[&range_begins, &range_ends], + ) + .await + .context("failed to clear ranges")?; + } + + // 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 ( + 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 + ), 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) + WHERE id = $8 AND epoch = $9 + RETURNING durable_version", + &[ + &point_deletes, + &upsert_keys, + &upsert_values, + &winner_dedup_nids, + &winner_dedup_seqs, + &winner_dedup_cvs, + &max_winner_cv, + &LEASE_ID, + &epoch, + ], + ) + .await + .context("failed to apply batch and 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")?; + + // 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); + + 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"), + } + } + + // 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::debug!( + 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) +} 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..5bb98fcbf2 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/shared.rs @@ -0,0 +1,207 @@ +use std::{ + sync::{ + Arc, + atomic::{AtomicI64, AtomicU64, Ordering}, + }, + time::Duration, +}; + +use deadpool_postgres::Pool; +use futures_util::StreamExt; +use tokio::sync::{Notify, watch}; + +use super::transport::Transport; + +/// The singleton row id of `udb_lease`. +pub const LEASE_ID: i32 = 1; + +/// 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); + +/// 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 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, in multi-node mode, a candidate leader. +pub struct PostgresShared { + pub pool: Pool, + /// Unique per-process id. Names this node's commit subject and is the dedup `client_node_id`. + pub node_id: String, + /// 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, transport: Transport) -> Arc { + let (lease_tx, lease_rx) = watch::channel(None); + let shared = Arc::new(Self { + pool, + node_id, + transport, + durable_version: AtomicI64::new(0), + watermark_notify: Notify::new(), + commit_seq: AtomicU64::new(0), + lease_tx, + lease_rx, + }); + + // 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); + 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 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::debug!( + 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)); + } + + /// 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 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! { + 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") + } + } + } + // The subscription ended (client closed). Fall back to lease polling only. + None => return shared.lease_poll_only().await, + } + } + _ = interval.tick() => { + shared.refresh_lease_row().await; + } + } + } + } + + /// 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, + 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/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 ec0175ee5a..718ced8836 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, @@ -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/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_conflict_tracker.rs b/engine/packages/universaldb/src/driver/rocksdb/transaction_conflict_tracker.rs deleted file mode 100644 index 370240760a..0000000000 --- a/engine/packages/universaldb/src/driver/rocksdb/transaction_conflict_tracker.rs +++ /dev/null @@ -1,109 +0,0 @@ -use std::{ - sync::{ - Arc, - atomic::{AtomicU64, Ordering}, - }, - time::{Duration, Instant}, -}; - -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 -const TXN_CONFLICT_TTL: Duration = Duration::from_secs(10); - -#[derive(Debug)] -struct PreviousTransaction { - insert_instant: Instant, - start_version: u64, - commit_version: u64, - conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, -} - -#[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>>, - global_version: Arc, -} - -impl TransactionConflictTracker { - pub fn new() -> Self { - TransactionConflictTracker { - txns: Arc::new(Mutex::new(Vec::new())), - global_version: Arc::new(AtomicU64::new(0)), - } - } - - /// Each number returned is unique. - pub fn next_global_version(&self) -> u64 { - self.global_version.fetch_add(1, Ordering::SeqCst) - } - - pub async fn check_and_insert( - &self, - txn1_start_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); - - for txn2 in &*txns { - // Check txn versions overlap (intersection or encapsulation) - if txn1_start_version < txn2.commit_version && 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 - if cr1_start < cr2_end && cr2_start < cr1_end && cr1_type != cr2_type { - tracing::debug!( - cr1_start=%hex::encode(cr1_start), - cr1_end=%hex::encode(cr1_end), - ?cr1_type, - cr2_start=%hex::encode(cr2_start), - cr2_end=%hex::encode(cr2_end), - ?cr2_type, - txn1_start_version, - txn1_commit_version, - txn2_start_version = txn2.start_version, - txn2_commit_version = txn2.commit_version, - "transaction conflict detected" - ); - return true; - } - } - } - } - } - - // 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, - }); - - false - } - - pub async fn remove(&self, txn_start_version: u64) { - let mut txns = self.txns.lock().await; - - if let Some(i) = txns - .iter() - .enumerate() - .find_map(|(i, txn)| (txn.start_version == txn_start_version).then_some(i)) - { - txns.remove(i); - } - } -} diff --git a/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs b/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs index 69835eb6a7..662a1a1f39 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()); @@ -425,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/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/failover.rs b/engine/packages/universaldb/tests/failover.rs new file mode 100644 index 0000000000..6293161760 --- /dev/null +++ b/engine/packages/universaldb/tests/failover.rs @@ -0,0 +1,317 @@ +use std::{sync::Arc, time::Duration}; + +use rivet_test_deps_docker::{TestDatabase, TestPubSub}; +use tokio_postgres::NoTls; +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"; + +/// 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)) +} + +/// 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 (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, &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, &nats_config).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); +} + +/// 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 (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, &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, &nats_config).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); +} 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/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(); +} 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/mod.rs b/engine/packages/universalpubsub/src/driver/postgres/mod.rs deleted file mode 100644 index 3ef8f02beb..0000000000 --- a/engine/packages/universalpubsub/src/driver/postgres/mod.rs +++ /dev/null @@ -1,859 +0,0 @@ -use anyhow::{Context, Result, anyhow}; -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::backoff::Backoff; -use scc::HashMap; -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; - -#[derive(Clone)] -struct Subscription { - // Channel to send messages to this subscription - tx: broadcast::Sender>, -} - -impl Subscription { - fn new(tx: broadcast::Sender>) -> Self { - Self { tx } - } -} - -/// > 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; - -/// Base64 encoding ratio -const BYTES_PER_BLOCK: usize = 3; -const CHARS_PER_BLOCK: usize = 4; - -/// 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; - -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 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; - -#[derive(Clone)] -pub struct PostgresDriver { - pool: Arc, - client: Arc>>, - subscriptions: Arc>, - /// Wakeup channels for queue subscriptions, keyed by queue channel name. - queue_subscriptions: 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 subscriptions: Arc> = Arc::new(HashMap::new()); - let queue_subscriptions: Arc> = Arc::new(HashMap::new()); - let client: Arc>> = Arc::new(Mutex::new(None)); - - // 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(), - subscriptions.clone(), - queue_subscriptions.clone(), - client.clone(), - ready_tx, - ssl_root_cert_path.clone(), - ssl_client_cert_path.clone(), - ssl_client_key_path.clone(), - )); - - let driver = Self { - pool: Arc::new(pool), - client, - subscriptions, - queue_subscriptions, - 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 - { - let conn = driver - .pool - .get() - .await - .context("failed to get connection for queue table creation")?; - conn.batch_execute( - "CREATE TABLE IF NOT EXISTS ups_queue_subs ( \ - id TEXT PRIMARY KEY, \ - subject_hash TEXT NOT NULL, \ - queue_hash TEXT NOT NULL, \ - heartbeat_at TIMESTAMPTZ NOT NULL DEFAULT NOW() \ - ); \ - 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 ( \ - 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 queue tables")?; - tracing::debug!("queue tables ready"); - } - - // 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, - subscriptions: Arc>, - queue_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 subscriptions_clone = subscriptions.clone(); - let queue_subscriptions_clone = queue_subscriptions.clone(); - let poll_handle = tokio::spawn(async move { - Self::poll_connection(conn, subscriptions_clone, queue_subscriptions_clone) - .await; - }); - - // Get regular channels to re-subscribe to - let mut channels = Vec::new(); - 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 { - tracing::debug!( - regular_channels = channels.len(), - queue_channels = queue_channels.len(), - "re-subscribing to channels after reconnection" - ); - } - - for channel in channels.iter().chain(queue_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, - subscriptions: Arc>, - queue_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()); - } else { - tracing::warn!(channel = %note.channel(), "received notification for unknown channel"); - } - } - 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 - // Hash the subject to ensure it fits - 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()) - } - - /// 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); - - 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], - ) - .await - .context("failed to query active queue subs")?; - - for row in rows { - let queue_hash: String = row.get(0); - let channel = self.queue_channel(&subject_hash, &queue_hash); - - 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")?; - - conn.execute(&format!("NOTIFY \"{}\"", channel), &[]) - .await - .context("failed to notify queue channel")?; - } - - Ok(()) - } - - fn spawn_subscription_cleanup_task( - &self, - subject_hash: 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 \"{}\"", subject_hash); - if let Err(err) = client.execute(sql.as_str(), &[]).await { - tracing::warn!(?err, %subject_hash, "failed to UNLISTEN channel"); - } else { - tracing::trace!(%subject_hash, "unlistened channel"); - } - } - driver.subscriptions.remove_async(&subject_hash).await; - metrics::POSTGRES_SUBSCRIPTION_COUNT.set(driver.subscriptions.len() as i64); - } - }); - - drop_guard - } - - fn spawn_queue_subscription_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 queue channel"); - } else { - tracing::trace!(%channel, "unlistened queue channel"); - } - } - driver.queue_subscriptions.remove_async(&channel).await; - } - }); - - drop_guard - } -} - -#[async_trait] -impl PubSubDriver for PostgresDriver { - async fn subscribe( - &self, - 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 drop_guard = self.spawn_subscription_cleanup_task(hashed.clone(), tx.clone()); - (rx, drop_guard) - } - }; - - Ok(Box::new(PostgresSubscriber { - subject: subject.to_string(), - rx: Some(rx), - _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 channel = self.queue_channel(&subject_hash, &queue_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, subject_hash, queue_hash) VALUES ($1, $2, $3)", - &[&sub_id, &subject_hash, &queue_hash], - ) - .await - .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) - } - }; - - // 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, - queue_hash, - sub_id, - pool: self.pool.clone(), - rx: Some(rx), - _drop_guard: drop_guard, - _heartbeat_token: heartbeat_token, - })) - } - - async fn publish( - &self, - subject: &str, - 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?; - } - - // Retry getting a connection from the pool with backoff in case the connection is - // currently disconnected - 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()); - } - } - } - Result::Err(e) => { - tracing::debug!(?e, "failed to get connection from pool, retrying"); - last_error = Some(e.into()); - } - } - - // Check if we should continue retrying - if !backoff.tick().await { - return Err( - last_error.unwrap_or_else(|| anyhow!("failed to publish after retries")) - ); - } - } - } - - async fn flush(&self) -> Result<()> { - Ok(()) - } - - fn max_message_size(&self) -> usize { - POSTGRES_MAX_MESSAGE_SIZE - } -} - -pub struct PostgresSubscriber { - subject: String, - rx: Option>>, - _drop_guard: tokio_util::sync::DropGuard, -} - -#[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 - } - } - } -} - -pub struct PostgresQueueSubscriber { - subject: String, - subject_hash: String, - queue_hash: String, - sub_id: String, - pool: Arc, - rx: Option>>, - _drop_guard: tokio_util::sync::DropGuard, - _heartbeat_token: tokio_util::sync::CancellationToken, -} - -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 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 - } - Err(tokio::sync::broadcast::error::RecvError::Closed) => { - return Ok(DriverOutput::Unsubscribed); - } - Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => { - // Notifications were dropped while lagged; loop back to claim in case - // messages are waiting - } - } - } - } -} - -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/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/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 35ace375d3..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; @@ -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,61 +77,22 @@ 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] -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, true, None, None, None) - .await - .unwrap(); - let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), true); - - test_all_inner(&pubsub, &docker).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, false, None, None, None) - .await - .unwrap(); - let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false); - - test_all_inner(&pubsub, &docker).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/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/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 d899e608c3..0c088b9df8 100644 --- a/engine/packages/util/src/lib.rs +++ b/engine/packages/util/src/lib.rs @@ -1,8 +1,9 @@ +use std::fmt::Display; + 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; @@ -12,10 +13,12 @@ pub mod future; pub mod geo; pub mod math; pub mod metric; +pub mod metrics; pub mod req; pub mod serde; pub mod size; pub mod sort; +pub mod throttle; pub mod timestamp; pub mod url; @@ -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..18a9823171 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 `rivet_util::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 `rivet_util::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; 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)); + } +} 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/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/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/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..804f533a64 --- /dev/null +++ b/engine/sdks/rust/universaldb-commit/src/versioned.rs @@ -0,0 +1,100 @@ +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), + } + } +} + +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/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), } } diff --git a/engine/sdks/schemas/universaldb-commit/v1.bare b/engine/sdks/schemas/universaldb-commit/v1.bare new file mode 100644 index 0000000000..c85d43f447 --- /dev/null +++ b/engine/sdks/schemas/universaldb-commit/v1.bare @@ -0,0 +1,98 @@ +# Commit wire format for the multi-node Postgres leader-resolver UDB driver. +# +# 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 + 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 + # 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/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/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( 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" + ); + } +} 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; + } + } } } } 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 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/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 ca2f6b1fc8..703559ced7 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; @@ -118,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, @@ -152,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", @@ -171,7 +185,11 @@ export function generateDockerCompose(context: TemplateContext) { ); services[postgresServiceName] = { restart: "unless-stopped", - image: "postgres:17-alpine", + image: "postgres:18-alpine", + // 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", "POSTGRES_PASSWORD=postgres", @@ -179,7 +197,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 +212,7 @@ export function generateDockerCompose(context: TemplateContext) { services[shellServiceName] = { build: { - context: "../../..", + context: "../..", dockerfile: "docker/engine/Dockerfile", target: "engine-full", args: { @@ -207,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: [ @@ -275,7 +295,7 @@ export function generateDockerCompose(context: TemplateContext) { services[serviceName] = { build: { - context: "../../..", + context: "../..", dockerfile: "docker/engine/Dockerfile", target: "engine-full", args: { @@ -286,6 +306,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`, @@ -293,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", }, @@ -335,30 +358,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..1c5b8247da 100644 --- a/self-host/compose/template/src/services/edge/rivet-engine.ts +++ b/self-host/compose/template/src/services/edge/rivet-engine.ts @@ -35,35 +35,37 @@ export function generateDatacenterRivetEngine( datacenters, }; - // Config structure matching Rust schema in packages/common/config/src/config/mod.rs - const config = { + // Config structure matching Rust schema in engine/packages/config/src/config/mod.rs. + // Values that match the engine's defaults are omitted. + const config: Record = { 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, }, }; + // 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/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 82% rename from self-host/compose/dev-host/docker-compose.yml rename to self-host/dev-host/docker-compose.yml index 3b54702d63..fc44c89a26 100644 --- a/self-host/compose/dev-host/docker-compose.yml +++ b/self-host/dev-host/docker-compose.yml @@ -72,14 +72,18 @@ services: network_mode: host postgres: restart: unless-stopped - image: postgres:17-alpine + 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: @@ -153,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 @@ -179,20 +184,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/1/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/1/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc index ab9cfa618a..29ca3ee404 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/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,15 @@ "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" + }, + "nats": { + "addresses": [ + "nats-dc-a:4222" + ] } } \ No newline at end of file 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/1/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/1/config.jsonc index ab9cfa618a..29ca3ee404 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/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,15 @@ "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" + }, + "nats": { + "addresses": [ + "nats-dc-a:4222" + ] } } \ 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/2/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/2/config.jsonc index ab9cfa618a..29ca3ee404 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/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,15 @@ "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" + }, + "nats": { + "addresses": [ + "nats-dc-a:4222" + ] } } \ 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..0fe743310a 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,15 @@ "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" + }, + "nats": { + "addresses": [ + "nats-dc-b:4222" + ] } } \ 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/1/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/1/config.jsonc index 2a34bc83ff..0fe743310a 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/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,15 @@ "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" + }, + "nats": { + "addresses": [ + "nats-dc-b:4222" + ] } } \ 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/2/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/2/config.jsonc index 2a34bc83ff..0fe743310a 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/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,15 @@ "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" + }, + "nats": { + "addresses": [ + "nats-dc-b:4222" + ] } } \ 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..5a639b2fd4 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,15 @@ "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" + }, + "nats": { + "addresses": [ + "nats-dc-c:4222" + ] } } \ 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..5a639b2fd4 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,15 @@ "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" + }, + "nats": { + "addresses": [ + "nats-dc-c:4222" + ] } } \ 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..5a639b2fd4 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,15 @@ "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" + }, + "nats": { + "addresses": [ + "nats-dc-c:4222" + ] } } \ 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 77% rename from self-host/compose/dev-multidc-multinode/docker-compose.yml rename to self-host/dev-multidc-multinode/docker-compose.yml index 75c500850b..81b39d39eb 100644 --- a/self-host/compose/dev-multidc-multinode/docker-compose.yml +++ b/self-host/dev-multidc-multinode/docker-compose.yml @@ -83,9 +83,30 @@ 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:17-alpine + image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -93,7 +114,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 +128,7 @@ services: - rivet-network-dc-a rivet-shell-dc-a: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -118,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: @@ -174,7 +197,7 @@ services: - '4317:4317' rivet-engine-dc-a-0: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -183,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: @@ -217,7 +243,7 @@ services: start_period: 30s rivet-engine-dc-a-1: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -226,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: @@ -258,7 +287,7 @@ services: start_period: 30s rivet-engine-dc-a-2: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -267,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: @@ -299,63 +331,97 @@ 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 + 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:17-alpine + image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -363,7 +429,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 +441,7 @@ services: - rivet-network-dc-b rivet-shell-dc-b: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -386,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: @@ -440,7 +508,7 @@ services: - rivet-network-dc-b-to-core rivet-engine-dc-b-0: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -449,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: @@ -481,7 +552,7 @@ services: start_period: 30s rivet-engine-dc-b-1: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -490,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: @@ -522,7 +596,7 @@ services: start_period: 30s rivet-engine-dc-b-2: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -531,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: @@ -563,61 +640,95 @@ 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 + 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:17-alpine + image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -625,7 +736,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 +748,7 @@ services: - rivet-network-dc-c rivet-shell-dc-c: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -648,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: @@ -702,7 +815,7 @@ services: - rivet-network-dc-c-to-core rivet-engine-dc-c-0: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -711,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: @@ -743,7 +859,7 @@ services: start_period: 30s rivet-engine-dc-c-1: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -752,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: @@ -784,7 +903,7 @@ services: start_period: 30s rivet-engine-dc-c-2: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -793,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: @@ -825,58 +947,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 85% rename from self-host/compose/dev-multidc/docker-compose.yml rename to self-host/dev-multidc/docker-compose.yml index f20db41ff3..77938647f1 100644 --- a/self-host/compose/dev-multidc/docker-compose.yml +++ b/self-host/dev-multidc/docker-compose.yml @@ -85,7 +85,11 @@ services: condition: service_healthy postgres-dc-a: restart: unless-stopped - image: postgres:17-alpine + 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: @@ -182,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 @@ -215,27 +220,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:17-alpine + image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -243,7 +257,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 +269,7 @@ services: - rivet-network-dc-b rivet-shell-dc-b: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -319,7 +333,7 @@ services: - rivet-network-dc-b-to-core rivet-engine-dc-b: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -328,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 @@ -359,25 +374,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:17-alpine + image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -385,7 +409,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 +421,7 @@ services: - rivet-network-dc-c rivet-shell-dc-c: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -461,7 +485,7 @@ services: - rivet-network-dc-c-to-core rivet-engine-dc-c: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -470,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 @@ -501,22 +526,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 76% rename from self-host/compose/dev-multinode/docker-compose.yml rename to self-host/dev-multinode/docker-compose.yml index 381ece396d..349b19edbc 100644 --- a/self-host/compose/dev-multinode/docker-compose.yml +++ b/self-host/dev-multinode/docker-compose.yml @@ -79,16 +79,37 @@ 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:17-alpine + 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 +123,7 @@ services: - rivet-network rivet-shell: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -113,6 +134,8 @@ services: command: infinity stop_grace_period: 0s depends_on: + nats: + condition: service_healthy postgres: condition: service_healthy volumes: @@ -166,7 +189,7 @@ services: - '4317:4317' rivet-engine-0: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -175,11 +198,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:4317 stop_grace_period: 0s depends_on: + nats: + condition: service_healthy vector-client: condition: service_started otel-collector: @@ -206,7 +232,7 @@ services: start_period: 30s rivet-engine-1: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -215,11 +241,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:4317 stop_grace_period: 0s depends_on: + nats: + condition: service_healthy vector-client: condition: service_started otel-collector: @@ -244,7 +273,7 @@ services: start_period: 30s rivet-engine-2: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -253,11 +282,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:4317 stop_grace_period: 0s depends_on: + nats: + condition: service_healthy vector-client: condition: service_started otel-collector: @@ -282,58 +314,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 80% 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..13e9996d6e 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,15 @@ "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" + }, + "nats": { + "addresses": [ + "nats:4222" + ] } } \ 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 80% 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..13e9996d6e 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,15 @@ "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" + }, + "nats": { + "addresses": [ + "nats:4222" + ] } } \ 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 80% 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..13e9996d6e 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,15 @@ "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" + }, + "nats": { + "addresses": [ + "nats:4222" + ] } } \ 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 83% rename from self-host/compose/dev/docker-compose.yml rename to self-host/dev/docker-compose.yml index 6c5dcc0ef2..2234f54130 100644 --- a/self-host/compose/dev/docker-compose.yml +++ b/self-host/dev/docker-compose.yml @@ -81,14 +81,18 @@ services: condition: service_healthy postgres: restart: unless-stopped - image: postgres:17-alpine + 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: @@ -175,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 @@ -206,22 +211,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 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/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/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 \ 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