diff --git a/Cargo.lock b/Cargo.lock index 12ad23aba3..7b0053c042 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,6 +8266,7 @@ version = "2.3.7" dependencies = [ "anyhow", "async-trait", + "base64 0.22.1", "deadpool-postgres", "foundationdb-tuple", "futures-util", @@ -8264,17 +8280,21 @@ dependencies = [ "rivet-postgres-util", "rivet-test-deps-docker", "rivet-tracing-utils", + "rivet-universaldb-commit", "rocksdb", + "scc", "serde", "tempfile", "thiserror 1.0.69", "tokio", "tokio-postgres", "tokio-postgres-rustls", + "tokio-util", "tracing", "tracing-subscriber", "url", "uuid", + "vbare", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 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/engine/artifacts/config-schema.json b/engine/artifacts/config-schema.json index ead08e9e6f..af2af30077 100644 --- a/engine/artifacts/config-schema.json +++ b/engine/artifacts/config-schema.json @@ -40,17 +40,6 @@ } ] }, - "api_public": { - "default": null, - "anyOf": [ - { - "$ref": "#/definitions/ApiPublic" - }, - { - "type": "null" - } - ] - }, "auth": { "default": null, "anyOf": [ @@ -214,27 +203,6 @@ }, "additionalProperties": false }, - "ApiPublic": { - "description": "Configuration for the public API service.", - "type": "object", - "properties": { - "respect_forwarded_for": { - "description": "Flag to respect the X-Forwarded-For header for client IP addresses.\n\nWill be ignored in favor of CF-Connecting-IP if DNS provider is configured as Cloudflare.", - "type": [ - "boolean", - "null" - ] - }, - "verbose_errors": { - "description": "Flag to enable verbose error reporting.", - "type": [ - "boolean", - "null" - ] - } - }, - "additionalProperties": false - }, "Auth": { "type": "object", "required": [ @@ -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": [ 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/mod.rs b/engine/packages/config/src/config/mod.rs index 3ee2aace1d..caa8bf196e 100644 --- a/engine/packages/config/src/config/mod.rs +++ b/engine/packages/config/src/config/mod.rs @@ -4,7 +4,6 @@ use serde::{Deserialize, Serialize}; use std::sync::LazyLock; pub mod api_peer; -pub mod api_public; pub mod auth; pub mod cache; pub mod clickhouse; @@ -21,7 +20,6 @@ pub mod telemetry; pub mod topology; pub use api_peer::*; -pub use api_public::*; pub use auth::*; pub use cache::*; pub use clickhouse::*; @@ -74,9 +72,6 @@ pub struct Root { #[serde(default)] pub guard: Option, - #[serde(default)] - pub api_public: Option, - #[serde(default)] pub api_peer: Option, @@ -122,7 +117,6 @@ impl Default for Root { Root { auth: None, guard: None, - api_public: None, api_peer: None, pegboard: None, logs: None, @@ -146,11 +140,6 @@ impl Root { self.guard.as_ref().unwrap_or(&DEFAULT) } - pub fn api_public(&self) -> &ApiPublic { - static DEFAULT: LazyLock = LazyLock::new(ApiPublic::default); - self.api_public.as_ref().unwrap_or(&DEFAULT) - } - pub fn api_peer(&self) -> &ApiPeer { static DEFAULT: LazyLock = LazyLock::new(ApiPeer::default); self.api_peer.as_ref().unwrap_or(&DEFAULT) diff --git a/engine/packages/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/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/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..955e4d7897 100644 --- a/engine/packages/test-deps-docker/src/database.rs +++ b/engine/packages/test-deps-docker/src/database.rs @@ -61,7 +61,7 @@ impl TestDatabase { }); let docker_config = DockerRunConfig { - image: "postgres:17".to_string(), + image: "postgres:18".to_string(), container_name: container_name.clone(), port_mapping: (port, 5432), env_vars: vec![ diff --git a/engine/packages/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..57dba4ab33 100644 --- a/engine/packages/universaldb/Cargo.toml +++ b/engine/packages/universaldb/Cargo.toml @@ -9,6 +9,7 @@ edition.workspace = true [dependencies] anyhow.workspace = true async-trait.workspace = true +base64.workspace = true deadpool-postgres.workspace = true foundationdb-tuple.workspace = true futures-util.workspace = true @@ -18,20 +19,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 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/driver/rocksdb/transaction_conflict_tracker.rs b/engine/packages/universaldb/src/conflict_tracker.rs similarity index 69% rename from engine/packages/universaldb/src/driver/rocksdb/transaction_conflict_tracker.rs rename to engine/packages/universaldb/src/conflict_tracker.rs index 370240760a..127e2bf615 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/transaction_conflict_tracker.rs +++ b/engine/packages/universaldb/src/conflict_tracker.rs @@ -11,7 +11,7 @@ use tokio::sync::Mutex; use crate::options::ConflictRangeType; // Transactions cannot live longer than 5 seconds so we don't need to store transaction conflicts longer than -// that +// that. const TXN_CONFLICT_TTL: Duration = Duration::from_secs(10); #[derive(Debug)] @@ -22,6 +22,15 @@ struct PreviousTransaction { conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, } +/// In-process FoundationDB-style resolver. Holds the last `TXN_CONFLICT_TTL` of committed +/// transactions and rejects a committing transaction if any retained transaction has both an +/// overlapping version window and an overlapping conflict range of a differing type. +/// +/// Used by the rocksdb driver (single process) and by the postgres leader-resolver. The two +/// differ only in where the commit version comes from: rocksdb generates it from the in-process +/// `global_version` counter, while the postgres leader assigns it from the durable +/// `udb_version_seq` so it survives leader failover and matches the versionstamp. For that reason +/// `check_and_insert` takes the commit version from the caller instead of generating it. #[derive(Clone)] pub struct TransactionConflictTracker { // NOTE: We use a mutex because we need to lock reads across all active txns. This could be optimized to @@ -40,18 +49,23 @@ impl TransactionConflictTracker { } } - /// Each number returned is unique. + /// Each number returned is unique. Used by the in-process rocksdb driver to assign both start + /// and commit versions. The postgres leader does not use this; it assigns versions from the + /// durable Postgres sequence. pub fn next_global_version(&self) -> u64 { self.global_version.fetch_add(1, Ordering::SeqCst) } + /// Returns `true` on conflict (same polarity as the original rocksdb tracker). The caller + /// supplies `commit_version` (e.g. `nextval('udb_version_seq')` on the postgres leader, or + /// `next_global_version()` on rocksdb) so version assignment stays the caller's responsibility. pub async fn check_and_insert( &self, txn1_start_version: u64, + txn1_commit_version: u64, txn1_conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, ) -> bool { let mut txns = self.txns.lock().await; - let txn1_commit_version = self.next_global_version(); // Prune old entries txns.retain(|txn| txn.insert_instant.elapsed() < TXN_CONFLICT_TTL); diff --git a/engine/packages/universaldb/src/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..432ede2f34 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/codec.rs @@ -0,0 +1,167 @@ +use anyhow::Result; +use rivet_universaldb_commit::{self as proto, versioned}; +use vbare::OwnedVersionedData; + +use crate::{ + options::{ConflictRangeType, MutationType}, + tx_ops::Operation, +}; + +/// Decoded form of a `udb_commit_requests.payload` blob. +/// +/// `read_version` is intentionally omitted: it is also denormalized into the `read_version` column, +/// which is what the leader's drain reads, so decoding it here would be dead. +pub struct DecodedCommit { + pub conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, + pub operations: Vec, +} + +/// Encode a follower's commit request to the versioned BARE wire format with an embedded version +/// header so a leader running older or newer code can still decode it during a rolling deploy. +pub fn encode_commit_request( + read_version: u64, + conflict_ranges: &[(Vec, Vec, ConflictRangeType)], + operations: &[Operation], +) -> Result> { + let request = proto::CommitRequest { + read_version, + conflict_ranges: conflict_ranges + .iter() + .map(|(begin, end, kind)| proto::ConflictRange { + begin: begin.clone(), + end: end.clone(), + kind: conflict_range_type_to_proto(*kind), + }) + .collect(), + operations: operations.iter().map(operation_to_proto).collect(), + }; + + versioned::CommitRequest::wrap_latest(request) + .serialize_with_embedded_version(proto::PROTOCOL_VERSION) +} + +/// Decode a `udb_commit_requests.payload` blob produced by [`encode_commit_request`]. +pub fn decode_commit_request(payload: &[u8]) -> Result { + let request = versioned::CommitRequest::deserialize_with_embedded_version(payload)?; + + let conflict_ranges = request + .conflict_ranges + .into_iter() + .map(|range| { + ( + range.begin, + range.end, + conflict_range_type_from_proto(range.kind), + ) + }) + .collect(); + + let operations = request + .operations + .into_iter() + .map(operation_from_proto) + .collect(); + + Ok(DecodedCommit { + conflict_ranges, + operations, + }) +} + +fn conflict_range_type_to_proto(kind: ConflictRangeType) -> proto::ConflictRangeType { + match kind { + ConflictRangeType::Read => proto::ConflictRangeType::Read, + ConflictRangeType::Write => proto::ConflictRangeType::Write, + } +} + +fn conflict_range_type_from_proto(kind: proto::ConflictRangeType) -> ConflictRangeType { + match kind { + proto::ConflictRangeType::Read => ConflictRangeType::Read, + proto::ConflictRangeType::Write => ConflictRangeType::Write, + } +} + +fn operation_to_proto(op: &Operation) -> proto::Operation { + match op { + Operation::SetValue { key, value } => proto::Operation::SetValue(proto::SetValue { + key: key.clone(), + value: value.clone(), + }), + Operation::Clear { key } => proto::Operation::Clear(proto::Clear { key: key.clone() }), + Operation::ClearRange { begin, end } => proto::Operation::ClearRange(proto::ClearRange { + begin: begin.clone(), + end: end.clone(), + }), + Operation::AtomicOp { + key, + param, + op_type, + } => proto::Operation::AtomicOp(proto::AtomicOp { + key: key.clone(), + param: param.clone(), + op_type: mutation_type_to_proto(*op_type), + }), + } +} + +fn operation_from_proto(op: proto::Operation) -> Operation { + match op { + proto::Operation::SetValue(proto::SetValue { key, value }) => { + Operation::SetValue { key, value } + } + proto::Operation::Clear(proto::Clear { key }) => Operation::Clear { key }, + proto::Operation::ClearRange(proto::ClearRange { begin, end }) => { + Operation::ClearRange { begin, end } + } + proto::Operation::AtomicOp(proto::AtomicOp { + key, + param, + op_type, + }) => Operation::AtomicOp { + key, + param, + op_type: mutation_type_from_proto(op_type), + }, + } +} + +fn mutation_type_to_proto(op_type: MutationType) -> proto::MutationType { + match op_type { + MutationType::Add => proto::MutationType::Add, + MutationType::And => proto::MutationType::And, + MutationType::BitAnd => proto::MutationType::BitAnd, + MutationType::Or => proto::MutationType::Or, + MutationType::BitOr => proto::MutationType::BitOr, + MutationType::Xor => proto::MutationType::Xor, + MutationType::BitXor => proto::MutationType::BitXor, + MutationType::AppendIfFits => proto::MutationType::AppendIfFits, + MutationType::Max => proto::MutationType::Max, + MutationType::Min => proto::MutationType::Min, + MutationType::SetVersionstampedKey => proto::MutationType::SetVersionstampedKey, + MutationType::SetVersionstampedValue => proto::MutationType::SetVersionstampedValue, + MutationType::ByteMin => proto::MutationType::ByteMin, + MutationType::ByteMax => proto::MutationType::ByteMax, + MutationType::CompareAndClear => proto::MutationType::CompareAndClear, + } +} + +fn mutation_type_from_proto(op_type: proto::MutationType) -> MutationType { + match op_type { + proto::MutationType::Add => MutationType::Add, + proto::MutationType::And => MutationType::And, + proto::MutationType::BitAnd => MutationType::BitAnd, + proto::MutationType::Or => MutationType::Or, + proto::MutationType::BitOr => MutationType::BitOr, + proto::MutationType::Xor => MutationType::Xor, + proto::MutationType::BitXor => MutationType::BitXor, + proto::MutationType::AppendIfFits => MutationType::AppendIfFits, + proto::MutationType::Max => MutationType::Max, + proto::MutationType::Min => MutationType::Min, + proto::MutationType::SetVersionstampedKey => MutationType::SetVersionstampedKey, + proto::MutationType::SetVersionstampedValue => MutationType::SetVersionstampedValue, + proto::MutationType::ByteMin => MutationType::ByteMin, + proto::MutationType::ByteMax => MutationType::ByteMax, + proto::MutationType::CompareAndClear => MutationType::CompareAndClear, + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/commit.rs b/engine/packages/universaldb/src/driver/postgres/commit.rs new file mode 100644 index 0000000000..e61714dd0e --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/commit.rs @@ -0,0 +1,185 @@ +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result}; + +use crate::{error::DatabaseError, options::ConflictRangeType, tx_ops::Operation}; + +use super::{ + codec, + shared::{LeaseInfo, PostgresShared, commit_channel, reply_channel}, +}; + +/// How long to wait for a leader to be elected before giving up a submit as retryable. +const LEADER_WAIT_TIMEOUT: Duration = Duration::from_secs(5); +/// Backstop poll cadence while waiting for a commit result, in case a reply NOTIFY is missed. +const RESULT_POLL_INTERVAL: Duration = Duration::from_millis(250); + +/// Submit a follower transaction's commit to the leader and await the result. +/// +/// `read_version` is the watermark captured when this transaction opened its read snapshot. A pure +/// snapshot read-only transaction (no operations and no read conflict ranges) submits nothing. +pub async fn submit( + shared: &Arc, + read_version: i64, + operations: Vec, + conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, +) -> Result<()> { + // A transaction with no writes and no serializable read ranges has nothing to order or + // validate; it never needs the leader. + if operations.is_empty() + && conflict_ranges + .iter() + .all(|(_, _, kind)| matches!(kind, ConflictRangeType::Write)) + { + return Ok(()); + } + + let lease = wait_for_leader(shared).await?; + let payload = + codec::encode_commit_request(read_version.max(0) as u64, &conflict_ranges, &operations) + .context("failed to encode commit request")?; + let reply_channel = reply_channel(&shared.node_id); + + // Subscribe to our reply channel before inserting so we cannot miss the leader's NOTIFY. + let mut reply_rx = shared.listener.listen(&reply_channel).await; + + let conn = shared + .pool + .get() + .await + .context("failed to get connection for commit submit")?; + + let id: i64 = conn + .query_one( + "INSERT INTO udb_commit_requests (epoch, read_version, payload, reply_channel) + VALUES ($1, $2, $3, $4) + RETURNING id", + &[&lease.epoch, &read_version, &payload, &reply_channel], + ) + .await + .context("failed to enqueue commit request")? + .get(0); + + // Wake the leader's drain loop. + if let Err(err) = conn + .execute( + "SELECT pg_notify($1, $2)", + &[&commit_channel(&lease.leader_addr), &id.to_string()], + ) + .await + { + tracing::debug!( + ?err, + "failed to notify leader; relying on its poll backstop" + ); + } + + // Release the connection before waiting so a long wait does not pin a pool slot. The request + // row is durable, so await_result re-acquires a connection per poll. + drop(conn); + + await_result(shared, id, lease.epoch, &mut reply_rx).await +} + +/// Wait for a known leader, returning a retryable error if none is elected in time. +async fn wait_for_leader(shared: &Arc) -> Result { + let deadline = Instant::now() + LEADER_WAIT_TIMEOUT; + loop { + if let Some(lease) = shared.current_lease() { + return Ok(lease); + } + if Instant::now() >= deadline { + return Err(DatabaseError::NotCommitted.into()); + } + tokio::time::sleep(RESULT_POLL_INTERVAL).await; + } +} + +/// Poll the request row until it reaches a terminal status, woken by reply NOTIFYs with a polling +/// backstop. Bails as retryable if the leader epoch advances (our request is now orphaned and will +/// never be applied, so it is definitively not committed). +async fn await_result( + shared: &Arc, + id: i64, + submit_epoch: i64, + reply_rx: &mut tokio::sync::broadcast::Receiver, +) -> Result<()> { + loop { + // Re-acquire a connection per poll: the request row is durable, so a transient pool/query + // error just means we retry the poll rather than failing a possibly-applied commit. + match read_status(shared, id).await { + Ok(Some(Status::Committed)) => return Ok(()), + Ok(Some(Status::Conflict)) => return Err(DatabaseError::NotCommitted.into()), + Ok(Some(Status::Pending)) => {} + Ok(None) => { + // The row was GC'd before we observed a terminal status. Treat as not committed + // and let the retry loop resubmit. + return Err(DatabaseError::NotCommitted.into()); + } + Err(err) => { + tracing::debug!(?err, "transient error polling commit status, retrying"); + } + } + + // If a new leader took over, our old-epoch request will never be claimed. + if let Some(current) = shared.current_lease() { + if current.epoch != submit_epoch { + return Err(DatabaseError::NotCommitted.into()); + } + } + + tokio::select! { + res = reply_rx.recv() => { + match res { + Ok(_) | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + *reply_rx = shared + .listener + .listen(&reply_channel(&shared.node_id)) + .await; + } + } + } + _ = tokio::time::sleep(RESULT_POLL_INTERVAL) => {} + } + } +} + +enum Status { + Pending, + Committed, + Conflict, +} + +/// Read the current status of a commit request. `Ok(None)` means the row no longer exists. +async fn read_status(shared: &Arc, id: i64) -> Result> { + let conn = shared + .pool + .get() + .await + .context("failed to get connection for commit status poll")?; + + let row = conn + .query_opt( + "SELECT status FROM udb_commit_requests WHERE id = $1", + &[&id], + ) + .await + .context("failed to read commit request status")?; + + let Some(row) = row else { + return Ok(None); + }; + + let status: String = row.get(0); + let status = match status.as_str() { + "committed" => Status::Committed, + "conflict" => Status::Conflict, + // 'pending' or any in-flight state. + _ => Status::Pending, + }; + Ok(Some(status)) +} diff --git a/engine/packages/universaldb/src/driver/postgres/database.rs b/engine/packages/universaldb/src/driver/postgres/database.rs index 4d50d08ab3..3ff891011f 100644 --- a/engine/packages/universaldb/src/driver/postgres/database.rs +++ b/engine/packages/universaldb/src/driver/postgres/database.rs @@ -13,6 +13,7 @@ use rivet_postgres_util::build_tls_config; use tokio::task::JoinHandle; use tokio_postgres_rustls::MakeRustlsConnect; use url::Url; +use uuid::Uuid; use crate::{ RetryableTransaction, Transaction, @@ -22,9 +23,15 @@ use crate::{ utils::{MaybeCommitted, calculate_tx_retry_backoff}, }; -use super::transaction::PostgresTransactionDriver; +use super::{ + listener::PgListener, resolver, shared::PostgresShared, transaction::PostgresTransactionDriver, +}; -const GC_INTERVAL: Duration = Duration::from_secs(5); +const GC_INTERVAL: Duration = Duration::from_secs(30); +/// Terminal and orphaned commit-request rows older than this are garbage collected. Must be well +/// beyond the longest a follower could spend awaiting a result, so a result is never deleted before +/// it is observed. +const COMMIT_ROW_MAX_AGE_SECS: i64 = 60; #[derive(Clone, Debug)] pub struct PostgresConfig { @@ -50,8 +57,9 @@ impl PostgresConfig { } pub struct PostgresDatabaseDriver { - pool: Pool, + shared: Arc, max_retries: AtomicI32, + resolver_handle: JoinHandle<()>, gc_handle: JoinHandle<()>, } @@ -63,7 +71,61 @@ impl PostgresDatabaseDriver { "creating PostgresDatabaseDriver" ); - // Create deadpool config from connection string + let ssl_disabled = if let Ok(url) = Url::parse(&config.connection_string) { + url.query_pairs() + .any(|(k, v)| k == "sslmode" && v == "disable") + } else { + false + }; + + let pool = Self::build_pool(&config, ssl_disabled)?; + + // Initialize the schema (idempotent). + { + let conn = pool + .get() + .await + .context("failed to get connection from postgres pool")?; + Self::init_schema(&conn).await?; + } + + // Unique per-process node id (no hyphens) used to name this node's NOTIFY channels. Kept + // short so `udb_commit_` stays within Postgres's 63-byte identifier limit. + let node_id = Uuid::new_v4().simple().to_string(); + + let listener = PgListener::new( + config.connection_string.clone(), + ssl_disabled, + config + .ssl_config + .as_ref() + .and_then(|c| c.ssl_root_cert_path.clone()), + config + .ssl_config + .as_ref() + .and_then(|c| c.ssl_client_cert_path.clone()), + config + .ssl_config + .as_ref() + .and_then(|c| c.ssl_client_key_path.clone()), + ); + + let shared = PostgresShared::new(pool, node_id, listener); + + // Every node runs the resolver; only the elected leader drains the commit queue. + let resolver_handle = resolver::spawn(shared.clone()); + + let gc_handle = Self::spawn_gc(shared.clone()); + + Ok(PostgresDatabaseDriver { + shared, + max_retries: AtomicI32::new(100), + resolver_handle, + gc_handle, + }) + } + + 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 +136,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 +155,87 @@ impl PostgresDatabaseDriver { .as_ref() .and_then(|c| c.ssl_client_key_path.as_ref()), )?; - let tls = MakeRustlsConnect::new(tls_config); - pool_config - .create_pool(Some(Runtime::Tokio1), tls) - .context("failed to create postgres connection pool")? - }; - - tracing::debug!("Getting Postgres connection from pool"); - // Get a connection from the pool to create the table - let conn = pool - .get() - .await - .context("failed to get connection from postgres pool")?; - - // Enable btree gist - conn.execute("CREATE EXTENSION IF NOT EXISTS btree_gist", &[]) - .await - .context("failed to create btree_gist extension")?; - - conn.execute("CREATE UNLOGGED SEQUENCE IF NOT EXISTS global_version_seq START WITH 1 INCREMENT BY 1 MINVALUE 1", &[]) - .await - .context("failed to create global version sequence")?; + .create_pool(Some(Runtime::Tokio1), MakeRustlsConnect::new(tls_config)) + .context("failed to create postgres connection pool") + } + } - // Create the KV table if it doesn't exist - conn.execute( + async fn init_schema(conn: &deadpool_postgres::Client) -> Result<()> { + // Durable latest-value store. + conn.batch_execute( "CREATE TABLE IF NOT EXISTS kv ( key BYTEA PRIMARY KEY, value BYTEA NOT NULL - )", - &[], + ); + + CREATE TABLE IF NOT EXISTS udb_lease ( + id INT PRIMARY KEY DEFAULT 1 CHECK (id = 1), + epoch BIGINT NOT NULL, + leader_addr TEXT NOT NULL, + durable_version BIGINT NOT NULL DEFAULT 0, + expires_at TIMESTAMPTZ NOT NULL + ); + + CREATE SEQUENCE IF NOT EXISTS udb_version_seq AS BIGINT + START WITH 1 INCREMENT BY 1 MINVALUE 1; + + CREATE TABLE IF NOT EXISTS udb_commit_requests ( + id BIGSERIAL PRIMARY KEY, + epoch BIGINT NOT NULL, + read_version BIGINT NOT NULL, + payload BYTEA NOT NULL, + reply_channel TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + commit_version BIGINT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + ); + + CREATE INDEX IF NOT EXISTS udb_commit_requests_pending + ON udb_commit_requests (id) WHERE status = 'pending';", ) .await - .context("failed to create 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, - commit_version BIGINT NOT NULL, - ts timestamp NOT NULL DEFAULT now(), - - EXCLUDE USING gist ( - -- Conflict if byte range overlaps... - range_data WITH &&, - -- And if conflict types are different... - conflict_type WITH <>, - -- And if the txn versions overlap... - int8range(start_version, commit_version, '[]') WITH &&, - -- But not if the start_version is the same (from the same txn) - start_version WITH <> - ) - )", - &[], - ) - .await - .context("failed to create conflict_ranges table")?; + .context("failed to initialize postgres schema")?; - // Create index on ts column for efficient garbage collection - conn.execute( - "CREATE INDEX IF NOT EXISTS idx_conflict_ranges_ts ON conflict_ranges (ts)", - &[], - ) - .await - .context("failed to create index on conflict_ranges ts column")?; + Ok(()) + } - let pool2 = pool.clone(); - let gc_handle = tokio::spawn(async move { + fn spawn_gc(shared: Arc) -> JoinHandle<()> { + tokio::spawn(async move { let mut interval = tokio::time::interval(GC_INTERVAL); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { interval.tick().await; - tracing::debug!(status=?pool2.status(), "postgres pool status"); + let conn = match shared.pool.get().await { + Ok(conn) => conn, + Err(err) => { + tracing::debug!(?err, "failed to get connection for commit gc"); + continue; + } + }; - // NOTE: Transactions have a max limit of 5 seconds, we delete after 10 seconds for extra padding - // Delete old conflict ranges if let Err(err) = conn .execute( - "DELETE FROM conflict_ranges where ts < now() - interval '10 seconds'", - &[], + "DELETE FROM udb_commit_requests + WHERE created_at < now() - ($1::bigint * interval '1 second')", + &[&COMMIT_ROW_MAX_AGE_SECS], ) .await { - tracing::error!(?err, "failed postgres gc task"); + tracing::error!(?err, "failed postgres commit-queue gc"); } } - }); - - Ok(PostgresDatabaseDriver { - pool, - max_retries: AtomicI32::new(100), - gc_handle, }) } } impl DatabaseDriver for PostgresDatabaseDriver { fn create_txn(&self) -> Result { - // Pass the connection pool and config to the transaction driver - Ok(Transaction::new(Arc::new( - PostgresTransactionDriver::with_config(self.pool.clone()), - ))) + Ok(Transaction::new(Arc::new(PostgresTransactionDriver::new( + self.shared.clone(), + )))) } fn run<'a>( @@ -291,10 +290,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/listener.rs b/engine/packages/universaldb/src/driver/postgres/listener.rs new file mode 100644 index 0000000000..e5a7919dfe --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/listener.rs @@ -0,0 +1,228 @@ +use std::{path::PathBuf, sync::Arc, time::Duration}; + +use futures_util::future::poll_fn; +use rivet_postgres_util::build_tls_config; +use scc::HashMap; +use tokio::{ + io::{AsyncRead, AsyncWrite}, + sync::{Mutex, broadcast}, +}; +use tokio_postgres::AsyncMessage; +use tokio_postgres_rustls::MakeRustlsConnect; + +/// How long to wait between reconnect attempts for the dedicated LISTEN connection. +const RECONNECT_BACKOFF: Duration = Duration::from_secs(1); +/// Capacity of each channel's broadcast buffer. Notifications are wakeup signals with a polling +/// backstop, so a lagged receiver only delays a wake, never drops a durable commit. +const BROADCAST_CAPACITY: usize = 1024; + +struct Subscription { + tx: broadcast::Sender, +} + +/// Owns a single dedicated Postgres connection used exclusively for `LISTEN`. Demultiplexes +/// incoming `NOTIFY` payloads to per-channel broadcast senders and re-`LISTEN`s every registered +/// channel after a reconnect. +/// +/// This is separate from the deadpool pool because deadpool recycles connections and drops the +/// async notification stream; LISTEN requires owning the connection's message stream directly. +pub struct PgListener { + conn_str: String, + ssl_disabled: bool, + ssl_root_cert_path: Option, + ssl_client_cert_path: Option, + ssl_client_key_path: Option, + channels: Arc>, + client: Arc>>, +} + +impl PgListener { + pub fn new( + conn_str: String, + ssl_disabled: bool, + ssl_root_cert_path: Option, + ssl_client_cert_path: Option, + ssl_client_key_path: Option, + ) -> Self { + let channels: Arc> = Arc::new(HashMap::new()); + let client: Arc>> = Arc::new(Mutex::new(None)); + + tokio::spawn(Self::connection_lifecycle( + conn_str.clone(), + ssl_disabled, + ssl_root_cert_path.clone(), + ssl_client_cert_path.clone(), + ssl_client_key_path.clone(), + channels.clone(), + client.clone(), + )); + + Self { + conn_str, + ssl_disabled, + ssl_root_cert_path, + ssl_client_cert_path, + ssl_client_key_path, + channels, + client, + } + } + + /// Subscribe to a channel, registering a `LISTEN` if this is the first subscriber. Returns a + /// broadcast receiver of notification payloads. Idempotent per channel. + pub async fn listen(&self, channel: &str) -> broadcast::Receiver { + match self.channels.entry_async(channel.to_string()).await { + scc::hash_map::Entry::Occupied(entry) => entry.get().tx.subscribe(), + scc::hash_map::Entry::Vacant(entry) => { + let (tx, rx) = broadcast::channel(BROADCAST_CAPACITY); + entry.insert_entry(Subscription { tx }); + + // Best-effort immediate LISTEN; the lifecycle task re-LISTENs on reconnect. + if let Some(client) = &*self.client.lock().await { + if let Err(err) = client.execute(&format!("LISTEN \"{channel}\""), &[]).await { + tracing::warn!(?err, %channel, "failed to LISTEN, will retry on reconnect"); + } + } + + rx + } + } + } + + async fn connection_lifecycle( + conn_str: String, + ssl_disabled: bool, + ssl_root_cert_path: Option, + ssl_client_cert_path: Option, + ssl_client_key_path: Option, + channels: Arc>, + client: Arc>>, + ) { + loop { + let connected = if ssl_disabled { + Self::connect_and_run(&conn_str, tokio_postgres::NoTls, &channels, &client).await + } else { + match build_tls_config( + ssl_root_cert_path.as_ref(), + ssl_client_cert_path.as_ref(), + ssl_client_key_path.as_ref(), + ) { + Ok(tls_config) => { + Self::connect_and_run( + &conn_str, + MakeRustlsConnect::new(tls_config), + &channels, + &client, + ) + .await + } + Err(err) => { + tracing::error!(?err, "failed to build listener TLS config"); + false + } + } + }; + + if !connected { + tokio::time::sleep(RECONNECT_BACKOFF).await; + } + } + } + + /// Connects, re-LISTENs all channels, then drives the notification poll loop until the + /// connection closes. Returns `true` if a connection was successfully established (so the caller + /// can skip the reconnect backoff). + async fn connect_and_run( + conn_str: &str, + tls: T, + channels: &Arc>, + client: &Arc>>, + ) -> bool + where + T: tokio_postgres::tls::MakeTlsConnect, + T::Stream: AsyncRead + AsyncWrite + Unpin + Send + 'static, + T::TlsConnect: Send, + >::Future: Send, + { + let (new_client, connection) = match tokio_postgres::connect(conn_str, tls).await { + Ok(pair) => pair, + Err(err) => { + tracing::error!(?err, "failed to connect postgres listener"); + return false; + } + }; + + let channels_poll = channels.clone(); + let poll_handle = + tokio::spawn(async move { Self::poll_connection(connection, channels_poll).await }); + + // Re-LISTEN all registered channels on the fresh connection. + let mut registered = Vec::new(); + channels + .iter_async(|k, _| { + registered.push(k.clone()); + true + }) + .await; + for channel in ®istered { + if let Err(err) = new_client + .execute(&format!("LISTEN \"{channel}\""), &[]) + .await + { + tracing::error!(?err, %channel, "failed to re-LISTEN channel after reconnect"); + } + } + + *client.lock().await = Some(new_client); + + // Block until the poll loop ends (connection closed or errored). + let _ = poll_handle.await; + + *client.lock().await = None; + + true + } + + async fn poll_connection( + mut connection: tokio_postgres::Connection, + channels: Arc>, + ) where + S: AsyncRead + AsyncWrite + Unpin, + T: AsyncRead + AsyncWrite + Unpin, + { + loop { + match poll_fn(|cx| connection.poll_message(cx)).await { + Some(Ok(AsyncMessage::Notification(note))) => { + if let Some(sub) = channels.get_async(note.channel()).await { + // Ignore send errors: no active receiver just means no one is waiting + // right now; the polling backstop covers them. + let _ = sub.tx.send(note.payload().to_string()); + } + } + Some(Ok(_)) => {} + Some(Err(err)) => { + tracing::warn!(?err, "postgres listener connection error"); + break; + } + None => { + tracing::warn!("postgres listener connection closed"); + break; + } + } + } + } +} + +impl Clone for PgListener { + fn clone(&self) -> Self { + Self { + conn_str: self.conn_str.clone(), + ssl_disabled: self.ssl_disabled, + ssl_root_cert_path: self.ssl_root_cert_path.clone(), + ssl_client_cert_path: self.ssl_client_cert_path.clone(), + ssl_client_key_path: self.ssl_client_key_path.clone(), + channels: self.channels.clone(), + client: self.client.clone(), + } + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/mod.rs b/engine/packages/universaldb/src/driver/postgres/mod.rs index 1c24f9cb94..64f4bbd1bf 100644 --- a/engine/packages/universaldb/src/driver/postgres/mod.rs +++ b/engine/packages/universaldb/src/driver/postgres/mod.rs @@ -1,4 +1,9 @@ +mod codec; +mod commit; mod database; +mod listener; +mod resolver; +mod shared; mod transaction; mod transaction_task; diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs b/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs new file mode 100644 index 0000000000..d9729de884 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/resolver/apply.rs @@ -0,0 +1,139 @@ +use anyhow::{Context, Result}; +use deadpool_postgres::Transaction; + +use crate::{ + atomic::apply_atomic_op, options::MutationType, tuple::Versionstamp, tx_ops::Operation, + versionstamp::substitute_raw_versionstamp, +}; + +/// Apply a winning transaction's operations to `kv` inside the leader's batch txn. +/// +/// `commit_version` is the Postgres-resolved version assigned to this commit (`nextval`). It is +/// substituted into the 8-byte committed-version slot of any versionstamped key/value so +/// versionstamps are globally monotonic with commit order across all follower processes. +pub async fn apply( + txn: &Transaction<'_>, + operations: Vec, + commit_version: u64, +) -> Result<()> { + // Distinguishes multiple versionstamped operations within a single commit so their 10-byte + // stamps stay unique (8-byte version shared, 2-byte counter incremented). + let mut versionstamp_counter: u16 = 0; + + for op in operations { + match op { + Operation::SetValue { key, value } => { + upsert(txn, &key, &value).await?; + } + Operation::Clear { key } => { + txn.execute("DELETE FROM kv WHERE key = $1", &[&key]) + .await + .context("failed to clear key")?; + } + Operation::ClearRange { begin, end } => { + txn.execute( + "DELETE FROM kv WHERE key >= $1 AND key < $2", + &[&begin, &end], + ) + .await + .context("failed to clear range")?; + } + Operation::AtomicOp { + key, + param, + op_type, + } => { + apply_atomic( + txn, + key, + param, + op_type, + commit_version, + &mut versionstamp_counter, + ) + .await?; + } + } + } + + Ok(()) +} + +async fn apply_atomic( + txn: &Transaction<'_>, + key: Vec, + param: Vec, + op_type: MutationType, + commit_version: u64, + versionstamp_counter: &mut u16, +) -> Result<()> { + match op_type { + MutationType::SetVersionstampedKey => { + let versionstamp = build_versionstamp(commit_version, versionstamp_counter); + let key = substitute_raw_versionstamp(key, &versionstamp) + .map_err(anyhow::Error::msg) + .context("failed substituting versionstamped key")?; + upsert(txn, &key, ¶m).await?; + } + MutationType::SetVersionstampedValue => { + let versionstamp = build_versionstamp(commit_version, versionstamp_counter); + let value = substitute_raw_versionstamp(param, &versionstamp) + .map_err(anyhow::Error::msg) + .context("failed substituting versionstamped value")?; + upsert(txn, &key, &value).await?; + } + // Read-modify-write atomics: the leader is the single writer, so reading the live value + // inside the apply txn and writing the result is serializable with no lost update. + MutationType::Add + | MutationType::And + | MutationType::BitAnd + | MutationType::Or + | MutationType::BitOr + | MutationType::Xor + | MutationType::BitXor + | MutationType::AppendIfFits + | MutationType::Max + | MutationType::Min + | MutationType::ByteMin + | MutationType::ByteMax + | MutationType::CompareAndClear => { + let current = txn + .query_opt("SELECT value FROM kv WHERE key = $1", &[&key]) + .await + .context("failed to read current value for atomic op")? + .map(|row| row.get::<_, Vec>(0)); + + let new_value = apply_atomic_op(current.as_deref(), ¶m, op_type); + + if let Some(new_value) = new_value { + upsert(txn, &key, &new_value).await?; + } else { + txn.execute("DELETE FROM kv WHERE key = $1", &[&key]) + .await + .context("failed to clear key after atomic op")?; + } + } + } + + Ok(()) +} + +async fn upsert(txn: &Transaction<'_>, key: &[u8], value: &[u8]) -> Result<()> { + txn.execute( + "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2", + &[&key, &value], + ) + .await + .context("failed to upsert kv")?; + Ok(()) +} + +/// Build a 10-byte versionstamp (plus the 2 user-version bytes the substitution helper ignores) +/// from the Postgres-resolved commit version and a per-commit counter. +fn build_versionstamp(commit_version: u64, counter: &mut u16) -> Versionstamp { + let mut bytes = [0u8; 12]; + bytes[0..8].copy_from_slice(&commit_version.to_be_bytes()); + bytes[8..10].copy_from_slice(&counter.to_be_bytes()); + *counter = counter.wrapping_add(1); + Versionstamp::from(bytes) +} diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs b/engine/packages/universaldb/src/driver/postgres/resolver/lease.rs new file mode 100644 index 0000000000..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..f9b6c45461 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs @@ -0,0 +1,396 @@ +mod apply; +mod lease; + +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result}; +use tokio::sync::broadcast; + +use crate::{conflict_tracker::TransactionConflictTracker, transaction::TXN_TIMEOUT}; + +use super::shared::{ + ELECTION_CHANNEL, LEASE_ID, LeaseInfo, PostgresShared, WATERMARK_CHANNEL, commit_channel, +}; + +/// Max commits resolved+applied per batch (group commit). Amortizes the resolver, Postgres +/// round-trips, and fsync across the batch. +const DRAIN_BATCH_SIZE: i64 = 256; + +/// How often a leader renews its lease. Must be comfortably under `LEASE_TTL_SECS`. +const RENEW_INTERVAL: Duration = Duration::from_secs(3); + +/// Backstop poll cadence so a missed `udb_commit` NOTIFY cannot stall the drain indefinitely. +const POLL_BACKSTOP: Duration = Duration::from_millis(50); + +/// How long a candidate waits before retrying election when another node holds the lease. +const ELECTION_RETRY: Duration = Duration::from_secs(2); + +enum DrainOutcome { + /// Processed zero or more requests; still leader. + Drained, + /// Lost the lease (epoch bumped by a new leader). Step down. + LostLease, +} + +/// Spawn the per-process resolver task. Every node runs this; only the elected leader drains the +/// commit queue. The returned handle is aborted when the owning driver drops, which stops lease +/// renewal so the lease expires and another node can take over (node-death / failover path). +pub fn spawn(shared: Arc) -> tokio::task::JoinHandle<()> { + tokio::spawn(run(shared)) +} + +async fn run(shared: Arc) { + // A departing leader NOTIFYs this channel after releasing its lease so we elect immediately + // rather than waiting out the full `ELECTION_RETRY` tick. + let mut election_rx = shared.listener.listen(ELECTION_CHANNEL).await; + + loop { + match lease::try_acquire(&shared.pool, &shared.node_id).await { + Ok(Some(acquired)) => { + tracing::info!(epoch = acquired.epoch, node_id = %shared.node_id, "acquired udb leader lease"); + if let Err(err) = lead(&shared, acquired.epoch).await { + tracing::error!(?err, "udb leader loop errored, stepping down"); + } + tracing::info!(epoch = acquired.epoch, "stepped down from udb leader"); + } + Ok(None) => { + wait_for_election_retry(&shared, &mut election_rx).await; + } + Err(err) => { + tracing::warn!(?err, "failed udb lease acquire attempt"); + wait_for_election_retry(&shared, &mut election_rx).await; + } + } + } +} + +/// Wait before retrying the election: either the `ELECTION_RETRY` backstop elapses, or a departing +/// leader wakes us via `ELECTION_CHANNEL` so handoff is near-instant. +async fn wait_for_election_retry( + shared: &Arc, + election_rx: &mut broadcast::Receiver, +) { + tokio::select! { + _ = tokio::time::sleep(ELECTION_RETRY) => {} + res = election_rx.recv() => { + if matches!(res, Err(broadcast::error::RecvError::Closed)) { + // The listener recreates the channel on reconnect; re-subscribe. + *election_rx = shared.listener.listen(ELECTION_CHANNEL).await; + } + } + } +} + +/// Best-effort graceful leadership handoff invoked on shutdown. If this node currently holds the +/// lease, expire it and wake a standby so it takes over immediately instead of waiting out the TTL. +/// Safe to call on a follower: the fenced release matches no row and nothing is notified. The +/// caller must already have stopped lease renewal before calling this. +pub async fn handoff(shared: &Arc) { + match lease::release(&shared.pool, &shared.node_id).await { + Ok(true) => { + tracing::info!(node_id = %shared.node_id, "released udb leader lease for graceful handoff"); + notify_election(shared).await; + } + Ok(false) => {} + Err(err) => { + tracing::warn!(?err, "failed to release udb lease on shutdown"); + } + } +} + +/// Wake standby candidates so the next election fires immediately after a graceful release. +async fn notify_election(shared: &Arc) { + let conn = match shared.pool.get().await { + Ok(conn) => conn, + Err(err) => { + tracing::debug!(?err, "failed to get connection for election notify"); + return; + } + }; + + if let Err(err) = conn + .execute("SELECT pg_notify($1, '')", &[&ELECTION_CHANNEL]) + .await + { + tracing::debug!(?err, "failed to notify election channel"); + } +} + +/// Leader main loop: hold the lease, drain the commit queue on wake or poll, and renew the lease. +async fn lead(shared: &Arc, epoch: i64) -> Result<()> { + // Publish our own lease into the cache immediately so our local commits route to us. + shared.set_lease(LeaseInfo { + epoch, + leader_addr: shared.node_id.clone(), + }); + + // The recovery floor: a freshly elected leader has a cold conflict window, so reject commits + // whose read_version predates the floor until the window warms (one TXN_TIMEOUT), forcing + // those followers to take a fresh read_version. + let recovery_version = recovery_floor(shared).await?; + let recovery_deadline = Instant::now() + TXN_TIMEOUT; + + let tracker = TransactionConflictTracker::new(); + + let mut wake_rx = shared + .listener + .listen(&commit_channel(&shared.node_id)) + .await; + + let mut renew_interval = tokio::time::interval(RENEW_INTERVAL); + renew_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut poll_interval = tokio::time::interval(POLL_BACKSTOP); + poll_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + // Drain anything already queued before our first wake. + if matches!( + drain(shared, epoch, &tracker, recovery_version, recovery_deadline).await?, + DrainOutcome::LostLease + ) { + return Ok(()); + } + + loop { + tokio::select! { + _ = renew_interval.tick() => { + if !lease::renew(&shared.pool, &shared.node_id, epoch).await? { + tracing::warn!(epoch, "lost udb lease on renew"); + return Ok(()); + } + } + res = wake_rx.recv() => { + match res { + Ok(_) | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + wake_rx = shared.listener.listen(&commit_channel(&shared.node_id)).await; + } + } + if matches!( + drain(shared, epoch, &tracker, recovery_version, recovery_deadline).await?, + DrainOutcome::LostLease + ) { + return Ok(()); + } + } + _ = poll_interval.tick() => { + if matches!( + drain(shared, epoch, &tracker, recovery_version, recovery_deadline).await?, + DrainOutcome::LostLease + ) { + return Ok(()); + } + } + } + } +} + +/// The version floor a freshly elected leader continues from: the higher of the durable watermark +/// and the sequence high-water. The LOGGED `udb_version_seq` is crash-safe, so this never regresses. +async fn recovery_floor(shared: &Arc) -> Result { + let durable = lease::current_durable_version(&shared.pool).await?; + + let conn = shared + .pool + .get() + .await + .context("failed to get connection for recovery floor")?; + let seq_high: i64 = conn + .query_one("SELECT last_value FROM udb_version_seq", &[]) + .await + .context("failed to read sequence high water")? + .get(0); + + Ok(durable.max(seq_high).max(0) as u64) +} + +/// Drain pending commit requests in id-ordered batches until none remain. Each batch resolves and +/// applies inside a single Postgres transaction (group commit), fenced on the leader's epoch. +async fn drain( + shared: &Arc, + epoch: i64, + tracker: &TransactionConflictTracker, + recovery_version: u64, + recovery_deadline: Instant, +) -> Result { + loop { + match drain_batch(shared, epoch, tracker, recovery_version, recovery_deadline).await? { + BatchOutcome::Empty => return Ok(DrainOutcome::Drained), + BatchOutcome::Processed => {} + BatchOutcome::LostLease => return Ok(DrainOutcome::LostLease), + } + } +} + +enum BatchOutcome { + Empty, + Processed, + LostLease, +} + +struct Reply { + channel: String, + id: i64, +} + +async fn drain_batch( + shared: &Arc, + epoch: i64, + tracker: &TransactionConflictTracker, + recovery_version: u64, + recovery_deadline: Instant, +) -> Result { + let mut conn = shared + .pool + .get() + .await + .context("failed to get connection for drain batch")?; + let txn = conn + .build_transaction() + .start() + .await + .context("failed to start drain batch txn")?; + + // Claim a batch in id order. FOR UPDATE SKIP LOCKED holds the rows for this txn so they are + // stamped terminal on COMMIT with no intermediate 'claimed' state to clean up. + let rows = txn + .query( + "SELECT id, read_version, payload, reply_channel + FROM udb_commit_requests + WHERE status = 'pending' AND epoch = $1 + ORDER BY id + LIMIT $2 + FOR UPDATE SKIP LOCKED", + &[&epoch, &DRAIN_BATCH_SIZE], + ) + .await + .context("failed to claim commit batch")?; + + if rows.is_empty() { + txn.rollback().await.ok(); + return Ok(BatchOutcome::Empty); + } + + let cold_window = Instant::now() < recovery_deadline; + let mut max_winner_cv: i64 = 0; + let mut replies = Vec::with_capacity(rows.len()); + + for row in &rows { + let id: i64 = row.get(0); + let read_version: i64 = row.get(1); + let payload: Vec = row.get(2); + let reply_channel: String = row.get(3); + + let decoded = super::codec::decode_commit_request(&payload) + .context("failed to decode commit payload")?; + + let commit_version: i64 = txn + .query_one("SELECT nextval('udb_version_seq')", &[]) + .await + .context("failed to get next commit version")? + .get(0); + + let start_version = read_version.max(0) as u64; + + // Cold-window guard: a commit whose read_version predates the recovery floor cannot be + // safely resolved against this leader's empty window. Reject it as retryable. + let conflicted = if cold_window && start_version < recovery_version { + true + } else { + tracker + .check_and_insert( + start_version, + commit_version.max(0) as u64, + decoded.conflict_ranges, + ) + .await + }; + + if conflicted { + txn.execute( + "UPDATE udb_commit_requests SET status = 'conflict' WHERE id = $1", + &[&id], + ) + .await + .context("failed to stamp conflict")?; + } else { + apply::apply(&txn, decoded.operations, commit_version.max(0) as u64) + .await + .context("failed to apply commit")?; + txn.execute( + "UPDATE udb_commit_requests SET status = 'committed', commit_version = $1 WHERE id = $2", + &[&commit_version, &id], + ) + .await + .context("failed to stamp committed")?; + max_winner_cv = max_winner_cv.max(commit_version); + } + + replies.push(Reply { + channel: reply_channel, + id, + }); + } + + // Advance the watermark, fenced on our epoch. A zombie old leader whose epoch was bumped sees + // zero rows updated and must step down before any of its writes become visible. + let new_durable: i64 = match txn + .query_opt( + "UPDATE udb_lease + SET durable_version = GREATEST(durable_version, $1) + WHERE id = $2 AND epoch = $3 + RETURNING durable_version", + &[&max_winner_cv, &LEASE_ID, &epoch], + ) + .await + .context("failed to advance watermark")? + { + Some(row) => row.get(0), + None => { + txn.rollback().await.ok(); + return Ok(BatchOutcome::LostLease); + } + }; + + txn.commit().await.context("failed to commit drain batch")?; + + // Watermark advances strictly after the apply txn is durably committed and visible, so a + // reader handed this read_version can never miss a write with commit_version <= read_version. + shared.advance_durable_version(new_durable); + + notify_after_commit(&conn, new_durable, &replies).await; + + Ok(BatchOutcome::Processed) +} + +/// Wake watermark listeners and the followers waiting on each processed request. Best-effort: a +/// missed NOTIFY is covered by the follower's polling backstop and the watermark refresh timer. +async fn notify_after_commit( + conn: &deadpool_postgres::Client, + new_durable: i64, + replies: &[Reply], +) { + if let Err(err) = conn + .execute( + "SELECT pg_notify($1, $2)", + &[&WATERMARK_CHANNEL, &new_durable.to_string()], + ) + .await + { + tracing::debug!(?err, "failed to notify watermark"); + } + + let channels: Vec<&str> = replies.iter().map(|r| r.channel.as_str()).collect(); + let ids: Vec = replies.iter().map(|r| r.id.to_string()).collect(); + if let Err(err) = conn + .execute( + "SELECT pg_notify(c, p) FROM unnest($1::text[], $2::text[]) AS t(c, p)", + &[&channels, &ids], + ) + .await + { + tracing::debug!(?err, "failed to notify commit replies"); + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/shared.rs b/engine/packages/universaldb/src/driver/postgres/shared.rs new file mode 100644 index 0000000000..5b14c4df08 --- /dev/null +++ b/engine/packages/universaldb/src/driver/postgres/shared.rs @@ -0,0 +1,165 @@ +use std::{ + sync::{ + Arc, + atomic::{AtomicI64, Ordering}, + }, + time::Duration, +}; + +use deadpool_postgres::Pool; +use tokio::sync::{Notify, watch}; + +use super::listener::PgListener; + +/// The singleton row id of `udb_lease`. +pub const LEASE_ID: i32 = 1; + +/// How often the follower refreshes its cached lease row (epoch, leader channel, watermark) as a +/// backstop to the `udb_watermark` NOTIFY. A stale-but-older watermark only widens the conflict +/// window, so this can be loose. +const LEASE_REFRESH_INTERVAL: Duration = Duration::from_millis(500); + +/// Channel a follower NOTIFYs (and the leader LISTENs) to wake the leader's drain loop. +pub fn commit_channel(node_id: &str) -> String { + format!("udb_commit_{node_id}") +} + +/// Channel the leader NOTIFYs (and a follower LISTENs) to deliver a commit result. +pub fn reply_channel(node_id: &str) -> String { + format!("udb_reply_{node_id}") +} + +/// Channel the leader NOTIFYs on every watermark advance; all nodes LISTEN. +pub const WATERMARK_CHANNEL: &str = "udb_watermark"; + +/// Channel a departing leader NOTIFYs after releasing its lease so a standby candidate elects +/// immediately instead of waiting out `ELECTION_RETRY`. All non-leader candidates LISTEN. +pub const ELECTION_CHANNEL: &str = "udb_election"; + +/// Cached view of the current leader lease, as seen by a follower. +#[derive(Clone, Debug)] +pub struct LeaseInfo { + pub epoch: i64, + /// Node id of the current leader, used to build its commit channel. + pub leader_addr: String, +} + +/// Process-wide state shared by the follower transaction tasks and the leader resolver. Every node +/// is both a follower (it submits its own commits) and a candidate leader. +pub struct PostgresShared { + pub pool: Pool, + /// Unique per-process id used to name this node's NOTIFY channels. + pub node_id: String, + pub listener: PgListener, + /// Highest durable commit version (`udb_lease.durable_version`); the follower read version. + durable_version: AtomicI64, + /// Pinged whenever `durable_version` advances. + watermark_notify: Notify, + lease_tx: watch::Sender>, + lease_rx: watch::Receiver>, +} + +impl PostgresShared { + pub fn new(pool: Pool, node_id: String, listener: PgListener) -> Arc { + let (lease_tx, lease_rx) = watch::channel(None); + let shared = Arc::new(Self { + pool, + node_id, + listener, + durable_version: AtomicI64::new(0), + watermark_notify: Notify::new(), + lease_tx, + lease_rx, + }); + + tokio::spawn(Self::cache_refresh_task(shared.clone())); + + shared + } + + /// The cached follower read version (`durable_version`). + pub fn read_version(&self) -> i64 { + self.durable_version.load(Ordering::SeqCst) + } + + /// Advance the cached watermark monotonically and wake any waiters. + pub fn advance_durable_version(&self, version: i64) { + let prev = self.durable_version.fetch_max(version, Ordering::SeqCst); + if version > prev { + self.watermark_notify.notify_waiters(); + } + } + + /// Current cached lease, if known. + pub fn current_lease(&self) -> Option { + self.lease_rx.borrow().clone() + } + + /// Publish a freshly observed/elected lease into the cache. + pub fn set_lease(&self, lease: LeaseInfo) { + let _ = self.lease_tx.send(Some(lease)); + } + + /// Background task: keep `durable_version` and the cached lease fresh via the `udb_watermark` + /// NOTIFY plus a periodic poll of `udb_lease`. + async fn cache_refresh_task(shared: Arc) { + let mut watermark_rx = shared.listener.listen(WATERMARK_CHANNEL).await; + let mut interval = tokio::time::interval(LEASE_REFRESH_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + tokio::select! { + notify = watermark_rx.recv() => { + match notify { + Ok(payload) => { + if let Ok(version) = payload.parse::() { + shared.advance_durable_version(version); + } + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + // Re-subscribe; the listener recreates the channel on reconnect. + watermark_rx = shared.listener.listen(WATERMARK_CHANNEL).await; + } + } + } + _ = interval.tick() => { + shared.refresh_lease_row().await; + } + } + } + } + + async fn refresh_lease_row(&self) { + let conn = match self.pool.get().await { + Ok(conn) => conn, + Err(err) => { + tracing::debug!(?err, "failed to get connection for lease refresh"); + return; + } + }; + + let row = conn + .query_opt( + "SELECT epoch, leader_addr, durable_version FROM udb_lease WHERE id = $1", + &[&LEASE_ID], + ) + .await; + + match row { + Ok(Some(row)) => { + let epoch: i64 = row.get(0); + let leader_addr: String = row.get(1); + let durable_version: i64 = row.get(2); + self.advance_durable_version(durable_version); + self.set_lease(LeaseInfo { epoch, leader_addr }); + } + Ok(None) => { + // No lease row yet; no leader elected. + } + Err(err) => { + tracing::debug!(?err, "failed to refresh lease row"); + } + } + } +} diff --git a/engine/packages/universaldb/src/driver/postgres/transaction.rs b/engine/packages/universaldb/src/driver/postgres/transaction.rs index f80cbba393..cc315feba4 100644 --- a/engine/packages/universaldb/src/driver/postgres/transaction.rs +++ b/engine/packages/universaldb/src/driver/postgres/transaction.rs @@ -1,11 +1,13 @@ use std::{ future::Future, pin::Pin, - sync::atomic::{AtomicBool, Ordering}, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, }; use anyhow::{Context, Result}; -use deadpool_postgres::Pool; use tokio::sync::{OnceCell, mpsc, oneshot}; use crate::{ @@ -18,33 +20,35 @@ use crate::{ value::{Slice, Value, Values}, }; -use super::transaction_task::{TransactionCommand, TransactionTask}; +use super::{ + shared::PostgresShared, + transaction_task::{TransactionCommand, TransactionTask}, +}; pub struct PostgresTransactionDriver { - pool: Pool, + shared: Arc, operations: TransactionOperations, committed: AtomicBool, tx_sender: OnceCell>, } impl PostgresTransactionDriver { - pub fn with_config(pool: Pool) -> Self { + pub fn new(shared: Arc) -> Self { PostgresTransactionDriver { - pool, + shared, operations: TransactionOperations::default(), committed: AtomicBool::new(false), tx_sender: OnceCell::new(), } } - /// Get or create the transaction task + /// Get or create the transaction task that owns this transaction's read snapshot. async fn ensure_transaction(&self) -> Result<&mpsc::UnboundedSender> { self.tx_sender .get_or_try_init(|| async { let (sender, receiver) = mpsc::unbounded_channel(); - // Spawn the transaction task with serializable isolation - let task = TransactionTask::new(self.pool.clone(), receiver); + let task = TransactionTask::new(self.shared.clone(), receiver); tokio::spawn(task.run()); anyhow::Ok(sender) diff --git a/engine/packages/universaldb/src/driver/postgres/transaction_task.rs b/engine/packages/universaldb/src/driver/postgres/transaction_task.rs index 297751f608..c4d3e3f9f6 100644 --- a/engine/packages/universaldb/src/driver/postgres/transaction_task.rs +++ b/engine/packages/universaldb/src/driver/postgres/transaction_task.rs @@ -1,17 +1,18 @@ -use anyhow::{Context, Result, anyhow, bail}; -use deadpool_postgres::{Pool, Transaction}; +use std::sync::Arc; + +use anyhow::{Result, anyhow, bail}; +use deadpool_postgres::Transaction; use tokio::sync::{mpsc, oneshot}; use tokio_postgres::IsolationLevel; use crate::{ - atomic::apply_atomic_op, - error::DatabaseError, - options::{ConflictRangeType, MutationType}, + options::ConflictRangeType, tx_ops::Operation, value::{KeyValue, Slice, Values}, - versionstamp::{generate_versionstamp, substitute_raw_versionstamp}, }; +use super::{commit, shared::PostgresShared}; + pub enum TransactionCommand { // Read operations Get { @@ -48,70 +49,57 @@ pub enum TransactionCommand { }, } -/// TransactionTask runs in a separate tokio task to manage a PostgreSQL transaction. -/// -/// This design is necessary because PostgreSQL transactions have lifetime constraints -/// that don't work well with the FoundationDB-style API. Specifically: -/// - The transaction must outlive all references to it -/// - We can't store the transaction in a mutex due to lifetime issues with the connection +/// TransactionTask runs in a separate tokio task to own a single pinned PostgreSQL `REPEATABLE READ` +/// snapshot connection for the lifetime of a follower transaction. /// -/// By running in a separate task and communicating via channels, we avoid these lifetime -/// issues while maintaining a single serializable transaction for all operations. +/// Reads go directly against this snapshot (they never involve the leader). Commits delegate to +/// [`commit::submit`], which enqueues the request on the leader and awaits the result. The +/// `read_version` is captured from the cached watermark before the snapshot is opened, so no write +/// with `commit_version <= read_version` can be invisible to the snapshot. pub struct TransactionTask { - pool: Pool, + shared: Arc, receiver: mpsc::UnboundedReceiver, } impl TransactionTask { - pub fn new(pool: Pool, receiver: mpsc::UnboundedReceiver) -> Self { - Self { pool, receiver } + pub fn new( + shared: Arc, + receiver: mpsc::UnboundedReceiver, + ) -> Self { + Self { shared, receiver } } pub async fn run(mut self) { - // Get connection from pool - let mut conn = match self.pool.get().await { + // Capture the read version BEFORE opening the snapshot so the snapshot reflects every write + // with commit_version <= read_version. + let read_version = self.shared.read_version(); + + let mut conn = match self.shared.pool.get().await { Ok(conn) => conn, Err(_) => { - // If we can't get a connection, respond to all pending commands with errors self.fail_receiver().await; return; } }; - // Start the read transaction let tx = match conn .build_transaction() .isolation_level(IsolationLevel::RepeatableRead) + .read_only(true) .start() .await { Ok(tx) => tx, Err(_) => { - // If we can't start a transaction, respond to all pending commands with errors self.fail_receiver().await; return; } }; - // TODO: Parallelize future - let start_version = match tx - .query_one("SELECT nextval('global_version_seq')", &[]) - .await - { - Ok(row) => row.get::<_, i64>(0), - Err(err) => { - tracing::error!(?err, "failed to get postgres txn start_version"); - self.fail_receiver().await; - return; - } - }; - - // Process commands while let Some(cmd) = self.receiver.recv().await { match cmd { TransactionCommand::Get { key, response } => { let result = self.handle_get(&tx, &key).await; - let _ = response.send(result); } TransactionCommand::GetKey { @@ -121,7 +109,6 @@ impl TransactionTask { response, } => { let result = self.handle_get_key(&tx, &key, or_equal, offset).await; - let _ = response.send(result); } TransactionCommand::GetRange { @@ -148,7 +135,6 @@ impl TransactionTask { reverse, ) .await; - let _ = response.send(result); } TransactionCommand::Commit { @@ -156,14 +142,12 @@ impl TransactionTask { conflict_ranges, response, } => { - let (_, result) = tokio::join!( - // Read-only txn, we don't care about the result - tx.commit(), - self.handle_commit(start_version, operations, conflict_ranges), - ); - + // The read snapshot is read-only; release it and submit the commit to the leader. + let _ = tx.commit().await; + let result = + commit::submit(&self.shared, read_version, operations, conflict_ranges) + .await; let _ = response.send(result); - // Exit after commit return; } TransactionCommand::GetEstimatedRangeSize { @@ -174,13 +158,12 @@ impl TransactionTask { let result = self .handle_get_estimated_range_size(&tx, &begin, &end) .await; - let _ = response.send(result); } } } - // If the channel is closed, the transaction will be rolled back when dropped + // If the channel is closed, the snapshot transaction is rolled back when dropped. } async fn handle_get(&mut self, tx: &Transaction<'_>, key: &[u8]) -> Result> { @@ -234,27 +217,18 @@ impl TransactionTask { reverse: bool, ) -> Result { // Determine SQL operators based on key selector types - // For begin selector: - // first_greater_or_equal: or_equal = false, offset = 1 -> ">=" - // first_greater_than: or_equal = true, offset = 1 -> ">" let begin_op = if begin_offset == 1 { if begin_or_equal { ">" } else { ">=" } } else { - // This shouldn't happen for begin in range queries ">=" }; - // For end selector: - // first_greater_than: or_equal = true, offset = 1 -> "<=" - // first_greater_or_equal: or_equal = false, offset = 1 -> "<" let end_op = if end_offset == 1 { if end_or_equal { "<=" } else { "<" } } else { - // This shouldn't happen for end in range queries "<" }; - // Build query with CTE that adds conflict range let query = if reverse { if let Some(limit) = limit { format!( @@ -301,22 +275,22 @@ impl TransactionTask { begin: &[u8], end: &[u8], ) -> Result { - // Sample's 1% of the range + // Sample 1% of the range. let query = " WITH range_stats AS ( - SELECT + SELECT COUNT(*) as estimated_count, COALESCE(SUM(pg_column_size(key) + pg_column_size(value)), 0) as sample_size - FROM kv TABLESAMPLE SYSTEM(1) + FROM kv TABLESAMPLE SYSTEM(1) WHERE key >= $1 AND key < $2 ), table_stats AS ( - SELECT reltuples::bigint as total_rows - FROM pg_class + SELECT reltuples::bigint as total_rows + FROM pg_class WHERE relname = 'kv' AND relkind = 'r' ) - SELECT - CASE + SELECT + CASE WHEN r.estimated_count = 0 THEN 0 ELSE (r.sample_size * 100)::bigint END as estimated_size @@ -329,165 +303,6 @@ impl TransactionTask { .map_err(map_postgres_error) } - async fn handle_commit( - &mut self, - start_version: i64, - operations: Vec, - conflict_ranges: Vec<(Vec, Vec, ConflictRangeType)>, - ) -> Result<()> { - // Get connection from pool - let mut conn = self.pool.get().await?; - - // Start write transaction - let tx = conn - .build_transaction() - .isolation_level(IsolationLevel::ReadCommitted) - .start() - .await - .context("failed to start write txn")?; - - let mut begins = Vec::with_capacity(conflict_ranges.len()); - let mut ends = Vec::with_capacity(conflict_ranges.len()); - let mut conflict_types = Vec::with_capacity(conflict_ranges.len()); - - for (begin, end, conflict_type) in conflict_ranges { - let conflict_type = match conflict_type { - ConflictRangeType::Read => "read", - ConflictRangeType::Write => "write", - }; - - begins.push(begin); - ends.push(end); - conflict_types.push(conflict_type); - } - - let query = " - WITH data AS ( - SELECT nextval('global_version_seq') AS commit_version - ) - INSERT INTO conflict_ranges (range_data, conflict_type, start_version, commit_version) - SELECT - bytearange(begin_key, end_key, '[)'), - conflict_type::range_type, - $4, - data.commit_version - FROM UNNEST($1::bytea[], $2::bytea[], $3::text[]) AS t(begin_key, end_key, conflict_type), data"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - // Insert all conflict ranges at once - tx.execute(&stmt, &[&begins, &ends, &conflict_types, &start_version]) - .await - .map_err(map_postgres_error)?; - - let transaction_versionstamp = generate_versionstamp(0); - - for op in operations { - match op { - Operation::SetValue { key, value } => { - let query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key, &value]) - .await - .map_err(map_postgres_error)?; - } - Operation::Clear { key } => { - let query = "DELETE FROM kv WHERE key = $1"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key]) - .await - .map_err(map_postgres_error)?; - } - Operation::ClearRange { begin, end } => { - let query = "DELETE FROM kv WHERE key >= $1 AND key < $2"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&begin, &end]) - .await - .map_err(map_postgres_error)?; - } - Operation::AtomicOp { - key, - param, - op_type, - } => { - if matches!(op_type, MutationType::SetVersionstampedKey) { - let key = substitute_raw_versionstamp(key, &transaction_versionstamp) - .map_err(anyhow::Error::msg) - .context("failed substituting versionstamped key")?; - let query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key, ¶m]) - .await - .map_err(map_postgres_error)?; - continue; - } - - if matches!(op_type, MutationType::SetVersionstampedValue) { - let value = substitute_raw_versionstamp(param, &transaction_versionstamp) - .map_err(anyhow::Error::msg) - .context("failed substituting versionstamped value")?; - let query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; - let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key, &value]) - .await - .map_err(map_postgres_error)?; - continue; - } - - // TODO: All operations need to be done on the sql side, not in rust - - // Get current value from database - let current_query = "SELECT value FROM kv WHERE key = $1"; - let stmt = tx - .prepare_cached(current_query) - .await - .map_err(map_postgres_error)?; - - let current_row = tx - .query_opt(&stmt, &[&key]) - .await - .map_err(map_postgres_error)?; - - // Extract current value or use None if key doesn't exist - let current_value = current_row.map(|row| row.get::<_, Vec>(0)); - let current_slice = current_value.as_deref(); - - // Apply atomic operation - let new_value = apply_atomic_op(current_slice, ¶m, op_type); - - // Store the result - if let Some(new_value) = new_value { - let update_query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; - let stmt = tx - .prepare_cached(update_query) - .await - .map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key, &new_value]) - .await - .map_err(map_postgres_error)?; - } else { - let update_query = "DELETE FROM kv WHERE key = $1"; - let stmt = tx - .prepare_cached(update_query) - .await - .map_err(map_postgres_error)?; - - tx.execute(&stmt, &[&key]) - .await - .map_err(map_postgres_error)?; - } - } - } - } - - tx.commit().await.map_err(map_postgres_error) - } - async fn fail_receiver(&mut self) { while let Some(cmd) = self.receiver.recv().await { match cmd { @@ -511,31 +326,19 @@ impl TransactionTask { } } -/// Maps PostgreSQL error to DatabaseError +/// Maps a PostgreSQL error from the read path to a `DatabaseError` where appropriate. fn map_postgres_error(err: tokio_postgres::Error) -> anyhow::Error { - let error_str = if let Some(err) = err.as_db_error() { - err.to_string() - } else { - err.to_string() - }; + let error_str = err.to_string(); - if error_str.contains("exclusion_violation") - || error_str.contains("violates exclusion constraint") - { - // Retryable - another transaction has a conflicting range - DatabaseError::NotCommitted.into() - } else if error_str.contains("serialization failure") + if error_str.contains("serialization failure") || error_str.contains("could not serialize") || error_str.contains("deadlock detected") { - // Retryable - transaction conflict - DatabaseError::NotCommitted.into() + crate::error::DatabaseError::NotCommitted.into() } else if error_str.contains("current transaction is aborted") { - // Returned by the rest of the commands in a txn if it failed for exclusion reasons - DatabaseError::NotCommitted.into() + crate::error::DatabaseError::NotCommitted.into() } else { tracing::error!(%err, "postgres error"); - // Non-retryable error anyhow::Error::new(err) } } diff --git a/engine/packages/universaldb/src/driver/rocksdb/database.rs b/engine/packages/universaldb/src/driver/rocksdb/database.rs index ec0175ee5a..ba186d872e 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/database.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/database.rs @@ -17,9 +17,9 @@ use crate::{ utils::{MaybeCommitted, calculate_tx_retry_backoff}, }; -use super::{ - transaction::RocksDbTransactionDriver, transaction_conflict_tracker::TransactionConflictTracker, -}; +use crate::conflict_tracker::TransactionConflictTracker; + +use super::transaction::RocksDbTransactionDriver; pub struct RocksDbDatabaseDriver { db: Arc, diff --git a/engine/packages/universaldb/src/driver/rocksdb/mod.rs b/engine/packages/universaldb/src/driver/rocksdb/mod.rs index a24bd72603..b18e28edca 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/mod.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/mod.rs @@ -1,6 +1,5 @@ mod database; mod transaction; -mod transaction_conflict_tracker; mod transaction_task; pub use database::RocksDbDatabaseDriver; diff --git a/engine/packages/universaldb/src/driver/rocksdb/transaction.rs b/engine/packages/universaldb/src/driver/rocksdb/transaction.rs index e62c3eda79..85cf5edf5c 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/transaction.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/transaction.rs @@ -21,10 +21,9 @@ use crate::{ value::{Slice, Value, Values}, }; -use super::{ - transaction_conflict_tracker::TransactionConflictTracker, - transaction_task::{TransactionCommand, TransactionTask}, -}; +use crate::conflict_tracker::TransactionConflictTracker; + +use super::transaction_task::{TransactionCommand, TransactionTask}; pub struct RocksDbTransactionDriver { db: Arc, diff --git a/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs b/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs index 69835eb6a7..986e6ba2bd 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs @@ -6,9 +6,9 @@ use rocksdb::{ }; use tokio::sync::{mpsc, oneshot}; -use super::transaction_conflict_tracker::TransactionConflictTracker; use crate::{ atomic::apply_atomic_op, + conflict_tracker::TransactionConflictTracker, error::DatabaseError, key_selector::KeySelector, options::{ConflictRangeType, MutationType}, @@ -412,9 +412,11 @@ impl TransactionTask { } } + // rocksdb generates both start and commit versions from the in-process counter. + let commit_version = self.txn_conflict_tracker.next_global_version(); if self .txn_conflict_tracker - .check_and_insert(start_version, conflict_ranges) + .check_and_insert(start_version, commit_version, conflict_ranges) .await { return Err(DatabaseError::NotCommitted.into()); diff --git a/engine/packages/universaldb/src/lib.rs b/engine/packages/universaldb/src/lib.rs index 96260a177d..1b6fb78b87 100644 --- a/engine/packages/universaldb/src/lib.rs +++ b/engine/packages/universaldb/src/lib.rs @@ -1,4 +1,5 @@ pub(crate) mod atomic; +pub(crate) mod conflict_tracker; mod database; pub mod driver; pub mod error; diff --git a/engine/packages/universaldb/tests/failover.rs b/engine/packages/universaldb/tests/failover.rs new file mode 100644 index 0000000000..5fb394453c --- /dev/null +++ b/engine/packages/universaldb/tests/failover.rs @@ -0,0 +1,289 @@ +use std::{sync::Arc, time::Duration}; + +use rivet_test_deps_docker::TestDatabase; +use tokio_postgres::NoTls; +use universaldb::{Database, utils::IsolationLevel::*}; +use uuid::Uuid; + +const ALPHA_KEY: &[u8] = b"failover/alpha"; +const BETA_KEY: &[u8] = b"failover/beta"; + +/// Build a fresh Postgres-backed `Database`. Each call spins up an independent driver (its own pool, +/// node id, listener, and resolver), so two of them against one Postgres model two engine nodes. +async fn make_db(connection_string: &str) -> Database { + let driver = universaldb::driver::PostgresDatabaseDriver::new_with_config( + universaldb::driver::postgres::PostgresConfig::new(connection_string.to_string()), + ) + .await + .unwrap(); + Database::new(Arc::new(driver)) +} + +/// Raw verification connection used to inspect leader/lease/version state out of band. +async fn connect_raw(connection_string: &str) -> tokio_postgres::Client { + let (client, connection) = tokio_postgres::connect(connection_string, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client +} + +struct LeaseRow { + epoch: i64, + leader_addr: String, + durable_version: i64, +} + +async fn read_lease(client: &tokio_postgres::Client) -> Option { + let row = client + .query_opt( + "SELECT epoch, leader_addr, durable_version FROM udb_lease WHERE id = 1", + &[], + ) + .await + .unwrap()?; + Some(LeaseRow { + epoch: row.get(0), + leader_addr: row.get(1), + durable_version: row.get(2), + }) +} + +/// High-water of the LOGGED version sequence. A freshly elected leader must continue from at least +/// this value, never regress below it. +async fn read_seq_high(client: &tokio_postgres::Client) -> i64 { + client + .query_one("SELECT last_value FROM udb_version_seq", &[]) + .await + .unwrap() + .get(0) +} + +/// Poll `udb_lease` until `pred` holds or the deadline passes. +async fn wait_for_lease bool>( + client: &tokio_postgres::Client, + timeout: Duration, + pred: F, +) -> LeaseRow { + let deadline = tokio::time::Instant::now() + timeout; + loop { + if let Some(lease) = read_lease(client).await { + if pred(&lease) { + return lease; + } + } + if tokio::time::Instant::now() >= deadline { + panic!("timed out waiting for lease condition"); + } + tokio::time::sleep(Duration::from_millis(200)).await; + } +} + +async fn write_key(db: &Database, key: &'static [u8], value: &'static [u8]) { + db.txn("test_failover", move |tx| async move { + tx.set(key, value); + Ok(()) + }) + .await + .unwrap(); +} + +async fn read_key(db: &Database, key: &'static [u8]) -> Option> { + db.txn("test_failover", move |tx| async move { + let val = tx.get(key, Serializable).await?; + Ok(val) + }) + .await + .unwrap() + .map(|slice| slice.to_vec()) +} + +/// Exercises leader failover: two nodes share one Postgres, the elected leader is killed, the +/// survivor must take over the lease (new epoch), continue the crash-safe version sequence without +/// regression, preserve the dead leader's committed data, and resume accepting commits. +#[tokio::test] +async fn test_postgres_leader_failover() { + let _ = tracing_subscriber::fmt() + .with_env_filter("info") + .with_test_writer() + .try_init(); + + let (db_config, docker_config) = TestDatabase::Postgres + .config(Uuid::new_v4(), 1) + .await + .unwrap(); + let mut docker_config = docker_config.unwrap(); + docker_config.start().await.unwrap(); + + tokio::time::sleep(Duration::from_secs(4)).await; + + let rivet_config::config::Database::Postgres(postgres_config) = db_config else { + unreachable!(); + }; + let connection_string = postgres_config.url.read().clone(); + + let raw = connect_raw(&connection_string).await; + + // Node 1 comes up first and deterministically wins the first election (epoch 1). + let db1 = make_db(&connection_string).await; + let lease1 = wait_for_lease(&raw, Duration::from_secs(15), |l| l.epoch == 1).await; + let leader1_addr = lease1.leader_addr.clone(); + + // Node 2 joins while node 1 holds a valid lease, so it loses the election and runs as a + // follower. + let db2 = make_db(&connection_string).await; + + // Leader (node 1) commits data. The version sequence and watermark advance. + write_key(&db1, ALPHA_KEY, b"1").await; + + // The follower (node 2) reads through its own snapshot and sees the leader's committed write, + // proving cross-node reads work before any failover. + assert_eq!( + read_key(&db2, ALPHA_KEY).await, + Some(b"1".to_vec()), + "follower must see the leader's committed write" + ); + + let lease_before = read_lease(&raw).await.unwrap(); + let seq_before = read_seq_high(&raw).await; + assert!( + lease_before.durable_version >= 1, + "durable_version must have advanced after the first commit" + ); + + // Kill node 1. Dropping the driver aborts its resolver, so it stops renewing the lease. + drop(db1); + + // Node 2 must take over once node 1's lease expires (TTL is 10s). The epoch is bumped and the + // leader address changes to node 2. + let lease_after = wait_for_lease(&raw, Duration::from_secs(40), |l| { + l.epoch > lease_before.epoch + }) + .await; + assert!( + lease_after.epoch > lease_before.epoch, + "new leader must bump the epoch (was {}, now {})", + lease_before.epoch, + lease_after.epoch + ); + assert_ne!( + lease_after.leader_addr, leader1_addr, + "the surviving node must become the new leader" + ); + + // The crash-safe LOGGED sequence continues from the prior high-water; it never regresses. + let seq_after_takeover = read_seq_high(&raw).await; + assert!( + seq_after_takeover >= seq_before, + "version sequence regressed across failover ({} -> {})", + seq_before, + seq_after_takeover + ); + assert!( + lease_after.durable_version >= lease_before.durable_version, + "durable_version regressed across failover ({} -> {})", + lease_before.durable_version, + lease_after.durable_version + ); + + // The data the dead leader committed survives the failover. + assert_eq!( + read_key(&db2, ALPHA_KEY).await, + Some(b"1".to_vec()), + "committed data must survive leader failover" + ); + + // The new leader resumes accepting commits. + write_key(&db2, BETA_KEY, b"2").await; + assert_eq!( + read_key(&db2, BETA_KEY).await, + Some(b"2".to_vec()), + "new leader must accept and durably apply commits" + ); + + // The new commit advanced the version sequence and watermark past the pre-failover floor, + // confirming the new leader sequences from a strictly higher version. + let lease_final = read_lease(&raw).await.unwrap(); + assert!( + read_seq_high(&raw).await > seq_before, + "a post-failover commit must advance the version sequence" + ); + assert!( + lease_final.durable_version > lease_before.durable_version, + "a post-failover commit must advance the durable watermark" + ); + + drop(db2); +} + +/// Exercises graceful leader handoff: a leader that is shut down cleanly (SIGTERM path) releases its +/// lease immediately instead of letting it expire, so a standby takes over well within the lease TTL +/// rather than after it. This is what turns a rolling deploy from a ~TTL commit stall into a +/// near-instant handoff. +#[tokio::test] +async fn test_postgres_graceful_handoff() { + let _ = tracing_subscriber::fmt() + .with_env_filter("info") + .with_test_writer() + .try_init(); + + let (db_config, docker_config) = TestDatabase::Postgres + .config(Uuid::new_v4(), 1) + .await + .unwrap(); + let mut docker_config = docker_config.unwrap(); + docker_config.start().await.unwrap(); + + tokio::time::sleep(Duration::from_secs(4)).await; + + let rivet_config::config::Database::Postgres(postgres_config) = db_config else { + unreachable!(); + }; + let connection_string = postgres_config.url.read().clone(); + + let raw = connect_raw(&connection_string).await; + + // Node 1 wins the first election; node 2 joins as a follower. + let db1 = make_db(&connection_string).await; + let lease1 = wait_for_lease(&raw, Duration::from_secs(15), |l| l.epoch == 1).await; + let leader1_addr = lease1.leader_addr.clone(); + let db2 = make_db(&connection_string).await; + + write_key(&db1, ALPHA_KEY, b"1").await; + let lease_before = read_lease(&raw).await.unwrap(); + + // Gracefully shut down the leader. Unlike a hard drop, this releases the lease in place and + // wakes the standby, so takeover must complete in well under the 10s TTL. + let handoff_start = tokio::time::Instant::now(); + db1.shutdown().await; + + // The lease TTL is 10s; a graceful handoff must take over well under that. The 5s deadline here + // is itself below the TTL, so reaching this line already proves the lease was not waited out. + let lease_after = wait_for_lease(&raw, Duration::from_secs(5), |l| { + l.epoch > lease_before.epoch + }) + .await; + let handoff_elapsed = handoff_start.elapsed(); + assert!( + handoff_elapsed < Duration::from_secs(8), + "graceful handoff must beat the lease TTL (took {handoff_elapsed:?})" + ); + assert_ne!( + lease_after.leader_addr, leader1_addr, + "the standby must become the new leader after a graceful handoff" + ); + + // The new leader serves the old leader's data and accepts fresh commits. + assert_eq!( + read_key(&db2, ALPHA_KEY).await, + Some(b"1".to_vec()), + "committed data must survive graceful handoff" + ); + write_key(&db2, BETA_KEY, b"2").await; + assert_eq!(read_key(&db2, BETA_KEY).await, Some(b"2".to_vec())); + + drop(db1); + drop(db2); +} 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/universalpubsub/src/driver/postgres/doorbell.rs b/engine/packages/universalpubsub/src/driver/postgres/doorbell.rs new file mode 100644 index 0000000000..f08ff7bbf3 --- /dev/null +++ b/engine/packages/universalpubsub/src/driver/postgres/doorbell.rs @@ -0,0 +1,133 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use deadpool_postgres::Pool; +use tokio::sync::Notify; +use tokio::time::Instant; + +/// Number of doorbell shards. A subject maps to a shard via `hash(subject_hash) % K`. +/// Subscribers LISTEN their subject's shard channel; publishers wake the local +/// doorbell task which NOTIFYs the shard. +pub const DOORBELL_SHARD_COUNT: usize = 32; + +/// Debounce window. Caps each (process, shard) NOTIFY rate at one per window, which +/// bounds how many backends are woken per shard over time. +const DOORBELL_WINDOW: Duration = Duration::from_millis(5); + +/// Returns the NOTIFY channel name for a doorbell shard. +pub fn shard_channel(shard: usize) -> String { + format!("ups_db_{shard}") +} + +/// Returns the doorbell shard for a subject hash. +pub fn shard_for(subject_hash: &str) -> usize { + use std::hash::{DefaultHasher, Hash, Hasher}; + let mut hasher = DefaultHasher::new(); + subject_hash.hash(&mut hasher); + (hasher.finish() as usize) % DOORBELL_SHARD_COUNT +} + +/// Coalesced, payload-free NOTIFY doorbell. +/// +/// Publishers call [`Doorbell::mark_dirty`] after committing a row. A single +/// per-process task drains dirty shards and emits at most one NOTIFY per shard per +/// debounce window using leading-edge fire plus a trailing-edge flush. The doorbell +/// is a latency optimization only. Correctness comes from the table plus the +/// subscriber poll backstop, so a dropped or failed NOTIFY only adds latency. +pub struct Doorbell { + dirty: [AtomicBool; DOORBELL_SHARD_COUNT], + notify: Notify, + pool: Arc, +} + +impl Doorbell { + pub fn new(pool: Arc) -> Arc { + let doorbell = Arc::new(Self { + dirty: std::array::from_fn(|_| AtomicBool::new(false)), + notify: Notify::new(), + pool, + }); + + let task_doorbell = doorbell.clone(); + tokio::spawn(async move { task_doorbell.run().await }); + + doorbell + } + + /// Marks a shard dirty and wakes the doorbell task. Never blocks. + pub fn mark_dirty(&self, shard: usize) { + self.dirty[shard].store(true, Ordering::Release); + self.notify.notify_one(); + } + + async fn run(self: Arc) { + // Per-shard timestamp of the last NOTIFY emitted by this process. + let mut last_notify: [Option; DOORBELL_SHARD_COUNT] = [None; DOORBELL_SHARD_COUNT]; + // Per-shard deadline for a pending trailing-edge NOTIFY, if any. + let mut trailing: [Option; DOORBELL_SHARD_COUNT] = [None; DOORBELL_SHARD_COUNT]; + + loop { + // Arm on the next pending trailing deadline so the trailing edge fires + // even with no further publishes. Wait on the notify permit otherwise. + let next_deadline = trailing.iter().filter_map(|x| *x).min(); + match next_deadline { + Some(deadline) => { + tokio::select! { + _ = self.notify.notified() => {} + _ = tokio::time::sleep_until(deadline) => {} + } + } + None => { + self.notify.notified().await; + } + } + + let now = Instant::now(); + for shard in 0..DOORBELL_SHARD_COUNT { + let is_dirty = self.dirty[shard].swap(false, Ordering::AcqRel); + if is_dirty { + match last_notify[shard] { + Some(last) if now.duration_since(last) < DOORBELL_WINDOW => { + // Within the window. Defer to a trailing-edge NOTIFY at + // window end so at most one NOTIFY fires per shard per W. + if trailing[shard].is_none() { + trailing[shard] = Some(last + DOORBELL_WINDOW); + } + } + _ => { + // Leading edge. Fire immediately for low idle latency. + self.notify_shard(shard).await; + last_notify[shard] = Some(now); + trailing[shard] = None; + } + } + } + + // Flush a trailing-edge NOTIFY whose window has elapsed. + if let Some(deadline) = trailing[shard] { + if now >= deadline { + self.notify_shard(shard).await; + last_notify[shard] = Some(now); + trailing[shard] = None; + } + } + } + } + } + + async fn notify_shard(&self, shard: usize) { + let channel = shard_channel(shard); + match self.pool.get().await { + Ok(conn) => { + // Payload-free doorbell. The payload lives in the table. + if let Err(err) = conn.execute("SELECT pg_notify($1, '')", &[&channel]).await { + tracing::warn!(?err, %channel, "failed to emit doorbell notify"); + } + } + Err(err) => { + tracing::warn!(?err, %channel, "failed to get connection for doorbell notify"); + } + } + } +} diff --git a/engine/packages/universalpubsub/src/driver/postgres/mod.rs b/engine/packages/universalpubsub/src/driver/postgres/mod.rs index 3ef8f02beb..2a492fbccf 100644 --- a/engine/packages/universalpubsub/src/driver/postgres/mod.rs +++ b/engine/packages/universalpubsub/src/driver/postgres/mod.rs @@ -1,12 +1,11 @@ -use anyhow::{Context, Result, anyhow}; +use anyhow::{Context, Result}; use async_trait::async_trait; -use base64::Engine; -use base64::engine::general_purpose::STANDARD_NO_PAD as BASE64; use deadpool_postgres::{Config, ManagerConfig, Pool, PoolConfig, RecyclingMethod, Runtime}; use futures_util::future::poll_fn; use rivet_postgres_util::build_tls_config; -use rivet_util::backoff::Backoff; +use rivet_util::throttle::Backoff; use scc::HashMap; +use std::collections::VecDeque; use std::hash::{DefaultHasher, Hash, Hasher}; use std::path::PathBuf; use std::sync::Arc; @@ -21,49 +20,69 @@ use crate::driver::{PubSubDriver, SubscriberDriver, SubscriberDriverHandle}; use crate::metrics; use crate::pubsub::DriverOutput; -#[derive(Clone)] -struct Subscription { - // Channel to send messages to this subscription - tx: broadcast::Sender>, -} +mod doorbell; -impl Subscription { - fn new(tx: broadcast::Sender>) -> Self { - Self { tx } - } -} +use doorbell::{Doorbell, shard_channel, shard_for}; + +/// The transport is the table, not the NOTIFY payload, so there is no per-message +/// size cap from the 8000-byte NOTIFY limit. Match the NATS ceiling so chunking +/// behaves identically across drivers. +pub const POSTGRES_MAX_MESSAGE_SIZE: usize = 1024 * 1024; + +/// Poll backstop interval. Every subscriber reads its table on this interval +/// regardless of doorbell wakeups. This is the correctness floor that makes delivery +/// independent of any NOTIFY arriving. +const POLL_INTERVAL: Duration = Duration::from_secs(1); + +/// Idle-in-transaction timeout applied to the LISTEN connection. A wedged listener +/// holding a transaction open would otherwise fill the shared notify queue and fail +/// NOTIFY cluster-wide. Bounding it keeps a stuck listener degrading to added latency +/// rather than a cluster outage. +const LISTEN_IDLE_IN_TRANSACTION_TIMEOUT_MS: i64 = 30_000; + +/// How often this process refreshes its node liveness heartbeat. One heartbeat per +/// process keeps all of its subscriber registrations alive at once. +const NODE_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10); +/// How recent a node's heartbeat must be for its subscribers to count as live +/// responders. +const NODE_TTL_SECS: i64 = 30; + +/// How often to GC expired broadcast messages. +const MESSAGE_GC_INTERVAL: Duration = Duration::from_secs(5); +/// Max age before a broadcast message row is garbage collected. Must exceed the poll +/// interval plus the reconnect gap. A subscriber that falls behind this misses +/// messages, matching NATS-core at-most-once semantics for slow consumers. +const MESSAGE_MAX_AGE_SECS: i64 = 10; + +/// How often to GC dead nodes and the subscriber rows orphaned by them. +const REGISTRY_GC_INTERVAL: Duration = Duration::from_secs(30); -/// > 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; +/// Per-shard signal carried over a subscriber's in-process wakeup channel. +#[derive(Clone)] +enum ShardSignal { + /// A doorbell NOTIFY landed for this shard. Poll the table. + Wakeup, + /// A local request found no responders for the given reply subject. The matching + /// reply subscriber surfaces a no-responders result. + NoResponders { subject: String }, +} + #[derive(Clone)] pub struct PostgresDriver { pool: Arc, client: Arc>>, - subscriptions: Arc>, - /// Wakeup channels for queue subscriptions, keyed by queue channel name. - queue_subscriptions: Arc>, + /// Identifies this process in the subscriber registry. A single heartbeat keeps + /// all of this node's registrations live. + node_id: String, + /// Wakeup channels keyed by doorbell shard channel name. Shared by broadcast and + /// queue subscribers whose subjects map to the same shard. + shard_subscriptions: Arc>>, + doorbell: Arc, client_ready: tokio::sync::watch::Receiver, } @@ -103,9 +122,11 @@ impl PostgresDriver { .context("failed to create postgres pool")?; tracing::debug!("postgres pool created successfully"); - let subscriptions: Arc> = Arc::new(HashMap::new()); - let queue_subscriptions: Arc> = Arc::new(HashMap::new()); + let pool = Arc::new(pool); + let shard_subscriptions: Arc>> = + Arc::new(HashMap::new()); let client: Arc>> = Arc::new(Mutex::new(None)); + let node_id = Uuid::new_v4().to_string(); // Create channel for client ready notifications let (ready_tx, client_ready) = tokio::sync::watch::channel(false); @@ -113,8 +134,7 @@ impl PostgresDriver { // Spawn connection lifecycle task tokio::spawn(Self::spawn_connection_lifecycle( conn_str.clone(), - subscriptions.clone(), - queue_subscriptions.clone(), + shard_subscriptions.clone(), client.clone(), ready_tx, ssl_root_cert_path.clone(), @@ -122,34 +142,63 @@ impl PostgresDriver { ssl_client_key_path.clone(), )); + let doorbell = Doorbell::new(pool.clone()); + let driver = Self { - pool: Arc::new(pool), + pool, client, - subscriptions, - queue_subscriptions, + node_id, + shard_subscriptions, + doorbell, client_ready, }; // Wait for initial connection to be established driver.wait_for_client().await?; - // Create queue tables eagerly so they exist before any publish or subscribe + // Create tables eagerly so they exist before any publish or subscribe. { + tracing::debug!("configuring postgres udb tables"); let conn = driver .pool .get() .await - .context("failed to get connection for queue table creation")?; + .context("failed to get connection for table creation")?; conn.batch_execute( - "CREATE TABLE IF NOT EXISTS ups_queue_subs ( \ - id TEXT PRIMARY KEY, \ + // Broadcast transport table. UNLOGGED gives at-most-once across a + // crash, matching NATS-core semantics, and avoids WAL fsync on every + // publish. The real subject is stored so receivers can verify it and + // reject DefaultHasher subject-hash collisions. + "CREATE UNLOGGED TABLE IF NOT EXISTS ups_messages ( \ + id BIGSERIAL PRIMARY KEY, \ subject_hash TEXT NOT NULL, \ - queue_hash TEXT NOT NULL, \ + subject TEXT NOT NULL, \ + payload BYTEA NOT NULL, \ + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() \ + ); \ + CREATE INDEX IF NOT EXISTS ups_messages_subject_id \ + ON ups_messages (subject_hash, id); \ + CREATE TABLE IF NOT EXISTS ups_nodes ( \ + node_id TEXT PRIMARY KEY, \ heartbeat_at TIMESTAMPTZ NOT NULL DEFAULT NOW() \ ); \ + CREATE TABLE IF NOT EXISTS ups_subs ( \ + id TEXT PRIMARY KEY, \ + node_id TEXT NOT NULL, \ + subject_hash TEXT NOT NULL, \ + subject TEXT NOT NULL \ + ); \ + CREATE INDEX IF NOT EXISTS ups_subs_subject \ + ON ups_subs (subject_hash); \ + CREATE TABLE IF NOT EXISTS ups_queue_subs ( \ + id TEXT PRIMARY KEY, \ + node_id TEXT NOT NULL, \ + subject_hash TEXT NOT NULL, \ + queue_hash TEXT NOT NULL \ + ); \ CREATE INDEX IF NOT EXISTS ups_queue_subs_subject_queue \ ON ups_queue_subs (subject_hash, queue_hash); \ - CREATE TABLE IF NOT EXISTS ups_queue_messages ( \ + CREATE UNLOGGED TABLE IF NOT EXISTS ups_queue_messages ( \ id BIGSERIAL PRIMARY KEY, \ subject_hash TEXT NOT NULL, \ queue_hash TEXT NOT NULL, \ @@ -160,10 +209,91 @@ impl PostgresDriver { ON ups_queue_messages (subject_hash, queue_hash, id);", ) .await - .context("failed to create queue tables")?; - tracing::debug!("queue tables ready"); + .context("failed to create tables")?; + tracing::debug!("postgres udb tables ready"); } + // Register this node and start its liveness heartbeat. + driver.heartbeat_node().await?; + let heartbeat_driver = driver.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(NODE_HEARTBEAT_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + interval.tick().await; + if let Err(e) = heartbeat_driver.heartbeat_node().await { + tracing::warn!(?e, "failed to heartbeat node"); + } + } + }); + + // Spawn GC task for expired broadcast messages + let message_gc_driver = driver.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(MESSAGE_GC_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + interval.tick().await; + if let Ok(conn) = message_gc_driver.pool.get().await { + let result = conn + .execute( + "DELETE FROM ups_messages \ + WHERE created_at < NOW() - ($1::bigint * INTERVAL '1 second')", + &[&MESSAGE_MAX_AGE_SECS], + ) + .await; + if let Err(e) = result { + tracing::warn!(?e, "failed to gc broadcast messages"); + } + } + } + }); + + // Spawn GC task for dead nodes and orphaned subscriber rows. + let registry_gc_driver = driver.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(REGISTRY_GC_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + interval.tick().await; + if let Ok(conn) = registry_gc_driver.pool.get().await { + if let Err(e) = conn + .execute( + "DELETE FROM ups_nodes \ + WHERE heartbeat_at < NOW() - ($1::bigint * INTERVAL '1 second')", + &[&NODE_TTL_SECS], + ) + .await + { + tracing::warn!(?e, "failed to gc dead nodes"); + } + if let Err(e) = conn + .execute( + "DELETE FROM ups_subs \ + WHERE node_id NOT IN (SELECT node_id FROM ups_nodes)", + &[], + ) + .await + { + tracing::warn!(?e, "failed to gc orphaned subs"); + } + if let Err(e) = conn + .execute( + "DELETE FROM ups_queue_subs \ + WHERE node_id NOT IN (SELECT node_id FROM ups_nodes)", + &[], + ) + .await + { + tracing::warn!(?e, "failed to gc orphaned queue subs"); + } + } + } + }); + // Spawn GC task for orphaned queue messages let gc_driver = driver.clone(); tokio::spawn(async move { @@ -193,8 +323,7 @@ impl PostgresDriver { /// Manages the connection lifecycle with automatic reconnection async fn spawn_connection_lifecycle( conn_str: String, - subscriptions: Arc>, - queue_subscriptions: Arc>, + shard_subscriptions: Arc>>, client: Arc>>, ready_tx: tokio::sync::watch::Sender, ssl_root_cert_path: Option, @@ -227,41 +356,42 @@ impl PostgresDriver { // Spawn the polling task immediately // This must be done before any operations on the client - let subscriptions_clone = subscriptions.clone(); - let queue_subscriptions_clone = queue_subscriptions.clone(); + let shard_subscriptions_clone = shard_subscriptions.clone(); let poll_handle = tokio::spawn(async move { - Self::poll_connection(conn, subscriptions_clone, queue_subscriptions_clone) - .await; + Self::poll_connection(conn, shard_subscriptions_clone).await; }); - // Get regular channels to re-subscribe to + // Bound a stuck listener so it cannot wedge the shared notify queue. + if let Result::Err(e) = new_client + .execute( + &format!( + "SET idle_in_transaction_session_timeout = '{}'", + LISTEN_IDLE_IN_TRANSACTION_TIMEOUT_MS + ), + &[], + ) + .await + { + tracing::warn!(?e, "failed to set idle_in_transaction_session_timeout"); + } + + // Get shard channels to re-subscribe to let mut channels = Vec::new(); - subscriptions + shard_subscriptions .iter_async(|k, _| { channels.push(k.clone()); true }) .await; - // Get queue wakeup channels to re-subscribe to - let mut queue_channels = Vec::new(); - queue_subscriptions - .iter_async(|k, _| { - queue_channels.push(k.clone()); - true - }) - .await; - - let needs_resubscribe = !channels.is_empty() || !queue_channels.is_empty(); - if needs_resubscribe { + if !channels.is_empty() { tracing::debug!( - regular_channels = channels.len(), - queue_channels = queue_channels.len(), - "re-subscribing to channels after reconnection" + channels = channels.len(), + "re-subscribing to doorbell shards after reconnection" ); } - for channel in channels.iter().chain(queue_channels.iter()) { + for channel in channels.iter() { tracing::debug!(?channel, "re-subscribing to channel"); if let Result::Err(e) = new_client .execute(&format!("LISTEN \"{}\"", channel), &[]) @@ -298,30 +428,20 @@ impl PostgresDriver { /// Polls the connection for notifications until it closes or errors async fn poll_connection( mut conn: tokio_postgres::Connection, - subscriptions: Arc>, - queue_subscriptions: Arc>, + shard_subscriptions: Arc>>, ) where T: tokio_postgres::tls::TlsStream + Unpin, { loop { match poll_fn(|cx| conn.poll_message(cx)).await { Some(std::result::Result::Ok(AsyncMessage::Notification(note))) => { - tracing::trace!(channel = %note.channel(), "received notification"); - if let Some(sub) = subscriptions.get_async(note.channel()).await { - let bytes = match BASE64.decode(note.payload()) { - std::result::Result::Ok(b) => b, - std::result::Result::Err(err) => { - tracing::error!(?err, "failed decoding base64"); - continue; - } - }; - tracing::trace!(channel = %note.channel(), bytes_len = bytes.len(), "sending to broadcast channel"); - let _ = sub.tx.send(bytes); - } else if let Some(sub) = queue_subscriptions.get_async(note.channel()).await { - // Queue notifications are wakeup signals only; payload lives in the table - let _ = sub.tx.send(Vec::new()); + tracing::trace!(channel = %note.channel(), "received doorbell wakeup"); + // Doorbell notifications are payload-free wakeup signals only. + // Subscribers read their payload from the table. + if let Some(sub) = shard_subscriptions.get_async(note.channel()).await { + let _ = sub.send(ShardSignal::Wakeup); } else { - tracing::warn!(channel = %note.channel(), "received notification for unknown channel"); + tracing::trace!(channel = %note.channel(), "wakeup for unknown shard"); } } Some(std::result::Result::Ok(_)) => { @@ -361,8 +481,9 @@ impl PostgresDriver { } fn hash_subject(&self, subject: &str) -> String { - // Postgres channel names have a 64 character limit - // Hash the subject to ensure it fits + // Postgres channel names have a 64 character limit, but this hash is also the + // table index key. Collisions are possible and resolved by verifying the real + // subject stored alongside each row. let mut hasher = DefaultHasher::new(); subject.hash(&mut hasher); format!("ups_{:x}", hasher.finish()) @@ -374,57 +495,89 @@ impl PostgresDriver { format!("{:x}", hasher.finish()) } - /// Returns the NOTIFY channel name for a (subject, queue) pair. - fn queue_channel(&self, subject_hash: &str, queue_hash: &str) -> String { - // Max length: "ups_q_" (6) + 16 + "_" (1) + 16 = 39 chars, well within 64 - format!("ups_q_{}_{}", subject_hash, queue_hash) + /// Upserts this node's liveness heartbeat. Re-inserts the row if a GC pass removed + /// it after a transient stall. + async fn heartbeat_node(&self) -> Result<()> { + let conn = self + .pool + .get() + .await + .context("failed to get connection for node heartbeat")?; + conn.execute( + "INSERT INTO ups_nodes (node_id, heartbeat_at) VALUES ($1, NOW()) \ + ON CONFLICT (node_id) DO UPDATE SET heartbeat_at = NOW()", + &[&self.node_id], + ) + .await + .context("failed to upsert node heartbeat")?; + Ok(()) } - /// Inserts messages into the queue table and notifies active queue subscribers. - async fn publish_to_queues(&self, subject: &str, payload: &[u8]) -> Result<()> { - let subject_hash = self.hash_subject(subject); - + /// Returns the current max broadcast message id, used as a subscriber's starting + /// cursor so it only sees future messages (NATS at-most-once, no replay). + async fn current_max_id(&self) -> Result { let conn = self .pool .get() .await - .context("failed to get connection for queue publish")?; - - // Find active queue groups for this subject - let rows = conn - .query( - "SELECT DISTINCT queue_hash FROM ups_queue_subs \ - WHERE subject_hash = $1 \ - AND heartbeat_at > NOW() - ($2::bigint * INTERVAL '1 second')", - &[&subject_hash, &QUEUE_SUB_TTL_SECS], - ) + .context("failed to get connection for cursor init")?; + let row = conn + .query_one("SELECT COALESCE(MAX(id), 0) FROM ups_messages", &[]) .await - .context("failed to query active queue subs")?; + .context("failed to read current max id")?; + Ok(row.get(0)) + } - for row in rows { - let queue_hash: String = row.get(0); - let channel = self.queue_channel(&subject_hash, &queue_hash); + /// Ensures this process is LISTENing on the given doorbell shard and returns a + /// wakeup receiver plus a drop guard that UNLISTENs once no receivers remain. + async fn ensure_shard_listen( + &self, + shard: usize, + ) -> ( + broadcast::Receiver, + tokio_util::sync::DropGuard, + ) { + let channel = shard_channel(shard); - conn.execute( - "INSERT INTO ups_queue_messages (subject_hash, queue_hash, payload) \ - VALUES ($1, $2, $3)", - &[&subject_hash, &queue_hash, &payload], - ) - .await - .context("failed to insert queue message")?; + match self.shard_subscriptions.entry_async(channel.clone()).await { + scc::hash_map::Entry::Occupied(existing) => { + let rx = existing.subscribe(); + let drop_guard = + self.spawn_shard_cleanup_task(channel.clone(), existing.get().clone()); + (rx, drop_guard) + } + scc::hash_map::Entry::Vacant(e) => { + let (tx, rx) = broadcast::channel(1024); + e.insert_entry(tx.clone()); + metrics::POSTGRES_SUBSCRIPTION_COUNT.set(self.shard_subscriptions.len() as i64); - conn.execute(&format!("NOTIFY \"{}\"", channel), &[]) - .await - .context("failed to notify queue channel")?; - } + if let Some(client) = &*self.client.lock().await { + match client + .execute(&format!("LISTEN \"{channel}\""), &[]) + .instrument(tracing::trace_span!("pg_listen")) + .await + { + Result::Ok(_) => { + tracing::debug!(%channel, "successfully subscribed to shard"); + } + Result::Err(e) => { + tracing::warn!(?e, %channel, "failed to LISTEN, will retry on reconnection"); + } + } + } else { + tracing::debug!(%channel, "client not connected, will LISTEN on reconnection"); + } - Ok(()) + let drop_guard = self.spawn_shard_cleanup_task(channel.clone(), tx.clone()); + (rx, drop_guard) + } + } } - fn spawn_subscription_cleanup_task( + fn spawn_shard_cleanup_task( &self, - subject_hash: String, - tx: broadcast::Sender>, + channel: String, + tx: broadcast::Sender, ) -> tokio_util::sync::DropGuard { let driver = self.clone(); let token = tokio_util::sync::CancellationToken::new(); @@ -434,47 +587,117 @@ impl PostgresDriver { token.cancelled().await; if tx.receiver_count() == 0 { if let Some(client) = &*driver.client.lock().await { - let sql = format!("UNLISTEN \"{}\"", subject_hash); + let sql = format!("UNLISTEN \"{}\"", channel); if let Err(err) = client.execute(sql.as_str(), &[]).await { - tracing::warn!(?err, %subject_hash, "failed to UNLISTEN channel"); + tracing::warn!(?err, %channel, "failed to UNLISTEN channel"); } else { - tracing::trace!(%subject_hash, "unlistened channel"); + tracing::trace!(%channel, "unlistened channel"); } } - driver.subscriptions.remove_async(&subject_hash).await; - metrics::POSTGRES_SUBSCRIPTION_COUNT.set(driver.subscriptions.len() as i64); + driver.shard_subscriptions.remove_async(&channel).await; + metrics::POSTGRES_SUBSCRIPTION_COUNT.set(driver.shard_subscriptions.len() as i64); } }); drop_guard } - fn spawn_queue_subscription_cleanup_task( + /// Inserts the broadcast row and any active queue-group rows in one transaction. + async fn try_publish_to_db( &self, - channel: String, - tx: broadcast::Sender>, - ) -> tokio_util::sync::DropGuard { - let driver = self.clone(); - let token = tokio_util::sync::CancellationToken::new(); - let drop_guard = token.clone().drop_guard(); + subject: &str, + subject_hash: &str, + payload: &[u8], + ) -> Result<()> { + let mut conn = self + .pool + .get() + .await + .context("failed to get connection for publish")?; + let tx = conn + .transaction() + .await + .context("failed to begin publish transaction")?; - tokio::spawn(async move { - token.cancelled().await; - if tx.receiver_count() == 0 { - if let Some(client) = &*driver.client.lock().await { - let sql = format!("UNLISTEN \"{}\"", channel); + // Broadcast row. + tx.execute( + "INSERT INTO ups_messages (subject_hash, subject, payload) VALUES ($1, $2, $3)", + &[&subject_hash, &subject, &payload], + ) + .await + .context("failed to insert broadcast message")?; - if let Err(err) = client.execute(sql.as_str(), &[]).await { - tracing::warn!(?err, %channel, "failed to UNLISTEN queue channel"); - } else { - tracing::trace!(%channel, "unlistened queue channel"); - } - } - driver.queue_subscriptions.remove_async(&channel).await; - } - }); + // Queue rows for every live queue group on this subject. Batched into the same + // transaction so a crash never strands a row mid-publish. + let rows = tx + .query( + "SELECT DISTINCT s.queue_hash FROM ups_queue_subs s \ + JOIN ups_nodes n ON s.node_id = n.node_id \ + WHERE s.subject_hash = $1 \ + AND n.heartbeat_at > NOW() - ($2::bigint * INTERVAL '1 second')", + &[&subject_hash, &NODE_TTL_SECS], + ) + .await + .context("failed to query active queue subs")?; - drop_guard + for row in rows { + let queue_hash: String = row.get(0); + tx.execute( + "INSERT INTO ups_queue_messages (subject_hash, queue_hash, payload) \ + VALUES ($1, $2, $3)", + &[&subject_hash, &queue_hash, &payload], + ) + .await + .context("failed to insert queue message")?; + } + + tx.commit().await.context("failed to commit publish")?; + + Ok(()) + } + + /// Returns whether any live subscriber (broadcast or queue) exists for the subject + /// anywhere in the fleet. Used to decide whether a request surfaces a no-responders + /// result instead of waiting out its timeout. + async fn has_responders(&self, subject_hash: &str, subject: &str) -> Result { + let conn = self + .pool + .get() + .await + .context("failed to get connection for responder check")?; + let row = conn + .query_one( + "SELECT \ + EXISTS( \ + SELECT 1 FROM ups_subs s \ + JOIN ups_nodes n ON s.node_id = n.node_id \ + WHERE s.subject_hash = $1 AND s.subject = $2 \ + AND n.heartbeat_at > NOW() - ($3::bigint * INTERVAL '1 second') \ + ) \ + OR EXISTS( \ + SELECT 1 FROM ups_queue_subs s \ + JOIN ups_nodes n ON s.node_id = n.node_id \ + WHERE s.subject_hash = $1 \ + AND n.heartbeat_at > NOW() - ($3::bigint * INTERVAL '1 second') \ + )", + &[&subject_hash, &subject, &NODE_TTL_SECS], + ) + .await + .context("failed to check responders")?; + Ok(row.get(0)) + } + + /// Delivers a no-responders result to the local reply subscriber. The requester is + /// always in this process, so the signal is routed in-memory over the reply + /// subject's shard channel rather than the table. + async fn signal_no_responders(&self, reply_subject: &str) { + let reply_hash = self.hash_subject(reply_subject); + let channel = shard_channel(shard_for(&reply_hash)); + if let Some(tx) = self.shard_subscriptions.get_async(&channel).await { + let _ = tx.send(ShardSignal::NoResponders { + subject: reply_subject.to_string(), + }); + } } } @@ -483,68 +706,48 @@ impl PubSubDriver for PostgresDriver { async fn subscribe( &self, subject: &str, - _reply_id: Option, + reply_id: Option, ) -> Result { - // TODO: To match NATS implementation, LISTEN must be pipelined (i.e. wait for the command - // to reach the server, but not wait for it to respond). However, this has to ensure that - // NOTIFY & LISTEN are called on the same connection (not diff connections in a pool) or - // else there will be race conditions where messages might be published before - // subscriptions are registered. - // - // tokio-postgres currently does not expose the API for pipelining, so we are SOL. - // - // We might be able to use a background tokio task in combination with flush if we use the - // same Postgres connection, but unsure if that will create a bottleneck. - - let hashed = self.hash_subject(subject); - - // Check if we already have a subscription for this channel - let (rx, drop_guard) = match self.subscriptions.entry_async(hashed.clone()).await { - scc::hash_map::Entry::Occupied(existing_sub) => { - // Reuse the existing broadcast channel - let rx = existing_sub.tx.subscribe(); - let drop_guard = - self.spawn_subscription_cleanup_task(hashed.clone(), existing_sub.tx.clone()); - (rx, drop_guard) - } - scc::hash_map::Entry::Vacant(e) => { - // Create a new broadcast channel for this subject - let (tx, rx) = tokio::sync::broadcast::channel(1024); - let subscription = Subscription::new(tx.clone()); - - // Register subscription - e.insert_entry(subscription.clone()); - metrics::POSTGRES_SUBSCRIPTION_COUNT.set(self.subscriptions.len() as i64); - - // Execute LISTEN command on the async client (for receiving notifications) - // This only needs to be done once per channel - // Try to LISTEN if client is available, but don't fail if disconnected - // The reconnection logic will handle re-subscribing - if let Some(client) = &*self.client.lock().await { - match client - .execute(&format!("LISTEN \"{hashed}\""), &[]) - .instrument(tracing::trace_span!("pg_listen")) - .await - { - Result::Ok(_) => { - tracing::debug!(%hashed, "successfully subscribed to channel"); - } - Result::Err(e) => { - tracing::warn!(?e, %hashed, "failed to LISTEN, will retry on reconnection"); - } - } - } else { - tracing::debug!(%hashed, "client not connected, will LISTEN on reconnection"); - } + let subject_hash = self.hash_subject(subject); + let shard = shard_for(&subject_hash); - let drop_guard = self.spawn_subscription_cleanup_task(hashed.clone(), tx.clone()); - (rx, drop_guard) - } + // Capture the cursor before LISTENing. Any message inserted after this point + // has a higher id and is delivered either by the doorbell wakeup or the poll + // backstop, so there is no subscribe/publish race. + let cursor = self.current_max_id().await?; + + let (rx, drop_guard) = self.ensure_shard_listen(shard).await; + + // Register in the responder registry so requests to this subject can detect + // responders. Reply inboxes are never request targets, so they skip the + // registry to keep request latency off this path. + let sub_id = if reply_id.is_none() { + let sub_id = Uuid::new_v4().to_string(); + let conn = self + .pool + .get() + .await + .context("failed to get connection for subscribe")?; + conn.execute( + "INSERT INTO ups_subs (id, node_id, subject_hash, subject) \ + VALUES ($1, $2, $3, $4)", + &[&sub_id, &self.node_id, &subject_hash, &subject], + ) + .await + .context("failed to register subscriber")?; + Some(sub_id) + } else { + None }; Ok(Box::new(PostgresSubscriber { subject: subject.to_string(), - rx: Some(rx), + subject_hash, + pool: self.pool.clone(), + cursor, + buffer: VecDeque::new(), + rx, + sub_id, _drop_guard: drop_guard, })) } @@ -552,7 +755,7 @@ impl PubSubDriver for PostgresDriver { async fn queue_subscribe(&self, subject: &str, queue: &str) -> Result { let subject_hash = self.hash_subject(subject); let queue_hash = self.hash_queue(queue); - let channel = self.queue_channel(&subject_hash, &queue_hash); + let shard = shard_for(&subject_hash); // Register this subscriber in the database so publishers know the queue exists let sub_id = Uuid::new_v4().to_string(); @@ -563,80 +766,15 @@ impl PubSubDriver for PostgresDriver { .await .context("failed to get connection for queue subscribe")?; conn.execute( - "INSERT INTO ups_queue_subs (id, subject_hash, queue_hash) VALUES ($1, $2, $3)", - &[&sub_id, &subject_hash, &queue_hash], + "INSERT INTO ups_queue_subs (id, node_id, subject_hash, queue_hash) \ + VALUES ($1, $2, $3, $4)", + &[&sub_id, &self.node_id, &subject_hash, &queue_hash], ) .await .context("failed to register queue subscriber")?; } - // 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"); - } - } - } - } - } - }); + let (rx, drop_guard) = self.ensure_shard_listen(shard).await; Ok(Box::new(PostgresQueueSubscriber { subject: subject.to_string(), @@ -644,9 +782,8 @@ impl PubSubDriver for PostgresDriver { queue_hash, sub_id, pool: self.pool.clone(), - rx: Some(rx), + rx, _drop_guard: drop_guard, - _heartbeat_token: heartbeat_token, })) } @@ -654,80 +791,53 @@ impl PubSubDriver for PostgresDriver { &self, subject: &str, payload: &[u8], - _reply_subject: Option<&str>, + reply_subject: Option<&str>, ) -> Result<()> { - // TODO: See `subscribe` about pipelining - - // Encode payload to base64 and send NOTIFY - let encoded = BASE64.encode(payload); - let hashed = self.hash_subject(subject); - - tracing::trace!("attempting to get connection for publish"); - - // Wait for listen connection to be ready first if this channel has subscribers - // This ensures that if we're reconnecting, the LISTEN is re-registered before NOTIFY - if self.subscriptions.contains_async(&hashed).await { - self.wait_for_client().await?; + let subject_hash = self.hash_subject(subject); + let shard = shard_for(&subject_hash); + + // Request semantics: if a reply is expected and no responder exists anywhere, + // surface a no-responders result immediately instead of persisting a message + // nobody will read. + if let Some(reply_subject) = reply_subject { + match self.has_responders(&subject_hash, subject).await { + Result::Ok(false) => { + self.signal_no_responders(reply_subject).await; + return Ok(()); + } + Result::Ok(true) => {} + Result::Err(e) => { + // On a failed check, fall through to a normal publish rather than + // risk a false no-responders result. + tracing::warn!(?e, %subject, "responder check failed, publishing anyway"); + } + } } - // Retry getting a connection from the pool with backoff in case the connection is - // currently disconnected + // Persist the message, retrying on transient connection errors. The row is + // committed before the doorbell rings so any wakeup observes it. let mut backoff = Backoff::default(); - let mut last_error; - loop { - match self.pool.get().await { - Result::Ok(conn) => { - // Test the connection with a simple query before using it - match conn.execute("SELECT 1", &[]).await { - Result::Ok(_) => { - // Connection is good; run NOTIFY and queue publish in parallel. - // publish_to_queues acquires its own pool connection so both - // can proceed concurrently. - let notify_sql = format!("NOTIFY \"{hashed}\", '{encoded}'"); - let (notify_result, queue_result) = tokio::join!( - conn.execute(notify_sql.as_str(), &[]) - .instrument(tracing::trace_span!("pg_notify")), - self.publish_to_queues(subject, payload), - ); - match notify_result { - Result::Ok(_) => { - if let Err(e) = queue_result { - tracing::warn!(?e, %subject, "failed to publish to queue subscribers"); - } - return Ok(()); - } - Result::Err(e) => { - tracing::debug!( - ?e, - "NOTIFY failed, retrying with new connection" - ); - last_error = Some(e.into()); - } - } - } - Result::Err(e) => { - tracing::debug!( - ?e, - "connection test failed, retrying with new connection" - ); - last_error = Some(e.into()); - } - } - } + match self + .try_publish_to_db(subject, &subject_hash, payload) + .await + { + Result::Ok(()) => break, Result::Err(e) => { - tracing::debug!(?e, "failed to get connection from pool, retrying"); - last_error = Some(e.into()); + if !backoff.tick().await { + tracing::warn!(?e, %subject, "failed to publish, cannot retry again"); + return Err(e); + } + tracing::debug!(?e, "publish failed, retrying"); } } - - // Check if we should continue retrying - if !backoff.tick().await { - return Err( - last_error.unwrap_or_else(|| anyhow!("failed to publish after retries")) - ); - } } + + // Ring the doorbell. Best-effort: the subscriber poll backstop covers a + // dropped or coalesced wakeup, so publish never blocks on NOTIFY. + self.doorbell.mark_dirty(shard); + + Ok(()) } async fn flush(&self) -> Result<()> { @@ -741,40 +851,120 @@ impl PubSubDriver for PostgresDriver { pub struct PostgresSubscriber { subject: String, - rx: Option>>, + subject_hash: String, + pool: Arc, + cursor: i64, + buffer: VecDeque>, + rx: broadcast::Receiver, + /// Responder-registry row id, present for non-inbox subscriptions. Deleted on drop. + sub_id: Option, _drop_guard: tokio_util::sync::DropGuard, } +impl PostgresSubscriber { + /// Reads new rows past the cursor into the buffer, advancing the cursor. Rows + /// whose stored subject does not match are skipped (DefaultHasher collisions) but + /// still advance the cursor. + async fn fetch(&mut self) -> Result<()> { + let conn = self + .pool + .get() + .await + .context("failed to get connection for poll")?; + let rows = conn + .query( + "SELECT id, subject, payload FROM ups_messages \ + WHERE subject_hash = $1 AND id > $2 ORDER BY id", + &[&self.subject_hash, &self.cursor], + ) + .await + .context("failed to poll broadcast messages")?; + + for row in rows { + let id: i64 = row.get(0); + let subject: String = row.get(1); + let payload: Vec = row.get(2); + self.cursor = id; + if subject == self.subject { + self.buffer.push_back(payload); + } + } + + Ok(()) + } +} + #[async_trait] impl SubscriberDriver for PostgresSubscriber { async fn next(&mut self) -> Result { - let rx = match self.rx.as_mut() { - Some(rx) => rx, - None => return Ok(DriverOutput::Unsubscribed), - }; - match rx.recv().await { - std::result::Result::Ok(payload) => Ok(DriverOutput::Message { - subject: self.subject.clone(), - payload, - }), - Err(tokio::sync::broadcast::error::RecvError::Closed) => Ok(DriverOutput::Unsubscribed), - Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => { - // Try again - self.next().await + loop { + if let Some(payload) = self.buffer.pop_front() { + return Ok(DriverOutput::Message { + subject: self.subject.clone(), + payload, + }); + } + + if let Err(e) = self.fetch().await { + // Transient DB errors must not kill the subscriber; the next poll + // tick retries. + tracing::warn!(?e, subject = %self.subject, "failed to poll, will retry"); + } + + if !self.buffer.is_empty() { + continue; + } + + // Wait for a doorbell wakeup, a no-responders signal, or the poll backstop. + tokio::select! { + res = self.rx.recv() => { + match res { + std::result::Result::Ok(ShardSignal::Wakeup) => {} + std::result::Result::Ok(ShardSignal::NoResponders { subject }) + if subject == self.subject => + { + return Ok(DriverOutput::NoResponders); + } + std::result::Result::Ok(ShardSignal::NoResponders { .. }) => {} + Err(broadcast::error::RecvError::Lagged(_)) => {} + Err(broadcast::error::RecvError::Closed) => { + return Ok(DriverOutput::Unsubscribed); + } + } + } + _ = tokio::time::sleep(POLL_INTERVAL) => {} } } } } +impl Drop for PostgresSubscriber { + fn drop(&mut self) { + let Some(sub_id) = self.sub_id.take() else { + return; + }; + let pool = self.pool.clone(); + tokio::spawn(async move { + if let Ok(conn) = pool.get().await { + if let Err(e) = conn + .execute("DELETE FROM ups_subs WHERE id = $1", &[&sub_id]) + .await + { + tracing::warn!(?e, %sub_id, "failed to deregister subscriber"); + } + } + }); + } +} + pub struct PostgresQueueSubscriber { subject: String, subject_hash: String, queue_hash: String, sub_id: String, pool: Arc, - rx: Option>>, + rx: broadcast::Receiver, _drop_guard: tokio_util::sync::DropGuard, - _heartbeat_token: tokio_util::sync::CancellationToken, } impl PostgresQueueSubscriber { @@ -811,31 +1001,32 @@ impl PostgresQueueSubscriber { impl SubscriberDriver for PostgresQueueSubscriber { async fn next(&mut self) -> Result { loop { - // Drain any messages that arrived before or between notifications. - // Do this before borrowing rx so claim_message can borrow self freely. - if let Some(payload) = self.claim_message().await? { - return Ok(DriverOutput::Message { - subject: self.subject.clone(), - payload, - }); - } - - // Wait for a wakeup notification, then loop back to claim. - let rx = match self.rx.as_mut() { - Some(rx) => rx, - None => return Ok(DriverOutput::Unsubscribed), - }; - match rx.recv().await { - std::result::Result::Ok(_) => { - // Wakeup received; loop back to claim + // Drain any messages that arrived before or between wakeups. + match self.claim_message().await { + Result::Ok(Some(payload)) => { + return Ok(DriverOutput::Message { + subject: self.subject.clone(), + payload, + }); } - Err(tokio::sync::broadcast::error::RecvError::Closed) => { - return Ok(DriverOutput::Unsubscribed); + Result::Ok(None) => {} + Result::Err(e) => { + tracing::warn!(?e, subject = %self.subject, "failed to claim, will retry"); } - Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => { - // Notifications were dropped while lagged; loop back to claim in case - // messages are waiting + } + + // Wait for any shard signal or the poll backstop, then loop back to claim. + tokio::select! { + res = self.rx.recv() => { + match res { + std::result::Result::Ok(_) => {} + Err(broadcast::error::RecvError::Lagged(_)) => {} + Err(broadcast::error::RecvError::Closed) => { + return Ok(DriverOutput::Unsubscribed); + } + } } + _ = tokio::time::sleep(POLL_INTERVAL) => {} } } } 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/reconnect.rs b/engine/packages/universalpubsub/tests/reconnect.rs index 35ace375d3..8c15f6ab54 100644 --- a/engine/packages/universalpubsub/tests/reconnect.rs +++ b/engine/packages/universalpubsub/tests/reconnect.rs @@ -43,7 +43,7 @@ async fn test_nats_driver_with_memory_reconnect() { .unwrap(); let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), true); - test_all_inner(&pubsub, &docker).await; + test_all_inner(&pubsub, &docker, true).await; } #[tokio::test] @@ -77,7 +77,7 @@ async fn test_nats_driver_without_memory_reconnect() { .unwrap(); let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false); - test_all_inner(&pubsub, &docker).await; + test_all_inner(&pubsub, &docker, true).await; } #[tokio::test] @@ -95,13 +95,12 @@ async fn test_postgres_driver_with_memory_reconnect() { }; let url = pg.url.read().clone(); - let driver = - universalpubsub::driver::postgres::PostgresDriver::connect(url, true, None, None, None) - .await - .unwrap(); + let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, None, None, None) + .await + .unwrap(); let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), true); - test_all_inner(&pubsub, &docker).await; + test_all_inner(&pubsub, &docker, false).await; } #[tokio::test] @@ -119,19 +118,27 @@ async fn test_postgres_driver_without_memory_reconnect() { }; let url = pg.url.read().clone(); - let driver = - universalpubsub::driver::postgres::PostgresDriver::connect(url, false, None, None, None) - .await - .unwrap(); + let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, None, None, None) + .await + .unwrap(); let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false); - test_all_inner(&pubsub, &docker).await; + test_all_inner(&pubsub, &docker, false).await; } -async fn test_all_inner(pubsub: &PubSub, docker: &rivet_test_deps_docker::DockerRunConfig) { +async fn test_all_inner( + pubsub: &PubSub, + docker: &rivet_test_deps_docker::DockerRunConfig, + supports_subscribe_while_stopped: bool, +) { test_reconnect_inner(&pubsub, &docker).await; test_publish_while_stopped(&pubsub, &docker).await; - test_subscribe_while_stopped(&pubsub, &docker).await; + // The table-backed Postgres driver must read its cursor and register in the + // responder table when subscribing, so it cannot subscribe while the backend is + // fully down. NATS buffers the subscribe and reconnects, so it can. + if supports_subscribe_while_stopped { + test_subscribe_while_stopped(&pubsub, &docker).await; + } } async fn test_reconnect_inner(pubsub: &PubSub, docker: &rivet_test_deps_docker::DockerRunConfig) { diff --git a/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..a9d637aa81 --- /dev/null +++ b/engine/sdks/rust/universaldb-commit/src/versioned.rs @@ -0,0 +1,38 @@ +use anyhow::{Ok, Result, bail}; +use vbare::OwnedVersionedData; + +use crate::generated::v1; + +// Only v1 exists today. When adding v2+, generate converters with +// `scripts/vbare-gen-converters` (see the envoy-protocol package for the +// resulting `versioned/` module layout) and wire them in here. +pub enum CommitRequest { + V1(v1::CommitRequest), +} + +impl OwnedVersionedData for CommitRequest { + type Latest = v1::CommitRequest; + + fn wrap_latest(latest: v1::CommitRequest) -> Self { + CommitRequest::V1(latest) + } + + fn unwrap_latest(self) -> Result { + match self { + CommitRequest::V1(data) => Ok(data), + } + } + + fn deserialize_version(payload: &[u8], version: u16) -> Result { + match version { + 1 => Ok(CommitRequest::V1(serde_bare::from_slice(payload)?)), + _ => bail!("invalid version: {version}"), + } + } + + fn serialize_version(self, _version: u16) -> Result> { + match self { + CommitRequest::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + } + } +} diff --git a/engine/sdks/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..2377312f84 --- /dev/null +++ b/engine/sdks/schemas/universaldb-commit/v1.bare @@ -0,0 +1,70 @@ +# Commit-queue wire format for the Postgres leader-resolver UDB driver. +# +# Followers encode a CommitRequest into the `payload` column of +# `udb_commit_requests`; the leader decodes it to resolve and apply. +# Rust-only (never leaves the engine), but versioned so rolling deploys can +# skew follower vs leader code. + +type ConflictRangeType enum { + READ + WRITE +} + +type ConflictRange struct { + begin: data + end: data + kind: ConflictRangeType +} + +# Order MUST match universaldb::options::MutationType declaration order so the +# enum tag round-trips. 15 variants today; never reorder, only append. +type MutationType enum { + ADD + AND + BIT_AND + OR + BIT_OR + XOR + BIT_XOR + APPEND_IF_FITS + MAX + MIN + SET_VERSIONSTAMPED_KEY + SET_VERSIONSTAMPED_VALUE + BYTE_MIN + BYTE_MAX + COMPARE_AND_CLEAR +} + +type SetValue struct { + key: data + value: data +} + +type Clear struct { + key: data +} + +type ClearRange struct { + begin: data + end: data +} + +type AtomicOp struct { + key: data + param: data + opType: MutationType +} + +type Operation union { + SetValue | + Clear | + ClearRange | + AtomicOp +} + +type CommitRequest struct { + readVersion: u64 + conflictRanges: list + operations: list +} diff --git a/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..4e3781aab2 100644 --- a/self-host/compose/template/src/docker-compose.ts +++ b/self-host/compose/template/src/docker-compose.ts @@ -1,6 +1,8 @@ import * as yaml from "js-yaml"; import { CORE_NETWORK_NAME, type TemplateContext } from "./context"; +const RUNNER_CONFIG_INIT_SERVICE = "runner-config-init"; + export function generateDockerCompose(context: TemplateContext) { const config = context.config; @@ -171,7 +173,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) plus a + // dedicated LISTEN connection and pubsub, so a multi-engine datacenter + // needs far more than the default max_connections of 100. + command: ["postgres", "-c", "max_connections=500"], environment: [ "POSTGRES_USER=postgres", "POSTGRES_PASSWORD=postgres", @@ -179,7 +185,7 @@ export function generateDockerCompose(context: TemplateContext) { ], volumes: [ `./${context.getDatacenterServicePath("postgres", datacenter.name)}/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh`, - `${postgresVolumeName}:/var/lib/postgresql/data`, + `${postgresVolumeName}:/var/lib/postgresql`, ], ports: isPrimary ? [`5432:5432`] : undefined, healthcheck: { @@ -194,7 +200,7 @@ export function generateDockerCompose(context: TemplateContext) { services[shellServiceName] = { build: { - context: "../../..", + context: "../..", dockerfile: "docker/engine/Dockerfile", target: "engine-full", args: { @@ -275,7 +281,7 @@ export function generateDockerCompose(context: TemplateContext) { services[serviceName] = { build: { - context: "../../..", + context: "../..", dockerfile: "docker/engine/Dockerfile", target: "engine-full", args: { @@ -335,30 +341,79 @@ export function generateDockerCompose(context: TemplateContext) { services[serviceName] = { build: { - context: "../../..", - dockerfile: "engine/sdks/rust/test-envoy/Dockerfile", + context: "../..", + dockerfile: "examples/kitchen-sink/Dockerfile", + // The runner copies the host-built napi binary, so the base + // image must have a glibc at least as new as the build host. + args: { + NODE_IMAGE: "node:22-trixie-slim", + }, }, platform: "linux/amd64", restart: "unless-stopped", environment: [ + `RIVET_KITCHEN_SINK_MODE=serverful`, `RIVET_ENDPOINT=http://${context.getServiceHost("rivet-engine", datacenter.name, 0)}:6420`, - `INTERNAL_SERVER_PORT=5050`, - `RIVET_POOL_NAME=test-envoy`, - `AUTOSTART_ENVOY=1`, - `AUTOCONFIGURE_SERVERLESS=0` + `RIVET_TOKEN=dev`, + `RIVET_NAMESPACE=default`, + `RIVET_POOL=default`, + `PORT=8080`, ], stop_grace_period: "4s", - ports: isPrimary && i === 0 ? [`5050:5050`] : undefined, + ports: isPrimary && i === 0 ? [`5050:8080`] : undefined, depends_on: { [engineServiceName]: { condition: "service_healthy", }, + [RUNNER_CONFIG_INIT_SERVICE]: { + condition: "service_completed_successfully", + }, }, networks: [dcNetworkName], }; } }); + // Serverful runners are rejected with `no_runner_config` until a runner + // config exists for their pool, so a one-shot init container upserts a + // `normal` runner config for the `default` pool across every datacenter once + // the leader engine is healthy. The runners wait for this to finish. + const primaryDc = config.datacenters[0]; + const primaryEngineHost = context.getServiceHost( + "rivet-engine", + primaryDc.name, + 0, + ); + const primaryEngineService = context.getServiceName( + "rivet-engine", + primaryDc.name, + 0, + ); + const runnerConfigBody = JSON.stringify({ + datacenters: Object.fromEntries( + config.datacenters.map((dc) => [dc.name, { normal: {} }]), + ), + }); + const runnerConfigScript = [ + `until curl -fsS -X PUT`, + `"http://${primaryEngineHost}:6420/runner-configs/default?namespace=default"`, + `-H "Authorization: Bearer dev"`, + `-H "Content-Type: application/json"`, + `-d '${runnerConfigBody}';`, + `do echo "waiting for engine to accept runner config"; sleep 2; done;`, + `echo "runner config upserted"`, + ].join(" "); + services[RUNNER_CONFIG_INIT_SERVICE] = { + image: "curlimages/curl:latest", + restart: "no", + depends_on: { + [primaryEngineService]: { condition: "service_healthy" }, + }, + entrypoint: ["sh", "-c"], + command: [runnerConfigScript], + networks: [context.getDatacenterNetworkName(primaryDc.name)], + }; + const dockerComposeConfig = { services, networks, diff --git a/self-host/compose/template/src/services/edge/rivet-engine.ts b/self-host/compose/template/src/services/edge/rivet-engine.ts index 25df2ce5ca..52d33dfc1b 100644 --- a/self-host/compose/template/src/services/edge/rivet-engine.ts +++ b/self-host/compose/template/src/services/edge/rivet-engine.ts @@ -35,32 +35,24 @@ export function generateDatacenterRivetEngine( datacenters, }; - // Config structure matching Rust schema in packages/common/config/src/config/mod.rs + // Config structure matching Rust schema in engine/packages/config/src/config/mod.rs. + // Values that match the engine's defaults are omitted. const config = { auth: { admin_token: "dev", }, - guard: { - port: GUARD_PORT, - // https is optional and not configured for local development - }, api_peer: { host: "0.0.0.0", - port: API_PEER_PORT, }, topology, postgres: { url: `postgresql://postgres:postgres@${context.getServiceHost("postgres", datacenter.name)}:5432/rivet_engine`, }, - cache: { - driver: "in_memory", - }, clickhouse: { - http_url: `http://${clickhouseHost}:9300`, // TODO: - native_url: `http://${clickhouseHost}:9301`, // TODO: + http_url: `http://${clickhouseHost}:9300`, + native_url: `http://${clickhouseHost}:9301`, username: "system", password: "default", - secure: false, }, }; diff --git a/self-host/compose/template/src/services/edge/runner.ts b/self-host/compose/template/src/services/edge/runner.ts index c04f936dd3..9a111e637c 100644 --- a/self-host/compose/template/src/services/edge/runner.ts +++ b/self-host/compose/template/src/services/edge/runner.ts @@ -1,6 +1,7 @@ import type { TemplateContext } from "../../context"; export function generateRunner(context: TemplateContext) { - // The test runner service now uses the Rust test-envoy binary. - // The docker-compose template points at the Rust Dockerfile directly. + // The runner service runs the kitchen-sink example in serverful mode, + // connecting to the engine as a long-lived runner. The docker-compose + // template builds examples/kitchen-sink/Dockerfile directly. } diff --git a/self-host/compose/dev-host/.gitattributes b/self-host/dev-host/.gitattributes similarity index 100% rename from self-host/compose/dev-host/.gitattributes rename to self-host/dev-host/.gitattributes diff --git a/self-host/compose/dev-host/README.md b/self-host/dev-host/README.md similarity index 100% rename from self-host/compose/dev-host/README.md rename to self-host/dev-host/README.md diff --git a/self-host/compose/dev-host/clickhouse/client-config.xml b/self-host/dev-host/clickhouse/client-config.xml similarity index 100% rename from self-host/compose/dev-host/clickhouse/client-config.xml rename to self-host/dev-host/clickhouse/client-config.xml diff --git a/self-host/compose/dev-host/clickhouse/config.xml b/self-host/dev-host/clickhouse/config.xml similarity index 100% rename from self-host/compose/dev-host/clickhouse/config.xml rename to self-host/dev-host/clickhouse/config.xml diff --git a/self-host/compose/dev-host/clickhouse/init/01-create-otel-table.sql b/self-host/dev-host/clickhouse/init/01-create-otel-table.sql similarity index 100% rename from self-host/compose/dev-host/clickhouse/init/01-create-otel-table.sql rename to self-host/dev-host/clickhouse/init/01-create-otel-table.sql diff --git a/self-host/compose/dev-host/clickhouse/users.xml b/self-host/dev-host/clickhouse/users.xml similarity index 100% rename from self-host/compose/dev-host/clickhouse/users.xml rename to self-host/dev-host/clickhouse/users.xml diff --git a/self-host/compose/dev-host/docker-compose.yml b/self-host/dev-host/docker-compose.yml similarity index 83% rename from self-host/compose/dev-host/docker-compose.yml rename to self-host/dev-host/docker-compose.yml index 3b54702d63..94619e9324 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: @@ -179,20 +183,42 @@ services: network_mode: host runner: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://127.0.0.1:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine: condition: service_healthy + runner-config-init: + condition: service_completed_successfully + network_mode: host + runner-config-init: + image: curlimages/curl:latest + restart: 'no' + depends_on: + rivet-engine: + condition: service_healthy + entrypoint: + - sh + - '-c' + command: + - >- + until curl -fsS -X PUT + "http://127.0.0.1:6420/runner-configs/default?namespace=default" -H + "Authorization: Bearer dev" -H "Content-Type: application/json" -d + '{"datacenters":{"default":{"normal":{}}}}'; do echo "waiting for engine + to accept runner config"; sleep 2; done; echo "runner config upserted" network_mode: host networks: rivet-core-network: diff --git a/self-host/compose/dev-host/grafana/dashboards/api.json b/self-host/dev-host/grafana/dashboards/api.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/api.json rename to self-host/dev-host/grafana/dashboards/api.json diff --git a/self-host/compose/dev-host/grafana/dashboards/cache.json b/self-host/dev-host/grafana/dashboards/cache.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/cache.json rename to self-host/dev-host/grafana/dashboards/cache.json diff --git a/self-host/compose/dev-host/grafana/dashboards/epoxy.json b/self-host/dev-host/grafana/dashboards/epoxy.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/epoxy.json rename to self-host/dev-host/grafana/dashboards/epoxy.json diff --git a/self-host/compose/dev-host/grafana/dashboards/futures.json b/self-host/dev-host/grafana/dashboards/futures.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/futures.json rename to self-host/dev-host/grafana/dashboards/futures.json diff --git a/self-host/compose/dev-host/grafana/dashboards/gasoline.json b/self-host/dev-host/grafana/dashboards/gasoline.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/gasoline.json rename to self-host/dev-host/grafana/dashboards/gasoline.json diff --git a/self-host/compose/dev-host/grafana/dashboards/guard.json b/self-host/dev-host/grafana/dashboards/guard.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/guard.json rename to self-host/dev-host/grafana/dashboards/guard.json diff --git a/self-host/compose/dev-host/grafana/dashboards/operation.json b/self-host/dev-host/grafana/dashboards/operation.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/operation.json rename to self-host/dev-host/grafana/dashboards/operation.json diff --git a/self-host/compose/dev-host/grafana/dashboards/pegboard.json b/self-host/dev-host/grafana/dashboards/pegboard.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/pegboard.json rename to self-host/dev-host/grafana/dashboards/pegboard.json diff --git a/self-host/compose/dev-host/grafana/dashboards/tokio.json b/self-host/dev-host/grafana/dashboards/tokio.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/tokio.json rename to self-host/dev-host/grafana/dashboards/tokio.json diff --git a/self-host/compose/dev-host/grafana/dashboards/traces.json b/self-host/dev-host/grafana/dashboards/traces.json similarity index 100% rename from self-host/compose/dev-host/grafana/dashboards/traces.json rename to self-host/dev-host/grafana/dashboards/traces.json diff --git a/self-host/compose/dev-host/grafana/grafana.ini b/self-host/dev-host/grafana/grafana.ini similarity index 100% rename from self-host/compose/dev-host/grafana/grafana.ini rename to self-host/dev-host/grafana/grafana.ini diff --git a/self-host/compose/dev-host/grafana/provisioning/dashboards/dashboards.yaml b/self-host/dev-host/grafana/provisioning/dashboards/dashboards.yaml similarity index 100% rename from self-host/compose/dev-host/grafana/provisioning/dashboards/dashboards.yaml rename to self-host/dev-host/grafana/provisioning/dashboards/dashboards.yaml diff --git a/self-host/compose/dev-host/grafana/provisioning/datasources/datasources.yaml b/self-host/dev-host/grafana/provisioning/datasources/datasources.yaml similarity index 100% rename from self-host/compose/dev-host/grafana/provisioning/datasources/datasources.yaml rename to self-host/dev-host/grafana/provisioning/datasources/datasources.yaml diff --git a/self-host/compose/dev-host/otel-collector/config.yaml b/self-host/dev-host/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-host/otel-collector/config.yaml rename to self-host/dev-host/otel-collector/config.yaml diff --git a/self-host/compose/dev-host/postgres/init-db.sh b/self-host/dev-host/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-host/postgres/init-db.sh rename to self-host/dev-host/postgres/init-db.sh diff --git a/self-host/compose/dev-host/prometheus/prometheus.yml b/self-host/dev-host/prometheus/prometheus.yml similarity index 100% rename from self-host/compose/dev-host/prometheus/prometheus.yml rename to self-host/dev-host/prometheus/prometheus.yml diff --git a/self-host/compose/dev-host/rivet-engine/config.jsonc b/self-host/dev-host/rivet-engine/config.jsonc similarity index 79% rename from self-host/compose/dev-host/rivet-engine/config.jsonc rename to self-host/dev-host/rivet-engine/config.jsonc index d5a3095269..ee79af406a 100644 --- a/self-host/compose/dev-host/rivet-engine/config.jsonc +++ b/self-host/dev-host/rivet-engine/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -28,14 +24,10 @@ "postgres": { "url": "postgresql://postgres:postgres@127.0.0.1:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://127.0.0.1:9300", "native_url": "http://127.0.0.1:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-host/vector-client/vector.yaml b/self-host/dev-host/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-host/vector-client/vector.yaml rename to self-host/dev-host/vector-client/vector.yaml diff --git a/self-host/compose/dev-host/vector-server/vector.yaml b/self-host/dev-host/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-host/vector-server/vector.yaml rename to self-host/dev-host/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/.gitattributes b/self-host/dev-multidc-multinode/.gitattributes similarity index 100% rename from self-host/compose/dev-multidc-multinode/.gitattributes rename to self-host/dev-multidc-multinode/.gitattributes diff --git a/self-host/compose/dev-multidc-multinode/README.md b/self-host/dev-multidc-multinode/README.md similarity index 100% rename from self-host/compose/dev-multidc-multinode/README.md rename to self-host/dev-multidc-multinode/README.md diff --git a/self-host/compose/dev-multidc-multinode/core/clickhouse/client-config.xml b/self-host/dev-multidc-multinode/core/clickhouse/client-config.xml similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/clickhouse/client-config.xml rename to self-host/dev-multidc-multinode/core/clickhouse/client-config.xml diff --git a/self-host/compose/dev-multidc-multinode/core/clickhouse/config.xml b/self-host/dev-multidc-multinode/core/clickhouse/config.xml similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/clickhouse/config.xml rename to self-host/dev-multidc-multinode/core/clickhouse/config.xml diff --git a/self-host/compose/dev-multidc-multinode/core/clickhouse/init/01-create-otel-table.sql b/self-host/dev-multidc-multinode/core/clickhouse/init/01-create-otel-table.sql similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/clickhouse/init/01-create-otel-table.sql rename to self-host/dev-multidc-multinode/core/clickhouse/init/01-create-otel-table.sql diff --git a/self-host/compose/dev-multidc-multinode/core/clickhouse/users.xml b/self-host/dev-multidc-multinode/core/clickhouse/users.xml similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/clickhouse/users.xml rename to self-host/dev-multidc-multinode/core/clickhouse/users.xml diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/api.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/api.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/api.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/api.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/cache.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/cache.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/cache.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/cache.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/epoxy.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/epoxy.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/epoxy.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/epoxy.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/futures.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/futures.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/futures.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/futures.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/gasoline.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/gasoline.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/gasoline.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/gasoline.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/guard.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/guard.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/guard.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/guard.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/operation.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/operation.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/operation.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/operation.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/pegboard.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/pegboard.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/pegboard.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/pegboard.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/tokio.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/tokio.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/tokio.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/tokio.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/dashboards/traces.json b/self-host/dev-multidc-multinode/core/grafana/dashboards/traces.json similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/dashboards/traces.json rename to self-host/dev-multidc-multinode/core/grafana/dashboards/traces.json diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/grafana.ini b/self-host/dev-multidc-multinode/core/grafana/grafana.ini similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/grafana.ini rename to self-host/dev-multidc-multinode/core/grafana/grafana.ini diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/provisioning/dashboards/dashboards.yaml b/self-host/dev-multidc-multinode/core/grafana/provisioning/dashboards/dashboards.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/provisioning/dashboards/dashboards.yaml rename to self-host/dev-multidc-multinode/core/grafana/provisioning/dashboards/dashboards.yaml diff --git a/self-host/compose/dev-multidc-multinode/core/grafana/provisioning/datasources/datasources.yaml b/self-host/dev-multidc-multinode/core/grafana/provisioning/datasources/datasources.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/grafana/provisioning/datasources/datasources.yaml rename to self-host/dev-multidc-multinode/core/grafana/provisioning/datasources/datasources.yaml diff --git a/self-host/compose/dev-multidc-multinode/core/prometheus/prometheus.yml b/self-host/dev-multidc-multinode/core/prometheus/prometheus.yml similarity index 100% rename from self-host/compose/dev-multidc-multinode/core/prometheus/prometheus.yml rename to self-host/dev-multidc-multinode/core/prometheus/prometheus.yml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/otel-collector/config.yaml b/self-host/dev-multidc-multinode/datacenters/dc-a/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-a/otel-collector/config.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-a/otel-collector/config.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/postgres/init-db.sh b/self-host/dev-multidc-multinode/datacenters/dc-a/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-a/postgres/init-db.sh rename to self-host/dev-multidc-multinode/datacenters/dc-a/postgres/init-db.sh diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/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..0532bce5c3 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-a:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc index ab9cfa618a..0532bce5c3 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-a:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/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/2/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc index ab9cfa618a..0532bce5c3 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-a:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/vector-client/vector.yaml b/self-host/dev-multidc-multinode/datacenters/dc-a/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-a/vector-client/vector.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-a/vector-client/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-a/vector-server/vector.yaml b/self-host/dev-multidc-multinode/datacenters/dc-a/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-a/vector-server/vector.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-a/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/otel-collector/config.yaml b/self-host/dev-multidc-multinode/datacenters/dc-b/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/otel-collector/config.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-b/otel-collector/config.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/postgres/init-db.sh b/self-host/dev-multidc-multinode/datacenters/dc-b/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/postgres/init-db.sh rename to self-host/dev-multidc-multinode/datacenters/dc-b/postgres/init-db.sh diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc index 2a34bc83ff..e75a22f77c 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 2, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-b:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc index 2a34bc83ff..e75a22f77c 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 2, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-b:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc index 2a34bc83ff..e75a22f77c 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 2, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-b:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/vector-client/vector.yaml b/self-host/dev-multidc-multinode/datacenters/dc-b/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/vector-client/vector.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-b/vector-client/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-b/vector-server/vector.yaml b/self-host/dev-multidc-multinode/datacenters/dc-b/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-b/vector-server/vector.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-b/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/otel-collector/config.yaml b/self-host/dev-multidc-multinode/datacenters/dc-c/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/otel-collector/config.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-c/otel-collector/config.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/postgres/init-db.sh b/self-host/dev-multidc-multinode/datacenters/dc-c/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/postgres/init-db.sh rename to self-host/dev-multidc-multinode/datacenters/dc-c/postgres/init-db.sh diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc index 970454f7f6..cb4c587a19 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 3, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-c:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc index 970454f7f6..cb4c587a19 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 3, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-c:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc rename to self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc index 970454f7f6..cb4c587a19 100644 --- a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc +++ b/self-host/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 3, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-c:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/vector-client/vector.yaml b/self-host/dev-multidc-multinode/datacenters/dc-c/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/vector-client/vector.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-c/vector-client/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/datacenters/dc-c/vector-server/vector.yaml b/self-host/dev-multidc-multinode/datacenters/dc-c/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc-multinode/datacenters/dc-c/vector-server/vector.yaml rename to self-host/dev-multidc-multinode/datacenters/dc-c/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc-multinode/docker-compose.yml b/self-host/dev-multidc-multinode/docker-compose.yml similarity index 83% rename from self-host/compose/dev-multidc-multinode/docker-compose.yml rename to self-host/dev-multidc-multinode/docker-compose.yml index 75c500850b..e4ccd4fc70 100644 --- a/self-host/compose/dev-multidc-multinode/docker-compose.yml +++ b/self-host/dev-multidc-multinode/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: @@ -174,7 +178,7 @@ services: - '4317:4317' rivet-engine-dc-a-0: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -217,7 +221,7 @@ services: start_period: 30s rivet-engine-dc-a-1: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -258,7 +262,7 @@ services: start_period: 30s rivet-engine-dc-a-2: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -299,63 +303,82 @@ services: start_period: 30s runner-dc-a-0: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-a-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s ports: - - '5050:5050' + - '5050:8080' depends_on: rivet-engine-dc-a-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-a runner-dc-a-1: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-a-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-a-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-a runner-dc-a-2: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-a-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-a-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-a postgres-dc-b: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -363,7 +386,7 @@ services: volumes: - >- ./datacenters/dc-b/postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data-dc-b:/var/lib/postgresql/data + - postgres-data-dc-b:/var/lib/postgresql healthcheck: test: - CMD-SHELL @@ -375,7 +398,7 @@ services: - rivet-network-dc-b rivet-shell-dc-b: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -440,7 +463,7 @@ services: - rivet-network-dc-b-to-core rivet-engine-dc-b-0: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -481,7 +504,7 @@ services: start_period: 30s rivet-engine-dc-b-1: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -522,7 +545,7 @@ services: start_period: 30s rivet-engine-dc-b-2: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -563,61 +586,80 @@ services: start_period: 30s runner-dc-b-0: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-b-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-b-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-b runner-dc-b-1: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-b-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-b-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-b runner-dc-b-2: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-b-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-b-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-b postgres-dc-c: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -625,7 +667,7 @@ services: volumes: - >- ./datacenters/dc-c/postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data-dc-c:/var/lib/postgresql/data + - postgres-data-dc-c:/var/lib/postgresql healthcheck: test: - CMD-SHELL @@ -637,7 +679,7 @@ services: - rivet-network-dc-c rivet-shell-dc-c: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -702,7 +744,7 @@ services: - rivet-network-dc-c-to-core rivet-engine-dc-c-0: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -743,7 +785,7 @@ services: start_period: 30s rivet-engine-dc-c-1: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -784,7 +826,7 @@ services: start_period: 30s rivet-engine-dc-c-2: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -825,58 +867,92 @@ services: start_period: 30s runner-dc-c-0: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-c-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-c-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-c runner-dc-c-1: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-c-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-c-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-c runner-dc-c-2: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-c-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-c-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-c + runner-config-init: + image: curlimages/curl:latest + restart: 'no' + depends_on: + rivet-engine-dc-a-0: + condition: service_healthy + entrypoint: + - sh + - '-c' + command: + - >- + until curl -fsS -X PUT + "http://rivet-engine-dc-a-0:6420/runner-configs/default?namespace=default" + -H "Authorization: Bearer dev" -H "Content-Type: application/json" -d + '{"datacenters":{"dc-a":{"normal":{}},"dc-b":{"normal":{}},"dc-c":{"normal":{}}}}'; + do echo "waiting for engine to accept runner config"; sleep 2; done; + echo "runner config upserted" + networks: + - rivet-network-dc-a networks: rivet-core-network: driver: bridge diff --git a/self-host/compose/dev-multidc/.gitattributes b/self-host/dev-multidc/.gitattributes similarity index 100% rename from self-host/compose/dev-multidc/.gitattributes rename to self-host/dev-multidc/.gitattributes diff --git a/self-host/compose/dev-multidc/README.md b/self-host/dev-multidc/README.md similarity index 100% rename from self-host/compose/dev-multidc/README.md rename to self-host/dev-multidc/README.md diff --git a/self-host/compose/dev-multidc/core/clickhouse/client-config.xml b/self-host/dev-multidc/core/clickhouse/client-config.xml similarity index 100% rename from self-host/compose/dev-multidc/core/clickhouse/client-config.xml rename to self-host/dev-multidc/core/clickhouse/client-config.xml diff --git a/self-host/compose/dev-multidc/core/clickhouse/config.xml b/self-host/dev-multidc/core/clickhouse/config.xml similarity index 100% rename from self-host/compose/dev-multidc/core/clickhouse/config.xml rename to self-host/dev-multidc/core/clickhouse/config.xml diff --git a/self-host/compose/dev-multidc/core/clickhouse/init/01-create-otel-table.sql b/self-host/dev-multidc/core/clickhouse/init/01-create-otel-table.sql similarity index 100% rename from self-host/compose/dev-multidc/core/clickhouse/init/01-create-otel-table.sql rename to self-host/dev-multidc/core/clickhouse/init/01-create-otel-table.sql diff --git a/self-host/compose/dev-multidc/core/clickhouse/users.xml b/self-host/dev-multidc/core/clickhouse/users.xml similarity index 100% rename from self-host/compose/dev-multidc/core/clickhouse/users.xml rename to self-host/dev-multidc/core/clickhouse/users.xml diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/api.json b/self-host/dev-multidc/core/grafana/dashboards/api.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/api.json rename to self-host/dev-multidc/core/grafana/dashboards/api.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/cache.json b/self-host/dev-multidc/core/grafana/dashboards/cache.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/cache.json rename to self-host/dev-multidc/core/grafana/dashboards/cache.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/epoxy.json b/self-host/dev-multidc/core/grafana/dashboards/epoxy.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/epoxy.json rename to self-host/dev-multidc/core/grafana/dashboards/epoxy.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/futures.json b/self-host/dev-multidc/core/grafana/dashboards/futures.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/futures.json rename to self-host/dev-multidc/core/grafana/dashboards/futures.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/gasoline.json b/self-host/dev-multidc/core/grafana/dashboards/gasoline.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/gasoline.json rename to self-host/dev-multidc/core/grafana/dashboards/gasoline.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/guard.json b/self-host/dev-multidc/core/grafana/dashboards/guard.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/guard.json rename to self-host/dev-multidc/core/grafana/dashboards/guard.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/operation.json b/self-host/dev-multidc/core/grafana/dashboards/operation.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/operation.json rename to self-host/dev-multidc/core/grafana/dashboards/operation.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/pegboard.json b/self-host/dev-multidc/core/grafana/dashboards/pegboard.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/pegboard.json rename to self-host/dev-multidc/core/grafana/dashboards/pegboard.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/tokio.json b/self-host/dev-multidc/core/grafana/dashboards/tokio.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/tokio.json rename to self-host/dev-multidc/core/grafana/dashboards/tokio.json diff --git a/self-host/compose/dev-multidc/core/grafana/dashboards/traces.json b/self-host/dev-multidc/core/grafana/dashboards/traces.json similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/dashboards/traces.json rename to self-host/dev-multidc/core/grafana/dashboards/traces.json diff --git a/self-host/compose/dev-multidc/core/grafana/grafana.ini b/self-host/dev-multidc/core/grafana/grafana.ini similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/grafana.ini rename to self-host/dev-multidc/core/grafana/grafana.ini diff --git a/self-host/compose/dev-multidc/core/grafana/provisioning/dashboards/dashboards.yaml b/self-host/dev-multidc/core/grafana/provisioning/dashboards/dashboards.yaml similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/provisioning/dashboards/dashboards.yaml rename to self-host/dev-multidc/core/grafana/provisioning/dashboards/dashboards.yaml diff --git a/self-host/compose/dev-multidc/core/grafana/provisioning/datasources/datasources.yaml b/self-host/dev-multidc/core/grafana/provisioning/datasources/datasources.yaml similarity index 100% rename from self-host/compose/dev-multidc/core/grafana/provisioning/datasources/datasources.yaml rename to self-host/dev-multidc/core/grafana/provisioning/datasources/datasources.yaml diff --git a/self-host/compose/dev-multidc/core/prometheus/prometheus.yml b/self-host/dev-multidc/core/prometheus/prometheus.yml similarity index 100% rename from self-host/compose/dev-multidc/core/prometheus/prometheus.yml rename to self-host/dev-multidc/core/prometheus/prometheus.yml diff --git a/self-host/compose/dev-multidc/datacenters/dc-a/otel-collector/config.yaml b/self-host/dev-multidc/datacenters/dc-a/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-a/otel-collector/config.yaml rename to self-host/dev-multidc/datacenters/dc-a/otel-collector/config.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-a/postgres/init-db.sh b/self-host/dev-multidc/datacenters/dc-a/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-a/postgres/init-db.sh rename to self-host/dev-multidc/datacenters/dc-a/postgres/init-db.sh diff --git a/self-host/compose/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc b/self-host/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc rename to self-host/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc index f989acaae9..ac3c4bc3c6 100644 --- a/self-host/compose/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc +++ b/self-host/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-a:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc/datacenters/dc-a/vector-client/vector.yaml b/self-host/dev-multidc/datacenters/dc-a/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-a/vector-client/vector.yaml rename to self-host/dev-multidc/datacenters/dc-a/vector-client/vector.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-a/vector-server/vector.yaml b/self-host/dev-multidc/datacenters/dc-a/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-a/vector-server/vector.yaml rename to self-host/dev-multidc/datacenters/dc-a/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-b/otel-collector/config.yaml b/self-host/dev-multidc/datacenters/dc-b/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-b/otel-collector/config.yaml rename to self-host/dev-multidc/datacenters/dc-b/otel-collector/config.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-b/postgres/init-db.sh b/self-host/dev-multidc/datacenters/dc-b/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-b/postgres/init-db.sh rename to self-host/dev-multidc/datacenters/dc-b/postgres/init-db.sh diff --git a/self-host/compose/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc b/self-host/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc rename to self-host/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc index fa082ce527..8e5c6aaa45 100644 --- a/self-host/compose/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc +++ b/self-host/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 2, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-b:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc/datacenters/dc-b/vector-client/vector.yaml b/self-host/dev-multidc/datacenters/dc-b/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-b/vector-client/vector.yaml rename to self-host/dev-multidc/datacenters/dc-b/vector-client/vector.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-b/vector-server/vector.yaml b/self-host/dev-multidc/datacenters/dc-b/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-b/vector-server/vector.yaml rename to self-host/dev-multidc/datacenters/dc-b/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-c/otel-collector/config.yaml b/self-host/dev-multidc/datacenters/dc-c/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-c/otel-collector/config.yaml rename to self-host/dev-multidc/datacenters/dc-c/otel-collector/config.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-c/postgres/init-db.sh b/self-host/dev-multidc/datacenters/dc-c/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-c/postgres/init-db.sh rename to self-host/dev-multidc/datacenters/dc-c/postgres/init-db.sh diff --git a/self-host/compose/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc b/self-host/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc similarity index 88% rename from self-host/compose/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc rename to self-host/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc index f3c3c6ae38..6ba3be6d65 100644 --- a/self-host/compose/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc +++ b/self-host/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 3, @@ -50,14 +46,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-c:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multidc/datacenters/dc-c/vector-client/vector.yaml b/self-host/dev-multidc/datacenters/dc-c/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-c/vector-client/vector.yaml rename to self-host/dev-multidc/datacenters/dc-c/vector-client/vector.yaml diff --git a/self-host/compose/dev-multidc/datacenters/dc-c/vector-server/vector.yaml b/self-host/dev-multidc/datacenters/dc-c/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multidc/datacenters/dc-c/vector-server/vector.yaml rename to self-host/dev-multidc/datacenters/dc-c/vector-server/vector.yaml diff --git a/self-host/compose/dev-multidc/docker-compose.yml b/self-host/dev-multidc/docker-compose.yml similarity index 86% rename from self-host/compose/dev-multidc/docker-compose.yml rename to self-host/dev-multidc/docker-compose.yml index f20db41ff3..437ad5bdee 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: @@ -215,27 +219,36 @@ services: start_period: 30s runner-dc-a: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-a:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s ports: - - '5050:5050' + - '5050:8080' depends_on: rivet-engine-dc-a: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-a postgres-dc-b: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -243,7 +256,7 @@ services: volumes: - >- ./datacenters/dc-b/postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data-dc-b:/var/lib/postgresql/data + - postgres-data-dc-b:/var/lib/postgresql healthcheck: test: - CMD-SHELL @@ -255,7 +268,7 @@ services: - rivet-network-dc-b rivet-shell-dc-b: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -319,7 +332,7 @@ services: - rivet-network-dc-b-to-core rivet-engine-dc-b: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -359,25 +372,34 @@ services: start_period: 30s runner-dc-b: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-b:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-b: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-b postgres-dc-c: restart: unless-stopped - image: postgres:17-alpine + image: postgres:18-alpine + command: + - postgres + - '-c' + - max_connections=500 environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres @@ -385,7 +407,7 @@ services: volumes: - >- ./datacenters/dc-c/postgres/init-db.sh:/docker-entrypoint-initdb.d/init-db.sh - - postgres-data-dc-c:/var/lib/postgresql/data + - postgres-data-dc-c:/var/lib/postgresql healthcheck: test: - CMD-SHELL @@ -397,7 +419,7 @@ services: - rivet-network-dc-c rivet-shell-dc-c: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -461,7 +483,7 @@ services: - rivet-network-dc-c-to-core rivet-engine-dc-c: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -501,22 +523,46 @@ services: start_period: 30s runner-dc-c: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-dc-c:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-dc-c: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network-dc-c + runner-config-init: + image: curlimages/curl:latest + restart: 'no' + depends_on: + rivet-engine-dc-a: + condition: service_healthy + entrypoint: + - sh + - '-c' + command: + - >- + until curl -fsS -X PUT + "http://rivet-engine-dc-a:6420/runner-configs/default?namespace=default" + -H "Authorization: Bearer dev" -H "Content-Type: application/json" -d + '{"datacenters":{"dc-a":{"normal":{}},"dc-b":{"normal":{}},"dc-c":{"normal":{}}}}'; + do echo "waiting for engine to accept runner config"; sleep 2; done; + echo "runner config upserted" + networks: + - rivet-network-dc-a networks: rivet-core-network: driver: bridge diff --git a/self-host/compose/dev-multinode/.gitattributes b/self-host/dev-multinode/.gitattributes similarity index 100% rename from self-host/compose/dev-multinode/.gitattributes rename to self-host/dev-multinode/.gitattributes diff --git a/self-host/compose/dev-multinode/README.md b/self-host/dev-multinode/README.md similarity index 100% rename from self-host/compose/dev-multinode/README.md rename to self-host/dev-multinode/README.md diff --git a/self-host/compose/dev-multinode/clickhouse/client-config.xml b/self-host/dev-multinode/clickhouse/client-config.xml similarity index 100% rename from self-host/compose/dev-multinode/clickhouse/client-config.xml rename to self-host/dev-multinode/clickhouse/client-config.xml diff --git a/self-host/compose/dev-multinode/clickhouse/config.xml b/self-host/dev-multinode/clickhouse/config.xml similarity index 100% rename from self-host/compose/dev-multinode/clickhouse/config.xml rename to self-host/dev-multinode/clickhouse/config.xml diff --git a/self-host/compose/dev-multinode/clickhouse/init/01-create-otel-table.sql b/self-host/dev-multinode/clickhouse/init/01-create-otel-table.sql similarity index 100% rename from self-host/compose/dev-multinode/clickhouse/init/01-create-otel-table.sql rename to self-host/dev-multinode/clickhouse/init/01-create-otel-table.sql diff --git a/self-host/compose/dev-multinode/clickhouse/users.xml b/self-host/dev-multinode/clickhouse/users.xml similarity index 100% rename from self-host/compose/dev-multinode/clickhouse/users.xml rename to self-host/dev-multinode/clickhouse/users.xml diff --git a/self-host/compose/dev-multinode/docker-compose.yml b/self-host/dev-multinode/docker-compose.yml similarity index 82% rename from self-host/compose/dev-multinode/docker-compose.yml rename to self-host/dev-multinode/docker-compose.yml index 381ece396d..389510f7b3 100644 --- a/self-host/compose/dev-multinode/docker-compose.yml +++ b/self-host/dev-multinode/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-0: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -206,7 +210,7 @@ services: start_period: 30s rivet-engine-1: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -244,7 +248,7 @@ services: start_period: 30s rivet-engine-2: build: - context: ../../.. + context: ../.. dockerfile: docker/engine/Dockerfile target: engine-full args: @@ -282,58 +286,91 @@ services: start_period: 30s runner-0: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s ports: - - '5050:5050' + - '5050:8080' depends_on: rivet-engine-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network runner-1: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully networks: - rivet-network runner-2: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine-0:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s depends_on: rivet-engine-0: condition: service_healthy + runner-config-init: + condition: service_completed_successfully + networks: + - rivet-network + runner-config-init: + image: curlimages/curl:latest + restart: 'no' + depends_on: + rivet-engine-0: + condition: service_healthy + entrypoint: + - sh + - '-c' + command: + - >- + until curl -fsS -X PUT + "http://rivet-engine-0:6420/runner-configs/default?namespace=default" -H + "Authorization: Bearer dev" -H "Content-Type: application/json" -d + '{"datacenters":{"default":{"normal":{}}}}'; do echo "waiting for engine + to accept runner config"; sleep 2; done; echo "runner config upserted" networks: - rivet-network networks: diff --git a/self-host/compose/dev-multinode/grafana/dashboards/api.json b/self-host/dev-multinode/grafana/dashboards/api.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/api.json rename to self-host/dev-multinode/grafana/dashboards/api.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/cache.json b/self-host/dev-multinode/grafana/dashboards/cache.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/cache.json rename to self-host/dev-multinode/grafana/dashboards/cache.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/epoxy.json b/self-host/dev-multinode/grafana/dashboards/epoxy.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/epoxy.json rename to self-host/dev-multinode/grafana/dashboards/epoxy.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/futures.json b/self-host/dev-multinode/grafana/dashboards/futures.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/futures.json rename to self-host/dev-multinode/grafana/dashboards/futures.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/gasoline.json b/self-host/dev-multinode/grafana/dashboards/gasoline.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/gasoline.json rename to self-host/dev-multinode/grafana/dashboards/gasoline.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/guard.json b/self-host/dev-multinode/grafana/dashboards/guard.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/guard.json rename to self-host/dev-multinode/grafana/dashboards/guard.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/operation.json b/self-host/dev-multinode/grafana/dashboards/operation.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/operation.json rename to self-host/dev-multinode/grafana/dashboards/operation.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/pegboard.json b/self-host/dev-multinode/grafana/dashboards/pegboard.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/pegboard.json rename to self-host/dev-multinode/grafana/dashboards/pegboard.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/tokio.json b/self-host/dev-multinode/grafana/dashboards/tokio.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/tokio.json rename to self-host/dev-multinode/grafana/dashboards/tokio.json diff --git a/self-host/compose/dev-multinode/grafana/dashboards/traces.json b/self-host/dev-multinode/grafana/dashboards/traces.json similarity index 100% rename from self-host/compose/dev-multinode/grafana/dashboards/traces.json rename to self-host/dev-multinode/grafana/dashboards/traces.json diff --git a/self-host/compose/dev-multinode/grafana/grafana.ini b/self-host/dev-multinode/grafana/grafana.ini similarity index 100% rename from self-host/compose/dev-multinode/grafana/grafana.ini rename to self-host/dev-multinode/grafana/grafana.ini diff --git a/self-host/compose/dev-multinode/grafana/provisioning/dashboards/dashboards.yaml b/self-host/dev-multinode/grafana/provisioning/dashboards/dashboards.yaml similarity index 100% rename from self-host/compose/dev-multinode/grafana/provisioning/dashboards/dashboards.yaml rename to self-host/dev-multinode/grafana/provisioning/dashboards/dashboards.yaml diff --git a/self-host/compose/dev-multinode/grafana/provisioning/datasources/datasources.yaml b/self-host/dev-multinode/grafana/provisioning/datasources/datasources.yaml similarity index 100% rename from self-host/compose/dev-multinode/grafana/provisioning/datasources/datasources.yaml rename to self-host/dev-multinode/grafana/provisioning/datasources/datasources.yaml diff --git a/self-host/compose/dev-multinode/otel-collector/config.yaml b/self-host/dev-multinode/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev-multinode/otel-collector/config.yaml rename to self-host/dev-multinode/otel-collector/config.yaml diff --git a/self-host/compose/dev-multinode/postgres/init-db.sh b/self-host/dev-multinode/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev-multinode/postgres/init-db.sh rename to self-host/dev-multinode/postgres/init-db.sh diff --git a/self-host/compose/dev-multinode/prometheus/prometheus.yml b/self-host/dev-multinode/prometheus/prometheus.yml similarity index 100% rename from self-host/compose/dev-multinode/prometheus/prometheus.yml rename to self-host/dev-multinode/prometheus/prometheus.yml diff --git a/self-host/compose/dev-multinode/rivet-engine/0/config.jsonc b/self-host/dev-multinode/rivet-engine/0/config.jsonc similarity index 79% rename from self-host/compose/dev-multinode/rivet-engine/0/config.jsonc rename to self-host/dev-multinode/rivet-engine/0/config.jsonc index 655fc83ed7..b25680c399 100644 --- a/self-host/compose/dev-multinode/rivet-engine/0/config.jsonc +++ b/self-host/dev-multinode/rivet-engine/0/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -28,14 +24,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multinode/rivet-engine/2/config.jsonc b/self-host/dev-multinode/rivet-engine/1/config.jsonc similarity index 79% rename from self-host/compose/dev-multinode/rivet-engine/2/config.jsonc rename to self-host/dev-multinode/rivet-engine/1/config.jsonc index 655fc83ed7..b25680c399 100644 --- a/self-host/compose/dev-multinode/rivet-engine/2/config.jsonc +++ b/self-host/dev-multinode/rivet-engine/1/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -28,14 +24,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multinode/rivet-engine/1/config.jsonc b/self-host/dev-multinode/rivet-engine/2/config.jsonc similarity index 79% rename from self-host/compose/dev-multinode/rivet-engine/1/config.jsonc rename to self-host/dev-multinode/rivet-engine/2/config.jsonc index 655fc83ed7..b25680c399 100644 --- a/self-host/compose/dev-multinode/rivet-engine/1/config.jsonc +++ b/self-host/dev-multinode/rivet-engine/2/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -28,14 +24,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev-multinode/vector-client/vector.yaml b/self-host/dev-multinode/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev-multinode/vector-client/vector.yaml rename to self-host/dev-multinode/vector-client/vector.yaml diff --git a/self-host/compose/dev-multinode/vector-server/vector.yaml b/self-host/dev-multinode/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev-multinode/vector-server/vector.yaml rename to self-host/dev-multinode/vector-server/vector.yaml diff --git a/self-host/compose/dev/.gitattributes b/self-host/dev/.gitattributes similarity index 100% rename from self-host/compose/dev/.gitattributes rename to self-host/dev/.gitattributes diff --git a/self-host/compose/dev/README.md b/self-host/dev/README.md similarity index 100% rename from self-host/compose/dev/README.md rename to self-host/dev/README.md diff --git a/self-host/compose/dev/clickhouse/client-config.xml b/self-host/dev/clickhouse/client-config.xml similarity index 100% rename from self-host/compose/dev/clickhouse/client-config.xml rename to self-host/dev/clickhouse/client-config.xml diff --git a/self-host/compose/dev/clickhouse/config.xml b/self-host/dev/clickhouse/config.xml similarity index 100% rename from self-host/compose/dev/clickhouse/config.xml rename to self-host/dev/clickhouse/config.xml diff --git a/self-host/compose/dev/clickhouse/init/01-create-otel-table.sql b/self-host/dev/clickhouse/init/01-create-otel-table.sql similarity index 100% rename from self-host/compose/dev/clickhouse/init/01-create-otel-table.sql rename to self-host/dev/clickhouse/init/01-create-otel-table.sql diff --git a/self-host/compose/dev/clickhouse/users.xml b/self-host/dev/clickhouse/users.xml similarity index 100% rename from self-host/compose/dev/clickhouse/users.xml rename to self-host/dev/clickhouse/users.xml diff --git a/self-host/compose/dev/docker-compose.yml b/self-host/dev/docker-compose.yml similarity index 84% rename from self-host/compose/dev/docker-compose.yml rename to self-host/dev/docker-compose.yml index 6c5dcc0ef2..7511952a55 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: @@ -206,22 +210,45 @@ services: start_period: 30s runner: build: - context: ../../.. - dockerfile: engine/sdks/rust/test-envoy/Dockerfile + context: ../.. + dockerfile: examples/kitchen-sink/Dockerfile + args: + NODE_IMAGE: node:22-trixie-slim platform: linux/amd64 restart: unless-stopped environment: + - RIVET_KITCHEN_SINK_MODE=serverful - RIVET_ENDPOINT=http://rivet-engine:6420 - - INTERNAL_SERVER_PORT=5050 - - RIVET_POOL_NAME=test-envoy - - AUTOSTART_ENVOY=1 - - AUTOCONFIGURE_SERVERLESS=0 + - RIVET_TOKEN=dev + - RIVET_NAMESPACE=default + - RIVET_POOL=default + - PORT=8080 stop_grace_period: 4s ports: - - '5050:5050' + - '5050:8080' + depends_on: + rivet-engine: + condition: service_healthy + runner-config-init: + condition: service_completed_successfully + networks: + - rivet-network + runner-config-init: + image: curlimages/curl:latest + restart: 'no' depends_on: rivet-engine: condition: service_healthy + entrypoint: + - sh + - '-c' + command: + - >- + until curl -fsS -X PUT + "http://rivet-engine:6420/runner-configs/default?namespace=default" -H + "Authorization: Bearer dev" -H "Content-Type: application/json" -d + '{"datacenters":{"default":{"normal":{}}}}'; do echo "waiting for engine + to accept runner config"; sleep 2; done; echo "runner config upserted" networks: - rivet-network networks: diff --git a/self-host/compose/dev/grafana/dashboards/api.json b/self-host/dev/grafana/dashboards/api.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/api.json rename to self-host/dev/grafana/dashboards/api.json diff --git a/self-host/compose/dev/grafana/dashboards/cache.json b/self-host/dev/grafana/dashboards/cache.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/cache.json rename to self-host/dev/grafana/dashboards/cache.json diff --git a/self-host/compose/dev/grafana/dashboards/epoxy.json b/self-host/dev/grafana/dashboards/epoxy.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/epoxy.json rename to self-host/dev/grafana/dashboards/epoxy.json diff --git a/self-host/compose/dev/grafana/dashboards/futures.json b/self-host/dev/grafana/dashboards/futures.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/futures.json rename to self-host/dev/grafana/dashboards/futures.json diff --git a/self-host/compose/dev/grafana/dashboards/gasoline.json b/self-host/dev/grafana/dashboards/gasoline.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/gasoline.json rename to self-host/dev/grafana/dashboards/gasoline.json diff --git a/self-host/compose/dev/grafana/dashboards/guard.json b/self-host/dev/grafana/dashboards/guard.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/guard.json rename to self-host/dev/grafana/dashboards/guard.json diff --git a/self-host/compose/dev/grafana/dashboards/operation.json b/self-host/dev/grafana/dashboards/operation.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/operation.json rename to self-host/dev/grafana/dashboards/operation.json diff --git a/self-host/compose/dev/grafana/dashboards/pegboard.json b/self-host/dev/grafana/dashboards/pegboard.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/pegboard.json rename to self-host/dev/grafana/dashboards/pegboard.json diff --git a/self-host/compose/dev/grafana/dashboards/tokio.json b/self-host/dev/grafana/dashboards/tokio.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/tokio.json rename to self-host/dev/grafana/dashboards/tokio.json diff --git a/self-host/compose/dev/grafana/dashboards/traces.json b/self-host/dev/grafana/dashboards/traces.json similarity index 100% rename from self-host/compose/dev/grafana/dashboards/traces.json rename to self-host/dev/grafana/dashboards/traces.json diff --git a/self-host/compose/dev/grafana/grafana.ini b/self-host/dev/grafana/grafana.ini similarity index 100% rename from self-host/compose/dev/grafana/grafana.ini rename to self-host/dev/grafana/grafana.ini diff --git a/self-host/compose/dev/grafana/provisioning/dashboards/dashboards.yaml b/self-host/dev/grafana/provisioning/dashboards/dashboards.yaml similarity index 100% rename from self-host/compose/dev/grafana/provisioning/dashboards/dashboards.yaml rename to self-host/dev/grafana/provisioning/dashboards/dashboards.yaml diff --git a/self-host/compose/dev/grafana/provisioning/datasources/datasources.yaml b/self-host/dev/grafana/provisioning/datasources/datasources.yaml similarity index 100% rename from self-host/compose/dev/grafana/provisioning/datasources/datasources.yaml rename to self-host/dev/grafana/provisioning/datasources/datasources.yaml diff --git a/self-host/compose/dev/otel-collector/config.yaml b/self-host/dev/otel-collector/config.yaml similarity index 100% rename from self-host/compose/dev/otel-collector/config.yaml rename to self-host/dev/otel-collector/config.yaml diff --git a/self-host/compose/dev/postgres/init-db.sh b/self-host/dev/postgres/init-db.sh similarity index 100% rename from self-host/compose/dev/postgres/init-db.sh rename to self-host/dev/postgres/init-db.sh diff --git a/self-host/compose/dev/prometheus/prometheus.yml b/self-host/dev/prometheus/prometheus.yml similarity index 100% rename from self-host/compose/dev/prometheus/prometheus.yml rename to self-host/dev/prometheus/prometheus.yml diff --git a/self-host/compose/dev/rivet-engine/config.jsonc b/self-host/dev/rivet-engine/config.jsonc similarity index 79% rename from self-host/compose/dev/rivet-engine/config.jsonc rename to self-host/dev/rivet-engine/config.jsonc index 74135c72fa..4c6312b6d3 100644 --- a/self-host/compose/dev/rivet-engine/config.jsonc +++ b/self-host/dev/rivet-engine/config.jsonc @@ -2,12 +2,8 @@ "auth": { "admin_token": "dev" }, - "guard": { - "port": 6420 - }, "api_peer": { - "host": "0.0.0.0", - "port": 6421 + "host": "0.0.0.0" }, "topology": { "datacenter_label": 1, @@ -28,14 +24,10 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" }, - "cache": { - "driver": "in_memory" - }, "clickhouse": { "http_url": "http://clickhouse:9300", "native_url": "http://clickhouse:9301", "username": "system", - "password": "default", - "secure": false + "password": "default" } } \ No newline at end of file diff --git a/self-host/compose/dev/vector-client/vector.yaml b/self-host/dev/vector-client/vector.yaml similarity index 100% rename from self-host/compose/dev/vector-client/vector.yaml rename to self-host/dev/vector-client/vector.yaml diff --git a/self-host/compose/dev/vector-server/vector.yaml b/self-host/dev/vector-server/vector.yaml similarity index 100% rename from self-host/compose/dev/vector-server/vector.yaml rename to self-host/dev/vector-server/vector.yaml 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