From 9d88491f2f5fd581aef1e82f391e653a887fea34 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Sat, 6 Sep 2025 23:44:02 -0700 Subject: [PATCH 01/17] chore(pegboard): consolidate to single subscriber per gateway --- Cargo.lock | 6 +- packages/common/universalpubsub/Cargo.toml | 2 +- .../src/driver/postgres/mod.rs | 106 ++- packages/core/guard/server/Cargo.toml | 1 + packages/core/guard/server/src/lib.rs | 7 +- packages/core/guard/server/src/routing/mod.rs | 17 +- .../server/src/routing/pegboard_gateway.rs | 26 +- .../server/src/routing/pegboard_tunnel.rs | 5 +- .../core/guard/server/src/shared_state.rs | 32 + packages/core/pegboard-gateway/Cargo.toml | 6 +- packages/core/pegboard-gateway/src/lib.rs | 427 ++++------ .../core/pegboard-gateway/src/shared_state.rs | 286 +++++++ packages/core/pegboard-runner-ws/src/lib.rs | 45 +- packages/core/pegboard-tunnel/Cargo.toml | 3 +- packages/core/pegboard-tunnel/src/lib.rs | 779 ++++-------------- .../core/pegboard-tunnel/tests/integration.rs | 13 +- .../infra/engine/tests/actors_lifecycle.rs | 9 +- .../pegboard/src/ops/runner/get_by_key.rs | 51 ++ .../services/pegboard/src/ops/runner/mod.rs | 1 + .../services/pegboard/src/pubsub_subjects.rs | 68 +- pnpm-lock.yaml | 9 + sdks/rust/runner-protocol/src/protocol.rs | 2 - sdks/rust/runner-protocol/src/versioned.rs | 2 - sdks/rust/tunnel-protocol/build.rs | 89 +- sdks/rust/tunnel-protocol/src/versioned.rs | 66 +- sdks/schemas/runner-protocol/v1.bare | 2 - sdks/schemas/tunnel-protocol/v1.bare | 66 +- sdks/typescript/runner-protocol/src/index.ts | 157 ++-- sdks/typescript/runner/package.json | 1 + sdks/typescript/runner/src/mod.ts | 91 +- sdks/typescript/runner/src/tunnel.ts | 774 ++++++++++------- .../runner/src/websocket-tunnel-adapter.ts | 4 +- sdks/typescript/tunnel-protocol/src/index.ts | 293 +++---- 33 files changed, 1783 insertions(+), 1663 deletions(-) create mode 100644 packages/core/guard/server/src/shared_state.rs create mode 100644 packages/core/pegboard-gateway/src/shared_state.rs create mode 100644 packages/services/pegboard/src/ops/runner/get_by_key.rs diff --git a/Cargo.lock b/Cargo.lock index bcf2da44ee..800d4ca464 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3270,9 +3270,11 @@ dependencies = [ "rivet-guard-core", "rivet-tunnel-protocol", "rivet-util", + "thiserror 1.0.69", "tokio", "tokio-tungstenite", "universalpubsub", + "versioned-data-util", ] [[package]] @@ -3304,6 +3306,7 @@ version = "0.0.1" dependencies = [ "anyhow", "async-trait", + "bytes", "futures", "gasoline", "http-body-util", @@ -4340,6 +4343,7 @@ dependencies = [ "tracing", "udb-util", "universaldb", + "universalpubsub", "url", "uuid", ] @@ -6329,7 +6333,6 @@ dependencies = [ "base64 0.22.1", "deadpool-postgres", "futures-util", - "moka", "rivet-config", "rivet-env", "rivet-error", @@ -6342,6 +6345,7 @@ dependencies = [ "tempfile", "tokio", "tokio-postgres", + "tokio-util", "tracing", "tracing-subscriber", "uuid", diff --git a/packages/common/universalpubsub/Cargo.toml b/packages/common/universalpubsub/Cargo.toml index 37028f9995..3cdb8509b6 100644 --- a/packages/common/universalpubsub/Cargo.toml +++ b/packages/common/universalpubsub/Cargo.toml @@ -12,7 +12,6 @@ async-trait.workspace = true base64.workspace = true deadpool-postgres.workspace = true futures-util.workspace = true -moka.workspace = true rivet-error.workspace = true rivet-ups-protocol.workspace = true serde_json.workspace = true @@ -21,6 +20,7 @@ serde.workspace = true sha2.workspace = true tokio-postgres.workspace = true tokio.workspace = true +tokio-util.workspace = true tracing.workspace = true uuid.workspace = true diff --git a/packages/common/universalpubsub/src/driver/postgres/mod.rs b/packages/common/universalpubsub/src/driver/postgres/mod.rs index 99cb703706..c2df9a36b4 100644 --- a/packages/common/universalpubsub/src/driver/postgres/mod.rs +++ b/packages/common/universalpubsub/src/driver/postgres/mod.rs @@ -1,5 +1,6 @@ +use std::collections::HashMap; use std::hash::{DefaultHasher, Hash, Hasher}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use anyhow::*; use async_trait::async_trait; @@ -7,7 +8,6 @@ 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 moka::future::Cache; use tokio_postgres::{AsyncMessage, NoTls}; use tracing::Instrument; @@ -18,6 +18,15 @@ use crate::pubsub::DriverOutput; struct Subscription { // Channel to send messages to this subscription tx: tokio::sync::broadcast::Sender>, + // Cancellation token shared by all subscribers of this subject + token: tokio_util::sync::CancellationToken, +} + +impl Subscription { + fn new(tx: tokio::sync::broadcast::Sender>) -> Self { + let token = tokio_util::sync::CancellationToken::new(); + Self { tx, token } + } } /// > In the default configuration it must be shorter than 8000 bytes @@ -40,7 +49,7 @@ pub const POSTGRES_MAX_MESSAGE_SIZE: usize = pub struct PostgresDriver { pool: Arc, client: Arc, - subscriptions: Cache, + subscriptions: Arc>>, } impl PostgresDriver { @@ -65,8 +74,8 @@ impl PostgresDriver { .context("failed to create postgres pool")?; tracing::debug!("postgres pool created successfully"); - let subscriptions: Cache = - Cache::builder().initial_capacity(5).build(); + let subscriptions: Arc>> = + Arc::new(Mutex::new(HashMap::new())); let subscriptions2 = subscriptions.clone(); let (client, mut conn) = tokio_postgres::connect(&conn_str, tokio_postgres::NoTls).await?; @@ -75,7 +84,9 @@ impl PostgresDriver { loop { match poll_fn(|cx| conn.poll_message(cx)).await { Some(std::result::Result::Ok(AsyncMessage::Notification(note))) => { - if let Some(sub) = subscriptions2.get(note.channel()).await { + if let Some(sub) = + subscriptions2.lock().unwrap().get(note.channel()).cloned() + { let bytes = match BASE64.decode(note.payload()) { std::result::Result::Ok(b) => b, std::result::Result::Err(err) => { @@ -121,7 +132,7 @@ impl PostgresDriver { #[async_trait] impl PubSubDriver for PostgresDriver { async fn subscribe(&self, subject: &str) -> Result { - // TODO: To match NATS implementation, LIST must be pipelined (i.e. wait for the command + // 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 @@ -135,33 +146,57 @@ impl PubSubDriver for PostgresDriver { let hashed = self.hash_subject(subject); // Check if we already have a subscription for this channel - let rx = if let Some(existing_sub) = self.subscriptions.get(&hashed).await { - // Reuse the existing broadcast channel - existing_sub.tx.subscribe() - } else { - // Create a new broadcast channel for this subject - let (tx, rx) = tokio::sync::broadcast::channel(1024); - let subscription = Subscription { tx: tx.clone() }; - - // Register subscription - self.subscriptions - .insert(hashed.clone(), subscription) - .await; - - // Execute LISTEN command on the async client (for receiving notifications) - // This only needs to be done once per channel - let span = tracing::trace_span!("pg_listen"); - self.client - .execute(&format!("LISTEN \"{hashed}\""), &[]) - .instrument(span) - .await?; - - rx - }; + let (rx, drop_guard) = + if let Some(existing_sub) = self.subscriptions.lock().unwrap().get(&hashed).cloned() { + // Reuse the existing broadcast channel + let rx = existing_sub.tx.subscribe(); + let drop_guard = existing_sub.token.clone().drop_guard(); + (rx, drop_guard) + } else { + // Create a new broadcast channel for this subject + let (tx, rx) = tokio::sync::broadcast::channel(1024); + let subscription = Subscription::new(tx.clone()); + + // Register subscription + self.subscriptions + .lock() + .unwrap() + .insert(hashed.clone(), subscription.clone()); + + // Execute LISTEN command on the async client (for receiving notifications) + // This only needs to be done once per channel + let span = tracing::trace_span!("pg_listen"); + self.client + .execute(&format!("LISTEN \"{hashed}\""), &[]) + .instrument(span) + .await?; + + // Spawn a single cleanup task for this subscription waiting on its token + let driver = self.clone(); + let hashed_clone = hashed.clone(); + let tx_clone = tx.clone(); + let token_clone = subscription.token.clone(); + tokio::spawn(async move { + token_clone.cancelled().await; + if tx_clone.receiver_count() == 0 { + let sql = format!("UNLISTEN \"{}\"", hashed_clone); + if let Err(err) = driver.client.execute(sql.as_str(), &[]).await { + tracing::warn!(?err, %hashed_clone, "failed to UNLISTEN channel"); + } else { + tracing::trace!(%hashed_clone, "unlistened channel"); + } + driver.subscriptions.lock().unwrap().remove(&hashed_clone); + } + }); + + let drop_guard = subscription.token.clone().drop_guard(); + (rx, drop_guard) + }; Ok(Box::new(PostgresSubscriber { subject: subject.to_string(), - rx, + rx: Some(rx), + _drop_guard: drop_guard, })) } @@ -191,13 +226,18 @@ impl PubSubDriver for PostgresDriver { pub struct PostgresSubscriber { subject: String, - rx: tokio::sync::broadcast::Receiver>, + rx: Option>>, + _drop_guard: tokio_util::sync::DropGuard, } #[async_trait] impl SubscriberDriver for PostgresSubscriber { async fn next(&mut self) -> Result { - match self.rx.recv().await { + 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, diff --git a/packages/core/guard/server/Cargo.toml b/packages/core/guard/server/Cargo.toml index fbca49ce9f..ef8695c6c2 100644 --- a/packages/core/guard/server/Cargo.toml +++ b/packages/core/guard/server/Cargo.toml @@ -20,6 +20,7 @@ hyper-tungstenite.workspace = true tower.workspace = true udb-util.workspace = true universaldb.workspace = true +universalpubsub.workspace = true futures.workspace = true # TODO: Make this use workspace version hyper = "1.6.0" diff --git a/packages/core/guard/server/src/lib.rs b/packages/core/guard/server/src/lib.rs index 7f3c6e4626..f7533372e3 100644 --- a/packages/core/guard/server/src/lib.rs +++ b/packages/core/guard/server/src/lib.rs @@ -5,6 +5,7 @@ pub mod cache; pub mod errors; pub mod middleware; pub mod routing; +pub mod shared_state; pub mod tls; #[tracing::instrument(skip_all)] @@ -26,8 +27,12 @@ pub async fn start(config: rivet_config::Config, pools: rivet_pools::Pools) -> R tracing::warn!("crypto provider already installed in this process"); } + // Share shared context + let shared_state = shared_state::SharedState::new(ctx.ups()?); + shared_state.start().await?; + // Create handlers - let routing_fn = routing::create_routing_function(ctx.clone()); + let routing_fn = routing::create_routing_function(ctx.clone(), shared_state.clone()); let cache_key_fn = cache::create_cache_key_function(ctx.clone()); let middleware_fn = middleware::create_middleware_function(ctx.clone()); let cert_resolver = tls::create_cert_resolver(&ctx).await?; diff --git a/packages/core/guard/server/src/routing/mod.rs b/packages/core/guard/server/src/routing/mod.rs index b8b0d757f6..524d63075d 100644 --- a/packages/core/guard/server/src/routing/mod.rs +++ b/packages/core/guard/server/src/routing/mod.rs @@ -5,7 +5,7 @@ use gas::prelude::*; use hyper::header::HeaderName; use rivet_guard_core::RoutingFn; -use crate::errors; +use crate::{errors, shared_state::SharedState}; mod api_peer; mod api_public; @@ -17,13 +17,14 @@ pub(crate) const X_RIVET_TARGET: HeaderName = HeaderName::from_static("x-rivet-t /// Creates the main routing function that handles all incoming requests #[tracing::instrument(skip_all)] -pub fn create_routing_function(ctx: StandaloneCtx) -> RoutingFn { +pub fn create_routing_function(ctx: StandaloneCtx, shared_state: SharedState) -> RoutingFn { Arc::new( move |hostname: &str, path: &str, port_type: rivet_guard_core::proxy_service::PortType, headers: &hyper::HeaderMap| { let ctx = ctx.clone(); + let shared_state = shared_state.clone(); Box::pin( async move { @@ -41,9 +42,15 @@ pub fn create_routing_function(ctx: StandaloneCtx) -> RoutingFn { return Ok(routing_output); } - if let Some(routing_output) = - pegboard_gateway::route_request(&ctx, target, host, path, headers) - .await? + if let Some(routing_output) = pegboard_gateway::route_request( + &ctx, + &shared_state, + target, + host, + path, + headers, + ) + .await? { return Ok(routing_output); } diff --git a/packages/core/guard/server/src/routing/pegboard_gateway.rs b/packages/core/guard/server/src/routing/pegboard_gateway.rs index fe8a5b2fe5..58088a1ac4 100644 --- a/packages/core/guard/server/src/routing/pegboard_gateway.rs +++ b/packages/core/guard/server/src/routing/pegboard_gateway.rs @@ -6,7 +6,7 @@ use hyper::header::HeaderName; use rivet_guard_core::proxy_service::{RouteConfig, RouteTarget, RoutingOutput, RoutingTimeout}; use udb_util::{SERIALIZABLE, TxnExt}; -use crate::errors; +use crate::{errors, shared_state::SharedState}; const ACTOR_READY_TIMEOUT: Duration = Duration::from_secs(10); pub const X_RIVET_ACTOR: HeaderName = HeaderName::from_static("x-rivet-actor"); @@ -16,6 +16,7 @@ pub const X_RIVET_PORT: HeaderName = HeaderName::from_static("x-rivet-port"); #[tracing::instrument(skip_all)] pub async fn route_request( ctx: &StandaloneCtx, + shared_state: &SharedState, target: &str, _host: &str, path: &str, @@ -73,7 +74,7 @@ pub async fn route_request( let port_name = port_name.to_str()?; // Lookup actor - find_actor(ctx, actor_id, port_name, path).await + find_actor(ctx, shared_state, actor_id, port_name, path).await } struct FoundActor { @@ -86,6 +87,7 @@ struct FoundActor { #[tracing::instrument(skip_all, fields(%actor_id, %port_name, %path))] async fn find_actor( ctx: &StandaloneCtx, + shared_state: &SharedState, actor_id: Id, port_name: &str, path: &str, @@ -158,10 +160,10 @@ async fn find_actor( actor_ids: vec![actor_id], }); let res = tokio::time::timeout(Duration::from_secs(5), get_runner_fut).await??; - let runner_info = res.actors.into_iter().next().filter(|x| x.is_connectable); + let actor = res.actors.into_iter().next().filter(|x| x.is_connectable); - let runner_id = if let Some(runner_info) = runner_info { - runner_info.runner_id + let runner_id = if let Some(actor) = actor { + actor.runner_id } else { tracing::info!(?actor_id, "waiting for actor to become ready"); @@ -185,11 +187,23 @@ async fn find_actor( tracing::debug!(?actor_id, ?runner_id, "actor ready"); + // Get runner key from runner_id + let runner_key = ctx + .udb()? + .run(|tx, _mc| async move { + let txs = tx.subspace(pegboard::keys::subspace()); + let key_key = pegboard::keys::runner::KeyKey::new(runner_id); + txs.read_opt(&key_key, SERIALIZABLE).await + }) + .await? + .context("runner key not found")?; + // Return pegboard-gateway instance let gateway = pegboard_gateway::PegboardGateway::new( ctx.clone(), + shared_state.pegboard_gateway.clone(), actor_id, - runner_id, + runner_key, port_name.to_string(), ); Ok(Some(RoutingOutput::CustomServe(std::sync::Arc::new( diff --git a/packages/core/guard/server/src/routing/pegboard_tunnel.rs b/packages/core/guard/server/src/routing/pegboard_tunnel.rs index a3b6b0eaad..7d69150a03 100644 --- a/packages/core/guard/server/src/routing/pegboard_tunnel.rs +++ b/packages/core/guard/server/src/routing/pegboard_tunnel.rs @@ -12,13 +12,10 @@ pub async fn route_request( _host: &str, _path: &str, ) -> Result> { - // Check target if target != "tunnel" { return Ok(None); } - // Create pegboard-tunnel service instance - let tunnel = pegboard_tunnel::PegboardTunnelCustomServe::new(ctx.clone()).await?; - + let tunnel = pegboard_tunnel::PegboardTunnelCustomServe::new(ctx.clone()); Ok(Some(RoutingOutput::CustomServe(Arc::new(tunnel)))) } diff --git a/packages/core/guard/server/src/shared_state.rs b/packages/core/guard/server/src/shared_state.rs new file mode 100644 index 0000000000..71bc25939e --- /dev/null +++ b/packages/core/guard/server/src/shared_state.rs @@ -0,0 +1,32 @@ +use anyhow::*; +use gas::prelude::*; +use std::{ops::Deref, sync::Arc}; +use universalpubsub::PubSub; + +#[derive(Clone)] +pub struct SharedState(Arc); + +impl SharedState { + pub fn new(pubsub: PubSub) -> SharedState { + SharedState(Arc::new(SharedStateInner { + pegboard_gateway: pegboard_gateway::shared_state::SharedState::new(pubsub), + })) + } + + pub async fn start(&self) -> Result<()> { + self.pegboard_gateway.start().await?; + Ok(()) + } +} + +impl Deref for SharedState { + type Target = SharedStateInner; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +pub struct SharedStateInner { + pub pegboard_gateway: pegboard_gateway::shared_state::SharedState, +} diff --git a/packages/core/pegboard-gateway/Cargo.toml b/packages/core/pegboard-gateway/Cargo.toml index aeeb32bc53..6a3ddc446a 100644 --- a/packages/core/pegboard-gateway/Cargo.toml +++ b/packages/core/pegboard-gateway/Cargo.toml @@ -14,12 +14,14 @@ gas.workspace = true http-body-util.workspace = true hyper = "1.6" hyper-tungstenite.workspace = true -pegboard = { path = "../../services/pegboard" } +pegboard.workspace = true rand.workspace = true rivet-error.workspace = true -rivet-guard-core = { path = "../guard/core" } +rivet-guard-core.workspace = true rivet-tunnel-protocol.workspace = true rivet-util.workspace = true tokio-tungstenite.workspace = true tokio.workspace = true universalpubsub.workspace = true +versioned-data-util.workspace = true +thiserror.workspace = true diff --git a/packages/core/pegboard-gateway/src/lib.rs b/packages/core/pegboard-gateway/src/lib.rs index 09a960fe88..8d752eba42 100644 --- a/packages/core/pegboard-gateway/src/lib.rs +++ b/packages/core/pegboard-gateway/src/lib.rs @@ -4,56 +4,48 @@ use bytes::Bytes; use futures_util::{SinkExt, StreamExt}; use gas::prelude::*; use http_body_util::{BodyExt, Full}; -use hyper::{Request, Response, StatusCode, body::Incoming as BodyIncoming}; +use hyper::{Request, Response, StatusCode}; use hyper_tungstenite::HyperWebsocket; -use pegboard::pubsub_subjects::{ - TunnelHttpResponseSubject, TunnelHttpRunnerSubject, TunnelHttpWebSocketSubject, -}; -use rivet_error::*; use rivet_guard_core::{ custom_serve::CustomServeTrait, proxy_service::{ResponseBody, X_RIVET_ERROR}, request_context::RequestContext, }; use rivet_tunnel_protocol::{ - MessageBody, StreamFinishReason, ToServerRequestFinish, ToServerRequestStart, - ToServerWebSocketClose, ToServerWebSocketMessage, ToServerWebSocketOpen, TunnelMessage, - versioned, + MessageKind, ToServerRequestStart, ToServerWebSocketClose, ToServerWebSocketMessage, + ToServerWebSocketOpen, }; use rivet_util::serde::HashableMap; -use std::result::Result::Ok as ResultOk; -use std::{ - collections::HashMap, - sync::{ - Arc, - atomic::{AtomicU64, Ordering}, - }, - time::Duration, -}; -use tokio::{ - sync::{Mutex, oneshot}, - time::timeout, -}; +use std::time::Duration; use tokio_tungstenite::tungstenite::Message; -use universalpubsub::NextOutput; + +use crate::shared_state::{SharedState, TunnelMessageData}; + +pub mod shared_state; const UPS_REQ_TIMEOUT: Duration = Duration::from_secs(2); pub struct PegboardGateway { ctx: StandaloneCtx, - request_counter: AtomicU64, + shared_state: SharedState, actor_id: Id, - runner_id: Id, + runner_key: String, port_name: String, } impl PegboardGateway { - pub fn new(ctx: StandaloneCtx, actor_id: Id, runner_id: Id, port_name: String) -> Self { + pub fn new( + ctx: StandaloneCtx, + shared_state: SharedState, + actor_id: Id, + runner_key: String, + port_name: String, + ) -> Self { Self { ctx, - request_counter: AtomicU64::new(0), + shared_state, actor_id, - runner_id, + runner_key, port_name, } } @@ -70,7 +62,7 @@ impl CustomServeTrait for PegboardGateway { match res { Result::Ok(x) => Ok(x), Err(err) => { - if is_tunnel_closed_error(&err) { + if is_tunnel_service_unavailable(&err) { // This will force the request to be retried with a new tunnel Ok(Response::builder() .status(StatusCode::SERVICE_UNAVAILABLE) @@ -96,7 +88,7 @@ impl CustomServeTrait for PegboardGateway { { Result::Ok(()) => std::result::Result::<(), (HyperWebsocket, anyhow::Error)>::Ok(()), Result::Err((client_ws, err)) => { - if is_tunnel_closed_error(&err) { + if is_tunnel_service_unavailable(&err) { Err(( client_ws, rivet_guard_core::errors::WebSocketServiceUnavailable.build(), @@ -124,13 +116,10 @@ impl PegboardGateway { .map_err(|_| anyhow!("invalid x-rivet-actor header"))? .to_string(); - // Generate request ID using atomic counter - let request_id = self.request_counter.fetch_add(1, Ordering::SeqCst); - // Extract request parts let mut headers = HashableMap::new(); for (name, value) in req.headers() { - if let ResultOk(value_str) = value.to_str() { + if let Result::Ok(value_str) = value.to_str() { headers.insert(name.to_string(), value_str.to_string()); } } @@ -149,9 +138,21 @@ impl PegboardGateway { .map_err(|e| anyhow!("failed to read body: {}", e))? .to_bytes(); - // Create tunnel message - let request_start = ToServerRequestStart { - request_id, + // Build subject to publish to + let tunnel_subject = pegboard::pubsub_subjects::TunnelRunnerReceiverSubject::new( + &self.runner_key, + &self.port_name, + ) + .to_string(); + + // Start listening for request responses + let (request_id, mut msg_rx) = self + .shared_state + .start_in_flight_request(tunnel_subject) + .await; + + // Start request + let message = MessageKind::ToServerRequestStart(ToServerRequestStart { actor_id: actor_id.clone(), method, path, @@ -162,117 +163,33 @@ impl PegboardGateway { Some(body_bytes.to_vec()) }, stream: false, - }; - - let message = TunnelMessage { - body: MessageBody::ToServerRequestStart(request_start), - }; + }); + self.shared_state.send_message(request_id, message).await?; + + // Wait for response + tracing::info!("starting response handler task"); + let response_start = loop { + let Some(msg) = msg_rx.recv().await else { + tracing::warn!("received no message response"); + return Err(RequestError::ServiceUnavailable.into()); + }; - // Serialize message - let serialized = versioned::TunnelMessage::serialize(versioned::TunnelMessage::V1(message)) - .map_err(|e| anyhow!("failed to serialize message: {}", e))?; - - // Build pubsub topic - let tunnel_subject = TunnelHttpRunnerSubject::new(self.runner_id, &self.port_name); - let topic = tunnel_subject.to_string(); - - tracing::info!( - %topic, - ?self.runner_id, - %self.port_name, - ?request_id, - "publishing request to pubsub" - ); - - // Create response channel - let (response_tx, response_rx) = oneshot::channel(); - let response_map = Arc::new(Mutex::new(HashMap::new())); - response_map.lock().await.insert(request_id, response_tx); - - // Subscribe to response topic - let response_subject = - TunnelHttpResponseSubject::new(self.runner_id, &self.port_name, request_id); - let response_topic = response_subject.to_string(); - - tracing::info!( - ?response_topic, - ?request_id, - "subscribing to response topic" - ); - - let mut subscriber = self.ctx.ups()?.subscribe(&response_topic).await?; - - // Spawn task to handle response - let response_map_clone = response_map.clone(); - tokio::spawn(async move { - tracing::info!("starting response handler task"); - while let ResultOk(NextOutput::Message(msg)) = subscriber.next().await { - // Ack message - match msg.reply(&[]).await { - Result::Ok(_) => {} - Err(err) => { - tracing::warn!(?err, "failed to ack gateway request response message") + match msg { + TunnelMessageData::Message(msg) => match msg { + MessageKind::ToClientResponseStart(response_start) => { + break response_start; } - }; - - tracing::info!( - payload_len = msg.payload.len(), - "received response from pubsub" - ); - if let ResultOk(tunnel_msg) = versioned::TunnelMessage::deserialize(&msg.payload) { - match tunnel_msg.body { - MessageBody::ToClientResponseStart(response_start) => { - tracing::info!(request_id = ?response_start.request_id, status = response_start.status, "received response from tunnel"); - if let Some(tx) = response_map_clone - .lock() - .await - .remove(&response_start.request_id) - { - tracing::info!(request_id = ?response_start.request_id, "sending response to handler"); - let _ = tx.send(response_start); - } else { - tracing::warn!(request_id = ?response_start.request_id, "no handler found for response"); - } - } - _ => { - tracing::warn!("received non-response message from pubsub"); - } + _ => { + tracing::warn!("received non-response message from pubsub"); } - } else { - tracing::error!("failed to deserialize response from pubsub"); + }, + TunnelMessageData::Timeout => { + tracing::warn!("tunnel message timeout"); + return Err(RequestError::ServiceUnavailable.into()); } } - tracing::info!("response handler task ended"); - }); - - // Publish request - self.ctx - .ups()? - .request_with_timeout(&topic, &serialized, UPS_REQ_TIMEOUT) - .await - .map_err(|e| anyhow!("failed to publish request: {}", e))?; - - // Send finish message - let finish_message = TunnelMessage { - body: MessageBody::ToServerRequestFinish(ToServerRequestFinish { - request_id, - reason: StreamFinishReason::Complete, - }), - }; - let finish_serialized = - versioned::TunnelMessage::serialize(versioned::TunnelMessage::V1(finish_message)) - .map_err(|e| anyhow!("failed to serialize finish message: {}", e))?; - self.ctx - .ups()? - .request_with_timeout(&topic, &finish_serialized, UPS_REQ_TIMEOUT) - .await - .map_err(|e| anyhow!("failed to publish finish message: {}", e))?; - - // Wait for response with timeout - let response_start = match timeout(Duration::from_secs(30), response_rx).await { - ResultOk(ResultOk(response)) => response, - _ => return Err(anyhow!("request timed out")), }; + tracing::info!("response handler task ended"); // Build HTTP response let mut response_builder = @@ -309,62 +226,69 @@ impl PegboardGateway { Err(err) => return Err((client_ws, err)), }; - // Generate WebSocket ID using atomic counter - let websocket_id = self.request_counter.fetch_add(1, Ordering::SeqCst); - // Extract headers let mut request_headers = HashableMap::new(); for (name, value) in headers { - if let ResultOk(value_str) = value.to_str() { + if let Result::Ok(value_str) = value.to_str() { request_headers.insert(name.to_string(), value_str.to_string()); } } - let ups = match self.ctx.ups() { - Result::Ok(u) => u, - Err(err) => return Err((client_ws, err.into())), - }; - - // Subscribe to messages from server before informing server that a client websocket is connecting to - // prevent race conditions. - let ws_subject = - TunnelHttpWebSocketSubject::new(self.runner_id, &self.port_name, websocket_id); - let response_topic = ws_subject.to_string(); - let mut subscriber = match ups.subscribe(&response_topic).await { - Result::Ok(sub) => sub, - Err(err) => return Err((client_ws, err.into())), - }; + // Build subject to publish to + let tunnel_subject = pegboard::pubsub_subjects::TunnelRunnerReceiverSubject::new( + &self.runner_key, + &self.port_name, + ) + .to_string(); - // Build pubsub topic - let tunnel_subject = TunnelHttpRunnerSubject::new(self.runner_id, &self.port_name); - let topic = tunnel_subject.to_string(); + // Start listening for WebSocket messages + let (request_id, mut msg_rx) = self + .shared_state + .start_in_flight_request(tunnel_subject.clone()) + .await; // Send WebSocket open message - let open_message = TunnelMessage { - body: MessageBody::ToServerWebSocketOpen(ToServerWebSocketOpen { - actor_id: actor_id.clone(), - web_socket_id: websocket_id, - path: path.to_string(), - headers: request_headers, - }), - }; - - let serialized = - match versioned::TunnelMessage::serialize(versioned::TunnelMessage::V1(open_message)) { - Result::Ok(s) => s, - Err(e) => { - return Err(( - client_ws, - anyhow!("failed to serialize websocket open: {}", e), - )); - } - }; + let open_message = MessageKind::ToServerWebSocketOpen(ToServerWebSocketOpen { + actor_id: actor_id.clone(), + path: path.to_string(), + headers: request_headers, + }); - if let Err(err) = ups - .request_with_timeout(&topic, &serialized, UPS_REQ_TIMEOUT) + if let Err(err) = self + .shared_state + .send_message(request_id, open_message) .await { - return Err((client_ws, err.into())); + return Err((client_ws, err)); + } + + // Wait for WebSocket open acknowledgment + let open_ack_received = loop { + let Some(msg) = msg_rx.recv().await else { + tracing::warn!("received no websocket open response"); + return Err((client_ws, RequestError::ServiceUnavailable.into())); + }; + + match msg { + TunnelMessageData::Message(MessageKind::ToClientWebSocketOpen) => { + break true; + } + TunnelMessageData::Message(MessageKind::ToClientWebSocketClose(close)) => { + tracing::info!(?close, "websocket closed before opening"); + return Err((client_ws, RequestError::ServiceUnavailable.into())); + } + TunnelMessageData::Timeout => { + tracing::warn!("websocket open timeout"); + return Err((client_ws, RequestError::ServiceUnavailable.into())); + } + _ => { + tracing::warn!("received unexpected message while waiting for websocket open"); + } + } + }; + + if !open_ack_received { + return Err((client_ws, anyhow!("failed to open websocket"))); } // Accept the WebSocket @@ -379,33 +303,30 @@ impl PegboardGateway { let (mut ws_sink, mut ws_stream) = ws_stream.split(); // Spawn task to forward messages from server to client + let mut msg_rx_for_task = msg_rx; tokio::spawn(async move { - while let ResultOk(NextOutput::Message(msg)) = subscriber.next().await { - // Ack message - match msg.reply(&[]).await { - Result::Ok(_) => {} - Err(err) => { - tracing::warn!(?err, "failed to ack gateway websocket message") - } - }; - - if let ResultOk(tunnel_msg) = versioned::TunnelMessage::deserialize(&msg.payload) { - match tunnel_msg.body { - MessageBody::ToClientWebSocketMessage(ws_msg) => { - if ws_msg.web_socket_id == websocket_id { - let msg = if ws_msg.binary { - Message::Binary(ws_msg.data.into()) - } else { - Message::Text( - String::from_utf8_lossy(&ws_msg.data).into_owned().into(), - ) - }; - let _ = ws_sink.send(msg).await; - } + while let Some(msg) = msg_rx_for_task.recv().await { + match msg { + TunnelMessageData::Message(MessageKind::ToClientWebSocketMessage(ws_msg)) => { + let msg = if ws_msg.binary { + Message::Binary(ws_msg.data.into()) + } else { + Message::Text(String::from_utf8_lossy(&ws_msg.data).into_owned().into()) + }; + if let Err(e) = ws_sink.send(msg).await { + tracing::warn!(?e, "failed to send websocket message to client"); + break; } - MessageBody::ToClientWebSocketClose(_) => break, - _ => {} } + TunnelMessageData::Message(MessageKind::ToClientWebSocketClose(close)) => { + tracing::info!(?close, "server closed websocket"); + break; + } + TunnelMessageData::Timeout => { + tracing::warn!("websocket message timeout"); + break; + } + _ => {} } } }); @@ -414,25 +335,14 @@ impl PegboardGateway { let mut close_reason = None; while let Some(msg) = ws_stream.next().await { match msg { - ResultOk(Message::Binary(data)) => { - let ws_message = TunnelMessage { - body: MessageBody::ToServerWebSocketMessage(ToServerWebSocketMessage { - web_socket_id: websocket_id, + Result::Ok(Message::Binary(data)) => { + let ws_message = + MessageKind::ToServerWebSocketMessage(ToServerWebSocketMessage { data: data.into(), binary: true, - }), - }; - let serialized = match versioned::TunnelMessage::serialize( - versioned::TunnelMessage::V1(ws_message), - ) { - Result::Ok(s) => s, - Err(_) => break, - }; - if let Err(err) = ups - .request_with_timeout(&topic, &serialized, UPS_REQ_TIMEOUT) - .await - { - if is_tunnel_closed_error(&err) { + }); + if let Err(err) = self.shared_state.send_message(request_id, ws_message).await { + if is_tunnel_service_unavailable(&err) { tracing::warn!("tunnel closed sending binary message"); close_reason = Some("Tunnel closed".to_string()); break; @@ -441,25 +351,14 @@ impl PegboardGateway { } } } - ResultOk(Message::Text(text)) => { - let ws_message = TunnelMessage { - body: MessageBody::ToServerWebSocketMessage(ToServerWebSocketMessage { - web_socket_id: websocket_id, + Result::Ok(Message::Text(text)) => { + let ws_message = + MessageKind::ToServerWebSocketMessage(ToServerWebSocketMessage { data: text.as_bytes().to_vec(), binary: false, - }), - }; - let serialized = match versioned::TunnelMessage::serialize( - versioned::TunnelMessage::V1(ws_message), - ) { - Result::Ok(s) => s, - Err(_) => break, - }; - if let Err(err) = ups - .request_with_timeout(&topic, &serialized, UPS_REQ_TIMEOUT) - .await - { - if is_tunnel_closed_error(&err) { + }); + if let Err(err) = self.shared_state.send_message(request_id, ws_message).await { + if is_tunnel_service_unavailable(&err) { tracing::warn!("tunnel closed sending text message"); close_reason = Some("Tunnel closed".to_string()); break; @@ -468,32 +367,23 @@ impl PegboardGateway { } } } - ResultOk(Message::Close(_)) | Err(_) => break, + Result::Ok(Message::Close(_)) | Err(_) => break, _ => {} } } // Send WebSocket close message - let close_message = TunnelMessage { - body: MessageBody::ToServerWebSocketClose(ToServerWebSocketClose { - web_socket_id: websocket_id, - code: None, - reason: close_reason, - }), - }; - - let serialized = match versioned::TunnelMessage::serialize(versioned::TunnelMessage::V1( - close_message, - )) { - Result::Ok(s) => s, - Err(_) => Vec::new(), - }; + let close_message = MessageKind::ToServerWebSocketClose(ToServerWebSocketClose { + code: None, + reason: close_reason, + }); - if let Err(err) = ups - .request_with_timeout(&topic, &serialized, UPS_REQ_TIMEOUT) + if let Err(err) = self + .shared_state + .send_message(request_id, close_message) .await { - if is_tunnel_closed_error(&err) { + if is_tunnel_service_unavailable(&err) { tracing::warn!("tunnel closed sending close message"); } else { tracing::error!(?err, "error sending close message"); @@ -504,14 +394,13 @@ impl PegboardGateway { } } +#[derive(thiserror::Error, Debug)] +enum RequestError { + #[error("service unavailable")] + ServiceUnavailable, +} + /// Determines if the tunnel is closed by if the UPS service is no longer responding. -fn is_tunnel_closed_error(err: &anyhow::Error) -> bool { - if let Some(err) = err.chain().find_map(|x| x.downcast_ref::()) - && err.group() == "ups" - && err.code() == "request_timeout" - { - true - } else { - false - } +fn is_tunnel_service_unavailable(err: &anyhow::Error) -> bool { + err.chain().any(|x| x.is::()) } diff --git a/packages/core/pegboard-gateway/src/shared_state.rs b/packages/core/pegboard-gateway/src/shared_state.rs new file mode 100644 index 0000000000..75b485679b --- /dev/null +++ b/packages/core/pegboard-gateway/src/shared_state.rs @@ -0,0 +1,286 @@ +use anyhow::*; +use gas::prelude::*; +use rivet_tunnel_protocol::{ + MessageId, MessageKind, PROTOCOL_VERSION, PubSubMessage, RequestId, versioned, +}; +use std::{ + collections::HashMap, + ops::Deref, + sync::Arc, + time::{Duration, Instant}, +}; +use tokio::sync::{Mutex, mpsc}; +use universalpubsub::{NextOutput, PubSub, PublishOpts, Subscriber}; +use versioned_data_util::OwnedVersionedData as _; + +const GC_INTERVAL: Duration = Duration::from_secs(60); +const MESSAGE_ACK_TIMEOUT: Duration = Duration::from_secs(5); + +struct InFlightRequest { + /// UPS subject to send messages to for this request. + receiver_subject: String, + /// Sender for incoming messages to this request. + msg_tx: mpsc::Sender, + /// True once first message for this request has been sent (so runner learned reply_to). + opened: bool, +} + +struct PendingMessage { + request_id: RequestId, + send_instant: Instant, +} + +pub enum TunnelMessageData { + Message(MessageKind), + Timeout, +} + +pub struct SharedStateInner { + ups: PubSub, + receiver_subject: String, + requests_in_flight: Mutex>, + pending_messages: Mutex>, +} + +#[derive(Clone)] +pub struct SharedState(Arc); + +impl SharedState { + pub fn new(ups: PubSub) -> Self { + let gateway_id = Uuid::new_v4(); + let receiver_subject = + pegboard::pubsub_subjects::TunnelGatewayReceiverSubject::new(gateway_id).to_string(); + + Self(Arc::new(SharedStateInner { + ups, + receiver_subject, + requests_in_flight: Mutex::new(HashMap::new()), + pending_messages: Mutex::new(HashMap::new()), + })) + } + + pub async fn start(&self) -> Result<()> { + let sub = self.ups.subscribe(&self.receiver_subject).await?; + + let self_clone = self.clone(); + tokio::spawn(async move { self_clone.receiver(sub).await }); + + let self_clone = self.clone(); + tokio::spawn(async move { self_clone.gc().await }); + + Ok(()) + } + + pub async fn send_message( + &self, + request_id: RequestId, + message_kind: MessageKind, + ) -> Result<()> { + let message_id = Uuid::new_v4().as_bytes().clone(); + + // Get subject and whether this is the first message for this request + let (tunnel_receiver_subject, include_reply_to) = { + let mut requests_in_flight = self.requests_in_flight.lock().await; + if let Some(req) = requests_in_flight.get_mut(&request_id) { + let receiver_subject = req.receiver_subject.clone(); + let include_reply_to = !req.opened; + if include_reply_to { + // Mark as opened so subsequent messages skip reply_to + req.opened = true; + } + (receiver_subject, include_reply_to) + } else { + bail!("request not in flight") + } + }; + + // Save pending message + { + let mut pending_messages = self.pending_messages.lock().await; + pending_messages.insert( + message_id, + PendingMessage { + request_id, + send_instant: Instant::now(), + }, + ); + } + + // Send message + let message = PubSubMessage { + request_id, + message_id, + // Only send reply to subject on the first message for this request. This reduces + // overhead of subsequent messages. + reply_to: if include_reply_to { + Some(self.receiver_subject.clone()) + } else { + None + }, + message_kind, + }; + let message_serialized = versioned::PubSubMessage::latest(message) + .serialize_with_embedded_version(PROTOCOL_VERSION)?; + self.ups + .publish( + &tunnel_receiver_subject, + &message_serialized, + PublishOpts::one(), + ) + .await?; + + Ok(()) + } + + pub async fn start_in_flight_request( + &self, + receiver_subject: String, + ) -> (RequestId, mpsc::Receiver) { + let id = Uuid::new_v4().into_bytes(); + let (msg_tx, msg_rx) = mpsc::channel(128); + self.requests_in_flight.lock().await.insert( + id, + InFlightRequest { + receiver_subject, + msg_tx, + opened: false, + }, + ); + (id, msg_rx) + } + + async fn receiver(&self, mut sub: Subscriber) { + while let Result::Ok(NextOutput::Message(msg)) = sub.next().await { + tracing::info!( + payload_len = msg.payload.len(), + "received message from pubsub" + ); + + match versioned::PubSubMessage::deserialize_with_embedded_version(&msg.payload) { + Result::Ok(msg) => { + tracing::debug!( + ?msg.request_id, + ?msg.message_id, + "successfully deserialized message" + ); + if let MessageKind::Ack = &msg.message_kind { + // Handle ack message + + let mut pending_messages = self.pending_messages.lock().await; + if pending_messages.remove(&msg.message_id).is_none() { + tracing::warn!( + "pending message does not exist or ack received after message body" + ); + } + } else { + // Forward message to receiver + + // Send message to sender using request_id directly + let requests_in_flight = self.requests_in_flight.lock().await; + let Some(in_flight) = requests_in_flight.get(&msg.request_id) else { + tracing::debug!( + ?msg.request_id, + "in flight has already been disconnected" + ); + continue; + }; + tracing::debug!( + ?msg.request_id, + "forwarding message to request handler" + ); + let _ = in_flight + .msg_tx + .send(TunnelMessageData::Message(msg.message_kind)) + .await; + + // Send ack + let ups_clone = self.ups.clone(); + let receiver_subject = in_flight.receiver_subject.clone(); + let ack_message = PubSubMessage { + request_id: msg.request_id, + message_id: Uuid::new_v4().into_bytes(), + reply_to: None, + message_kind: MessageKind::Ack, + }; + let ack_message_serialized = + match versioned::PubSubMessage::latest(ack_message) + .serialize_with_embedded_version(PROTOCOL_VERSION) + { + Result::Ok(x) => x, + Err(err) => { + tracing::error!(?err, "failed to serialize ack"); + continue; + } + }; + tokio::spawn(async move { + if let Result::Err(err) = ups_clone + .publish( + &receiver_subject, + &ack_message_serialized, + PublishOpts::one(), + ) + .await + { + tracing::warn!(?err, "failed to ack message") + } + }); + } + } + Result::Err(err) => { + tracing::error!(?err, "failed to parse message"); + } + } + } + } + + async fn gc(&self) { + let mut interval = tokio::time::interval(GC_INTERVAL); + loop { + interval.tick().await; + + let now = Instant::now(); + + // Purge unacked messages + { + let mut pending_messages = self.pending_messages.lock().await; + let mut removed_req_ids = Vec::new(); + pending_messages.retain(|_k, v| { + if now.duration_since(v.send_instant) > MESSAGE_ACK_TIMEOUT { + // Expired + removed_req_ids.push(v.request_id.clone()); + false + } else { + true + } + }); + + // Close in-flight messages + let requests_in_flight = self.requests_in_flight.lock().await; + for req_id in removed_req_ids { + if let Some(x) = requests_in_flight.get(&req_id) { + let _ = x.msg_tx.send(TunnelMessageData::Timeout); + } else { + tracing::warn!( + ?req_id, + "message expired for in flight that does not exist" + ); + } + } + } + + // Purge no longer in flight + { + let mut requests_in_flight = self.requests_in_flight.lock().await; + requests_in_flight.retain(|_k, v| !v.msg_tx.is_closed()); + } + } + } +} + +impl Deref for SharedState { + type Target = SharedStateInner; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} diff --git a/packages/core/pegboard-runner-ws/src/lib.rs b/packages/core/pegboard-runner-ws/src/lib.rs index 32de54f51d..576c7d567e 100644 --- a/packages/core/pegboard-runner-ws/src/lib.rs +++ b/packages/core/pegboard-runner-ws/src/lib.rs @@ -270,6 +270,7 @@ async fn build_connection( UrlData { protocol_version, namespace, + runner_key, }: UrlData, ) -> Result<(Id, Arc)> { let namespace = ctx @@ -300,9 +301,7 @@ async fn build_connection( .map_err(|err: anyhow::Error| WsError::InvalidPacket(err.to_string()).build())?; let (runner_id, workflow_id) = if let protocol::ToServer::Init { - runner_id, name, - key, version, total_slots, addresses_http, @@ -311,7 +310,16 @@ async fn build_connection( .. } = &packet { - let runner_id = if let Some(runner_id) = runner_id { + // Look up existing runner by key + let existing_runner = ctx + .op(pegboard::ops::runner::get_by_key::Input { + namespace_id: namespace.namespace_id, + name: name.clone(), + key: runner_key.clone(), + }) + .await?; + + let runner_id = if let Some(runner) = existing_runner.runner { // IMPORTANT: Before we spawn/get the workflow, we try to update the runner's last ping ts. // This ensures if the workflow is currently checking for expiry that it will not expire // (because we are about to send signals to it) and if it is already expired (but not @@ -319,7 +327,7 @@ async fn build_connection( let update_ping_res = ctx .op(pegboard::ops::runner::update_alloc_idx::Input { runners: vec![pegboard::ops::runner::update_alloc_idx::Runner { - runner_id: *runner_id, + runner_id: runner.runner_id, action: Action::UpdatePing { rtt: 0 }, }], }) @@ -332,11 +340,14 @@ async fn build_connection( .map(|notif| matches!(notif.eligibility, RunnerEligibility::Expired)) .unwrap_or_default() { + // Runner expired, create a new one Id::new_v1(ctx.config().dc_label()) } else { - *runner_id + // Use existing runner + runner.runner_id } } else { + // No existing runner for this key, create a new one Id::new_v1(ctx.config().dc_label()) }; @@ -346,7 +357,7 @@ async fn build_connection( runner_id, namespace_id: namespace.namespace_id, name: name.clone(), - key: key.clone(), + key: runner_key.clone(), version: version.clone(), total_slots: *total_slots, @@ -790,19 +801,17 @@ async fn msg_thread_inner(ctx: &StandaloneCtx, conns: Arc>) struct UrlData { protocol_version: u16, namespace: String, + runner_key: String, } fn parse_url(addr: SocketAddr, uri: hyper::Uri) -> Result { let url = url::Url::parse(&format!("ws://{addr}{uri}"))?; - // Get protocol version from last path segment - let last_segment = url - .path_segments() - .context("invalid url")? - .last() - .context("no path segments")?; - ensure!(last_segment.starts_with('v'), "invalid protocol version"); - let protocol_version = last_segment[1..] + // Read protocol version from query parameters (required) + let protocol_version = url + .query_pairs() + .find_map(|(n, v)| (n == "protocol_version").then_some(v)) + .context("missing `protocol_version` query parameter")? .parse::() .context("invalid protocol version")?; @@ -813,9 +822,17 @@ fn parse_url(addr: SocketAddr, uri: hyper::Uri) -> Result { .context("missing `namespace` query parameter")? .to_string(); + // Read runner key from query parameters (required) + let runner_key = url + .query_pairs() + .find_map(|(n, v)| (n == "runner_key").then_some(v)) + .context("missing `runner_key` query parameter")? + .to_string(); + Ok(UrlData { protocol_version, namespace, + runner_key, }) } diff --git a/packages/core/pegboard-tunnel/Cargo.toml b/packages/core/pegboard-tunnel/Cargo.toml index ad8747ffd8..ef927c3559 100644 --- a/packages/core/pegboard-tunnel/Cargo.toml +++ b/packages/core/pegboard-tunnel/Cargo.toml @@ -8,6 +8,7 @@ edition.workspace = true [dependencies] anyhow.workspace = true async-trait.workspace = true +bytes.workspace = true futures.workspace = true gas.workspace = true http-body-util = "0.1" @@ -37,4 +38,4 @@ versioned-data-util.workspace = true rand.workspace = true rivet-cache.workspace = true rivet-pools.workspace = true -rivet-util.workspace = true \ No newline at end of file +rivet-util.workspace = true diff --git a/packages/core/pegboard-tunnel/src/lib.rs b/packages/core/pegboard-tunnel/src/lib.rs index 37ddc97cb8..a71d3dc62f 100644 --- a/packages/core/pegboard-tunnel/src/lib.rs +++ b/packages/core/pegboard-tunnel/src/lib.rs @@ -1,54 +1,29 @@ -use std::collections::HashMap; -use std::sync::Arc; -use std::time::Duration; - use anyhow::*; use async_trait::async_trait; +use bytes::Bytes; use futures::{SinkExt, StreamExt}; use gas::prelude::*; use http_body_util::Full; -use hyper::body::{Bytes, Incoming as BodyIncoming}; -use hyper::{Request, Response, StatusCode}; -use hyper_tungstenite::tungstenite::Utf8Bytes as WsUtf8Bytes; -use hyper_tungstenite::tungstenite::protocol::frame::CloseFrame as WsCloseFrame; -use hyper_tungstenite::tungstenite::protocol::frame::coding::CloseCode as WsCloseCode; +use hyper::{Response, StatusCode}; use hyper_tungstenite::{HyperWebsocket, tungstenite::Message as WsMessage}; -use pegboard::pubsub_subjects::{ - TunnelHttpResponseSubject, TunnelHttpRunnerSubject, TunnelHttpWebSocketSubject, +use rivet_guard_core::{ + custom_serve::CustomServeTrait, proxy_service::ResponseBody, request_context::RequestContext, +}; +use rivet_tunnel_protocol::{ + MessageKind, PROTOCOL_VERSION, PubSubMessage, RequestId, RunnerMessage, versioned, }; -use rivet_guard_core::custom_serve::CustomServeTrait; -use rivet_guard_core::proxy_service::ResponseBody; -use rivet_guard_core::request_context::RequestContext; -use rivet_pools::Pools; -use rivet_tunnel_protocol::{MessageBody, TunnelMessage, versioned}; -use rivet_util::Id; -use std::net::SocketAddr; -use tokio::net::TcpListener; +use std::{collections::HashMap, sync::Arc}; use tokio::sync::Mutex; -use tokio::sync::RwLock; -use tokio_tungstenite::accept_async; -use tracing::{error, info}; -use universalpubsub::pubsub::NextOutput; - -const UPS_REQ_TIMEOUT: Duration = Duration::from_secs(2); - -struct RunnerConnection { - _runner_id: Id, - _port_name: String, -} - -type Connections = Arc>>>; +use universalpubsub::{PublishOpts, pubsub::NextOutput}; +use versioned_data_util::OwnedVersionedData as _; pub struct PegboardTunnelCustomServe { ctx: StandaloneCtx, - connections: Connections, } impl PegboardTunnelCustomServe { - pub async fn new(ctx: StandaloneCtx) -> Result { - let connections = Arc::new(RwLock::new(HashMap::new())); - - Ok(Self { ctx, connections }) + pub fn new(ctx: StandaloneCtx) -> Self { + Self { ctx } } } @@ -82,654 +57,258 @@ impl CustomServeTrait for PegboardTunnelCustomServe { Result::Ok(u) => u, Err(e) => return Err((client_ws, e.into())), }; - let connections = self.connections.clone(); - // Extract runner_id from query parameters - let runner_id = if let std::result::Result::Ok(url) = - url::Url::parse(&format!("ws://placeholder/{path}")) + // Parse URL to extract runner_id and protocol version + let url = match url::Url::parse(&format!("ws://placeholder/{path}")) { + Result::Ok(u) => u, + Err(e) => return Err((client_ws, e.into())), + }; + + // Extract runner_key from query parameters (required) + let runner_key = match url + .query_pairs() + .find_map(|(n, v)| (n == "runner_key").then_some(v)) + { + Some(key) => key.to_string(), + None => { + return Err((client_ws, anyhow!("runner_key query parameter is required"))); + } + }; + + // Extract protocol version from query parameters (required) + let protocol_version = match url + .query_pairs() + .find_map(|(n, v)| (n == "protocol_version").then_some(v)) + .as_ref() + .and_then(|v| v.parse::().ok()) { - url.query_pairs() - .find_map(|(n, v)| (n == "runner_id").then_some(v)) - .as_ref() - .and_then(|id| Id::parse(id).ok()) - .unwrap_or(Id::nil()) - } else { - Id::nil() + Some(version) => version, + None => { + return Err(( + client_ws, + anyhow!("protocol_version query parameter is required and must be a valid u16"), + )); + } }; let port_name = "main".to_string(); // Use "main" as default port name - info!( - ?runner_id, + tracing::info!( + ?runner_key, ?port_name, + ?protocol_version, ?path, "tunnel WebSocket connection established" ); - let connection_id = Id::nil(); - // Subscribe to pubsub topic for this runner before accepting the client websocket so // that failures can be retried by the proxy. - let topic = TunnelHttpRunnerSubject::new(runner_id, &port_name).to_string(); - info!(%topic, ?runner_id, "subscribing to pubsub topic"); - + let topic = + pegboard::pubsub_subjects::TunnelRunnerReceiverSubject::new(&runner_key, &port_name) + .to_string(); + tracing::info!(%topic, ?runner_key, "subscribing to runner receiver topic"); let mut sub = match ups.subscribe(&topic).await { Result::Ok(s) => s, Err(e) => return Err((client_ws, e.into())), }; + // Accept WS let ws_stream = match client_ws.await { Result::Ok(ws) => ws, Err(e) => { // Handshake already in progress; cannot retry safely here - error!(error=?e, "client websocket await failed"); + tracing::error!(error=?e, "client websocket await failed"); return std::result::Result::<(), (HyperWebsocket, anyhow::Error)>::Ok(()); } }; - // Split WebSocket stream into read and write halves let (ws_write, mut ws_read) = ws_stream.split(); let ws_write = Arc::new(tokio::sync::Mutex::new(ws_write)); - // Store connection - let connection = Arc::new(RunnerConnection { - _runner_id: runner_id, - _port_name: port_name.clone(), - }); + struct ActiveRequest { + /// Subject to send replies to. + reply_to: String, + } - connections - .write() - .await - .insert(connection_id, connection.clone()); + // Active HTTP & WebSocket requests. They are separate but use the same mechanism to + // maintain state. + let active_requests = Arc::new(Mutex::new(HashMap::::new())); - // Handle bidirectional message forwarding + // Forward pubsub -> WebSocket let ws_write_pubsub_to_ws = ws_write.clone(); - let connections_clone = connections.clone(); let ups_clone = ups.clone(); - - // Task for forwarding pubsub -> WebSocket + let active_requests_clone = active_requests.clone(); let pubsub_to_ws = tokio::spawn(async move { - info!("starting pubsub to WebSocket forwarding task"); - while let ::std::result::Result::Ok(NextOutput::Message(msg)) = sub.next().await { - // Ack message - match msg.reply(&[]).await { - Result::Ok(_) => {} + while let Result::Ok(NextOutput::Message(ups_msg)) = sub.next().await { + tracing::info!( + payload_len = ups_msg.payload.len(), + "received message from pubsub, forwarding to WebSocket" + ); + + // Parse message + let msg = match versioned::PubSubMessage::deserialize_with_embedded_version( + &ups_msg.payload, + ) { + Result::Ok(x) => x, Err(err) => { - tracing::warn!(?err, "failed to ack gateway request response message") + tracing::error!(?err, "failed to parse tunnel message"); + continue; } }; - info!( - payload_len = msg.payload.len(), - "received message from pubsub, forwarding to WebSocket" - ); + // Save active request + if let Some(reply_to) = msg.reply_to { + let mut active_requests = active_requests_clone.lock().await; + active_requests.insert(msg.request_id, ActiveRequest { reply_to }); + } + + // If terminal, remove active request tracking + if is_message_kind_request_close(&msg.message_kind) { + let mut active_requests = active_requests_clone.lock().await; + active_requests.remove(&msg.request_id); + } + // Forward raw message to WebSocket - let ws_msg = WsMessage::Binary(msg.payload.to_vec().into()); + let tunnel_msg = match versioned::RunnerMessage::latest(RunnerMessage { + request_id: msg.request_id, + message_id: msg.message_id, + message_kind: msg.message_kind, + }) + .serialize_version(protocol_version) + { + Result::Ok(x) => x, + Err(err) => { + tracing::error!(?err, "failed to serialize tunnel message"); + continue; + } + }; + let ws_msg = WsMessage::Binary(tunnel_msg.into()); { let mut stream = ws_write_pubsub_to_ws.lock().await; if let Err(e) = stream.send(ws_msg).await { - error!(?e, "failed to send message to WebSocket"); + tracing::error!(?e, "failed to send message to WebSocket"); break; } } } - info!("pubsub to WebSocket forwarding task ended"); + tracing::info!("pubsub to WebSocket forwarding task ended"); }); - // Task for forwarding WebSocket -> pubsub - let ws_write_ws_to_pubsub = ws_write.clone(); + // Forward WebSocket -> pubsub + let active_requests_clone = active_requests.clone(); + let runner_key_clone = runner_key.clone(); let ws_to_pubsub = tokio::spawn(async move { - info!("starting WebSocket to pubsub forwarding task"); + tracing::info!("starting WebSocket to pubsub forwarding task"); while let Some(msg) = ws_read.next().await { match msg { - ::std::result::Result::Ok(WsMessage::Binary(data)) => { - info!( + Result::Ok(WsMessage::Binary(data)) => { + tracing::info!( data_len = data.len(), "received binary message from WebSocket" ); - // Parse the tunnel message to extract request_id - match versioned::TunnelMessage::deserialize(&data) { - ::std::result::Result::Ok(tunnel_msg) => { - // Handle different message types - match &tunnel_msg.body { - MessageBody::ToClientResponseStart(resp) => { - info!(?resp.request_id, status = resp.status, "forwarding HTTP response to pubsub"); - let response_topic = TunnelHttpResponseSubject::new( - runner_id, - &port_name, - resp.request_id, - ) - .to_string(); - - info!(%response_topic, ?resp.request_id, "publishing HTTP response to pubsub"); - if let Err(e) = ups_clone - .request_with_timeout( - &response_topic, - &data.to_vec(), - UPS_REQ_TIMEOUT, - ) - .await - { - let err_any: anyhow::Error = e.into(); - if is_tunnel_closed_error(&err_any) { - info!( - "tunnel closed while publishing HTTP response; closing client websocket" - ); - // Close client websocket with reason - send_tunnel_closed_close_hyper( - &ws_write_ws_to_pubsub, - ) - .await; - break; - } else { - error!(?err_any, ?resp.request_id, "failed to publish HTTP response to pubsub"); - } - } else { - info!(?resp.request_id, "successfully published HTTP response to pubsub"); - } - } - MessageBody::ToClientWebSocketMessage(ws_msg) => { - info!(?ws_msg.web_socket_id, "forwarding WebSocket message to pubsub"); - // Forward WebSocket messages to the topic that pegboard-gateway subscribes to - let ws_topic = TunnelHttpWebSocketSubject::new( - runner_id, - &port_name, - ws_msg.web_socket_id, - ) - .to_string(); - - info!(%ws_topic, ?ws_msg.web_socket_id, "publishing WebSocket message to pubsub"); - - if let Err(e) = ups_clone - .request_with_timeout( - &ws_topic, - &data.to_vec(), - UPS_REQ_TIMEOUT, - ) - .await - { - let err_any: anyhow::Error = e.into(); - if is_tunnel_closed_error(&err_any) { - info!( - "tunnel closed while publishing WebSocket message; closing client websocket" - ); - // Close client websocket with reason - send_tunnel_closed_close_hyper( - &ws_write_ws_to_pubsub, - ) - .await; - break; - } else { - error!(?err_any, ?ws_msg.web_socket_id, "failed to publish WebSocket message to pubsub"); - } - } else { - info!(?ws_msg.web_socket_id, "successfully published WebSocket message to pubsub"); - } - } - MessageBody::ToClientWebSocketOpen(ws_open) => { - info!(?ws_open.web_socket_id, "forwarding WebSocket open to pubsub"); - let ws_topic = TunnelHttpWebSocketSubject::new( - runner_id, - &port_name, - ws_open.web_socket_id, - ) - .to_string(); + // Parse message + let msg = match versioned::RunnerMessage::deserialize_version( + &data, + protocol_version, + ) + .and_then(|x| x.into_latest()) + { + Result::Ok(x) => x, + Err(err) => { + tracing::error!(?err, "failed to deserialize message"); + continue; + } + }; + + // Determine reply to subject + let request_id = msg.request_id; + let reply_to = { + let active_requests = active_requests_clone.lock().await; + if let Some(req) = active_requests.get(&request_id) { + req.reply_to.clone() + } else { + tracing::warn!( + "no active request for tunnel message, may have timed out" + ); + continue; + } + }; - if let Err(e) = ups_clone - .request_with_timeout( - &ws_topic, - &data.to_vec(), - UPS_REQ_TIMEOUT, - ) - .await - { - let err_any: anyhow::Error = e.into(); - if is_tunnel_closed_error(&err_any) { - info!( - "tunnel closed while publishing WebSocket open; closing client websocket" - ); - // Close client websocket with reason - send_tunnel_closed_close_hyper( - &ws_write_ws_to_pubsub, - ) - .await; - break; - } else { - error!(?err_any, ?ws_open.web_socket_id, "failed to publish WebSocket open to pubsub"); - } - } else { - info!(?ws_open.web_socket_id, "successfully published WebSocket open to pubsub"); - } - } - MessageBody::ToClientWebSocketClose(ws_close) => { - info!(?ws_close.web_socket_id, "forwarding WebSocket close to pubsub"); - let ws_topic = TunnelHttpWebSocketSubject::new( - runner_id, - &port_name, - ws_close.web_socket_id, - ) - .to_string(); + // Remove active request entries when terminal + if is_message_kind_request_close(&msg.message_kind) { + let mut active_requests = active_requests_clone.lock().await; + active_requests.remove(&request_id); + } - if let Err(e) = ups_clone - .request_with_timeout( - &ws_topic, - &data.to_vec(), - UPS_REQ_TIMEOUT, - ) - .await - { - let err_any: anyhow::Error = e.into(); - if is_tunnel_closed_error(&err_any) { - info!( - "tunnel closed while publishing WebSocket close; closing client websocket" - ); - // Close client websocket with reason - send_tunnel_closed_close_hyper( - &ws_write_ws_to_pubsub, - ) - .await; - break; - } else { - error!(?err_any, ?ws_close.web_socket_id, "failed to publish WebSocket close to pubsub"); - } - } else { - info!(?ws_close.web_socket_id, "successfully published WebSocket close to pubsub"); - } - } - _ => { - // For other message types, we might not need to forward to pubsub - info!( - "Received non-response message from WebSocket, skipping pubsub forward" - ); - continue; - } + // Publish message to UPS + let message_serialized = + match versioned::PubSubMessage::latest(PubSubMessage { + request_id: msg.request_id, + message_id: msg.message_id, + reply_to: None, + message_kind: msg.message_kind, + }) + .serialize_with_embedded_version(PROTOCOL_VERSION) + { + Result::Ok(x) => x, + Err(err) => { + tracing::error!(?err, "failed to serialize tunnel to gateway"); + continue; } - } - ::std::result::Result::Err(e) => { - error!(?e, "failed to deserialize tunnel message from WebSocket"); + }; + match ups_clone + .publish(&reply_to, &message_serialized, PublishOpts::one()) + .await + { + Result::Ok(_) => {} + Err(err) => { + tracing::error!(?err, "error publishing ups message"); } } } - ::std::result::Result::Ok(WsMessage::Close(_)) => { - info!(?runner_id, "WebSocket closed"); + Result::Ok(WsMessage::Close(_)) => { + tracing::info!(?runner_key_clone, "WebSocket closed"); break; } - ::std::result::Result::Ok(_) => { + Result::Ok(_) => { // Ignore other message types } Err(e) => { - error!(?e, "WebSocket error"); + tracing::error!(?e, "WebSocket error"); break; } } } - info!("WebSocket to pubsub forwarding task ended"); - - // Clean up connection - connections_clone.write().await.remove(&connection_id); + tracing::info!("WebSocket to pubsub forwarding task ended"); }); // Wait for either task to complete tokio::select! { _ = pubsub_to_ws => { - info!("pubsub to WebSocket task completed"); + tracing::info!("pubsub to WebSocket task completed"); } _ = ws_to_pubsub => { - info!("WebSocket to pubsub task completed"); + tracing::info!("WebSocket to pubsub task completed"); } } // Clean up - connections.write().await.remove(&connection_id); - info!(?runner_id, "connection closed"); + tracing::info!(?runner_key, "connection closed"); std::result::Result::<(), (HyperWebsocket, anyhow::Error)>::Ok(()) } } -// Keep the old start function for backward compatibility in tests -pub async fn start(config: rivet_config::Config, pools: Pools) -> Result<()> { - let cache = rivet_cache::CacheInner::from_env(&config, pools.clone())?; - let ctx = StandaloneCtx::new( - gas::db::DatabaseKv::from_pools(pools.clone()).await?, - config.clone(), - pools.clone(), - cache, - "pegboard-tunnel", - Id::new_v1(config.dc_label()), - Id::new_v1(config.dc_label()), - )?; - - main_loop(ctx).await -} - -async fn main_loop(ctx: gas::prelude::StandaloneCtx) -> Result<()> { - let connections: Connections = Arc::new(RwLock::new(HashMap::new())); - - // Start WebSocket server - // Use pegboard config since pegboard_tunnel doesn't exist - let server_addr = SocketAddr::new( - ctx.config().pegboard().host(), - ctx.config().pegboard().port(), - ); - - info!(?server_addr, "starting pegboard-tunnel"); - - let listener = TcpListener::bind(&server_addr).await?; - - // Accept connections - loop { - let (tcp_stream, addr) = listener.accept().await?; - let connections = connections.clone(); - let ctx = ctx.clone(); - - tokio::spawn(async move { - if let Err(e) = handle_connection(ctx, tcp_stream, addr, connections).await { - error!(?e, ?addr, "connection handler error"); - } - }); +fn is_message_kind_request_close(kind: &MessageKind) -> bool { + match kind { + // HTTP terminal states + MessageKind::ToClientResponseStart(resp) => !resp.stream, + MessageKind::ToClientResponseChunk(chunk) => chunk.finish, + MessageKind::ToClientResponseAbort => true, + // WebSocket terminal states (either side closes) + MessageKind::ToClientWebSocketClose(_) => true, + MessageKind::ToServerWebSocketClose(_) => true, + _ => false, } } - -async fn handle_connection( - ctx: gas::prelude::StandaloneCtx, - tcp_stream: tokio::net::TcpStream, - addr: std::net::SocketAddr, - connections: Connections, -) -> Result<()> { - info!(?addr, "new connection"); - - // Parse WebSocket upgrade request - let ws_stream = accept_async(tcp_stream).await?; - - // For now, we'll expect the runner to send an initial message with its ID - // In production, this would be parsed from the URL path or headers - let runner_id = rivet_util::Id::nil(); // Placeholder - should be extracted from connection - let port_name = "default".to_string(); // Placeholder - should be extracted - - let connection_id = rivet_util::Id::nil(); - - // Subscribe to pubsub topic for this runner using raw pubsub client - let topic = TunnelHttpRunnerSubject::new(runner_id, &port_name).to_string(); - info!(%topic, ?runner_id, "subscribing to pubsub topic"); - - // Get UPS (UniversalPubSub) client - let ups = ctx.pools().ups()?; - let mut sub = ups.subscribe(&topic).await?; - - // Split WebSocket stream into read and write halves - let (ws_write, mut ws_read) = ws_stream.split(); - let ws_write = Arc::new(Mutex::new(ws_write)); - - // Store connection - let connection = Arc::new(RunnerConnection { - _runner_id: runner_id, - _port_name: port_name.clone(), - }); - - connections - .write() - .await - .insert(connection_id, connection.clone()); - - // Handle bidirectional message forwarding - let ws_write_clone = ws_write.clone(); - let connections_clone = connections.clone(); - let ups_clone = ups.clone(); - - // Task for forwarding pubsub -> WebSocket - let pubsub_to_ws = tokio::spawn(async move { - while let ::std::result::Result::Ok(NextOutput::Message(msg)) = sub.next().await { - // Ack message - match msg.reply(&[]).await { - Result::Ok(_) => {} - Err(err) => { - tracing::warn!(?err, "failed to ack gateway request response message") - } - }; - - // Forward raw message to WebSocket - let ws_msg = - tokio_tungstenite::tungstenite::Message::Binary(msg.payload.to_vec().into()); - { - let mut stream = ws_write_clone.lock().await; - if let Err(e) = stream.send(ws_msg).await { - error!(?e, "failed to send message to WebSocket"); - break; - } - } - } - }); - - // Task for forwarding WebSocket -> pubsub - let ws_write_ws_to_pubsub = ws_write.clone(); - let ws_to_pubsub = tokio::spawn(async move { - while let Some(msg) = ws_read.next().await { - match msg { - ::std::result::Result::Ok(tokio_tungstenite::tungstenite::Message::Binary( - data, - )) => { - // Parse the tunnel message to extract request_id - match versioned::TunnelMessage::deserialize(&data) { - ::std::result::Result::Ok(tunnel_msg) => { - // Handle different message types - match &tunnel_msg.body { - MessageBody::ToClientResponseStart(resp) => { - let response_topic = TunnelHttpResponseSubject::new( - runner_id, - &port_name, - resp.request_id, - ) - .to_string(); - - if let Err(e) = ups_clone - .request_with_timeout( - &response_topic, - &data.to_vec(), - UPS_REQ_TIMEOUT, - ) - .await - { - let err_any: anyhow::Error = e.into(); - if is_tunnel_closed_error(&err_any) { - info!( - "tunnel closed while publishing HTTP response; closing client websocket" - ); - // Close client websocket with reason - send_tunnel_closed_close_tokio(&ws_write_ws_to_pubsub) - .await; - break; - } else { - error!(?err_any, ?resp.request_id, "failed to publish HTTP response to pubsub"); - } - } - } - MessageBody::ToClientWebSocketMessage(ws_msg) => { - let ws_topic = TunnelHttpWebSocketSubject::new( - runner_id, - &port_name, - ws_msg.web_socket_id, - ) - .to_string(); - - if let Err(e) = ups_clone - .request_with_timeout( - &ws_topic, - &data.to_vec(), - UPS_REQ_TIMEOUT, - ) - .await - { - let err_any: anyhow::Error = e.into(); - if is_tunnel_closed_error(&err_any) { - info!( - "tunnel closed while publishing WebSocket message; closing client websocket" - ); - // Close client websocket with reason - send_tunnel_closed_close_tokio(&ws_write_ws_to_pubsub) - .await; - break; - } else { - error!(?err_any, ?ws_msg.web_socket_id, "failed to publish WebSocket message to pubsub"); - } - } - } - MessageBody::ToClientWebSocketOpen(ws_open) => { - let ws_topic = TunnelHttpWebSocketSubject::new( - runner_id, - &port_name, - ws_open.web_socket_id, - ) - .to_string(); - - if let Err(e) = ups_clone - .request_with_timeout( - &ws_topic, - &data.to_vec(), - UPS_REQ_TIMEOUT, - ) - .await - { - let err_any: anyhow::Error = e.into(); - if is_tunnel_closed_error(&err_any) { - info!( - "tunnel closed while publishing WebSocket open; closing client websocket" - ); - // Close client websocket with reason - send_tunnel_closed_close_tokio(&ws_write_ws_to_pubsub) - .await; - break; - } else { - error!(?err_any, ?ws_open.web_socket_id, "failed to publish WebSocket open to pubsub"); - } - } - } - MessageBody::ToClientWebSocketClose(ws_close) => { - let ws_topic = TunnelHttpWebSocketSubject::new( - runner_id, - &port_name, - ws_close.web_socket_id, - ) - .to_string(); - - if let Err(e) = ups_clone - .request_with_timeout( - &ws_topic, - &data.to_vec(), - UPS_REQ_TIMEOUT, - ) - .await - { - let err_any: anyhow::Error = e.into(); - if is_tunnel_closed_error(&err_any) { - info!( - "tunnel closed while publishing WebSocket close; closing client websocket" - ); - // Close client websocket with reason - send_tunnel_closed_close_tokio(&ws_write_ws_to_pubsub) - .await; - break; - } else { - error!(?err_any, ?ws_close.web_socket_id, "failed to publish WebSocket close to pubsub"); - } - } - } - _ => { - // For other message types, we might not need to forward to pubsub - info!( - "Received non-response message from WebSocket, skipping pubsub forward" - ); - continue; - } - } - } - ::std::result::Result::Err(e) => { - error!(?e, "failed to deserialize tunnel message from WebSocket"); - } - } - } - ::std::result::Result::Ok(tokio_tungstenite::tungstenite::Message::Close(_)) => { - info!(?runner_id, "WebSocket closed"); - break; - } - ::std::result::Result::Ok(_) => { - // Ignore other message types - } - Err(e) => { - error!(?e, "WebSocket error"); - break; - } - } - } - - // Clean up connection - connections_clone.write().await.remove(&connection_id); - }); - - // Wait for either task to complete - tokio::select! { - _ = pubsub_to_ws => { - info!("pubsub to WebSocket task completed"); - } - _ = ws_to_pubsub => { - info!("WebSocket to pubsub task completed"); - } - } - - // Clean up - connections.write().await.remove(&connection_id); - info!(?runner_id, "connection closed"); - - Ok(()) -} - -/// Determines if the tunnel is closed by if the UPS service is no longer responding. -fn is_tunnel_closed_error(err: &anyhow::Error) -> bool { - if let Some(err) = err - .chain() - .find_map(|x| x.downcast_ref::()) - && err.group() == "ups" - && err.code() == "request_timeout" - { - true - } else { - false - } -} - -// Helper: Build and send a standard tunnel-closed Close frame (hyper-tungstenite) -fn tunnel_closed_close_msg_hyper() -> WsMessage { - WsMessage::Close(Some(WsCloseFrame { - code: WsCloseCode::Error, - reason: WsUtf8Bytes::from_static("Tunnel closed"), - })) -} - -// Helper: Build and send a standard tunnel-closed Close frame (tokio-tungstenite) -fn tunnel_closed_close_msg_tokio() -> tokio_tungstenite::tungstenite::Message { - tokio_tungstenite::tungstenite::Message::Close(Some( - tokio_tungstenite::tungstenite::protocol::frame::CloseFrame { - code: tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::Error, - reason: tokio_tungstenite::tungstenite::Utf8Bytes::from_static("Tunnel closed"), - }, - )) -} - -// Helper: Send the tunnel-closed Close frame on a hyper-tungstenite sink -async fn send_tunnel_closed_close_hyper(ws_write: &tokio::sync::Mutex) -where - S: futures::Sink + Unpin, -{ - let mut stream = ws_write.lock().await; - let _ = stream.send(tunnel_closed_close_msg_hyper()).await; -} - -// Helper: Send the tunnel-closed Close frame on a tokio-tungstenite sink -async fn send_tunnel_closed_close_tokio(ws_write: &tokio::sync::Mutex) -where - S: futures::Sink + Unpin, -{ - let mut stream = ws_write.lock().await; - let _ = stream.send(tunnel_closed_close_msg_tokio()).await; -} diff --git a/packages/core/pegboard-tunnel/tests/integration.rs b/packages/core/pegboard-tunnel/tests/integration.rs index 2c20bfe638..70051af60a 100644 --- a/packages/core/pegboard-tunnel/tests/integration.rs +++ b/packages/core/pegboard-tunnel/tests/integration.rs @@ -91,7 +91,9 @@ async fn test_pubsub_to_websocket( }; // Serialize the message - let serialized = versioned::TunnelMessage::serialize(versioned::TunnelMessage::V1(message))?; + let serialized = versioned::RunnerMessage::serialize_with_embedded_version( + versioned::RunnerMessage::V1(message), + )?; // Publish to pubsub topic using proper subject let topic = TunnelHttpRunnerSubject::new(&runner_id.to_string(), port_name).to_string(); @@ -105,7 +107,7 @@ async fn test_pubsub_to_websocket( match received? { WsMessage::Binary(data) => { // Deserialize and verify the message - let tunnel_msg = versioned::TunnelMessage::deserialize(&data)?; + let tunnel_msg = versioned::RunnerMessage::deserialize_with_embedded_version(&data)?; match tunnel_msg.body { MessageBody::ToServerRequestStart(req) => { assert_eq!(req.request_id, request_id); @@ -150,7 +152,9 @@ async fn test_websocket_to_pubsub( }; // Serialize and send via WebSocket - let serialized = versioned::TunnelMessage::serialize(versioned::TunnelMessage::V1(message))?; + let serialized = versioned::RunnerMessage::serialize_with_embedded_version( + versioned::RunnerMessage::V1(message), + )?; ws_stream.send(WsMessage::Binary(serialized.into())).await?; // Wait for message on pubsub @@ -159,7 +163,8 @@ async fn test_websocket_to_pubsub( match received { universalpubsub::pubsub::NextOutput::Message(msg) => { // Deserialize and verify the message - let tunnel_msg = versioned::TunnelMessage::deserialize(&msg.payload)?; + let tunnel_msg = + versioned::RunnerMessage::deserialize_with_embedded_version(&msg.payload)?; match tunnel_msg.body { MessageBody::ToClientResponseStart(resp) => { assert_eq!(resp.request_id, request_id); diff --git a/packages/infra/engine/tests/actors_lifecycle.rs b/packages/infra/engine/tests/actors_lifecycle.rs index 85b991eb4a..46aeff7e9d 100644 --- a/packages/infra/engine/tests/actors_lifecycle.rs +++ b/packages/infra/engine/tests/actors_lifecycle.rs @@ -4,7 +4,7 @@ use std::time::Duration; #[test] fn actor_lifecycle_single_dc() { - common::run(common::TestOpts::new(2), |ctx| async move { + common::run(common::TestOpts::new(1), |ctx| async move { actor_lifecycle_inner(&ctx, false).await; }); } @@ -29,10 +29,6 @@ async fn actor_lifecycle_inner(ctx: &common::TestCtx, multi_dc: bool) { let actor_id = common::create_actor(&namespace, target_dc.guard_port()).await; - // TODO: This is a race condition. we might need to move this after the guard ping since guard - // correctly waits for the actor to start. - tokio::time::sleep(Duration::from_millis(500)).await; - // Test ping via guard let ping_response = common::ping_actor_via_guard(ctx.leader_dc().guard_port(), &actor_id, "main").await; @@ -52,11 +48,10 @@ async fn actor_lifecycle_inner(ctx: &common::TestCtx, multi_dc: bool) { // Destroy tracing::info!("destroying actor"); - tokio::time::sleep(Duration::from_millis(500)).await; common::destroy_actor(&actor_id, &namespace, target_dc.guard_port()).await; - tokio::time::sleep(Duration::from_millis(500)).await; // Validate runner state + tokio::time::sleep(Duration::from_millis(500)).await; assert!( !runner.has_actor(&actor_id).await, "Runner should not have the actor after destroy" diff --git a/packages/services/pegboard/src/ops/runner/get_by_key.rs b/packages/services/pegboard/src/ops/runner/get_by_key.rs new file mode 100644 index 0000000000..5a6ea63bf1 --- /dev/null +++ b/packages/services/pegboard/src/ops/runner/get_by_key.rs @@ -0,0 +1,51 @@ +use anyhow::*; +use gas::prelude::*; +use rivet_types::runners::Runner; +use udb_util::{SERIALIZABLE, TxnExt}; +use universaldb as udb; + +use crate::keys; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Input { + pub namespace_id: Id, + pub name: String, + pub key: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Output { + pub runner: Option, +} + +#[operation] +pub async fn pegboard_runner_get_by_key(ctx: &OperationCtx, input: &Input) -> Result { + let dc_name = ctx.config().dc_name()?.to_string(); + + let runner = ctx + .udb()? + .run(|tx, _mc| { + let dc_name = dc_name.to_string(); + let input = input.clone(); + async move { + let txs = tx.subspace(keys::subspace()); + + // Look up runner by key + let runner_by_key_key = + keys::ns::RunnerByKeyKey::new(input.namespace_id, input.name, input.key); + + let runner_data = txs.read_opt(&runner_by_key_key, SERIALIZABLE).await?; + + if let Some(data) = runner_data { + // Get full runner details using the runner_id + let runner = super::get::get_inner(&dc_name, &tx, data.runner_id).await?; + std::result::Result::<_, udb::FdbBindingError>::Ok(runner) + } else { + std::result::Result::<_, udb::FdbBindingError>::Ok(None) + } + } + }) + .await?; + + Ok(Output { runner }) +} diff --git a/packages/services/pegboard/src/ops/runner/mod.rs b/packages/services/pegboard/src/ops/runner/mod.rs index cd56c69267..1885c60f23 100644 --- a/packages/services/pegboard/src/ops/runner/mod.rs +++ b/packages/services/pegboard/src/ops/runner/mod.rs @@ -1,4 +1,5 @@ pub mod get; +pub mod get_by_key; pub mod list_for_ns; pub mod list_names; pub mod update_alloc_idx; diff --git a/packages/services/pegboard/src/pubsub_subjects.rs b/packages/services/pegboard/src/pubsub_subjects.rs index 13d39b1704..ae9e9438b6 100644 --- a/packages/services/pegboard/src/pubsub_subjects.rs +++ b/packages/services/pegboard/src/pubsub_subjects.rs @@ -1,77 +1,41 @@ -use rivet_util::Id; +use gas::prelude::*; -pub struct TunnelHttpRunnerSubject<'a> { - runner_id: Id, +pub struct TunnelRunnerReceiverSubject<'a> { + runner_key: &'a str, port_name: &'a str, } -impl<'a> TunnelHttpRunnerSubject<'a> { - pub fn new(runner_id: Id, port_name: &'a str) -> Self { +impl<'a> TunnelRunnerReceiverSubject<'a> { + pub fn new(runner_key: &'a str, port_name: &'a str) -> Self { Self { - runner_id, + runner_key, port_name, } } } -impl std::fmt::Display for TunnelHttpRunnerSubject<'_> { +impl std::fmt::Display for TunnelRunnerReceiverSubject<'_> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, - "pegboard.tunnel.http.runner.{}.{}", - self.runner_id, self.port_name + "pegboard.tunnel.runner_receiver.{}.{}", + self.runner_key, self.port_name ) } } -pub struct TunnelHttpResponseSubject<'a> { - runner_id: Id, - port_name: &'a str, - request_id: u64, +pub struct TunnelGatewayReceiverSubject { + gateway_id: Uuid, } -impl<'a> TunnelHttpResponseSubject<'a> { - pub fn new(runner_id: Id, port_name: &'a str, request_id: u64) -> Self { - Self { - runner_id, - port_name, - request_id, - } +impl<'a> TunnelGatewayReceiverSubject { + pub fn new(gateway_id: Uuid) -> Self { + Self { gateway_id } } } -impl std::fmt::Display for TunnelHttpResponseSubject<'_> { +impl std::fmt::Display for TunnelGatewayReceiverSubject { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "pegboard.tunnel.http.request.{}.{}.{}", - self.runner_id, self.port_name, self.request_id - ) - } -} - -pub struct TunnelHttpWebSocketSubject<'a> { - runner_id: Id, - port_name: &'a str, - websocket_id: u64, -} - -impl<'a> TunnelHttpWebSocketSubject<'a> { - pub fn new(runner_id: Id, port_name: &'a str, websocket_id: u64) -> Self { - Self { - runner_id, - port_name, - websocket_id, - } - } -} - -impl std::fmt::Display for TunnelHttpWebSocketSubject<'_> { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "pegboard.tunnel.http.websocket.{}.{}.{}", - self.runner_id, self.port_name, self.websocket_id - ) + write!(f, "pegboard.gateway.receiver.{}", self.gateway_id) } } diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 62a80f8ad4..91185e953a 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -685,6 +685,9 @@ importers: '@rivetkit/engine-tunnel-protocol': specifier: workspace:* version: link:../tunnel-protocol + uuid: + specifier: ^12.0.0 + version: 12.0.0 ws: specifier: ^8.18.3 version: 8.18.3 @@ -7241,6 +7244,10 @@ packages: util-deprecate@1.0.2: resolution: {integrity: sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==} + uuid@12.0.0: + resolution: {integrity: sha512-USe1zesMYh4fjCA8ZH5+X5WIVD0J4V1Jksm1bFTVBX2F/cwSXt0RO5w/3UXbdLKmZX65MiWV+hwhSS8p6oBTGA==} + hasBin: true + validator@13.15.15: resolution: {integrity: sha512-BgWVbCI72aIQy937xbawcs+hrVaN/CZ2UwutgaJ36hGqRrLNM+f5LUT/YPRbo8IV/ASeFzXszezV+y2+rq3l8A==} engines: {node: '>= 0.10'} @@ -14748,6 +14755,8 @@ snapshots: util-deprecate@1.0.2: {} + uuid@12.0.0: {} + validator@13.15.15: {} vary@1.1.2: {} diff --git a/sdks/rust/runner-protocol/src/protocol.rs b/sdks/rust/runner-protocol/src/protocol.rs index 4ca049f634..a037941e2a 100644 --- a/sdks/rust/runner-protocol/src/protocol.rs +++ b/sdks/rust/runner-protocol/src/protocol.rs @@ -20,9 +20,7 @@ pub enum ToClient { #[serde(rename_all = "snake_case")] pub enum ToServer { Init { - runner_id: Option, name: String, - key: String, version: u32, total_slots: u32, diff --git a/sdks/rust/runner-protocol/src/versioned.rs b/sdks/rust/runner-protocol/src/versioned.rs index 2b61c2c2dc..34c54168b6 100644 --- a/sdks/rust/runner-protocol/src/versioned.rs +++ b/sdks/rust/runner-protocol/src/versioned.rs @@ -285,9 +285,7 @@ impl TryFrom for protocol::ToServer { fn try_from(value: v1::ToServer) -> Result { match value { v1::ToServer::ToServerInit(init) => Ok(protocol::ToServer::Init { - runner_id: init.runner_id.map(|id| util::Id::parse(&id)).transpose()?, name: init.name, - key: init.key, version: init.version, total_slots: init.total_slots, addresses_http: init diff --git a/sdks/rust/tunnel-protocol/build.rs b/sdks/rust/tunnel-protocol/build.rs index 453be2c763..f43ed92983 100644 --- a/sdks/rust/tunnel-protocol/build.rs +++ b/sdks/rust/tunnel-protocol/build.rs @@ -1,4 +1,8 @@ -use std::{env, fs, path::Path}; +use std::{ + env, fs, + path::{Path, PathBuf}, + process::Command, +}; use indoc::formatdoc; @@ -52,6 +56,78 @@ mod rust { } } +mod typescript { + use super::*; + + pub fn generate_sdk(schema_dir: &Path) { + let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap(); + let workspace_root = Path::new(&manifest_dir) + .parent() + .and_then(|p| p.parent()) + .and_then(|p| p.parent()) + .expect("Failed to find workspace root"); + + let sdk_dir = workspace_root + .join("sdks") + .join("typescript") + .join("tunnel-protocol"); + let src_dir = sdk_dir.join("src"); + + let highest_version_path = super::find_highest_version(schema_dir); + + let _ = fs::remove_dir_all(&src_dir); + if let Err(e) = fs::create_dir_all(&src_dir) { + panic!("Failed to create SDK directory: {}", e); + } + + let output = + Command::new(workspace_root.join("node_modules/@bare-ts/tools/dist/bin/cli.js")) + .arg("compile") + .arg("--generator") + .arg("ts") + .arg(highest_version_path) + .arg("-o") + .arg(src_dir.join("index.ts")) + .output() + .expect("Failed to execute bare compiler for TypeScript"); + + if !output.status.success() { + panic!( + "BARE TypeScript generation failed: {}", + String::from_utf8_lossy(&output.stderr), + ); + } + } +} + +fn find_highest_version(schema_dir: &Path) -> 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_path +} + fn main() { let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap(); let workspace_root = Path::new(&manifest_dir) @@ -68,4 +144,15 @@ fn main() { println!("cargo:rerun-if-changed={}", schema_dir.display()); rust::generate_sdk(&schema_dir); + + // Check if cli.js exists before attempting TypeScript generation + let cli_js_path = workspace_root.join("node_modules/@bare-ts/tools/dist/bin/cli.js"); + if cli_js_path.exists() { + typescript::generate_sdk(&schema_dir); + } else { + println!( + "cargo:warning=TypeScript SDK generation skipped: cli.js not found at {}. Run `pnpm install` to install.", + cli_js_path.display() + ); + } } diff --git a/sdks/rust/tunnel-protocol/src/versioned.rs b/sdks/rust/tunnel-protocol/src/versioned.rs index d37e5d5447..f5208f27fb 100644 --- a/sdks/rust/tunnel-protocol/src/versioned.rs +++ b/sdks/rust/tunnel-protocol/src/versioned.rs @@ -3,20 +3,20 @@ use versioned_data_util::OwnedVersionedData; use crate::{PROTOCOL_VERSION, generated::v1}; -pub enum TunnelMessage { - V1(v1::TunnelMessage), +pub enum RunnerMessage { + V1(v1::RunnerMessage), } -impl OwnedVersionedData for TunnelMessage { - type Latest = v1::TunnelMessage; +impl OwnedVersionedData for RunnerMessage { + type Latest = v1::RunnerMessage; - fn latest(latest: v1::TunnelMessage) -> Self { - TunnelMessage::V1(latest) + fn latest(latest: v1::RunnerMessage) -> Self { + RunnerMessage::V1(latest) } fn into_latest(self) -> Result { #[allow(irrefutable_let_patterns)] - if let TunnelMessage::V1(data) = self { + if let RunnerMessage::V1(data) = self { Ok(data) } else { bail!("version not latest"); @@ -25,20 +25,64 @@ impl OwnedVersionedData for TunnelMessage { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(TunnelMessage::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(RunnerMessage::V1(serde_bare::from_slice(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - TunnelMessage::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + RunnerMessage::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), } } } -impl TunnelMessage { - pub fn deserialize(buf: &[u8]) -> Result { +impl RunnerMessage { + pub fn deserialize(buf: &[u8]) -> Result { + ::deserialize(buf, PROTOCOL_VERSION) + } + + pub fn serialize(self) -> Result> { + ::serialize(self, PROTOCOL_VERSION) + } +} + +pub enum PubSubMessage { + V1(v1::PubSubMessage), +} + +impl OwnedVersionedData for PubSubMessage { + type Latest = v1::PubSubMessage; + + fn latest(latest: v1::PubSubMessage) -> Self { + PubSubMessage::V1(latest) + } + + fn into_latest(self) -> Result { + #[allow(irrefutable_let_patterns)] + if let PubSubMessage::V1(data) = self { + Ok(data) + } else { + bail!("version not latest"); + } + } + + fn deserialize_version(payload: &[u8], version: u16) -> Result { + match version { + 1 => Ok(PubSubMessage::V1(serde_bare::from_slice(payload)?)), + _ => bail!("invalid version: {version}"), + } + } + + fn serialize_version(self, _version: u16) -> Result> { + match self { + PubSubMessage::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + } + } +} + +impl PubSubMessage { + pub fn deserialize(buf: &[u8]) -> Result { ::deserialize(buf, PROTOCOL_VERSION) } diff --git a/sdks/schemas/runner-protocol/v1.bare b/sdks/schemas/runner-protocol/v1.bare index ed99e59afb..5db5afcea4 100644 --- a/sdks/schemas/runner-protocol/v1.bare +++ b/sdks/schemas/runner-protocol/v1.bare @@ -133,9 +133,7 @@ type CommandWrapper struct { } type ToServerInit struct { - runnerId: optional name: str - key: str version: u32 totalSlots: u32 addressesHttp: optional> diff --git a/sdks/schemas/tunnel-protocol/v1.bare b/sdks/schemas/tunnel-protocol/v1.bare index 5405587c1b..f9e0e9c63f 100644 --- a/sdks/schemas/tunnel-protocol/v1.bare +++ b/sdks/schemas/tunnel-protocol/v1.bare @@ -1,15 +1,13 @@ -type RequestId u64 -type WebSocketId u64 +type RequestId data[16] # UUIDv4 +type MessageId data[16] # UUIDv4 type Id str -type StreamFinishReason enum { - COMPLETE - ABORT -} -# MARK: HTTP Request Forwarding +# MARK: Ack +type Ack void + +# MARK: HTTP type ToServerRequestStart struct { - requestId: RequestId actorId: Id method: str path: str @@ -19,17 +17,13 @@ type ToServerRequestStart struct { } type ToServerRequestChunk struct { - requestId: RequestId body: data + finish: bool } -type ToServerRequestFinish struct { - requestId: RequestId - reason: StreamFinishReason -} +type ToServerRequestAbort void type ToClientResponseStart struct { - requestId: RequestId status: u16 headers: map body: optional @@ -37,60 +31,52 @@ type ToClientResponseStart struct { } type ToClientResponseChunk struct { - requestId: RequestId body: data + finish: bool } -type ToClientResponseFinish struct { - requestId: RequestId - reason: StreamFinishReason -} +type ToClientResponseAbort void -# MARK: WebSocket Forwarding +# MARK: WebSocket type ToServerWebSocketOpen struct { actorId: Id - webSocketId: WebSocketId path: str headers: map } type ToServerWebSocketMessage struct { - webSocketId: WebSocketId data: data binary: bool } type ToServerWebSocketClose struct { - webSocketId: WebSocketId code: optional reason: optional } -type ToClientWebSocketOpen struct { - webSocketId: WebSocketId -} +type ToClientWebSocketOpen void type ToClientWebSocketMessage struct { - webSocketId: WebSocketId data: data binary: bool } type ToClientWebSocketClose struct { - webSocketId: WebSocketId code: optional reason: optional } # MARK: Message -type MessageBody union { +type MessageKind union { + Ack | + # HTTP ToServerRequestStart | ToServerRequestChunk | - ToServerRequestFinish | + ToServerRequestAbort | ToClientResponseStart | ToClientResponseChunk | - ToClientResponseFinish | + ToClientResponseAbort | # WebSocket ToServerWebSocketOpen | @@ -101,7 +87,19 @@ type MessageBody union { ToClientWebSocketClose } -# Main tunnel message -type TunnelMessage struct { - body: MessageBody +# MARK: Message sent over tunnel WebSocket +type RunnerMessage struct { + requestId: RequestId + messageId: MessageId + messageKind: MessageKind +} + +# MARK: Message sent over UPS +type PubSubMessage struct { + requestId: RequestId + messageId: MessageId + # Subject to send replies to. Only sent when opening a new request from gateway -> runner. + replyTo: optional + messageKind: MessageKind } + diff --git a/sdks/typescript/runner-protocol/src/index.ts b/sdks/typescript/runner-protocol/src/index.ts index 45a947c051..96c00212c6 100644 --- a/sdks/typescript/runner-protocol/src/index.ts +++ b/sdks/typescript/runner-protocol/src/index.ts @@ -590,18 +590,7 @@ export function writeCommandWrapper(bc: bare.ByteCursor, x: CommandWrapper): voi writeCommand(bc, x.inner) } -function read3(bc: bare.ByteCursor): Id | null { - return bare.readBool(bc) ? readId(bc) : null -} - -function write3(bc: bare.ByteCursor, x: Id | null): void { - bare.writeBool(bc, x != null) - if (x != null) { - writeId(bc, x) - } -} - -function read4(bc: bare.ByteCursor): ReadonlyMap { +function read3(bc: bare.ByteCursor): ReadonlyMap { const len = bare.readUintSafe(bc) const result = new Map() for (let i = 0; i < len; i++) { @@ -616,7 +605,7 @@ function read4(bc: bare.ByteCursor): ReadonlyMap { return result } -function write4(bc: bare.ByteCursor, x: ReadonlyMap): void { +function write3(bc: bare.ByteCursor, x: ReadonlyMap): void { bare.writeUintSafe(bc, x.size) for (const kv of x) { bare.writeString(bc, kv[0]) @@ -624,18 +613,18 @@ function write4(bc: bare.ByteCursor, x: ReadonlyMap): } } -function read5(bc: bare.ByteCursor): ReadonlyMap | null { - return bare.readBool(bc) ? read4(bc) : null +function read4(bc: bare.ByteCursor): ReadonlyMap | null { + return bare.readBool(bc) ? read3(bc) : null } -function write5(bc: bare.ByteCursor, x: ReadonlyMap | null): void { +function write4(bc: bare.ByteCursor, x: ReadonlyMap | null): void { bare.writeBool(bc, x != null) if (x != null) { - write4(bc, x) + write3(bc, x) } } -function read6(bc: bare.ByteCursor): ReadonlyMap { +function read5(bc: bare.ByteCursor): ReadonlyMap { const len = bare.readUintSafe(bc) const result = new Map() for (let i = 0; i < len; i++) { @@ -650,7 +639,7 @@ function read6(bc: bare.ByteCursor): ReadonlyMap { return result } -function write6(bc: bare.ByteCursor, x: ReadonlyMap): void { +function write5(bc: bare.ByteCursor, x: ReadonlyMap): void { bare.writeUintSafe(bc, x.size) for (const kv of x) { bare.writeString(bc, kv[0]) @@ -658,18 +647,18 @@ function write6(bc: bare.ByteCursor, x: ReadonlyMap): } } -function read7(bc: bare.ByteCursor): ReadonlyMap | null { - return bare.readBool(bc) ? read6(bc) : null +function read6(bc: bare.ByteCursor): ReadonlyMap | null { + return bare.readBool(bc) ? read5(bc) : null } -function write7(bc: bare.ByteCursor, x: ReadonlyMap | null): void { +function write6(bc: bare.ByteCursor, x: ReadonlyMap | null): void { bare.writeBool(bc, x != null) if (x != null) { - write6(bc, x) + write5(bc, x) } } -function read8(bc: bare.ByteCursor): ReadonlyMap { +function read7(bc: bare.ByteCursor): ReadonlyMap { const len = bare.readUintSafe(bc) const result = new Map() for (let i = 0; i < len; i++) { @@ -684,7 +673,7 @@ function read8(bc: bare.ByteCursor): ReadonlyMap { return result } -function write8(bc: bare.ByteCursor, x: ReadonlyMap): void { +function write7(bc: bare.ByteCursor, x: ReadonlyMap): void { bare.writeUintSafe(bc, x.size) for (const kv of x) { bare.writeString(bc, kv[0]) @@ -692,18 +681,18 @@ function write8(bc: bare.ByteCursor, x: ReadonlyMap): } } -function read9(bc: bare.ByteCursor): ReadonlyMap | null { - return bare.readBool(bc) ? read8(bc) : null +function read8(bc: bare.ByteCursor): ReadonlyMap | null { + return bare.readBool(bc) ? read7(bc) : null } -function write9(bc: bare.ByteCursor, x: ReadonlyMap | null): void { +function write8(bc: bare.ByteCursor, x: ReadonlyMap | null): void { bare.writeBool(bc, x != null) if (x != null) { - write8(bc, x) + write7(bc, x) } } -function read10(bc: bare.ByteCursor): ReadonlyMap { +function read9(bc: bare.ByteCursor): ReadonlyMap { const len = bare.readUintSafe(bc) const result = new Map() for (let i = 0; i < len; i++) { @@ -718,7 +707,7 @@ function read10(bc: bare.ByteCursor): ReadonlyMap { return result } -function write10(bc: bare.ByteCursor, x: ReadonlyMap): void { +function write9(bc: bare.ByteCursor, x: ReadonlyMap): void { bare.writeUintSafe(bc, x.size) for (const kv of x) { bare.writeString(bc, kv[0]) @@ -726,22 +715,22 @@ function write10(bc: bare.ByteCursor, x: ReadonlyMap): void { } } -function read11(bc: bare.ByteCursor): ReadonlyMap | null { - return bare.readBool(bc) ? read10(bc) : null +function read10(bc: bare.ByteCursor): ReadonlyMap | null { + return bare.readBool(bc) ? read9(bc) : null } -function write11(bc: bare.ByteCursor, x: ReadonlyMap | null): void { +function write10(bc: bare.ByteCursor, x: ReadonlyMap | null): void { bare.writeBool(bc, x != null) if (x != null) { - write10(bc, x) + write9(bc, x) } } -function read12(bc: bare.ByteCursor): Json | null { +function read11(bc: bare.ByteCursor): Json | null { return bare.readBool(bc) ? readJson(bc) : null } -function write12(bc: bare.ByteCursor, x: Json | null): void { +function write11(bc: bare.ByteCursor, x: Json | null): void { bare.writeBool(bc, x != null) if (x != null) { writeJson(bc, x) @@ -749,9 +738,7 @@ function write12(bc: bare.ByteCursor, x: Json | null): void { } export type ToServerInit = { - readonly runnerId: Id | null readonly name: string - readonly key: string readonly version: u32 readonly totalSlots: u32 readonly addressesHttp: ReadonlyMap | null @@ -764,32 +751,28 @@ export type ToServerInit = { export function readToServerInit(bc: bare.ByteCursor): ToServerInit { return { - runnerId: read3(bc), name: bare.readString(bc), - key: bare.readString(bc), version: bare.readU32(bc), totalSlots: bare.readU32(bc), - addressesHttp: read5(bc), - addressesTcp: read7(bc), - addressesUdp: read9(bc), + addressesHttp: read4(bc), + addressesTcp: read6(bc), + addressesUdp: read8(bc), lastCommandIdx: read1(bc), - prepopulateActorNames: read11(bc), - metadata: read12(bc), + prepopulateActorNames: read10(bc), + metadata: read11(bc), } } export function writeToServerInit(bc: bare.ByteCursor, x: ToServerInit): void { - write3(bc, x.runnerId) bare.writeString(bc, x.name) - bare.writeString(bc, x.key) bare.writeU32(bc, x.version) bare.writeU32(bc, x.totalSlots) - write5(bc, x.addressesHttp) - write7(bc, x.addressesTcp) - write9(bc, x.addressesUdp) + write4(bc, x.addressesHttp) + write6(bc, x.addressesTcp) + write8(bc, x.addressesUdp) write1(bc, x.lastCommandIdx) - write11(bc, x.prepopulateActorNames) - write12(bc, x.metadata) + write10(bc, x.prepopulateActorNames) + write11(bc, x.metadata) } export type ToServerEvents = readonly EventWrapper[] @@ -843,7 +826,7 @@ export function writeToServerPing(bc: bare.ByteCursor, x: ToServerPing): void { bare.writeI64(bc, x.ts) } -function read13(bc: bare.ByteCursor): readonly KvKey[] { +function read12(bc: bare.ByteCursor): readonly KvKey[] { const len = bare.readUintSafe(bc) if (len === 0) { return [] @@ -855,7 +838,7 @@ function read13(bc: bare.ByteCursor): readonly KvKey[] { return result } -function write13(bc: bare.ByteCursor, x: readonly KvKey[]): void { +function write12(bc: bare.ByteCursor, x: readonly KvKey[]): void { bare.writeUintSafe(bc, x.length) for (let i = 0; i < x.length; i++) { writeKvKey(bc, x[i]) @@ -868,30 +851,30 @@ export type KvGetRequest = { export function readKvGetRequest(bc: bare.ByteCursor): KvGetRequest { return { - keys: read13(bc), + keys: read12(bc), } } export function writeKvGetRequest(bc: bare.ByteCursor, x: KvGetRequest): void { - write13(bc, x.keys) + write12(bc, x.keys) } -function read14(bc: bare.ByteCursor): boolean | null { +function read13(bc: bare.ByteCursor): boolean | null { return bare.readBool(bc) ? bare.readBool(bc) : null } -function write14(bc: bare.ByteCursor, x: boolean | null): void { +function write13(bc: bare.ByteCursor, x: boolean | null): void { bare.writeBool(bc, x != null) if (x != null) { bare.writeBool(bc, x) } } -function read15(bc: bare.ByteCursor): u64 | null { +function read14(bc: bare.ByteCursor): u64 | null { return bare.readBool(bc) ? bare.readU64(bc) : null } -function write15(bc: bare.ByteCursor, x: u64 | null): void { +function write14(bc: bare.ByteCursor, x: u64 | null): void { bare.writeBool(bc, x != null) if (x != null) { bare.writeU64(bc, x) @@ -907,18 +890,18 @@ export type KvListRequest = { export function readKvListRequest(bc: bare.ByteCursor): KvListRequest { return { query: readKvListQuery(bc), - reverse: read14(bc), - limit: read15(bc), + reverse: read13(bc), + limit: read14(bc), } } export function writeKvListRequest(bc: bare.ByteCursor, x: KvListRequest): void { writeKvListQuery(bc, x.query) - write14(bc, x.reverse) - write15(bc, x.limit) + write13(bc, x.reverse) + write14(bc, x.limit) } -function read16(bc: bare.ByteCursor): readonly KvValue[] { +function read15(bc: bare.ByteCursor): readonly KvValue[] { const len = bare.readUintSafe(bc) if (len === 0) { return [] @@ -930,7 +913,7 @@ function read16(bc: bare.ByteCursor): readonly KvValue[] { return result } -function write16(bc: bare.ByteCursor, x: readonly KvValue[]): void { +function write15(bc: bare.ByteCursor, x: readonly KvValue[]): void { bare.writeUintSafe(bc, x.length) for (let i = 0; i < x.length; i++) { writeKvValue(bc, x[i]) @@ -944,14 +927,14 @@ export type KvPutRequest = { export function readKvPutRequest(bc: bare.ByteCursor): KvPutRequest { return { - keys: read13(bc), - values: read16(bc), + keys: read12(bc), + values: read15(bc), } } export function writeKvPutRequest(bc: bare.ByteCursor, x: KvPutRequest): void { - write13(bc, x.keys) - write16(bc, x.values) + write12(bc, x.keys) + write15(bc, x.values) } export type KvDeleteRequest = { @@ -960,12 +943,12 @@ export type KvDeleteRequest = { export function readKvDeleteRequest(bc: bare.ByteCursor): KvDeleteRequest { return { - keys: read13(bc), + keys: read12(bc), } } export function writeKvDeleteRequest(bc: bare.ByteCursor, x: KvDeleteRequest): void { - write13(bc, x.keys) + write12(bc, x.keys) } export type KvDropRequest = null @@ -1214,7 +1197,7 @@ export function writeKvErrorResponse(bc: bare.ByteCursor, x: KvErrorResponse): v bare.writeString(bc, x.message) } -function read17(bc: bare.ByteCursor): readonly KvMetadata[] { +function read16(bc: bare.ByteCursor): readonly KvMetadata[] { const len = bare.readUintSafe(bc) if (len === 0) { return [] @@ -1226,7 +1209,7 @@ function read17(bc: bare.ByteCursor): readonly KvMetadata[] { return result } -function write17(bc: bare.ByteCursor, x: readonly KvMetadata[]): void { +function write16(bc: bare.ByteCursor, x: readonly KvMetadata[]): void { bare.writeUintSafe(bc, x.length) for (let i = 0; i < x.length; i++) { writeKvMetadata(bc, x[i]) @@ -1241,16 +1224,16 @@ export type KvGetResponse = { export function readKvGetResponse(bc: bare.ByteCursor): KvGetResponse { return { - keys: read13(bc), - values: read16(bc), - metadata: read17(bc), + keys: read12(bc), + values: read15(bc), + metadata: read16(bc), } } export function writeKvGetResponse(bc: bare.ByteCursor, x: KvGetResponse): void { - write13(bc, x.keys) - write16(bc, x.values) - write17(bc, x.metadata) + write12(bc, x.keys) + write15(bc, x.values) + write16(bc, x.metadata) } export type KvListResponse = { @@ -1261,16 +1244,16 @@ export type KvListResponse = { export function readKvListResponse(bc: bare.ByteCursor): KvListResponse { return { - keys: read13(bc), - values: read16(bc), - metadata: read17(bc), + keys: read12(bc), + values: read15(bc), + metadata: read16(bc), } } export function writeKvListResponse(bc: bare.ByteCursor, x: KvListResponse): void { - write13(bc, x.keys) - write16(bc, x.values) - write17(bc, x.metadata) + write12(bc, x.keys) + write15(bc, x.values) + write16(bc, x.metadata) } export type KvPutResponse = null diff --git a/sdks/typescript/runner/package.json b/sdks/typescript/runner/package.json index 0b7bfdf4c3..984f065374 100644 --- a/sdks/typescript/runner/package.json +++ b/sdks/typescript/runner/package.json @@ -22,6 +22,7 @@ "dependencies": { "@rivetkit/engine-runner-protocol": "workspace:*", "@rivetkit/engine-tunnel-protocol": "workspace:*", + "uuid": "^12.0.0", "ws": "^8.18.3" }, "devDependencies": { diff --git a/sdks/typescript/runner/src/mod.ts b/sdks/typescript/runner/src/mod.ts index 92be4cda50..7d26a37925 100644 --- a/sdks/typescript/runner/src/mod.ts +++ b/sdks/typescript/runner/src/mod.ts @@ -1,16 +1,18 @@ import WebSocket from "ws"; import { importWebSocket } from "./websocket.js"; import * as protocol from "@rivetkit/engine-runner-protocol"; -import { unreachable, calculateBackoff } from "./utils.js"; -import { Tunnel } from "./tunnel.js"; -import { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter.js"; +import { unreachable, calculateBackoff } from "./utils"; +import { Tunnel } from "./tunnel"; +import { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter"; const KV_EXPIRE: number = 30_000; -interface ActorInstance { +export interface ActorInstance { actorId: string; generation: number; config: ActorConfig; + requests: Set; // Track active request IDs + webSockets: Set; // Track active WebSocket IDs } export interface ActorConfig { @@ -60,6 +62,11 @@ interface KvRequestEntry { export class Runner { #config: RunnerConfig; + + get config(): RunnerConfig { + return this.#config; + } + #actors: Map = new Map(); #actorWebSockets: Map> = new Map(); @@ -110,7 +117,7 @@ export class Runner { // MARK: Manage actors sleepActor(actorId: string, generation?: number) { - const actor = this.#getActor(actorId, generation); + const actor = this.getActor(actorId, generation); if (!actor) return; // Keep the actor instance in memory during sleep @@ -126,7 +133,7 @@ export class Runner { // Unregister actor from tunnel if (this.#tunnel) { - this.#tunnel.unregisterActor(actorId); + this.#tunnel.unregisterActor(actor); } this.#sendActorStateUpdate(actorId, actor.generation, "stopped"); @@ -147,7 +154,7 @@ export class Runner { } } - #getActor(actorId: string, generation?: number): ActorInstance | undefined { + getActor(actorId: string, generation?: number): ActorInstance | undefined { const actor = this.#actors.get(actorId); if (!actor) { console.error(`Actor ${actorId} not found`); @@ -363,10 +370,10 @@ export class Runner { const wsEndpoint = endpoint .replace("http://", "ws://") .replace("https://", "wss://"); - return `${wsEndpoint}/v1?namespace=${encodeURIComponent(this.#config.namespace)}`; + return `${wsEndpoint}?protocol_version=1&namespace=${encodeURIComponent(this.#config.namespace)}&runner_key=${encodeURIComponent(this.#config.runnerKey)}`; } - get pegboardRelayUrl() { + get pegboardTunnelUrl() { const endpoint = this.#config.pegboardRelayEndpoint || this.#config.pegboardEndpoint || @@ -374,26 +381,19 @@ export class Runner { const wsEndpoint = endpoint .replace("http://", "ws://") .replace("https://", "wss://"); - // Include runner ID if we have it - if (this.runnerId) { - return `${wsEndpoint}/tunnel?namespace=${encodeURIComponent(this.#config.namespace)}&runner_id=${this.runnerId}`; - } - return `${wsEndpoint}/tunnel?namespace=${encodeURIComponent(this.#config.namespace)}`; + return `${wsEndpoint}?protocol_version=1&namespace=${encodeURIComponent(this.#config.namespace)}&runner_key=${this.#config.runnerKey}`; } async #openTunnelAndWait(): Promise { return new Promise((resolve, reject) => { - const url = this.pegboardRelayUrl; + const url = this.pegboardTunnelUrl; //console.log("[RUNNER] Opening tunnel to:", url); //console.log("[RUNNER] Current runner ID:", this.runnerId || "none"); //console.log("[RUNNER] Active actors count:", this.#actors.size); let connected = false; - this.#tunnel = new Tunnel(url); - this.#tunnel.setCallbacks({ - fetch: this.#config.fetch, - websocket: this.#config.websocket, + this.#tunnel = new Tunnel(this, url, { onConnected: () => { if (!connected) { connected = true; @@ -410,35 +410,9 @@ export class Runner { }, }); this.#tunnel.start(); - - // Re-register all active actors with the new tunnel - for (const actorId of this.#actors.keys()) { - //console.log("[RUNNER] Re-registering actor with tunnel:", actorId); - this.#tunnel.registerActor(actorId); - } }); } - #openTunnel() { - const url = this.pegboardRelayUrl; - //console.log("[RUNNER] Opening tunnel to:", url); - //console.log("[RUNNER] Current runner ID:", this.runnerId || "none"); - //console.log("[RUNNER] Active actors count:", this.#actors.size); - - this.#tunnel = new Tunnel(url); - this.#tunnel.setCallbacks({ - fetch: this.#config.fetch, - websocket: this.#config.websocket, - }); - this.#tunnel.start(); - - // Re-register all active actors with the new tunnel - for (const actorId of this.#actors.keys()) { - //console.log("[RUNNER] Re-registering actor with tunnel:", actorId); - this.#tunnel.registerActor(actorId); - } - } - // MARK: Runner protocol async #openPegboardWebSocket() { const WS = await importWebSocket(); @@ -469,9 +443,7 @@ export class Runner { // Send init message const init: protocol.ToServerInit = { - runnerId: this.runnerId || null, name: this.#config.runnerName, - key: this.#config.runnerKey, version: this.#config.version, totalSlots: this.#config.totalSlots, addressesHttp: new Map(), // No addresses needed with tunnel @@ -560,18 +532,6 @@ export class Runner { // runnerLostThreshold: this.#runnerLostThreshold, //}); - // Only reopen tunnel if we didn't have a runner ID before - // This happens on reconnection after losing connection - if (!hadRunnerId && this.runnerId) { - // Reopen tunnel with runner ID - //console.log("[RUNNER] Received runner ID, reopening tunnel"); - if (this.#tunnel) { - //console.log("[RUNNER] Shutting down existing tunnel"); - this.#tunnel.shutdown(); - } - this.#openTunnel(); - } - // Resend events that haven't been acknowledged this.#resendUnacknowledgedEvents(init.lastEventIdx); @@ -664,21 +624,12 @@ export class Runner { actorId, generation, config: actorConfig, + requests: new Set(), + webSockets: new Set(), }; this.#actors.set(actorId, instance); - // Register actor with tunnel - if (this.#tunnel) { - //console.log("[RUNNER] Registering new actor with tunnel:", actorId); - this.#tunnel.registerActor(actorId); - } else { - console.error( - "[RUNNER] WARNING: No tunnel available to register actor:", - actorId, - ); - } - this.#sendActorStateUpdate(actorId, generation, "running"); // TODO: Add timeout to onActorStart diff --git a/sdks/typescript/runner/src/tunnel.ts b/sdks/typescript/runner/src/tunnel.ts index 7c868c44ce..4a00c878e2 100644 --- a/sdks/typescript/runner/src/tunnel.ts +++ b/sdks/typescript/runner/src/tunnel.ts @@ -1,162 +1,269 @@ import WebSocket from "ws"; import * as tunnel from "@rivetkit/engine-tunnel-protocol"; -import { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter.js"; -import { calculateBackoff } from "./utils.js"; +import { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter"; +import { calculateBackoff } from "./utils"; +import type { Runner, ActorInstance } from "./mod"; +import { v4 as uuidv4 } from "uuid"; + +const GC_INTERVAL = 60000; // 60 seconds +const MESSAGE_ACK_TIMEOUT = 5000; // 5 seconds + +interface PendingRequest { + resolve: (response: Response) => void; + reject: (error: Error) => void; + streamController?: ReadableStreamDefaultController; + actorId?: string; +} + +interface TunnelCallbacks { + onConnected(): void; + onDisconnected(): void; +} + +interface PendingMessage { + sentAt: number; + requestIdStr: string; +} export class Tunnel { #pegboardTunnelUrl: string; - #ws?: WebSocket; - #pendingRequests: Map void; - reject: (error: Error) => void; - streamController?: ReadableStreamDefaultController; - actorId?: string; - }> = new Map(); - #webSockets: Map = new Map(); + + #runner: Runner; + + #tunnelWs?: WebSocket; #shutdown = false; #reconnectTimeout?: NodeJS.Timeout; #reconnectAttempt = 0; - - // Track actors and their connections - #activeActors: Set = new Set(); - #actorRequests: Map> = new Map(); - #actorWebSockets: Map> = new Map(); - - // Callbacks - #onConnected?: () => void; - #onDisconnected?: () => void; - #fetchHandler?: (actorId: string, request: Request) => Promise; - #websocketHandler?: (actorId: string, ws: any, request: Request) => Promise; - - constructor(pegboardTunnelUrl: string) { - this.#pegboardTunnelUrl = pegboardTunnelUrl; - } - setCallbacks(options: { - onConnected?: () => void; - onDisconnected?: () => void; - fetch?: (actorId: string, request: Request) => Promise; - websocket?: (actorId: string, ws: any, request: Request) => Promise; - }) { - this.#onConnected = options.onConnected; - this.#onDisconnected = options.onDisconnected; - this.#fetchHandler = options.fetch; - this.#websocketHandler = options.websocket; + #actorPendingRequests: Map = new Map(); + #actorWebSockets: Map = new Map(); + + #pendingMessages: Map = new Map(); + #gcInterval?: NodeJS.Timeout; + + #callbacks: TunnelCallbacks; + + constructor( + runner: Runner, + pegboardTunnelUrl: string, + callbacks: TunnelCallbacks, + ) { + this.#pegboardTunnelUrl = pegboardTunnelUrl; + this.#runner = runner; + this.#callbacks = callbacks; } start(): void { - if (this.#ws?.readyState === WebSocket.OPEN) { + if (this.#tunnelWs?.readyState === WebSocket.OPEN) { return; } - + this.#connect(); + this.#startGarbageCollector(); } shutdown() { this.#shutdown = true; - + if (this.#reconnectTimeout) { clearTimeout(this.#reconnectTimeout); this.#reconnectTimeout = undefined; } - if (this.#ws) { - this.#ws.close(); - this.#ws = undefined; + if (this.#gcInterval) { + clearInterval(this.#gcInterval); + this.#gcInterval = undefined; + } + + if (this.#tunnelWs) { + this.#tunnelWs.close(); + this.#tunnelWs = undefined; } + // TODO: Should we use unregisterActor instead + // Reject all pending requests - for (const [_, request] of this.#pendingRequests) { + for (const [_, request] of this.#actorPendingRequests) { request.reject(new Error("Tunnel shutting down")); } - this.#pendingRequests.clear(); + this.#actorPendingRequests.clear(); // Close all WebSockets - for (const [_, ws] of this.#webSockets) { + for (const [_, ws] of this.#actorWebSockets) { ws.close(); } - this.#webSockets.clear(); - - // Clear actor tracking - this.#activeActors.clear(); - this.#actorRequests.clear(); this.#actorWebSockets.clear(); } - registerActor(actorId: string) { - this.#activeActors.add(actorId); - this.#actorRequests.set(actorId, new Set()); - this.#actorWebSockets.set(actorId, new Set()); + #sendMessage(requestId: tunnel.RequestId, messageKind: tunnel.MessageKind) { + if (!this.#tunnelWs || this.#tunnelWs.readyState !== WebSocket.OPEN) { + console.warn("Cannot send tunnel message, WebSocket not connected"); + return; + } + + // Build message + const messageId = generateUuidBuffer(); + + const requestIdStr = bufferToString(requestId); + this.#pendingMessages.set(bufferToString(messageId), { + sentAt: Date.now(), + requestIdStr, + }); + + // Send message + const message: tunnel.RunnerMessage = { + requestId, + messageId, + messageKind, + }; + + const encoded = tunnel.encodeRunnerMessage(message); + this.#tunnelWs.send(encoded); } - unregisterActor(actorId: string) { - this.#activeActors.delete(actorId); - - // Terminate all requests for this actor - const requests = this.#actorRequests.get(actorId); - if (requests) { - for (const requestId of requests) { - const pending = this.#pendingRequests.get(requestId); - if (pending) { - pending.reject(new Error(`Actor ${actorId} stopped`)); - this.#pendingRequests.delete(requestId); + #sendAck(requestId: tunnel.RequestId, messageId: tunnel.MessageId) { + if (!this.#tunnelWs || this.#tunnelWs.readyState !== WebSocket.OPEN) { + return; + } + + const message: tunnel.RunnerMessage = { + requestId, + messageId, + messageKind: { tag: "Ack", val: null }, + }; + + const encoded = tunnel.encodeRunnerMessage(message); + this.#tunnelWs.send(encoded); + } + + #startGarbageCollector() { + if (this.#gcInterval) { + clearInterval(this.#gcInterval); + } + + this.#gcInterval = setInterval(() => { + this.#gc(); + }, GC_INTERVAL); + } + + #gc() { + const now = Date.now(); + const messagesToDelete: string[] = []; + + for (const [messageId, pendingMessage] of this.#pendingMessages) { + // Check if message is older than timeout + if ( + now - pendingMessage.sentAt > MESSAGE_ACK_TIMEOUT + ) { + messagesToDelete.push(messageId); + + const requestIdStr = pendingMessage.requestIdStr; + + // Check if this is an HTTP request + const pendingRequest = + this.#actorPendingRequests.get(requestIdStr); + if (pendingRequest) { + // Reject the pending HTTP request + pendingRequest.reject( + new Error("Message acknowledgment timeout"), + ); + + // Close stream controller if it exists + if (pendingRequest.streamController) { + pendingRequest.streamController.error( + new Error("Message acknowledgment timeout"), + ); + } + + // Clean up from actorPendingRequests map + this.#actorPendingRequests.delete(requestIdStr); + } + + // Check if this is a WebSocket + const webSocket = this.#actorWebSockets.get(requestIdStr); + if (webSocket) { + // Close the WebSocket connection + webSocket.close(1000, "Message acknowledgment timeout"); + + // Clean up from actorWebSockets map + this.#actorWebSockets.delete(requestIdStr); } } - this.#actorRequests.delete(actorId); } - + + // Remove timed out messages + for (const messageId of messagesToDelete) { + this.#pendingMessages.delete(messageId); + console.warn(`Purged unacked message: ${messageId}`); + } + } + + unregisterActor(actor: ActorInstance) { + const actorId = actor.actorId; + + // Terminate all requests for this actor + for (const requestId of actor.requests) { + const pending = this.#actorPendingRequests.get(requestId); + if (pending) { + pending.reject(new Error(`Actor ${actorId} stopped`)); + this.#actorPendingRequests.delete(requestId); + } + } + actor.requests.clear(); + // Close all WebSockets for this actor - const webSockets = this.#actorWebSockets.get(actorId); - if (webSockets) { - for (const webSocketId of webSockets) { - const ws = this.#webSockets.get(webSocketId); - if (ws) { - ws.close(1000, "Actor stopped"); - this.#webSockets.delete(webSocketId); - } + for (const webSocketId of actor.webSockets) { + const ws = this.#actorWebSockets.get(webSocketId); + if (ws) { + ws.close(1000, "Actor stopped"); + this.#actorWebSockets.delete(webSocketId); } - this.#actorWebSockets.delete(actorId); } + actor.webSockets.clear(); } async #fetch(actorId: string, request: Request): Promise { // Validate actor exists - if (!this.#activeActors.has(actorId)) { - console.warn(`[TUNNEL] Ignoring request for unknown actor: ${actorId}`); + if (!this.#runner.hasActor(actorId)) { + console.warn( + `[TUNNEL] Ignoring request for unknown actor: ${actorId}`, + ); return new Response("Actor not found", { status: 404 }); } - - if (!this.#fetchHandler) { + + const fetchHandler = this.#runner.config.fetch(actorId, request); + + if (!fetchHandler) { return new Response("Not Implemented", { status: 501 }); } - - return this.#fetchHandler(actorId, request); + + return fetchHandler; } #connect() { if (this.#shutdown) return; try { - this.#ws = new WebSocket(this.#pegboardTunnelUrl, { + this.#tunnelWs = new WebSocket(this.#pegboardTunnelUrl, { headers: { "x-rivet-target": "tunnel", }, }); - this.#ws.binaryType = "arraybuffer"; + this.#tunnelWs.binaryType = "arraybuffer"; - this.#ws.addEventListener("open", () => { + this.#tunnelWs.addEventListener("open", () => { this.#reconnectAttempt = 0; - + if (this.#reconnectTimeout) { clearTimeout(this.#reconnectTimeout); this.#reconnectTimeout = undefined; } - this.#onConnected?.(); + this.#callbacks.onConnected(); }); - this.#ws.addEventListener("message", async (event) => { + this.#tunnelWs.addEventListener("message", async (event) => { try { await this.#handleMessage(event.data as ArrayBuffer); } catch (error) { @@ -164,12 +271,12 @@ export class Tunnel { } }); - this.#ws.addEventListener("error", (event) => { + this.#tunnelWs.addEventListener("error", (event) => { console.error("Tunnel WebSocket error:", event); }); - this.#ws.addEventListener("close", () => { - this.#onDisconnected?.(); + this.#tunnelWs.addEventListener("close", () => { + this.#callbacks.onDisconnected(); if (!this.#shutdown) { this.#scheduleReconnect(); @@ -192,9 +299,8 @@ export class Tunnel { multiplier: 2, jitter: true, }); - - this.#reconnectAttempt++; + this.#reconnectAttempt++; this.#reconnectTimeout = setTimeout(() => { this.#connect(); @@ -202,53 +308,97 @@ export class Tunnel { } async #handleMessage(data: ArrayBuffer) { - const message = tunnel.decodeTunnelMessage(new Uint8Array(data)); - - switch (message.body.tag) { - case "ToServerRequestStart": - await this.#handleRequestStart(message.body.val); - break; - case "ToServerRequestChunk": - await this.#handleRequestChunk(message.body.val); - break; - case "ToServerRequestFinish": - await this.#handleRequestFinish(message.body.val); - break; - case "ToServerWebSocketOpen": - await this.#handleWebSocketOpen(message.body.val); - break; - case "ToServerWebSocketMessage": - await this.#handleWebSocketMessage(message.body.val); - break; - case "ToServerWebSocketClose": - await this.#handleWebSocketClose(message.body.val); - break; - case "ToClientResponseStart": - this.#handleResponseStart(message.body.val); - break; - case "ToClientResponseChunk": - this.#handleResponseChunk(message.body.val); - break; - case "ToClientResponseFinish": - this.#handleResponseFinish(message.body.val); - break; - case "ToClientWebSocketOpen": - this.#handleWebSocketOpenResponse(message.body.val); - break; - case "ToClientWebSocketMessage": - this.#handleWebSocketMessageResponse(message.body.val); - break; - case "ToClientWebSocketClose": - this.#handleWebSocketCloseResponse(message.body.val); - break; + const message = tunnel.decodeRunnerMessage(new Uint8Array(data)); + + if (message.messageKind.tag === "Ack") { + // Mark pending message as acknowledged and remove it + const msgIdStr = bufferToString(message.messageId); + const pending = this.#pendingMessages.get(msgIdStr); + if (pending) { + this.#pendingMessages.delete(msgIdStr); + } + } else { + this.#sendAck(message.requestId, message.messageId); + switch (message.messageKind.tag) { + case "ToServerRequestStart": + await this.#handleRequestStart( + message.requestId, + message.messageKind.val, + ); + break; + case "ToServerRequestChunk": + await this.#handleRequestChunk( + message.requestId, + message.messageKind.val, + ); + break; + case "ToServerRequestAbort": + await this.#handleRequestAbort(message.requestId); + break; + case "ToServerWebSocketOpen": + await this.#handleWebSocketOpen( + message.requestId, + message.messageKind.val, + ); + break; + case "ToServerWebSocketMessage": + await this.#handleWebSocketMessage( + message.requestId, + message.messageKind.val, + ); + break; + case "ToServerWebSocketClose": + await this.#handleWebSocketClose( + message.requestId, + message.messageKind.val, + ); + break; + case "ToClientResponseStart": + this.#handleResponseStart( + message.requestId, + message.messageKind.val, + ); + break; + case "ToClientResponseChunk": + this.#handleResponseChunk( + message.requestId, + message.messageKind.val, + ); + break; + case "ToClientResponseAbort": + this.#handleResponseAbort(message.requestId); + break; + case "ToClientWebSocketOpen": + this.#handleWebSocketOpenResponse( + message.requestId, + message.messageKind.val, + ); + break; + case "ToClientWebSocketMessage": + this.#handleWebSocketMessageResponse( + message.requestId, + message.messageKind.val, + ); + break; + case "ToClientWebSocketClose": + this.#handleWebSocketCloseResponse( + message.requestId, + message.messageKind.val, + ); + break; + } } } - async #handleRequestStart(req: tunnel.ToServerRequestStart) { + async #handleRequestStart( + requestId: ArrayBuffer, + req: tunnel.ToServerRequestStart, + ) { // Track this request for the actor - const requests = this.#actorRequests.get(req.actorId); - if (requests) { - requests.add(req.requestId); + const requestIdStr = bufferToString(requestId); + const actor = this.#runner.getActor(req.actorId); + if (actor) { + actor.requests.add(requestIdStr); } try { @@ -271,12 +421,13 @@ export class Tunnel { const stream = new ReadableStream({ start: (controller) => { // Store controller for chunks - const existing = this.#pendingRequests.get(req.requestId); + const existing = + this.#actorPendingRequests.get(requestIdStr); if (existing) { existing.streamController = controller; existing.actorId = req.actorId; } else { - this.#pendingRequests.set(req.requestId, { + this.#actorPendingRequests.set(requestIdStr, { resolve: () => {}, reject: () => {}, streamController: controller, @@ -293,194 +444,193 @@ export class Tunnel { } as any); // Call fetch handler with validation - const response = await this.#fetch(req.actorId, streamingRequest); - await this.#sendResponse(req.requestId, response); + const response = await this.#fetch( + req.actorId, + streamingRequest, + ); + await this.#sendResponse(requestId, response); } else { // Non-streaming request const response = await this.#fetch(req.actorId, request); - await this.#sendResponse(req.requestId, response); + await this.#sendResponse(requestId, response); } } catch (error) { console.error("Error handling request:", error); - this.#sendResponseError(req.requestId, 500, "Internal Server Error"); + this.#sendResponseError(requestId, 500, "Internal Server Error"); } finally { // Clean up request tracking - const requests = this.#actorRequests.get(req.actorId); - if (requests) { - requests.delete(req.requestId); + const actor = this.#runner.getActor(req.actorId); + if (actor) { + actor.requests.delete(requestIdStr); } } } - async #handleRequestChunk(chunk: tunnel.ToServerRequestChunk) { - const pending = this.#pendingRequests.get(chunk.requestId); + async #handleRequestChunk( + requestId: ArrayBuffer, + chunk: tunnel.ToServerRequestChunk, + ) { + const requestIdStr = bufferToString(requestId); + const pending = this.#actorPendingRequests.get(requestIdStr); if (pending?.streamController) { pending.streamController.enqueue(new Uint8Array(chunk.body)); + if (chunk.finish) { + pending.streamController.close(); + this.#actorPendingRequests.delete(requestIdStr); + } } } - async #handleRequestFinish(finish: tunnel.ToServerRequestFinish) { - const pending = this.#pendingRequests.get(finish.requestId); + async #handleRequestAbort(requestId: ArrayBuffer) { + const requestIdStr = bufferToString(requestId); + const pending = this.#actorPendingRequests.get(requestIdStr); if (pending?.streamController) { - if (finish.reason === tunnel.StreamFinishReason.Complete) { - pending.streamController.close(); - } else { - pending.streamController.error(new Error("Request aborted")); - } + pending.streamController.error(new Error("Request aborted")); } - this.#pendingRequests.delete(finish.requestId); + this.#actorPendingRequests.delete(requestIdStr); } - async #sendResponse(requestId: bigint, response: Response) { + async #sendResponse(requestId: ArrayBuffer, response: Response) { // Always treat responses as non-streaming for now // In the future, we could detect streaming responses based on: // - Transfer-Encoding: chunked // - Content-Type: text/event-stream // - Explicit stream flag from the handler - + // Read the body first to get the actual content const body = response.body ? await response.arrayBuffer() : null; - + // Convert headers to map and add Content-Length if not present const headers = new Map(); response.headers.forEach((value, key) => { headers.set(key, value); }); - + // Add Content-Length header if we have a body and it's not already set if (body && !headers.has("content-length")) { headers.set("content-length", String(body.byteLength)); } // Send as non-streaming response - this.#send({ - body: { - tag: "ToClientResponseStart", - val: { - requestId, - status: response.status as tunnel.u16, - headers, - body: body || null, - stream: false, - }, + this.#sendMessage(requestId, { + tag: "ToClientResponseStart", + val: { + status: response.status as tunnel.u16, + headers, + body: body || null, + stream: false, }, }); } - #sendResponseError(requestId: bigint, status: number, message: string) { + #sendResponseError( + requestId: ArrayBuffer, + status: number, + message: string, + ) { const headers = new Map(); headers.set("content-type", "text/plain"); - this.#send({ - body: { - tag: "ToClientResponseStart", - val: { - requestId, - status: status as tunnel.u16, - headers, - body: new TextEncoder().encode(message).buffer as ArrayBuffer, - stream: false, - }, + this.#sendMessage(requestId, { + tag: "ToClientResponseStart", + val: { + status: status as tunnel.u16, + headers, + body: new TextEncoder().encode(message).buffer as ArrayBuffer, + stream: false, }, }); } - async #handleWebSocketOpen(open: tunnel.ToServerWebSocketOpen) { + async #handleWebSocketOpen( + requestId: ArrayBuffer, + open: tunnel.ToServerWebSocketOpen, + ) { + const webSocketId = bufferToString(requestId); // Validate actor exists - if (!this.#activeActors.has(open.actorId)) { - console.warn(`Ignoring WebSocket for unknown actor: ${open.actorId}`); + const actor = this.#runner.getActor(open.actorId); + if (!actor) { + console.warn( + `Ignoring WebSocket for unknown actor: ${open.actorId}`, + ); // Send close immediately - this.#send({ - body: { - tag: "ToClientWebSocketClose", - val: { - webSocketId: open.webSocketId, - code: 1011, - reason: "Actor not found", - }, + this.#sendMessage(requestId, { + tag: "ToClientWebSocketClose", + val: { + code: 1011, + reason: "Actor not found", }, }); return; } - if (!this.#websocketHandler) { + const websocketHandler = this.#runner.config.websocket; + + if (!websocketHandler) { console.error("No websocket handler configured for tunnel"); // Send close immediately - this.#send({ - body: { - tag: "ToClientWebSocketClose", - val: { - webSocketId: open.webSocketId, - code: 1011, - reason: "Not Implemented", - }, + this.#sendMessage(requestId, { + tag: "ToClientWebSocketClose", + val: { + code: 1011, + reason: "Not Implemented", }, }); return; } // Track this WebSocket for the actor - const webSockets = this.#actorWebSockets.get(open.actorId); - if (webSockets) { - webSockets.add(open.webSocketId); + if (actor) { + actor.webSockets.add(webSocketId); } try { // Create WebSocket adapter const adapter = new WebSocketTunnelAdapter( - open.webSocketId, + webSocketId, (data: ArrayBuffer | string, isBinary: boolean) => { // Send message through tunnel - const dataBuffer = typeof data === "string" - ? new TextEncoder().encode(data).buffer as ArrayBuffer - : data; - - this.#send({ - body: { - tag: "ToClientWebSocketMessage", - val: { - webSocketId: open.webSocketId, - data: dataBuffer, - binary: isBinary, - }, + const dataBuffer = + typeof data === "string" + ? (new TextEncoder().encode(data) + .buffer as ArrayBuffer) + : data; + + this.#sendMessage(requestId, { + tag: "ToClientWebSocketMessage", + val: { + data: dataBuffer, + binary: isBinary, }, }); }, (code?: number, reason?: string) => { // Send close through tunnel - this.#send({ - body: { - tag: "ToClientWebSocketClose", - val: { - webSocketId: open.webSocketId, - code: code || null, - reason: reason || null, - }, + this.#sendMessage(requestId, { + tag: "ToClientWebSocketClose", + val: { + code: code || null, + reason: reason || null, }, }); - + // Remove from map - this.#webSockets.delete(open.webSocketId); - + this.#actorWebSockets.delete(webSocketId); + // Clean up actor tracking - const webSockets = this.#actorWebSockets.get(open.actorId); - if (webSockets) { - webSockets.delete(open.webSocketId); + if (actor) { + actor.webSockets.delete(webSocketId); } - } + }, ); // Store adapter - this.#webSockets.set(open.webSocketId, adapter); + this.#actorWebSockets.set(webSocketId, adapter); // Send open confirmation - this.#send({ - body: { - tag: "ToClientWebSocketOpen", - val: { - webSocketId: open.webSocketId, - }, - }, + this.#sendMessage(requestId, { + tag: "ToClientWebSocketOpen", + val: null, }); // Notify adapter that connection is open @@ -490,7 +640,10 @@ export class Tunnel { // Include original headers from the open message const headerInit: Record = {}; if (open.headers) { - for (const [k, v] of open.headers as ReadonlyMap) { + for (const [k, v] of open.headers as ReadonlyMap< + string, + string + >) { headerInit[k] = v; } } @@ -504,55 +657,68 @@ export class Tunnel { }); // Call websocket handler - await this.#websocketHandler(open.actorId, adapter, request); + await websocketHandler(open.actorId, adapter, request); } catch (error) { console.error("Error handling WebSocket open:", error); - + // Send close on error - this.#send({ - body: { - tag: "ToClientWebSocketClose", - val: { - webSocketId: open.webSocketId, - code: 1011, - reason: "Server Error", - }, + this.#sendMessage(requestId, { + tag: "ToClientWebSocketClose", + val: { + code: 1011, + reason: "Server Error", }, }); - - this.#webSockets.delete(open.webSocketId); - + + this.#actorWebSockets.delete(webSocketId); + // Clean up actor tracking - const webSockets = this.#actorWebSockets.get(open.actorId); - if (webSockets) { - webSockets.delete(open.webSocketId); + if (actor) { + actor.webSockets.delete(webSocketId); } } } - async #handleWebSocketMessage(msg: tunnel.ToServerWebSocketMessage) { - const adapter = this.#webSockets.get(msg.webSocketId); + async #handleWebSocketMessage( + requestId: ArrayBuffer, + msg: tunnel.ToServerWebSocketMessage, + ) { + const webSocketId = bufferToString(requestId); + const adapter = this.#actorWebSockets.get(webSocketId); if (adapter) { const data = msg.binary ? new Uint8Array(msg.data) : new TextDecoder().decode(new Uint8Array(msg.data)); - + adapter._handleMessage(data, msg.binary); } } - async #handleWebSocketClose(close: tunnel.ToServerWebSocketClose) { - const adapter = this.#webSockets.get(close.webSocketId); + async #handleWebSocketClose( + requestId: ArrayBuffer, + close: tunnel.ToServerWebSocketClose, + ) { + const webSocketId = bufferToString(requestId); + const adapter = this.#actorWebSockets.get(webSocketId); if (adapter) { - adapter._handleClose(close.code || undefined, close.reason || undefined); - this.#webSockets.delete(close.webSocketId); + adapter._handleClose( + close.code || undefined, + close.reason || undefined, + ); + this.#actorWebSockets.delete(webSocketId); } } - #handleResponseStart(resp: tunnel.ToClientResponseStart) { - const pending = this.#pendingRequests.get(resp.requestId); + #handleResponseStart( + requestId: ArrayBuffer, + resp: tunnel.ToClientResponseStart, + ) { + const requestIdStr = bufferToString(requestId); + const pending = this.#actorPendingRequests.get(requestIdStr); if (!pending) { - console.warn(`Received response for unknown request ${resp.requestId}`); + console.warn( + `Received response for unknown request ${requestIdStr}`, + ); return; } @@ -585,62 +751,84 @@ export class Tunnel { }); pending.resolve(response); - this.#pendingRequests.delete(resp.requestId); + this.#actorPendingRequests.delete(requestIdStr); } } - #handleResponseChunk(chunk: tunnel.ToClientResponseChunk) { - const pending = this.#pendingRequests.get(chunk.requestId); + #handleResponseChunk( + requestId: ArrayBuffer, + chunk: tunnel.ToClientResponseChunk, + ) { + const requestIdStr = bufferToString(requestId); + const pending = this.#actorPendingRequests.get(requestIdStr); if (pending?.streamController) { pending.streamController.enqueue(new Uint8Array(chunk.body)); + if (chunk.finish) { + pending.streamController.close(); + this.#actorPendingRequests.delete(requestIdStr); + } } } - #handleResponseFinish(finish: tunnel.ToClientResponseFinish) { - const pending = this.#pendingRequests.get(finish.requestId); + #handleResponseAbort(requestId: ArrayBuffer) { + const requestIdStr = bufferToString(requestId); + const pending = this.#actorPendingRequests.get(requestIdStr); if (pending?.streamController) { - if (finish.reason === tunnel.StreamFinishReason.Complete) { - pending.streamController.close(); - } else { - pending.streamController.error(new Error("Response aborted")); - } + pending.streamController.error(new Error("Response aborted")); } - this.#pendingRequests.delete(finish.requestId); + this.#actorPendingRequests.delete(requestIdStr); } - #handleWebSocketOpenResponse(open: tunnel.ToClientWebSocketOpen) { - const adapter = this.#webSockets.get(open.webSocketId); + #handleWebSocketOpenResponse( + requestId: ArrayBuffer, + open: tunnel.ToClientWebSocketOpen, + ) { + const webSocketId = bufferToString(requestId); + const adapter = this.#actorWebSockets.get(webSocketId); if (adapter) { adapter._handleOpen(); } } - #handleWebSocketMessageResponse(msg: tunnel.ToClientWebSocketMessage) { - const adapter = this.#webSockets.get(msg.webSocketId); + #handleWebSocketMessageResponse( + requestId: ArrayBuffer, + msg: tunnel.ToClientWebSocketMessage, + ) { + const webSocketId = bufferToString(requestId); + const adapter = this.#actorWebSockets.get(webSocketId); if (adapter) { const data = msg.binary ? new Uint8Array(msg.data) : new TextDecoder().decode(new Uint8Array(msg.data)); - + adapter._handleMessage(data, msg.binary); } } - #handleWebSocketCloseResponse(close: tunnel.ToClientWebSocketClose) { - const adapter = this.#webSockets.get(close.webSocketId); + #handleWebSocketCloseResponse( + requestId: ArrayBuffer, + close: tunnel.ToClientWebSocketClose, + ) { + const webSocketId = bufferToString(requestId); + const adapter = this.#actorWebSockets.get(webSocketId); if (adapter) { - adapter._handleClose(close.code || undefined, close.reason || undefined); - this.#webSockets.delete(close.webSocketId); + adapter._handleClose( + close.code || undefined, + close.reason || undefined, + ); + this.#actorWebSockets.delete(webSocketId); } } +} - #send(message: tunnel.TunnelMessage) { - if (!this.#ws || this.#ws.readyState !== WebSocket.OPEN) { - console.warn("Cannot send tunnel message, WebSocket not connected"); - return; - } +/** Converts a buffer to a string. Used for storing strings in a lookup map. */ +function bufferToString(buffer: ArrayBuffer): string { + return Buffer.from(buffer).toString("base64"); +} - const encoded = tunnel.encodeTunnelMessage(message); - this.#ws.send(encoded); - } +/** Generates a UUID as bytes. */ +function generateUuidBuffer(): ArrayBuffer { + const buffer = new Uint8Array(16); + uuidv4(undefined, buffer); + return buffer.buffer; } diff --git a/sdks/typescript/runner/src/websocket-tunnel-adapter.ts b/sdks/typescript/runner/src/websocket-tunnel-adapter.ts index bf097f8a5c..b2430d3238 100644 --- a/sdks/typescript/runner/src/websocket-tunnel-adapter.ts +++ b/sdks/typescript/runner/src/websocket-tunnel-adapter.ts @@ -2,7 +2,7 @@ // Implements a subset of the WebSocket interface for compatibility with runner code export class WebSocketTunnelAdapter { - #webSocketId: bigint; + #webSocketId: string; #readyState: number = 0; // CONNECTING #eventListeners: Map void>> = new Map(); #onopen: ((this: any, ev: any) => any) | null = null; @@ -24,7 +24,7 @@ export class WebSocketTunnelAdapter { }> = []; constructor( - webSocketId: bigint, + webSocketId: string, sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void, closeCallback: (code?: number, reason?: string) => void ) { diff --git a/sdks/typescript/tunnel-protocol/src/index.ts b/sdks/typescript/tunnel-protocol/src/index.ts index 0ea9a5b8ce..9e4b6ac609 100644 --- a/sdks/typescript/tunnel-protocol/src/index.ts +++ b/sdks/typescript/tunnel-protocol/src/index.ts @@ -1,30 +1,38 @@ +import assert from "node:assert" import * as bare from "@bare-ts/lib" const DEFAULT_CONFIG = /* @__PURE__ */ bare.Config({}) export type u16 = number -export type u64 = bigint -export type RequestId = u64 +export type RequestId = ArrayBuffer export function readRequestId(bc: bare.ByteCursor): RequestId { - return bare.readU64(bc) + return bare.readFixedData(bc, 16) } export function writeRequestId(bc: bare.ByteCursor, x: RequestId): void { - bare.writeU64(bc, x) + assert(x.byteLength === 16) + bare.writeFixedData(bc, x) } -export type WebSocketId = u64 +/** + * UUIDv4 + */ +export type MessageId = ArrayBuffer -export function readWebSocketId(bc: bare.ByteCursor): WebSocketId { - return bare.readU64(bc) +export function readMessageId(bc: bare.ByteCursor): MessageId { + return bare.readFixedData(bc, 16) } -export function writeWebSocketId(bc: bare.ByteCursor, x: WebSocketId): void { - bare.writeU64(bc, x) +export function writeMessageId(bc: bare.ByteCursor, x: MessageId): void { + assert(x.byteLength === 16) + bare.writeFixedData(bc, x) } +/** + * UUIDv4 + */ export type Id = string export function readId(bc: bare.ByteCursor): Id { @@ -35,38 +43,10 @@ export function writeId(bc: bare.ByteCursor, x: Id): void { bare.writeString(bc, x) } -export enum StreamFinishReason { - Complete = "Complete", - Abort = "Abort", -} - -export function readStreamFinishReason(bc: bare.ByteCursor): StreamFinishReason { - const offset = bc.offset - const tag = bare.readU8(bc) - switch (tag) { - case 0: - return StreamFinishReason.Complete - case 1: - return StreamFinishReason.Abort - default: { - bc.offset = offset - throw new bare.BareError(offset, "invalid tag") - } - } -} - -export function writeStreamFinishReason(bc: bare.ByteCursor, x: StreamFinishReason): void { - switch (x) { - case StreamFinishReason.Complete: { - bare.writeU8(bc, 0) - break - } - case StreamFinishReason.Abort: { - bare.writeU8(bc, 1) - break - } - } -} +/** + * MARK: Ack + */ +export type Ack = null function read0(bc: bare.ByteCursor): ReadonlyMap { const len = bare.readUintSafe(bc) @@ -103,10 +83,9 @@ function write1(bc: bare.ByteCursor, x: ArrayBuffer | null): void { } /** - * MARK: HTTP Request Forwarding + * MARK: HTTP */ export type ToServerRequestStart = { - readonly requestId: RequestId readonly actorId: Id readonly method: string readonly path: string @@ -117,7 +96,6 @@ export type ToServerRequestStart = { export function readToServerRequestStart(bc: bare.ByteCursor): ToServerRequestStart { return { - requestId: readRequestId(bc), actorId: readId(bc), method: bare.readString(bc), path: bare.readString(bc), @@ -128,7 +106,6 @@ export function readToServerRequestStart(bc: bare.ByteCursor): ToServerRequestSt } export function writeToServerRequestStart(bc: bare.ByteCursor, x: ToServerRequestStart): void { - writeRequestId(bc, x.requestId) writeId(bc, x.actorId) bare.writeString(bc, x.method) bare.writeString(bc, x.path) @@ -138,41 +115,25 @@ export function writeToServerRequestStart(bc: bare.ByteCursor, x: ToServerReques } export type ToServerRequestChunk = { - readonly requestId: RequestId readonly body: ArrayBuffer + readonly finish: boolean } export function readToServerRequestChunk(bc: bare.ByteCursor): ToServerRequestChunk { return { - requestId: readRequestId(bc), body: bare.readData(bc), + finish: bare.readBool(bc), } } export function writeToServerRequestChunk(bc: bare.ByteCursor, x: ToServerRequestChunk): void { - writeRequestId(bc, x.requestId) bare.writeData(bc, x.body) + bare.writeBool(bc, x.finish) } -export type ToServerRequestFinish = { - readonly requestId: RequestId - readonly reason: StreamFinishReason -} - -export function readToServerRequestFinish(bc: bare.ByteCursor): ToServerRequestFinish { - return { - requestId: readRequestId(bc), - reason: readStreamFinishReason(bc), - } -} - -export function writeToServerRequestFinish(bc: bare.ByteCursor, x: ToServerRequestFinish): void { - writeRequestId(bc, x.requestId) - writeStreamFinishReason(bc, x.reason) -} +export type ToServerRequestAbort = null export type ToClientResponseStart = { - readonly requestId: RequestId readonly status: u16 readonly headers: ReadonlyMap readonly body: ArrayBuffer | null @@ -181,7 +142,6 @@ export type ToClientResponseStart = { export function readToClientResponseStart(bc: bare.ByteCursor): ToClientResponseStart { return { - requestId: readRequestId(bc), status: bare.readU16(bc), headers: read0(bc), body: read1(bc), @@ -190,7 +150,6 @@ export function readToClientResponseStart(bc: bare.ByteCursor): ToClientResponse } export function writeToClientResponseStart(bc: bare.ByteCursor, x: ToClientResponseStart): void { - writeRequestId(bc, x.requestId) bare.writeU16(bc, x.status) write0(bc, x.headers) write1(bc, x.body) @@ -198,45 +157,29 @@ export function writeToClientResponseStart(bc: bare.ByteCursor, x: ToClientRespo } export type ToClientResponseChunk = { - readonly requestId: RequestId readonly body: ArrayBuffer + readonly finish: boolean } export function readToClientResponseChunk(bc: bare.ByteCursor): ToClientResponseChunk { return { - requestId: readRequestId(bc), body: bare.readData(bc), + finish: bare.readBool(bc), } } export function writeToClientResponseChunk(bc: bare.ByteCursor, x: ToClientResponseChunk): void { - writeRequestId(bc, x.requestId) bare.writeData(bc, x.body) + bare.writeBool(bc, x.finish) } -export type ToClientResponseFinish = { - readonly requestId: RequestId - readonly reason: StreamFinishReason -} - -export function readToClientResponseFinish(bc: bare.ByteCursor): ToClientResponseFinish { - return { - requestId: readRequestId(bc), - reason: readStreamFinishReason(bc), - } -} - -export function writeToClientResponseFinish(bc: bare.ByteCursor, x: ToClientResponseFinish): void { - writeRequestId(bc, x.requestId) - writeStreamFinishReason(bc, x.reason) -} +export type ToClientResponseAbort = null /** - * MARK: WebSocket Forwarding + * MARK: WebSocket */ export type ToServerWebSocketOpen = { readonly actorId: Id - readonly webSocketId: WebSocketId readonly path: string readonly headers: ReadonlyMap } @@ -244,7 +187,6 @@ export type ToServerWebSocketOpen = { export function readToServerWebSocketOpen(bc: bare.ByteCursor): ToServerWebSocketOpen { return { actorId: readId(bc), - webSocketId: readWebSocketId(bc), path: bare.readString(bc), headers: read0(bc), } @@ -252,27 +194,23 @@ export function readToServerWebSocketOpen(bc: bare.ByteCursor): ToServerWebSocke export function writeToServerWebSocketOpen(bc: bare.ByteCursor, x: ToServerWebSocketOpen): void { writeId(bc, x.actorId) - writeWebSocketId(bc, x.webSocketId) bare.writeString(bc, x.path) write0(bc, x.headers) } export type ToServerWebSocketMessage = { - readonly webSocketId: WebSocketId readonly data: ArrayBuffer readonly binary: boolean } export function readToServerWebSocketMessage(bc: bare.ByteCursor): ToServerWebSocketMessage { return { - webSocketId: readWebSocketId(bc), data: bare.readData(bc), binary: bare.readBool(bc), } } export function writeToServerWebSocketMessage(bc: bare.ByteCursor, x: ToServerWebSocketMessage): void { - writeWebSocketId(bc, x.webSocketId) bare.writeData(bc, x.data) bare.writeBool(bc, x.binary) } @@ -300,75 +238,54 @@ function write3(bc: bare.ByteCursor, x: string | null): void { } export type ToServerWebSocketClose = { - readonly webSocketId: WebSocketId readonly code: u16 | null readonly reason: string | null } export function readToServerWebSocketClose(bc: bare.ByteCursor): ToServerWebSocketClose { return { - webSocketId: readWebSocketId(bc), code: read2(bc), reason: read3(bc), } } export function writeToServerWebSocketClose(bc: bare.ByteCursor, x: ToServerWebSocketClose): void { - writeWebSocketId(bc, x.webSocketId) write2(bc, x.code) write3(bc, x.reason) } -export type ToClientWebSocketOpen = { - readonly webSocketId: WebSocketId -} - -export function readToClientWebSocketOpen(bc: bare.ByteCursor): ToClientWebSocketOpen { - return { - webSocketId: readWebSocketId(bc), - } -} - -export function writeToClientWebSocketOpen(bc: bare.ByteCursor, x: ToClientWebSocketOpen): void { - writeWebSocketId(bc, x.webSocketId) -} +export type ToClientWebSocketOpen = null export type ToClientWebSocketMessage = { - readonly webSocketId: WebSocketId readonly data: ArrayBuffer readonly binary: boolean } export function readToClientWebSocketMessage(bc: bare.ByteCursor): ToClientWebSocketMessage { return { - webSocketId: readWebSocketId(bc), data: bare.readData(bc), binary: bare.readBool(bc), } } export function writeToClientWebSocketMessage(bc: bare.ByteCursor, x: ToClientWebSocketMessage): void { - writeWebSocketId(bc, x.webSocketId) bare.writeData(bc, x.data) bare.writeBool(bc, x.binary) } export type ToClientWebSocketClose = { - readonly webSocketId: WebSocketId readonly code: u16 | null readonly reason: string | null } export function readToClientWebSocketClose(bc: bare.ByteCursor): ToClientWebSocketClose { return { - webSocketId: readWebSocketId(bc), code: read2(bc), reason: read3(bc), } } export function writeToClientWebSocketClose(bc: bare.ByteCursor, x: ToClientWebSocketClose): void { - writeWebSocketId(bc, x.webSocketId) write2(bc, x.code) write3(bc, x.reason) } @@ -376,16 +293,17 @@ export function writeToClientWebSocketClose(bc: bare.ByteCursor, x: ToClientWebS /** * MARK: Message */ -export type MessageBody = +export type MessageKind = + | { readonly tag: "Ack"; readonly val: Ack } /** * HTTP */ | { readonly tag: "ToServerRequestStart"; readonly val: ToServerRequestStart } | { readonly tag: "ToServerRequestChunk"; readonly val: ToServerRequestChunk } - | { readonly tag: "ToServerRequestFinish"; readonly val: ToServerRequestFinish } + | { readonly tag: "ToServerRequestAbort"; readonly val: ToServerRequestAbort } | { readonly tag: "ToClientResponseStart"; readonly val: ToClientResponseStart } | { readonly tag: "ToClientResponseChunk"; readonly val: ToClientResponseChunk } - | { readonly tag: "ToClientResponseFinish"; readonly val: ToClientResponseFinish } + | { readonly tag: "ToClientResponseAbort"; readonly val: ToClientResponseAbort } /** * WebSocket */ @@ -396,33 +314,35 @@ export type MessageBody = | { readonly tag: "ToClientWebSocketMessage"; readonly val: ToClientWebSocketMessage } | { readonly tag: "ToClientWebSocketClose"; readonly val: ToClientWebSocketClose } -export function readMessageBody(bc: bare.ByteCursor): MessageBody { +export function readMessageKind(bc: bare.ByteCursor): MessageKind { const offset = bc.offset const tag = bare.readU8(bc) switch (tag) { case 0: - return { tag: "ToServerRequestStart", val: readToServerRequestStart(bc) } + return { tag: "Ack", val: null } case 1: - return { tag: "ToServerRequestChunk", val: readToServerRequestChunk(bc) } + return { tag: "ToServerRequestStart", val: readToServerRequestStart(bc) } case 2: - return { tag: "ToServerRequestFinish", val: readToServerRequestFinish(bc) } + return { tag: "ToServerRequestChunk", val: readToServerRequestChunk(bc) } case 3: - return { tag: "ToClientResponseStart", val: readToClientResponseStart(bc) } + return { tag: "ToServerRequestAbort", val: null } case 4: - return { tag: "ToClientResponseChunk", val: readToClientResponseChunk(bc) } + return { tag: "ToClientResponseStart", val: readToClientResponseStart(bc) } case 5: - return { tag: "ToClientResponseFinish", val: readToClientResponseFinish(bc) } + return { tag: "ToClientResponseChunk", val: readToClientResponseChunk(bc) } case 6: - return { tag: "ToServerWebSocketOpen", val: readToServerWebSocketOpen(bc) } + return { tag: "ToClientResponseAbort", val: null } case 7: - return { tag: "ToServerWebSocketMessage", val: readToServerWebSocketMessage(bc) } + return { tag: "ToServerWebSocketOpen", val: readToServerWebSocketOpen(bc) } case 8: - return { tag: "ToServerWebSocketClose", val: readToServerWebSocketClose(bc) } + return { tag: "ToServerWebSocketMessage", val: readToServerWebSocketMessage(bc) } case 9: - return { tag: "ToClientWebSocketOpen", val: readToClientWebSocketOpen(bc) } + return { tag: "ToServerWebSocketClose", val: readToServerWebSocketClose(bc) } case 10: - return { tag: "ToClientWebSocketMessage", val: readToClientWebSocketMessage(bc) } + return { tag: "ToClientWebSocketOpen", val: null } case 11: + return { tag: "ToClientWebSocketMessage", val: readToClientWebSocketMessage(bc) } + case 12: return { tag: "ToClientWebSocketClose", val: readToClientWebSocketClose(bc) } default: { bc.offset = offset @@ -431,65 +351,66 @@ export function readMessageBody(bc: bare.ByteCursor): MessageBody { } } -export function writeMessageBody(bc: bare.ByteCursor, x: MessageBody): void { +export function writeMessageKind(bc: bare.ByteCursor, x: MessageKind): void { switch (x.tag) { - case "ToServerRequestStart": { + case "Ack": { bare.writeU8(bc, 0) + break + } + case "ToServerRequestStart": { + bare.writeU8(bc, 1) writeToServerRequestStart(bc, x.val) break } case "ToServerRequestChunk": { - bare.writeU8(bc, 1) + bare.writeU8(bc, 2) writeToServerRequestChunk(bc, x.val) break } - case "ToServerRequestFinish": { - bare.writeU8(bc, 2) - writeToServerRequestFinish(bc, x.val) + case "ToServerRequestAbort": { + bare.writeU8(bc, 3) break } case "ToClientResponseStart": { - bare.writeU8(bc, 3) + bare.writeU8(bc, 4) writeToClientResponseStart(bc, x.val) break } case "ToClientResponseChunk": { - bare.writeU8(bc, 4) + bare.writeU8(bc, 5) writeToClientResponseChunk(bc, x.val) break } - case "ToClientResponseFinish": { - bare.writeU8(bc, 5) - writeToClientResponseFinish(bc, x.val) + case "ToClientResponseAbort": { + bare.writeU8(bc, 6) break } case "ToServerWebSocketOpen": { - bare.writeU8(bc, 6) + bare.writeU8(bc, 7) writeToServerWebSocketOpen(bc, x.val) break } case "ToServerWebSocketMessage": { - bare.writeU8(bc, 7) + bare.writeU8(bc, 8) writeToServerWebSocketMessage(bc, x.val) break } case "ToServerWebSocketClose": { - bare.writeU8(bc, 8) + bare.writeU8(bc, 9) writeToServerWebSocketClose(bc, x.val) break } case "ToClientWebSocketOpen": { - bare.writeU8(bc, 9) - writeToClientWebSocketOpen(bc, x.val) + bare.writeU8(bc, 10) break } case "ToClientWebSocketMessage": { - bare.writeU8(bc, 10) + bare.writeU8(bc, 11) writeToClientWebSocketMessage(bc, x.val) break } case "ToClientWebSocketClose": { - bare.writeU8(bc, 11) + bare.writeU8(bc, 12) writeToClientWebSocketClose(bc, x.val) break } @@ -497,35 +418,89 @@ export function writeMessageBody(bc: bare.ByteCursor, x: MessageBody): void { } /** - * Main tunnel message + * MARK: Message sent over tunnel WebSocket */ -export type TunnelMessage = { - readonly body: MessageBody +export type RunnerMessage = { + readonly requestId: RequestId + readonly messageId: MessageId + readonly messageKind: MessageKind } -export function readTunnelMessage(bc: bare.ByteCursor): TunnelMessage { +export function readRunnerMessage(bc: bare.ByteCursor): RunnerMessage { return { - body: readMessageBody(bc), + requestId: readRequestId(bc), + messageId: readMessageId(bc), + messageKind: readMessageKind(bc), } } -export function writeTunnelMessage(bc: bare.ByteCursor, x: TunnelMessage): void { - writeMessageBody(bc, x.body) +export function writeRunnerMessage(bc: bare.ByteCursor, x: RunnerMessage): void { + writeRequestId(bc, x.requestId) + writeMessageId(bc, x.messageId) + writeMessageKind(bc, x.messageKind) +} + +export function encodeRunnerMessage(x: RunnerMessage, config?: Partial): Uint8Array { + const fullConfig = config != null ? bare.Config(config) : DEFAULT_CONFIG + const bc = new bare.ByteCursor( + new Uint8Array(fullConfig.initialBufferLength), + fullConfig, + ) + writeRunnerMessage(bc, x) + return new Uint8Array(bc.view.buffer, bc.view.byteOffset, bc.offset) +} + +export function decodeRunnerMessage(bytes: Uint8Array): RunnerMessage { + const bc = new bare.ByteCursor(bytes, DEFAULT_CONFIG) + const result = readRunnerMessage(bc) + if (bc.offset < bc.view.byteLength) { + throw new bare.BareError(bc.offset, "remaining bytes") + } + return result +} + +/** + * MARK: Message sent over UPS + */ +export type PubSubMessage = { + readonly requestId: RequestId + readonly messageId: MessageId + /** + * Subject to send replies to. Only sent when opening a new request from gateway -> runner. + */ + readonly replyTo: string | null + readonly messageKind: MessageKind +} + +export function readPubSubMessage(bc: bare.ByteCursor): PubSubMessage { + return { + requestId: readRequestId(bc), + messageId: readMessageId(bc), + replyTo: read3(bc), + messageKind: readMessageKind(bc), + } +} + +export function writePubSubMessage(bc: bare.ByteCursor, x: PubSubMessage): void { + writeRequestId(bc, x.requestId) + writeMessageId(bc, x.messageId) + write3(bc, x.replyTo) + writeMessageKind(bc, x.messageKind) } -export function encodeTunnelMessage(x: TunnelMessage, config?: Partial): Uint8Array { +export function encodePubSubMessage(x: PubSubMessage, config?: Partial): Uint8Array { const fullConfig = config != null ? bare.Config(config) : DEFAULT_CONFIG const bc = new bare.ByteCursor( new Uint8Array(fullConfig.initialBufferLength), fullConfig, ) - writeTunnelMessage(bc, x) + writePubSubMessage(bc, x) return new Uint8Array(bc.view.buffer, bc.view.byteOffset, bc.offset) } -export function decodeTunnelMessage(bytes: Uint8Array): TunnelMessage { +export function decodePubSubMessage(bytes: Uint8Array): PubSubMessage { const bc = new bare.ByteCursor(bytes, DEFAULT_CONFIG) - const result = readTunnelMessage(bc) + const result = readPubSubMessage(bc) if (bc.offset < bc.view.byteLength) { throw new bare.BareError(bc.offset, "remaining bytes") } From 74e5db56c2c121c60091e27807b8a85f1f56d07c Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Mon, 8 Sep 2025 17:58:08 -0700 Subject: [PATCH 02/17] chore(ups): add postgres auto-reconnect --- Cargo.lock | 1 + packages/common/test-deps-docker/src/lib.rs | 36 +++ packages/common/universalpubsub/Cargo.toml | 1 + .../src/driver/postgres/mod.rs | 222 ++++++++++++++---- .../common/universalpubsub/tests/reconnect.rs | 217 +++++++++++++++++ 5 files changed, 430 insertions(+), 47 deletions(-) create mode 100644 packages/common/universalpubsub/tests/reconnect.rs diff --git a/Cargo.lock b/Cargo.lock index 800d4ca464..44bd7cd6e0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6338,6 +6338,7 @@ dependencies = [ "rivet-error", "rivet-test-deps-docker", "rivet-ups-protocol", + "rivet-util", "serde", "serde_json", "sha2", diff --git a/packages/common/test-deps-docker/src/lib.rs b/packages/common/test-deps-docker/src/lib.rs index efa46a3152..2ed73f8600 100644 --- a/packages/common/test-deps-docker/src/lib.rs +++ b/packages/common/test-deps-docker/src/lib.rs @@ -73,6 +73,42 @@ impl DockerRunConfig { Ok(true) } + pub async fn restart(&self) -> Result<()> { + let container_id = self + .container_id + .as_ref() + .ok_or_else(|| anyhow!("No container ID found, container not started"))?; + + tracing::debug!( + container_name = %self.container_name, + container_id = %container_id, + "restarting docker container" + ); + + let output = Command::new("docker") + .arg("restart") + .arg(container_id) + .output() + .await?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + anyhow::bail!( + "Failed to restart container {}: {}", + self.container_name, + stderr + ); + } + + tracing::debug!( + container_name = %self.container_name, + container_id = %container_id, + "container restarted successfully" + ); + + Ok(()) + } + pub fn container_id(&self) -> Option<&str> { self.container_id.as_deref() } diff --git a/packages/common/universalpubsub/Cargo.toml b/packages/common/universalpubsub/Cargo.toml index 3cdb8509b6..4ffddeb77c 100644 --- a/packages/common/universalpubsub/Cargo.toml +++ b/packages/common/universalpubsub/Cargo.toml @@ -14,6 +14,7 @@ deadpool-postgres.workspace = true futures-util.workspace = true rivet-error.workspace = true rivet-ups-protocol.workspace = true +rivet-util.workspace = true serde_json.workspace = true versioned-data-util.workspace = true serde.workspace = true diff --git a/packages/common/universalpubsub/src/driver/postgres/mod.rs b/packages/common/universalpubsub/src/driver/postgres/mod.rs index c2df9a36b4..21cf96f5f5 100644 --- a/packages/common/universalpubsub/src/driver/postgres/mod.rs +++ b/packages/common/universalpubsub/src/driver/postgres/mod.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; use std::hash::{DefaultHasher, Hash, Hasher}; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use anyhow::*; use async_trait::async_trait; @@ -8,6 +8,8 @@ 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_util::backoff::Backoff; +use tokio::sync::{Mutex, broadcast}; use tokio_postgres::{AsyncMessage, NoTls}; use tracing::Instrument; @@ -17,13 +19,13 @@ use crate::pubsub::DriverOutput; #[derive(Clone)] struct Subscription { // Channel to send messages to this subscription - tx: tokio::sync::broadcast::Sender>, + tx: broadcast::Sender>, // Cancellation token shared by all subscribers of this subject token: tokio_util::sync::CancellationToken, } impl Subscription { - fn new(tx: tokio::sync::broadcast::Sender>) -> Self { + fn new(tx: broadcast::Sender>) -> Self { let token = tokio_util::sync::CancellationToken::new(); Self { tx, token } } @@ -48,8 +50,9 @@ pub const POSTGRES_MAX_MESSAGE_SIZE: usize = #[derive(Clone)] pub struct PostgresDriver { pool: Arc, - client: Arc, + client: Arc>>>, subscriptions: Arc>>, + client_ready: tokio::sync::watch::Receiver, } impl PostgresDriver { @@ -76,48 +79,168 @@ impl PostgresDriver { let subscriptions: Arc>> = Arc::new(Mutex::new(HashMap::new())); - let subscriptions2 = subscriptions.clone(); + let client: Arc>>> = Arc::new(Mutex::new(None)); - let (client, mut conn) = tokio_postgres::connect(&conn_str, tokio_postgres::NoTls).await?; - tokio::spawn(async move { - // NOTE: This loop will stop automatically when client is dropped - loop { - match poll_fn(|cx| conn.poll_message(cx)).await { - Some(std::result::Result::Ok(AsyncMessage::Notification(note))) => { - if let Some(sub) = - subscriptions2.lock().unwrap().get(note.channel()).cloned() - { - let bytes = match BASE64.decode(note.payload()) { - std::result::Result::Ok(b) => b, - std::result::Result::Err(err) => { - tracing::error!(?err, "failed decoding base64"); - break; - } - }; - let _ = sub.tx.send(bytes); - } - } - Some(std::result::Result::Ok(_)) => { - // Ignore other async messages + // Create channel for client ready notifications + let (ready_tx, client_ready) = tokio::sync::watch::channel(false); + + // Spawn connection lifecycle task + tokio::spawn(Self::spawn_connection_lifecycle( + conn_str.clone(), + subscriptions.clone(), + client.clone(), + ready_tx, + )); + + let driver = Self { + pool: Arc::new(pool), + client, + subscriptions, + client_ready, + }; + + // Wait for initial connection to be established + driver.wait_for_client().await?; + + Ok(driver) + } + + /// Manages the connection lifecycle with automatic reconnection + async fn spawn_connection_lifecycle( + conn_str: String, + subscriptions: Arc>>, + client: Arc>>>, + ready_tx: tokio::sync::watch::Sender, + ) { + let mut backoff = Backoff::new(8, None, 1_000, 1_000); + + loop { + match tokio_postgres::connect(&conn_str, tokio_postgres::NoTls).await { + Result::Ok((new_client, conn)) => { + tracing::info!("postgres listen connection established"); + // Reset backoff on successful connection + backoff = Backoff::new(8, None, 1_000, 1_000); + + let new_client = Arc::new(new_client); + + // Update the client reference immediately + *client.lock().await = Some(new_client.clone()); + // Notify that client is ready + let _ = ready_tx.send(true); + + // Get channels to re-subscribe to + let channels: Vec = + subscriptions.lock().await.keys().cloned().collect(); + let needs_resubscribe = !channels.is_empty(); + + if needs_resubscribe { + tracing::debug!( + ?channels, + "will re-subscribe to channels after connection starts" + ); } - Some(std::result::Result::Err(err)) => { - tracing::error!(?err, "async postgres error"); - break; + + // Spawn a task to re-subscribe after a short delay + if needs_resubscribe { + let client_for_resub = new_client.clone(); + let channels_clone = channels.clone(); + tokio::spawn(async move { + tracing::debug!( + ?channels_clone, + "re-subscribing to channels after reconnection" + ); + for channel in &channels_clone { + if let Result::Err(e) = client_for_resub + .execute(&format!("LISTEN \"{}\"", channel), &[]) + .await + { + tracing::error!(?e, %channel, "failed to re-subscribe to channel"); + } else { + tracing::debug!(%channel, "successfully re-subscribed to channel"); + } + } + }); } - None => { - tracing::debug!("async postgres connection closed"); - break; + + // Poll the connection until it closes + Self::poll_connection(conn, subscriptions.clone()).await; + + // Clear the client reference on disconnect + *client.lock().await = None; + // Notify that client is disconnected + let _ = ready_tx.send(false); + } + Result::Err(e) => { + tracing::error!(?e, "failed to connect to postgres, retrying"); + backoff.tick().await; + } + } + } + } + + /// Polls the connection for notifications until it closes or errors + async fn poll_connection( + mut conn: tokio_postgres::Connection< + tokio_postgres::Socket, + tokio_postgres::tls::NoTlsStream, + >, + subscriptions: Arc>>, + ) { + 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.lock().await.get(note.channel()).cloned() { + 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 { + tracing::warn!(channel = %note.channel(), "received notification for unknown channel"); } } + Some(std::result::Result::Ok(_)) => { + // Ignore other async messages + } + Some(std::result::Result::Err(err)) => { + tracing::error!(?err, "postgres connection error, reconnecting"); + break; // Exit loop to reconnect + } + None => { + tracing::warn!("postgres connection closed, reconnecting"); + break; // Exit loop to reconnect + } } - tracing::debug!("listen connection closed"); - }); + } + } - Ok(Self { - pool: Arc::new(pool), - client: Arc::new(client), - subscriptions, + /// Wait for the client to be connected + async fn wait_for_client(&self) -> Result> { + let mut ready_rx = self.client_ready.clone(); + tokio::time::timeout(tokio::time::Duration::from_secs(5), async { + loop { + // Subscribe to changed before attempting to access client + let changed_fut = ready_rx.changed(); + + // Check if client is already available + if let Some(client) = self.client.lock().await.clone() { + return Ok(client); + } + + // Wait for change, will return client if exists on next iteration + changed_fut + .await + .map_err(|_| anyhow!("connection lifecycle task ended"))?; + tracing::debug!("client does not exist immediately after receive ready"); + } }) + .await + .map_err(|_| anyhow!("timeout waiting for postgres client connection"))? } fn hash_subject(&self, subject: &str) -> String { @@ -147,7 +270,7 @@ impl PubSubDriver for PostgresDriver { // Check if we already have a subscription for this channel let (rx, drop_guard) = - if let Some(existing_sub) = self.subscriptions.lock().unwrap().get(&hashed).cloned() { + if let Some(existing_sub) = self.subscriptions.lock().await.get(&hashed).cloned() { // Reuse the existing broadcast channel let rx = existing_sub.tx.subscribe(); let drop_guard = existing_sub.token.clone().drop_guard(); @@ -160,13 +283,15 @@ impl PubSubDriver for PostgresDriver { // Register subscription self.subscriptions .lock() - .unwrap() + .await .insert(hashed.clone(), subscription.clone()); // Execute LISTEN command on the async client (for receiving notifications) // This only needs to be done once per channel + // Wait for client to be connected with retry logic + let client = self.wait_for_client().await?; let span = tracing::trace_span!("pg_listen"); - self.client + client .execute(&format!("LISTEN \"{hashed}\""), &[]) .instrument(span) .await?; @@ -179,13 +304,16 @@ impl PubSubDriver for PostgresDriver { tokio::spawn(async move { token_clone.cancelled().await; if tx_clone.receiver_count() == 0 { - let sql = format!("UNLISTEN \"{}\"", hashed_clone); - if let Err(err) = driver.client.execute(sql.as_str(), &[]).await { - tracing::warn!(?err, %hashed_clone, "failed to UNLISTEN channel"); - } else { - tracing::trace!(%hashed_clone, "unlistened channel"); + let client = driver.client.lock().await.clone(); + if let Some(client) = client { + let sql = format!("UNLISTEN \"{}\"", hashed_clone); + if let Err(err) = client.execute(sql.as_str(), &[]).await { + tracing::warn!(?err, %hashed_clone, "failed to UNLISTEN channel"); + } else { + tracing::trace!(%hashed_clone, "unlistened channel"); + } } - driver.subscriptions.lock().unwrap().remove(&hashed_clone); + driver.subscriptions.lock().await.remove(&hashed_clone); } }); diff --git a/packages/common/universalpubsub/tests/reconnect.rs b/packages/common/universalpubsub/tests/reconnect.rs new file mode 100644 index 0000000000..fa98767e95 --- /dev/null +++ b/packages/common/universalpubsub/tests/reconnect.rs @@ -0,0 +1,217 @@ +use anyhow::*; +use rivet_test_deps_docker::{TestDatabase, TestPubSub}; +use std::{sync::Arc, time::Duration}; +use universalpubsub::{NextOutput, PubSub, PublishOpts}; +use uuid::Uuid; + +fn setup_logging() { + let _ = tracing_subscriber::fmt() + .with_env_filter("debug") + .with_ansi(false) + .with_test_writer() + .try_init(); +} + +#[tokio::test] +async fn test_nats_driver_with_memory_reconnect() { + setup_logging(); + + let test_id = Uuid::new_v4(); + let (pubsub_config, docker_config) = TestPubSub::Nats.config(test_id, 1).await.unwrap(); + let mut docker = docker_config.unwrap(); + docker.start().await.unwrap(); + tokio::time::sleep(Duration::from_secs(1)).await; + + let rivet_config::config::PubSub::Nats(nats) = pubsub_config else { + unreachable!(); + }; + + use std::str::FromStr; + let server_addrs = nats + .addresses + .iter() + .map(|addr| format!("nats://{addr}")) + .map(|url| async_nats::ServerAddr::from_str(url.as_ref())) + .collect::, _>>() + .unwrap(); + + let driver = universalpubsub::driver::nats::NatsDriver::connect( + async_nats::ConnectOptions::new(), + &server_addrs[..], + ) + .await + .unwrap(); + let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), true); + + test_reconnect_inner(&pubsub, &docker).await; +} + +#[tokio::test] +async fn test_nats_driver_without_memory_reconnect() { + setup_logging(); + + let test_id = Uuid::new_v4(); + let (pubsub_config, docker_config) = TestPubSub::Nats.config(test_id, 1).await.unwrap(); + let mut docker = docker_config.unwrap(); + docker.start().await.unwrap(); + tokio::time::sleep(Duration::from_secs(1)).await; + + let rivet_config::config::PubSub::Nats(nats) = pubsub_config else { + unreachable!(); + }; + + use std::str::FromStr; + let server_addrs = nats + .addresses + .iter() + .map(|addr| format!("nats://{addr}")) + .map(|url| async_nats::ServerAddr::from_str(url.as_ref())) + .collect::, _>>() + .unwrap(); + + let driver = universalpubsub::driver::nats::NatsDriver::connect( + async_nats::ConnectOptions::new(), + &server_addrs[..], + ) + .await + .unwrap(); + let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false); + + test_reconnect_inner(&pubsub, &docker).await; +} + +#[tokio::test] +async fn test_postgres_driver_with_memory_reconnect() { + setup_logging(); + + let test_id = Uuid::new_v4(); + let (db_config, docker_config) = TestDatabase::Postgres.config(test_id, 1).await.unwrap(); + let mut docker = docker_config.unwrap(); + docker.start().await.unwrap(); + tokio::time::sleep(Duration::from_secs(5)).await; + + let rivet_config::config::Database::Postgres(pg) = db_config else { + unreachable!(); + }; + let url = pg.url.read().clone(); + + let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, true) + .await + .unwrap(); + let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), true); + + test_reconnect_inner(&pubsub, &docker).await; +} + +#[tokio::test] +async fn test_postgres_driver_without_memory_reconnect() { + setup_logging(); + + let test_id = Uuid::new_v4(); + let (db_config, docker_config) = TestDatabase::Postgres.config(test_id, 1).await.unwrap(); + let mut docker = docker_config.unwrap(); + docker.start().await.unwrap(); + tokio::time::sleep(Duration::from_secs(5)).await; + + let rivet_config::config::Database::Postgres(pg) = db_config else { + unreachable!(); + }; + let url = pg.url.read().clone(); + + let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, false) + .await + .unwrap(); + let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false); + + test_reconnect_inner(&pubsub, &docker).await; +} + +async fn test_reconnect_inner(pubsub: &PubSub, docker: &rivet_test_deps_docker::DockerRunConfig) { + tracing::info!("testing reconnect functionality"); + + // Open subscription + let mut subscriber = pubsub.subscribe("test.reconnect").await.unwrap(); + tracing::info!("opened initial subscription"); + + // Test publish/receive message before restart + let message_before = b"message before restart"; + pubsub + .publish("test.reconnect", message_before, PublishOpts::broadcast()) + .await + .unwrap(); + pubsub.flush().await.unwrap(); + + match subscriber.next().await.unwrap() { + NextOutput::Message(msg) => { + assert_eq!( + msg.payload, message_before, + "message before restart should match" + ); + tracing::info!("received message before restart"); + } + NextOutput::Unsubscribed => { + panic!("unexpected unsubscribe before restart"); + } + } + + // Restart container + tracing::info!("restarting docker container"); + docker.restart().await.unwrap(); + + // Give the service time to come back up + tokio::time::sleep(Duration::from_secs(3)).await; + tracing::info!("docker container restarted"); + + // Test publish/receive message after restart + let message_after = b"message after restart"; + + // Retry logic for publish after restart since connection might need to reconnect + let mut retries = 0; + const MAX_RETRIES: u32 = 10; + loop { + match pubsub + .publish("test.reconnect", message_after, PublishOpts::broadcast()) + .await + { + Result::Ok(_) => { + tracing::info!("published message after restart"); + break; + } + Result::Err(e) if retries < MAX_RETRIES => { + retries += 1; + tracing::debug!(?e, retries, "failed to publish after restart, retrying"); + tokio::time::sleep(Duration::from_millis(500)).await; + } + Result::Err(e) => { + panic!("failed to publish after {} retries: {}", MAX_RETRIES, e); + } + } + } + + pubsub.flush().await.unwrap(); + + // Try to receive with timeout to handle reconnection delays + let receive_timeout = Duration::from_secs(10); + let receive_result = tokio::time::timeout(receive_timeout, subscriber.next()).await; + + match receive_result { + Result::Ok(Result::Ok(NextOutput::Message(msg))) => { + assert_eq!( + msg.payload, message_after, + "message after restart should match" + ); + tracing::info!("received message after restart - reconnection successful"); + } + Result::Ok(Result::Ok(NextOutput::Unsubscribed)) => { + panic!("unexpected unsubscribe after restart"); + } + Result::Ok(Result::Err(e)) => { + panic!("error receiving message after restart: {}", e); + } + Result::Err(_) => { + panic!("timeout receiving message after restart"); + } + } + + tracing::info!("reconnect test completed successfully"); +} From e29523a2c49904f573638386a1038549be9e368d Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Mon, 8 Sep 2025 20:18:37 -0700 Subject: [PATCH 03/17] chore(ups): handle edge cases with postgres listen/unlisten/notify when disconnected/reconnecting --- out/errors/ups.publish_failed.json | 5 + packages/common/test-deps-docker/src/lib.rs | 72 +++++++ .../src/driver/postgres/mod.rs | 166 +++++++++----- packages/common/universalpubsub/src/errors.rs | 2 + packages/common/universalpubsub/src/pubsub.rs | 31 ++- .../common/universalpubsub/tests/reconnect.rs | 202 +++++++++++++++--- 6 files changed, 395 insertions(+), 83 deletions(-) create mode 100644 out/errors/ups.publish_failed.json diff --git a/out/errors/ups.publish_failed.json b/out/errors/ups.publish_failed.json new file mode 100644 index 0000000000..b4c1ac0b42 --- /dev/null +++ b/out/errors/ups.publish_failed.json @@ -0,0 +1,5 @@ +{ + "code": "publish_failed", + "group": "ups", + "message": "Failed to publish message after retries" +} \ No newline at end of file diff --git a/packages/common/test-deps-docker/src/lib.rs b/packages/common/test-deps-docker/src/lib.rs index 2ed73f8600..644e57ad53 100644 --- a/packages/common/test-deps-docker/src/lib.rs +++ b/packages/common/test-deps-docker/src/lib.rs @@ -109,6 +109,78 @@ impl DockerRunConfig { Ok(()) } + pub async fn stop_container(&self) -> Result<()> { + let container_id = self + .container_id + .as_ref() + .ok_or_else(|| anyhow!("No container ID found, container not started"))?; + + tracing::debug!( + container_name = %self.container_name, + container_id = %container_id, + "stopping docker container" + ); + + let output = Command::new("docker") + .arg("stop") + .arg(container_id) + .output() + .await?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + anyhow::bail!( + "Failed to stop container {}: {}", + self.container_name, + stderr + ); + } + + tracing::debug!( + container_name = %self.container_name, + container_id = %container_id, + "container stopped successfully" + ); + + Ok(()) + } + + pub async fn start_container(&self) -> Result<()> { + let container_id = self + .container_id + .as_ref() + .ok_or_else(|| anyhow!("No container ID found, container not started"))?; + + tracing::debug!( + container_name = %self.container_name, + container_id = %container_id, + "starting docker container" + ); + + let output = Command::new("docker") + .arg("start") + .arg(container_id) + .output() + .await?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + anyhow::bail!( + "Failed to start container {}: {}", + self.container_name, + stderr + ); + } + + tracing::debug!( + container_name = %self.container_name, + container_id = %container_id, + "container started successfully" + ); + + Ok(()) + } + pub fn container_id(&self) -> Option<&str> { self.container_id.as_deref() } diff --git a/packages/common/universalpubsub/src/driver/postgres/mod.rs b/packages/common/universalpubsub/src/driver/postgres/mod.rs index 21cf96f5f5..3115a7ac72 100644 --- a/packages/common/universalpubsub/src/driver/postgres/mod.rs +++ b/packages/common/universalpubsub/src/driver/postgres/mod.rs @@ -112,21 +112,23 @@ impl PostgresDriver { client: Arc>>>, ready_tx: tokio::sync::watch::Sender, ) { - let mut backoff = Backoff::new(8, None, 1_000, 1_000); + let mut backoff = Backoff::default(); loop { match tokio_postgres::connect(&conn_str, tokio_postgres::NoTls).await { Result::Ok((new_client, conn)) => { tracing::info!("postgres listen connection established"); // Reset backoff on successful connection - backoff = Backoff::new(8, None, 1_000, 1_000); + backoff = Backoff::default(); let new_client = Arc::new(new_client); - // Update the client reference immediately - *client.lock().await = Some(new_client.clone()); - // Notify that client is ready - let _ = ready_tx.send(true); + // Spawn the polling task immediately + // This must be done before any operations on the client + let subscriptions_clone = subscriptions.clone(); + let poll_handle = tokio::spawn(async move { + Self::poll_connection(conn, subscriptions_clone).await; + }); // Get channels to re-subscribe to let channels: Vec = @@ -135,38 +137,41 @@ impl PostgresDriver { if needs_resubscribe { tracing::debug!( - ?channels, + channels=?channels.len(), "will re-subscribe to channels after connection starts" ); } - // Spawn a task to re-subscribe after a short delay + // Re-subscribe to channels if needs_resubscribe { - let client_for_resub = new_client.clone(); - let channels_clone = channels.clone(); - tokio::spawn(async move { - tracing::debug!( - ?channels_clone, - "re-subscribing to channels after reconnection" - ); - for channel in &channels_clone { - if let Result::Err(e) = client_for_resub - .execute(&format!("LISTEN \"{}\"", channel), &[]) - .await - { - tracing::error!(?e, %channel, "failed to re-subscribe to channel"); - } else { - tracing::debug!(%channel, "successfully re-subscribed to channel"); - } + tracing::debug!( + channels=?channels.len(), + "re-subscribing to channels after reconnection" + ); + for channel in &channels { + tracing::info!(?channel, "re-subscribing to channel"); + if let Result::Err(e) = new_client + .execute(&format!("LISTEN \"{}\"", channel), &[]) + .await + { + tracing::error!(?e, %channel, "failed to re-subscribe to channel"); + } else { + tracing::debug!(%channel, "successfully re-subscribed to channel"); } - }); + } } - // Poll the connection until it closes - Self::poll_connection(conn, subscriptions.clone()).await; + // Update the client reference and signal ready + // Do this AFTER re-subscribing to ensure LISTEN is complete + *client.lock().await = Some(new_client.clone()); + let _ = ready_tx.send(true); + + // Wait for the polling task to complete (when the connection closes) + let _ = poll_handle.await; // Clear the client reference on disconnect *client.lock().await = None; + // Notify that client is disconnected let _ = ready_tx.send(false); } @@ -208,12 +213,12 @@ impl PostgresDriver { // Ignore other async messages } Some(std::result::Result::Err(err)) => { - tracing::error!(?err, "postgres connection error, reconnecting"); - break; // Exit loop to reconnect + tracing::error!(?err, "postgres connection error"); + break; } None => { - tracing::warn!("postgres connection closed, reconnecting"); - break; // Exit loop to reconnect + tracing::warn!("postgres connection closed"); + break; } } } @@ -224,19 +229,16 @@ impl PostgresDriver { let mut ready_rx = self.client_ready.clone(); tokio::time::timeout(tokio::time::Duration::from_secs(5), async { loop { - // Subscribe to changed before attempting to access client - let changed_fut = ready_rx.changed(); - // Check if client is already available if let Some(client) = self.client.lock().await.clone() { return Ok(client); } - // Wait for change, will return client if exists on next iteration - changed_fut + // Wait for the ready signal to change + ready_rx + .changed() .await .map_err(|_| anyhow!("connection lifecycle task ended"))?; - tracing::debug!("client does not exist immediately after receive ready"); } }) .await @@ -288,13 +290,25 @@ impl PubSubDriver for PostgresDriver { // Execute LISTEN command on the async client (for receiving notifications) // This only needs to be done once per channel - // Wait for client to be connected with retry logic - let client = self.wait_for_client().await?; - let span = tracing::trace_span!("pg_listen"); - client - .execute(&format!("LISTEN \"{hashed}\""), &[]) - .instrument(span) - .await?; + // 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.clone() { + let span = tracing::trace_span!("pg_listen"); + match client + .execute(&format!("LISTEN \"{hashed}\""), &[]) + .instrument(span) + .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"); + } // Spawn a single cleanup task for this subscription waiting on its token let driver = self.clone(); @@ -333,14 +347,66 @@ impl PubSubDriver for PostgresDriver { // Encode payload to base64 and send NOTIFY let encoded = BASE64.encode(payload); - let conn = self.pool.get().await?; let hashed = self.hash_subject(subject); - let span = tracing::trace_span!("pg_notify"); - conn.execute(&format!("NOTIFY \"{hashed}\", '{encoded}'"), &[]) - .instrument(span) - .await?; - Ok(()) + tracing::debug!("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.lock().await.contains_key(&hashed) { + self.wait_for_client().await?; + } + + // Retry getting a connection from the pool with backoff in case the connection is + // currently disconnected + let mut backoff = Backoff::default(); + let mut last_error = None; + + 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, use it for NOTIFY + let span = tracing::trace_span!("pg_notify"); + match conn + .execute(&format!("NOTIFY \"{hashed}\", '{encoded}'"), &[]) + .instrument(span) + .await + { + Result::Ok(_) => return Ok(()), + Result::Err(e) => { + tracing::debug!( + ?e, + "NOTIFY failed, retrying with new connection" + ); + last_error = Some(e.into()); + } + } + } + Result::Err(e) => { + tracing::debug!( + ?e, + "connection test failed, retrying with new connection" + ); + last_error = Some(e.into()); + } + } + } + Result::Err(e) => { + tracing::debug!(?e, "failed to get connection from pool, retrying"); + last_error = Some(e.into()); + } + } + + // Check if we should continue retrying + if !backoff.tick().await { + return Err( + last_error.unwrap_or_else(|| anyhow!("failed to publish after retries")) + ); + } + } } async fn flush(&self) -> Result<()> { diff --git a/packages/common/universalpubsub/src/errors.rs b/packages/common/universalpubsub/src/errors.rs index afab64db4a..69849d4592 100644 --- a/packages/common/universalpubsub/src/errors.rs +++ b/packages/common/universalpubsub/src/errors.rs @@ -6,4 +6,6 @@ use serde::{Deserialize, Serialize}; pub enum Ups { #[error("request_timeout", "Request timeout.")] RequestTimeout, + #[error("publish_failed", "Failed to publish message after retries")] + PublishFailed, } diff --git a/packages/common/universalpubsub/src/pubsub.rs b/packages/common/universalpubsub/src/pubsub.rs index fd24e41ff2..ac6ff27696 100644 --- a/packages/common/universalpubsub/src/pubsub.rs +++ b/packages/common/universalpubsub/src/pubsub.rs @@ -8,6 +8,8 @@ use tokio::sync::broadcast; use tokio::sync::{RwLock, oneshot}; use uuid::Uuid; +use rivet_util::backoff::Backoff; + use crate::chunking::{ChunkTracker, encode_chunk, split_payload_into_chunks}; use crate::driver::{PubSubDriverHandle, PublishOpts, SubscriberDriverHandle}; @@ -131,7 +133,8 @@ impl PubSub { break; } } else { - self.driver.publish(subject, &encoded).await?; + // Use backoff when publishing through the driver + self.publish_with_backoff(subject, &encoded).await?; } } Ok(()) @@ -174,7 +177,26 @@ impl PubSub { break; } } else { - self.driver.publish(subject, &encoded).await?; + // Use backoff when publishing through the driver + self.publish_with_backoff(subject, &encoded).await?; + } + } + Ok(()) + } + + async fn publish_with_backoff(&self, subject: &str, encoded: &[u8]) -> Result<()> { + let mut backoff = Backoff::default(); + loop { + match self.driver.publish(subject, encoded).await { + Result::Ok(_) => break, + Err(err) if !backoff.tick().await => { + tracing::info!(?err, "error publishing, cannot retry again"); + return Err(crate::errors::Ups::PublishFailed.build().into()); + } + Err(err) => { + tracing::info!(?err, "error publishing, retrying"); + // Continue retrying + } } } Ok(()) @@ -293,7 +315,10 @@ impl Subscriber { pub async fn next(&mut self) -> Result { loop { match self.driver.next().await? { - DriverOutput::Message { subject, payload } => { + DriverOutput::Message { + subject: _, + payload, + } => { // Process chunks let mut tracker = self.pubsub.chunk_tracker.lock().unwrap(); match tracker.process_chunk(&payload) { diff --git a/packages/common/universalpubsub/tests/reconnect.rs b/packages/common/universalpubsub/tests/reconnect.rs index fa98767e95..230110789b 100644 --- a/packages/common/universalpubsub/tests/reconnect.rs +++ b/packages/common/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_reconnect_inner(&pubsub, &docker).await; + test_all_inner(&pubsub, &docker).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_reconnect_inner(&pubsub, &docker).await; + test_all_inner(&pubsub, &docker).await; } #[tokio::test] @@ -100,7 +100,7 @@ async fn test_postgres_driver_with_memory_reconnect() { .unwrap(); let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), true); - test_reconnect_inner(&pubsub, &docker).await; + test_all_inner(&pubsub, &docker).await; } #[tokio::test] @@ -123,7 +123,13 @@ async fn test_postgres_driver_without_memory_reconnect() { .unwrap(); let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false); + test_all_inner(&pubsub, &docker).await; +} + +async fn test_all_inner(pubsub: &PubSub, docker: &rivet_test_deps_docker::DockerRunConfig) { test_reconnect_inner(&pubsub, &docker).await; + test_publish_while_stopped(&pubsub, &docker).await; + test_subscribe_while_stopped(&pubsub, &docker).await; } async fn test_reconnect_inner(pubsub: &PubSub, docker: &rivet_test_deps_docker::DockerRunConfig) { @@ -158,35 +164,14 @@ async fn test_reconnect_inner(pubsub: &PubSub, docker: &rivet_test_deps_docker:: tracing::info!("restarting docker container"); docker.restart().await.unwrap(); - // Give the service time to come back up - tokio::time::sleep(Duration::from_secs(3)).await; - tracing::info!("docker container restarted"); - // Test publish/receive message after restart + // + // This should retry under the hood, since the container will still be starting let message_after = b"message after restart"; - - // Retry logic for publish after restart since connection might need to reconnect - let mut retries = 0; - const MAX_RETRIES: u32 = 10; - loop { - match pubsub - .publish("test.reconnect", message_after, PublishOpts::broadcast()) - .await - { - Result::Ok(_) => { - tracing::info!("published message after restart"); - break; - } - Result::Err(e) if retries < MAX_RETRIES => { - retries += 1; - tracing::debug!(?e, retries, "failed to publish after restart, retrying"); - tokio::time::sleep(Duration::from_millis(500)).await; - } - Result::Err(e) => { - panic!("failed to publish after {} retries: {}", MAX_RETRIES, e); - } - } - } + pubsub + .publish("test.reconnect", message_after, PublishOpts::broadcast()) + .await + .unwrap(); pubsub.flush().await.unwrap(); @@ -215,3 +200,160 @@ async fn test_reconnect_inner(pubsub: &PubSub, docker: &rivet_test_deps_docker:: tracing::info!("reconnect test completed successfully"); } + +async fn test_publish_while_stopped( + pubsub: &PubSub, + docker: &rivet_test_deps_docker::DockerRunConfig, +) { + tracing::info!("testing publish while container stopped"); + + // 1. Subscribe + let mut subscriber = pubsub.subscribe("test.publish_stopped").await.unwrap(); + tracing::info!("opened subscription"); + + // 2. Stop container + tracing::info!("stopping docker container"); + docker.stop_container().await.unwrap(); + tokio::time::sleep(Duration::from_secs(2)).await; + + // 3. Publish while stopped (should queue/retry) + let message = b"message while stopped"; + let publish_handle = tokio::spawn({ + let pubsub = pubsub.clone(); + let message = message.to_vec(); + async move { + pubsub + .publish("test.publish_stopped", &message, PublishOpts::broadcast()) + .await + } + }); + + // 4. Start container + tokio::time::sleep(Duration::from_secs(3)).await; + tracing::info!("starting docker container"); + docker.start_container().await.unwrap(); + tokio::time::sleep(Duration::from_secs(5)).await; + + // Wait for publish to complete + publish_handle.await.unwrap().unwrap(); + pubsub.flush().await.unwrap(); + + // 5. Receive message + tracing::info!("waiting for message"); + let receive_timeout = Duration::from_secs(5); + let receive_result = tokio::time::timeout(receive_timeout, subscriber.next()).await; + + match receive_result { + Result::Ok(Result::Ok(NextOutput::Message(msg))) => { + assert_eq!( + msg.payload, message, + "message published while stopped should be received" + ); + tracing::info!("received message published while stopped - reconnection successful"); + } + Result::Ok(Result::Ok(NextOutput::Unsubscribed)) => { + panic!("unexpected unsubscribe"); + } + Result::Ok(Result::Err(e)) => { + panic!("error receiving message: {}", e); + } + Result::Err(_) => { + panic!("timeout receiving message"); + } + } + + tracing::info!("publish while stopped test completed successfully"); +} + +async fn test_subscribe_while_stopped( + pubsub: &PubSub, + docker: &rivet_test_deps_docker::DockerRunConfig, +) { + tracing::info!("testing subscribe while container stopped"); + + // 1. Subscribe & test publish & unsubscribe + let mut subscriber = pubsub.subscribe("test.subscribe_stopped").await.unwrap(); + tracing::info!("opened initial subscription"); + + let test_message = b"test message"; + pubsub + .publish( + "test.subscribe_stopped", + test_message, + PublishOpts::broadcast(), + ) + .await + .unwrap(); + pubsub.flush().await.unwrap(); + + match subscriber.next().await.unwrap() { + NextOutput::Message(msg) => { + assert_eq!(msg.payload, test_message, "initial message should match"); + tracing::info!("received initial test message"); + } + NextOutput::Unsubscribed => { + panic!("unexpected unsubscribe"); + } + } + + drop(subscriber); // Drop to unsubscribe + tracing::info!("unsubscribed from initial subscription"); + + // 2. Stop container + tracing::info!("stopping docker container"); + docker.stop_container().await.unwrap(); + tokio::time::sleep(Duration::from_secs(2)).await; + + // 3. Subscribe while stopped + let subscribe_handle = tokio::spawn({ + let pubsub = pubsub.clone(); + async move { pubsub.subscribe("test.subscribe_stopped").await } + }); + + // 4. Start container + tokio::time::sleep(Duration::from_secs(3)).await; + tracing::info!("starting docker container"); + docker.start_container().await.unwrap(); + tokio::time::sleep(Duration::from_secs(5)).await; + + // Wait for subscription to complete + let mut new_subscriber = subscribe_handle.await.unwrap().unwrap(); + tracing::info!("new subscription established after reconnect"); + + // 5. Publish message + let final_message = b"message after reconnect"; + pubsub + .publish( + "test.subscribe_stopped", + final_message, + PublishOpts::broadcast(), + ) + .await + .unwrap(); + pubsub.flush().await.unwrap(); + + // 6. Receive + let receive_timeout = Duration::from_secs(10); + let receive_result = tokio::time::timeout(receive_timeout, new_subscriber.next()).await; + + match receive_result { + Result::Ok(Result::Ok(NextOutput::Message(msg))) => { + assert_eq!( + msg.payload, final_message, + "message after reconnect should match" + ); + tracing::info!("received message after reconnect - subscription successful"); + } + Result::Ok(Result::Ok(NextOutput::Unsubscribed)) => { + panic!("unexpected unsubscribe"); + } + Result::Ok(Result::Err(e)) => { + panic!("error receiving message: {}", e); + } + Result::Err(_) => { + panic!("timeout receiving message"); + } + } + + tracing::info!("subscribe while stopped test completed successfully"); +} From d4b39b34edd2b6e9d6703c602c51930520931e67 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Mon, 8 Sep 2025 20:42:18 -0700 Subject: [PATCH 04/17] ci: switch to depot --- .github/workflows/release.yaml | 14 +++++++------- .github/workflows/rust.yml | 6 +++--- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 961fbcbfb1..ce35afb4d2 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -79,22 +79,22 @@ jobs: matrix: include: - platform: linux - runner: ubuntu-latest + runner: depot-ubuntu-24.04 target: x86_64-unknown-linux-musl binary_ext: "" arch: x86_64 - platform: windows - runner: ubuntu-latest + runner: depot-ubuntu-24.04 target: x86_64-pc-windows-gnu binary_ext: ".exe" arch: x86_64 - platform: macos - runner: ubuntu-latest + runner: depot-ubuntu-24.04 target: x86_64-apple-darwin binary_ext: "" arch: x86_64 - platform: macos - runner: ubuntu-latest + runner: depot-ubuntu-24.04 target: aarch64-apple-darwin binary_ext: "" arch: aarch64 @@ -155,10 +155,10 @@ jobs: include: # TODO(RVT-4479): Add back ARM builder once manifest generation fixed # - platform: linux/arm64 - # runner: ubuntu-latest + # runner: depot-ubuntu-24.04 # arch_suffix: -arm64 - platform: linux/x86_64 - runner: ubuntu-latest + runner: depot-ubuntu-24.04 # TODO: Replace with appropriate arch_suffix when needed # arch_suffix: -amd64 arch_suffix: '' @@ -246,4 +246,4 @@ jobs: ./scripts/release/main.ts --version "${{ github.event.inputs.version }}" --completeCi else ./scripts/release/main.ts --version "${{ github.event.inputs.version }}" --no-latest --completeCi - fi \ No newline at end of file + fi diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 1671a3fc41..745f5ae9cd 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -43,7 +43,7 @@ jobs: # clippy: # name: Clippy - # runs-on: ubuntu-latest + # runs-on: depot-ubuntu-24.04 # steps: # - uses: actions/checkout@v4 @@ -59,7 +59,7 @@ jobs: check: name: Check - runs-on: ubuntu-latest + runs-on: depot-ubuntu-24.04 steps: - uses: actions/checkout@v4 @@ -77,7 +77,7 @@ jobs: test: name: Test - runs-on: ubuntu-latest + runs-on: depot-ubuntu-24.04 steps: - uses: actions/checkout@v4 From e0afd30b73ad2137e587f0412a6f80c30e19281d Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Tue, 9 Sep 2025 09:41:07 -0700 Subject: [PATCH 05/17] chore(runner): handle async actor stop --- sdks/typescript/runner/src/mod.ts | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/sdks/typescript/runner/src/mod.ts b/sdks/typescript/runner/src/mod.ts index 7d26a37925..2623d63848 100644 --- a/sdks/typescript/runner/src/mod.ts +++ b/sdks/typescript/runner/src/mod.ts @@ -127,7 +127,7 @@ export class Runner { // The server will send a StopActor command if it wants to fully stop } - stopActor(actorId: string, generation?: number) { + async stopActor(actorId: string, generation?: number) { const actor = this.#removeActor(actorId, generation); if (!actor) return; @@ -136,11 +136,14 @@ export class Runner { this.#tunnel.unregisterActor(actor); } - this.#sendActorStateUpdate(actorId, actor.generation, "stopped"); - - this.#config.onActorStop(actorId, actor.generation).catch((err) => { + // If onActorStop times out, Pegboard will handle this timeout with ACTOR_STOP_THRESHOLD_DURATION_MS + try { + await this.#config.onActorStop(actorId, actor.generation); + } catch (err) { console.error(`Error in onActorStop for actor ${actorId}:`, err); - }); + } + + this.#sendActorStateUpdate(actorId, actor.generation, "stopped"); } #stopAllActors() { From c70a66726194a3488bdf36fdcbb3d3f988e9a4c5 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Fri, 12 Sep 2025 13:22:24 -0700 Subject: [PATCH 06/17] fix(epoxy): fix Any quorum type not reaching any node --- packages/services/epoxy/src/http_client.rs | 15 ++++++++++++--- packages/services/epoxy/src/utils.rs | 2 +- packages/services/epoxy/tests/proposal.rs | 1 - 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/packages/services/epoxy/src/http_client.rs b/packages/services/epoxy/src/http_client.rs index bfac730d8a..51864305c6 100644 --- a/packages/services/epoxy/src/http_client.rs +++ b/packages/services/epoxy/src/http_client.rs @@ -54,12 +54,21 @@ where ) .collect::>() .await; + tracing::info!(?quorum_size, len = ?responses.len(), ?quorum_type, "fanout quorum size"); + + // Choow how many successful responses we need before considering a success + let target_responses = match quorum_type { + // Only require 1 response + utils::QuorumType::Any => 1, + // Include all responses + utils::QuorumType::All => responses.len(), + // Subtract 1 from quorum size since we're not counting ourselves + utils::QuorumType::Fast | utils::QuorumType::Slow => quorum_size - 1, + }; // Collect responses until we reach quorum or all futures complete - // - // Subtract 1 from quorum size since we're not counting ourselves let mut successful_responses = Vec::new(); - while successful_responses.len() < quorum_size - 1 { + while successful_responses.len() < target_responses { if let Some(response) = responses.next().await { match response { std::result::Result::Ok(result) => match result { diff --git a/packages/services/epoxy/src/utils.rs b/packages/services/epoxy/src/utils.rs index d7a6cf44e7..b6896f1a08 100644 --- a/packages/services/epoxy/src/utils.rs +++ b/packages/services/epoxy/src/utils.rs @@ -2,7 +2,7 @@ use anyhow::*; use epoxy_protocol::protocol::{self, ReplicaId}; use universaldb::Transaction; -#[derive(Clone, Copy)] +#[derive(Clone, Copy, Debug)] pub enum QuorumType { Fast, Slow, diff --git a/packages/services/epoxy/tests/proposal.rs b/packages/services/epoxy/tests/proposal.rs index f14ab6cd98..f297dfbc40 100644 --- a/packages/services/epoxy/tests/proposal.rs +++ b/packages/services/epoxy/tests/proposal.rs @@ -4,7 +4,6 @@ use common::THREE_REPLICAS; use epoxy::ops::propose::ProposalResult; use epoxy_protocol::protocol; use gas::prelude::*; -use rivet_acl::{Verifier, config::AclConfig}; use rivet_api_builder::{ApiCtx, GlobalApiCtx}; use rivet_util::Id; From 9549b26deb5716f56661b6ef7af2cca1d829e8a1 Mon Sep 17 00:00:00 2001 From: Kacper Wojciechowski <39823706+jog1t@users.noreply.github.com> Date: Thu, 11 Sep 2025 23:28:49 +0200 Subject: [PATCH 07/17] fix(engine/fe): remove addresses --- .../routes/_layout/ns.$namespace/runners.tsx | 39 ------------------- 1 file changed, 39 deletions(-) diff --git a/frontend/src/routes/_layout/ns.$namespace/runners.tsx b/frontend/src/routes/_layout/ns.$namespace/runners.tsx index 21474beec0..9969231100 100644 --- a/frontend/src/routes/_layout/ns.$namespace/runners.tsx +++ b/frontend/src/routes/_layout/ns.$namespace/runners.tsx @@ -71,7 +71,6 @@ function RouteComponent() { ID Name - HTTP Slots Last ping Created @@ -155,19 +154,11 @@ function RowSkeleton() { - - - ); } -const MAX_TO_SHOW = 2; - function Row(runner: Rivet.Runner) { - const [isExpanded, setExpanded] = useState(false); - const addresses = Object.values(runner.addressesHttp); - return ( @@ -186,36 +177,6 @@ function Row(runner: Rivet.Runner) { - -
- {addresses - .slice(0, isExpanded ? addresses.length : MAX_TO_SHOW) - .map((http) => { - const address = `${http.hostname}:${http.port}`; - return ( - - {address} - - ); - })} - - {addresses.length > MAX_TO_SHOW && !isExpanded ? ( - - ) : null} -
- - {runner.remainingSlots}/{runner.totalSlots} From 5ddcc8537f73ebac4e53f3a61182f16c54f2fc0f Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Fri, 5 Sep 2025 13:09:49 -0700 Subject: [PATCH 08/17] feat(pegboard): outbound runners --- Cargo.lock | 143 +++++-- Cargo.toml | 16 +- docker/dev/docker-compose.yml | 1 + out/openapi.json | 70 +++- packages/common/api-builder/Cargo.toml | 1 + packages/common/api-builder/src/wrappers.rs | 3 +- .../common/gasoline/core/src/utils/tags.rs | 10 + packages/common/pools/src/reqwest.rs | 7 + packages/common/types/Cargo.toml | 2 +- packages/common/types/src/runners.rs | 2 +- packages/common/udb-util/src/keys.rs | 4 +- packages/core/api-peer/src/namespaces.rs | 16 +- packages/core/guard/server/Cargo.toml | 2 +- packages/core/pegboard-outbound/Cargo.toml | 20 + packages/core/pegboard-outbound/src/lib.rs | 288 ++++++++++++++ packages/infra/engine/Cargo.toml | 9 +- packages/infra/engine/src/run_config.rs | 5 + packages/services/namespace/Cargo.toml | 6 +- packages/services/namespace/src/keys.rs | 50 +++ .../services/namespace/src/ops/get_global.rs | 58 ++- .../services/namespace/src/ops/get_local.rs | 39 +- packages/services/namespace/src/types.rs | 54 +++ .../namespace/src/workflows/namespace.rs | 57 +-- packages/services/pegboard/Cargo.toml | 2 +- .../services/pegboard/src/keys/datacenter.rs | 249 ------------ packages/services/pegboard/src/keys/mod.rs | 1 - packages/services/pegboard/src/keys/ns.rs | 357 +++++++++++++++++- packages/services/pegboard/src/keys/runner.rs | 16 +- packages/services/pegboard/src/lib.rs | 1 + packages/services/pegboard/src/messages.rs | 4 + .../services/pegboard/src/ops/actor/create.rs | 6 +- .../pegboard/src/ops/actor/get_for_key.rs | 4 +- .../pegboard/src/ops/actor/list_names.rs | 2 +- .../services/pegboard/src/ops/runner/get.rs | 2 +- .../src/ops/runner/update_alloc_idx.rs | 8 +- .../src/workflows/actor/actor_keys.rs | 2 +- .../pegboard/src/workflows/actor/destroy.rs | 21 +- .../pegboard/src/workflows/actor/mod.rs | 46 ++- .../pegboard/src/workflows/actor/runtime.rs | 95 +++-- .../pegboard/src/workflows/actor/setup.rs | 22 +- .../services/pegboard/src/workflows/runner.rs | 32 +- pnpm-lock.yaml | 3 + sdks/rust/{key-data => data}/Cargo.toml | 2 +- sdks/rust/{key-data => data}/build.rs | 2 +- sdks/rust/{key-data => data}/src/converted.rs | 8 +- sdks/rust/{key-data => data}/src/generated.rs | 0 sdks/rust/{key-data => data}/src/lib.rs | 2 +- sdks/rust/{key-data => data}/src/versioned.rs | 40 +- .../data/namespace.runner_kind.v1.bare | 14 + .../pegboard.namespace.actor_by_key.v1.bare | 0 .../pegboard.namespace.actor_name.v1.bare | 0 ...gboard.namespace.runner_alloc_idx.v1.bare} | 0 .../pegboard.namespace.runner_by_key.v1.bare | 0 .../pegboard.runner.address.v1.bare | 0 .../pegboard.runner.metadata.v1.bare | 0 sdks/typescript/runner/src/mod.ts | 7 +- sdks/typescript/test-runner/package.json | 3 +- sdks/typescript/test-runner/src/main.ts | 240 +++++++----- 58 files changed, 1432 insertions(+), 622 deletions(-) create mode 100644 packages/core/pegboard-outbound/Cargo.toml create mode 100644 packages/core/pegboard-outbound/src/lib.rs delete mode 100644 packages/services/pegboard/src/keys/datacenter.rs create mode 100644 packages/services/pegboard/src/messages.rs rename sdks/rust/{key-data => data}/Cargo.toml (95%) rename sdks/rust/{key-data => data}/build.rs (99%) rename sdks/rust/{key-data => data}/src/converted.rs (94%) rename sdks/rust/{key-data => data}/src/generated.rs (100%) rename sdks/rust/{key-data => data}/src/lib.rs (84%) rename sdks/rust/{key-data => data}/src/versioned.rs (83%) create mode 100644 sdks/schemas/data/namespace.runner_kind.v1.bare rename sdks/schemas/{key-data => data}/pegboard.namespace.actor_by_key.v1.bare (100%) rename sdks/schemas/{key-data => data}/pegboard.namespace.actor_name.v1.bare (100%) rename sdks/schemas/{key-data/pegboard.datacenter.runner_alloc_idx.v1.bare => data/pegboard.namespace.runner_alloc_idx.v1.bare} (100%) rename sdks/schemas/{key-data => data}/pegboard.namespace.runner_by_key.v1.bare (100%) rename sdks/schemas/{key-data => data}/pegboard.runner.address.v1.bare (100%) rename sdks/schemas/{key-data => data}/pegboard.runner.metadata.v1.bare (100%) diff --git a/Cargo.lock b/Cargo.lock index 44bd7cd6e0..f1f2a7a52d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -377,6 +377,31 @@ dependencies = [ "tracing", ] +[[package]] +name = "axum-extra" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45bf463831f5131b7d3c756525b305d40f1185b688565648a92e1392ca35713d" +dependencies = [ + "axum 0.8.4", + "axum-core 0.5.2", + "bytes", + "form_urlencoded", + "futures-util", + "http 1.3.1", + "http-body 1.0.1", + "http-body-util", + "mime", + "pin-project-lite", + "rustversion", + "serde", + "serde_html_form", + "serde_path_to_error", + "tower 0.5.2", + "tower-layer", + "tower-service", +] + [[package]] name = "axum-test" version = "17.3.0" @@ -1454,6 +1479,17 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "eventsource-stream" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74fef4569247a5f429d9156b9d0a2599914385dd189c539334c625d8099d90ab" +dependencies = [ + "futures-core", + "nom", + "pin-project-lite", +] + [[package]] name = "fallible-iterator" version = "0.2.0" @@ -2737,6 +2773,7 @@ dependencies = [ "gasoline", "rivet-api-builder", "rivet-api-util", + "rivet-data", "rivet-error", "rivet-util", "serde", @@ -2744,6 +2781,7 @@ dependencies = [ "udb-util", "universaldb", "utoipa", + "versioned-data-util", ] [[package]] @@ -3217,8 +3255,8 @@ dependencies = [ "rivet-api-client", "rivet-api-types", "rivet-api-util", + "rivet-data", "rivet-error", - "rivet-key-data", "rivet-metrics", "rivet-runner-protocol", "rivet-types", @@ -3277,6 +3315,23 @@ dependencies = [ "versioned-data-util", ] +[[package]] +name = "pegboard-outbound" +version = "0.0.1" +dependencies = [ + "anyhow", + "epoxy", + "gasoline", + "namespace", + "pegboard", + "reqwest-eventsource", + "rivet-config", + "rivet-runner-protocol", + "tracing", + "udb-util", + "universaldb", +] + [[package]] name = "pegboard-runner-ws" version = "0.0.1" @@ -3960,16 +4015,34 @@ dependencies = [ "tokio", "tokio-native-tls", "tokio-rustls 0.26.2", + "tokio-util", "tower 0.5.2", "tower-http", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", "webpki-roots 1.0.2", ] +[[package]] +name = "reqwest-eventsource" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "632c55746dbb44275691640e7b40c907c16a2dc1a5842aa98aaec90da6ec6bde" +dependencies = [ + "eventsource-stream", + "futures-core", + "futures-timer", + "mime", + "nom", + "pin-project-lite", + "reqwest", + "thiserror 1.0.69", +] + [[package]] name = "reserve-port" version = "2.3.0" @@ -3999,6 +4072,7 @@ version = "0.0.1" dependencies = [ "anyhow", "axum 0.8.4", + "axum-extra", "axum-test", "chrono", "gasoline", @@ -4204,6 +4278,24 @@ dependencies = [ "uuid", ] +[[package]] +name = "rivet-data" +version = "0.0.1" +dependencies = [ + "anyhow", + "bare_gen", + "gasoline", + "indoc", + "prettyplease", + "rivet-runner-protocol", + "rivet-util", + "serde", + "serde_bare", + "serde_json", + "syn 2.0.104", + "versioned-data-util", +] + [[package]] name = "rivet-dump-openapi" version = "0.0.1" @@ -4231,6 +4323,7 @@ dependencies = [ "lz4_flex", "namespace", "pegboard", + "pegboard-outbound", "pegboard-runner-ws", "portpicker", "rand 0.8.5", @@ -4327,9 +4420,9 @@ dependencies = [ "rivet-api-public", "rivet-cache", "rivet-config", + "rivet-data", "rivet-error", "rivet-guard-core", - "rivet-key-data", "rivet-logs", "rivet-metrics", "rivet-pools", @@ -4392,24 +4485,6 @@ dependencies = [ "uuid", ] -[[package]] -name = "rivet-key-data" -version = "0.0.1" -dependencies = [ - "anyhow", - "bare_gen", - "gasoline", - "indoc", - "prettyplease", - "rivet-runner-protocol", - "rivet-util", - "serde", - "serde_bare", - "serde_json", - "syn 2.0.104", - "versioned-data-util", -] - [[package]] name = "rivet-logs" version = "0.0.1" @@ -4615,7 +4690,7 @@ dependencies = [ "anyhow", "gasoline", "rivet-api-builder", - "rivet-key-data", + "rivet-data", "rivet-runner-protocol", "rivet-util", "serde", @@ -5215,6 +5290,19 @@ dependencies = [ "syn 2.0.104", ] +[[package]] +name = "serde_html_form" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d2de91cf02bbc07cde38891769ccd5d4f073d22a40683aa4bc7a95781aaa2c4" +dependencies = [ + "form_urlencoded", + "indexmap 2.10.0", + "itoa 1.0.15", + "ryu", + "serde", +] + [[package]] name = "serde_json" version = "1.0.141" @@ -6618,6 +6706,19 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "web-sys" version = "0.3.77" diff --git a/Cargo.toml b/Cargo.toml index 08114056cb..cc32db3c24 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [workspace] resolver = "2" -members = ["packages/common/api-builder","packages/common/api-client","packages/common/api-types","packages/common/api-util","packages/common/cache/build","packages/common/cache/result","packages/common/clickhouse-inserter","packages/common/clickhouse-user-query","packages/common/config","packages/common/env","packages/common/error/core","packages/common/error/macros","packages/common/gasoline/core","packages/common/gasoline/macros","packages/common/logs","packages/common/metrics","packages/common/pools","packages/common/runtime","packages/common/service-manager","packages/common/telemetry","packages/common/test-deps","packages/common/test-deps-docker","packages/common/types","packages/common/udb-util","packages/common/universaldb","packages/common/universalpubsub","packages/common/util/core","packages/common/util/id","packages/common/versioned-data-util","packages/core/actor-kv","packages/core/api-peer","packages/core/api-public","packages/core/bootstrap","packages/core/dump-openapi","packages/core/guard/core","packages/core/guard/server","packages/core/pegboard-gateway","packages/core/pegboard-runner-ws","packages/core/pegboard-tunnel","packages/core/workflow-worker","packages/infra/engine","packages/services/epoxy","packages/services/namespace","packages/services/pegboard","sdks/rust/api-full","sdks/rust/bare_gen","sdks/rust/epoxy-protocol","sdks/rust/key-data","sdks/rust/runner-protocol","sdks/rust/tunnel-protocol","sdks/rust/ups-protocol"] +members = ["packages/common/api-builder","packages/common/api-client","packages/common/api-types","packages/common/api-util","packages/common/cache/build","packages/common/cache/result","packages/common/clickhouse-inserter","packages/common/clickhouse-user-query","packages/common/config","packages/common/env","packages/common/error/core","packages/common/error/macros","packages/common/gasoline/core","packages/common/gasoline/macros","packages/common/logs","packages/common/metrics","packages/common/pools","packages/common/runtime","packages/common/service-manager","packages/common/telemetry","packages/common/test-deps","packages/common/test-deps-docker","packages/common/types","packages/common/udb-util","packages/common/universaldb","packages/common/universalpubsub","packages/common/util/core","packages/common/util/id","packages/common/versioned-data-util","packages/core/actor-kv","packages/core/api-peer","packages/core/api-public","packages/core/bootstrap","packages/core/dump-openapi","packages/core/guard/core","packages/core/guard/server","packages/core/pegboard-gateway","packages/core/pegboard-outbound","packages/core/pegboard-runner-ws","packages/core/pegboard-tunnel","packages/core/workflow-worker","packages/infra/engine","packages/services/epoxy","packages/services/namespace","packages/services/pegboard","sdks/rust/api-full","sdks/rust/bare_gen","sdks/rust/data","sdks/rust/epoxy-protocol","sdks/rust/runner-protocol","sdks/rust/tunnel-protocol","sdks/rust/ups-protocol"] [workspace.package] version = "0.0.1" @@ -79,6 +79,7 @@ tracing-core = "0.1" tracing-opentelemetry = "0.29" tracing-slog = "0.2" vergen = "9.0.4" +reqwest-eventsource = "0.6.0" [workspace.dependencies.sentry] version = "0.37.0" @@ -118,6 +119,10 @@ features = ["uuid"] version = "0.8" features = ["http2"] +[workspace.dependencies.axum-extra] +version = "0.10.1" +features = ["query"] + [workspace.dependencies.tower-http] version = "0.6" features = ["cors","trace"] @@ -359,6 +364,9 @@ path = "packages/core/guard/server" [workspace.dependencies.pegboard-gateway] path = "packages/core/pegboard-gateway" +[workspace.dependencies.pegboard-outbound] +path = "packages/core/pegboard-outbound" + [workspace.dependencies.pegboard-runner-ws] path = "packages/core/pegboard-runner-ws" @@ -386,12 +394,12 @@ path = "sdks/rust/api-full" [workspace.dependencies.bare_gen] path = "sdks/rust/bare_gen" +[workspace.dependencies.rivet-data] +path = "sdks/rust/data" + [workspace.dependencies.epoxy-protocol] path = "sdks/rust/epoxy-protocol" -[workspace.dependencies.rivet-key-data] -path = "sdks/rust/key-data" - [workspace.dependencies.rivet-runner-protocol] path = "sdks/rust/runner-protocol" diff --git a/docker/dev/docker-compose.yml b/docker/dev/docker-compose.yml index 6803272ebe..64de5f178d 100644 --- a/docker/dev/docker-compose.yml +++ b/docker/dev/docker-compose.yml @@ -187,6 +187,7 @@ services: environment: - RIVET_ENDPOINT=http://rivet-engine:6420 - RUNNER_HOST=runner + # - NO_AUTOSTART=1 stop_grace_period: 4s ports: - '5050:5050' diff --git a/out/openapi.json b/out/openapi.json index 87571399cc..ff97997a91 100644 --- a/out/openapi.json +++ b/out/openapi.json @@ -465,6 +465,17 @@ "schema": { "type": "string" } + }, + { + "name": "namespace_id", + "in": "query", + "required": true, + "schema": { + "type": "array", + "items": { + "$ref": "#/components/schemas/RivetId" + } + } } ], "responses": { @@ -1042,7 +1053,8 @@ "namespace_id", "name", "display_name", - "create_ts" + "create_ts", + "runner_kind" ], "properties": { "create_ts": { @@ -1057,6 +1069,9 @@ }, "namespace_id": { "$ref": "#/components/schemas/RivetId" + }, + "runner_kind": { + "$ref": "#/components/schemas/RunnerKind" } } }, @@ -1238,6 +1253,59 @@ }, "additionalProperties": false }, + "RunnerKind": { + "oneOf": [ + { + "type": "object", + "required": [ + "outbound" + ], + "properties": { + "outbound": { + "type": "object", + "required": [ + "url", + "slots_per_runner", + "min_runners", + "max_runners", + "runners_margin" + ], + "properties": { + "max_runners": { + "type": "integer", + "format": "int32", + "minimum": 0 + }, + "min_runners": { + "type": "integer", + "format": "int32", + "minimum": 0 + }, + "runners_margin": { + "type": "integer", + "format": "int32", + "minimum": 0 + }, + "slots_per_runner": { + "type": "integer", + "format": "int32", + "minimum": 0 + }, + "url": { + "type": "string" + } + } + } + } + }, + { + "type": "string", + "enum": [ + "custom" + ] + } + ] + }, "RunnersGetResponse": { "type": "object", "required": [ diff --git a/packages/common/api-builder/Cargo.toml b/packages/common/api-builder/Cargo.toml index f9ce906458..5655703737 100644 --- a/packages/common/api-builder/Cargo.toml +++ b/packages/common/api-builder/Cargo.toml @@ -8,6 +8,7 @@ license.workspace = true [dependencies] anyhow.workspace = true axum.workspace = true +axum-extra.workspace = true gas.workspace = true chrono.workspace = true hyper = { workspace = true, features = ["full"] } diff --git a/packages/common/api-builder/src/wrappers.rs b/packages/common/api-builder/src/wrappers.rs index e6d8193c2d..926b6ece48 100644 --- a/packages/common/api-builder/src/wrappers.rs +++ b/packages/common/api-builder/src/wrappers.rs @@ -1,13 +1,14 @@ use anyhow::Result; use axum::{ body::Bytes, - extract::{Extension, Path, Query}, + extract::{Extension, Path}, response::{IntoResponse, Json}, routing::{ delete as axum_delete, get as axum_get, patch as axum_patch, post as axum_post, put as axum_put, }, }; +use axum_extra::extract::Query; use serde::{Serialize, de::DeserializeOwned}; use std::future::Future; diff --git a/packages/common/gasoline/core/src/utils/tags.rs b/packages/common/gasoline/core/src/utils/tags.rs index 967c59d10f..92326b65bd 100644 --- a/packages/common/gasoline/core/src/utils/tags.rs +++ b/packages/common/gasoline/core/src/utils/tags.rs @@ -61,6 +61,16 @@ impl AsTags for serde_json::Value { } } +impl AsTags for () { + fn as_tags(&self) -> WorkflowResult { + Ok(serde_json::Value::Object(serde_json::Map::new())) + } + + fn as_cjson_tags(&self) -> WorkflowResult { + Ok(String::new()) + } +} + impl AsTags for &T { fn as_tags(&self) -> WorkflowResult { (*self).as_tags() diff --git a/packages/common/pools/src/reqwest.rs b/packages/common/pools/src/reqwest.rs index b3044041af..78f1f2e7cb 100644 --- a/packages/common/pools/src/reqwest.rs +++ b/packages/common/pools/src/reqwest.rs @@ -13,3 +13,10 @@ pub async fn client() -> Result { .await .cloned() } + +pub async fn client_no_timeout() -> Result { + CLIENT + .get_or_try_init(|| async { Client::builder().build() }) + .await + .cloned() +} diff --git a/packages/common/types/Cargo.toml b/packages/common/types/Cargo.toml index 6e429cbacf..9ae35c64aa 100644 --- a/packages/common/types/Cargo.toml +++ b/packages/common/types/Cargo.toml @@ -10,7 +10,7 @@ anyhow.workspace = true gas.workspace = true rivet-api-builder.workspace = true rivet-runner-protocol.workspace = true -rivet-key-data.workspace = true +rivet-data.workspace = true rivet-util.workspace = true serde.workspace = true utoipa.workspace = true diff --git a/packages/common/types/src/runners.rs b/packages/common/types/src/runners.rs index f56db3f670..d70ed49eda 100644 --- a/packages/common/types/src/runners.rs +++ b/packages/common/types/src/runners.rs @@ -1,5 +1,5 @@ use gas::prelude::*; -use rivet_key_data::generated::pegboard_runner_address_v1; +use rivet_data::generated::pegboard_runner_address_v1; use rivet_runner_protocol::protocol; use rivet_util::Id; use serde::{Deserialize, Serialize}; diff --git a/packages/common/udb-util/src/keys.rs b/packages/common/udb-util/src/keys.rs index 397c727c28..c0406fe67f 100644 --- a/packages/common/udb-util/src/keys.rs +++ b/packages/common/udb-util/src/keys.rs @@ -59,7 +59,7 @@ define_keys! { (31, DBS, "dbs"), (32, ACTOR, "actor"), (33, BY_NAME, "by_name"), - (34, DATACENTER, "datacenter"), + // 34 (35, REMAINING_MEMORY, "remaining_memory"), (36, REMAINING_CPU, "remaining_cpu"), (37, TOTAL_MEMORY, "total_memory"), @@ -119,4 +119,6 @@ define_keys! { (91, METRIC, "metric"), (92, CURRENT_BALLOT, "current_ballot"), (93, INSTANCE_BALLOT, "instance_ballot"), + (94, OUTBOUND, "outbound"), + (95, DESIRED_SLOTS, "desired_slots"), } diff --git a/packages/core/api-peer/src/namespaces.rs b/packages/core/api-peer/src/namespaces.rs index 196764252f..36b0a3c9fc 100644 --- a/packages/core/api-peer/src/namespaces.rs +++ b/packages/core/api-peer/src/namespaces.rs @@ -73,6 +73,7 @@ pub struct ListQuery { pub limit: Option, pub cursor: Option, pub name: Option, + pub namespace_id: Vec, } #[derive(Serialize, Deserialize, ToSchema)] @@ -85,7 +86,7 @@ pub struct ListResponse { #[utoipa::path( get, - operation_id = "actors_list", + operation_id = "namespaces_list", path = "/namespaces", params(ListQuery), responses( @@ -105,6 +106,17 @@ pub async fn list(ctx: ApiCtx, _path: (), query: ListQuery) -> Result, + shutdown_tx: oneshot::Sender<()>, + draining: Arc, +} + +#[tracing::instrument(skip_all)] +pub async fn start(config: rivet_config::Config, pools: rivet_pools::Pools) -> Result<()> { + let cache = rivet_cache::CacheInner::from_env(&config, pools.clone())?; + let ctx = StandaloneCtx::new( + db::DatabaseKv::from_pools(pools.clone()).await?, + config.clone(), + pools, + cache, + "pegboard-outbound", + Id::new_v1(config.dc_label()), + Id::new_v1(config.dc_label()), + )?; + + let mut sub = ctx + .subscribe::(()) + .await?; + let mut outbound_connections = HashMap::new(); + + loop { + tick(&ctx, &mut outbound_connections).await?; + + sub.next().await?; + } +} + +async fn tick( + ctx: &StandaloneCtx, + outbound_connections: &mut HashMap<(Id, String), Vec>, +) -> Result<()> { + let outbound_data = ctx + .udb()? + .run(|tx, _mc| async move { + let txs = tx.subspace(keys::subspace()); + let outbound_desired_subspace = + txs.subspace(&keys::ns::OutboundDesiredSlotsKey::subspace()); + + txs.get_ranges_keyvalues( + udb::RangeOption { + mode: StreamingMode::WantAll, + ..(&outbound_desired_subspace).into() + }, + // NOTE: This is a snapshot to prevent conflict with updates to this subspace + SNAPSHOT, + ) + .map(|res| match res { + Ok(entry) => { + let (key, desired_slots) = + txs.read_entry::(&entry)?; + + Ok((key.namespace_id, key.runner_name_selector, desired_slots)) + } + Err(err) => Err(err.into()), + }) + .try_collect::>() + .await + + // outbound/{ns_id}/{runner_name_selector}/desired_slots + }) + .await?; + + let mut namespace_ids = outbound_data + .iter() + .map(|(ns_id, _, _)| *ns_id) + .collect::>(); + namespace_ids.dedup(); + + let namespaces = ctx + .op(namespace::ops::get_global::Input { namespace_ids }) + .await?; + + for (ns_id, runner_name_selector, desired_slots) in &outbound_data { + let namespace = namespaces + .iter() + .find(|ns| ns.namespace_id == *ns_id) + .context("ns not found")?; + + let RunnerKind::Outbound { + url, + slots_per_runner, + min_runners, + max_runners, + runners_margin, + } = &namespace.runner_kind + else { + tracing::warn!( + ?ns_id, + "this namespace should not be in the outbound subspace (wrong runner kind)" + ); + continue; + }; + + let curr = outbound_connections + .entry((*ns_id, runner_name_selector.clone())) + .or_insert_with(Vec::new); + + // Remove finished and draining connections from list + curr.retain(|conn| !conn.handle.is_finished() && !conn.draining.load(Ordering::SeqCst)); + + let desired_count = (desired_slots + .div_ceil(*slots_per_runner) + .max(*min_runners) + .min(*max_runners) + + runners_margin) + .try_into()?; + + // Calculate diff + let drain_count = curr.len().saturating_sub(desired_count); + let start_count = desired_count.saturating_sub(curr.len()); + + if drain_count != 0 { + // TODO: Implement smart logic of draining runners with the lowest allocated actors + let draining_connections = curr.split_off(desired_count); + + for conn in draining_connections { + if conn.shutdown_tx.send(()).is_err() { + tracing::warn!( + "outbound connection shutdown channel dropped, likely already stopped" + ); + } + } + } + + let starting_connections = + std::iter::repeat_with(|| spawn_connection(ctx.clone(), url.clone())).take(start_count); + curr.extend(starting_connections); + } + + // Remove entries that aren't returned from udb + outbound_connections.retain(|(ns_id, runner_name_selector), _| { + outbound_data + .iter() + .any(|(ns_id2, runner_name_selector2, _)| { + ns_id == ns_id2 && runner_name_selector == runner_name_selector2 + }) + }); + + Ok(()) +} + +fn spawn_connection(ctx: StandaloneCtx, url: String) -> OutboundConnection { + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + let draining = Arc::new(AtomicBool::new(false)); + + let draining2 = draining.clone(); + let handle = tokio::spawn(async move { + if let Err(err) = outbound_handler(&ctx, url, shutdown_rx, draining2).await { + tracing::error!(?err, "outbound req failed"); + + // TODO: Add backoff + tokio::time::sleep(Duration::from_secs(1)).await; + + // On error, bump the autoscaler loop again + let _ = ctx + .msg(pegboard::messages::BumpOutboundAutoscaler {}) + .send() + .await; + } + }); + + OutboundConnection { + handle, + shutdown_tx, + draining, + } +} + +async fn outbound_handler( + ctx: &StandaloneCtx, + url: String, + shutdown_rx: oneshot::Receiver<()>, + draining: Arc, +) -> Result<()> { + let client = rivet_pools::reqwest::client_no_timeout().await?; + let mut es = sse::EventSource::new(client.get(url))?; + let mut runner_id = None; + + let stream_handler = async { + while let Some(event) = es.next().await { + match event { + Ok(sse::Event::Open) => {} + Ok(sse::Event::Message(msg)) => { + tracing::debug!(%msg.data, "received outbound req message"); + + if runner_id.is_none() { + runner_id = Some(Id::parse(&msg.data)?); + } + } + Err(sse::Error::StreamEnded) => { + tracing::debug!("outbound req stopped early"); + + return Ok(()); + } + Err(err) => return Err(err.into()), + } + } + + anyhow::Ok(()) + }; + + tokio::select! { + res = stream_handler => return res.map_err(Into::into), + _ = tokio::time::sleep(OUTBOUND_REQUEST_LIFESPAN) => {} + _ = shutdown_rx => {} + } + + draining.store(true, Ordering::SeqCst); + + ctx.msg(pegboard::messages::BumpOutboundAutoscaler {}) + .send() + .await?; + + if let Some(runner_id) = runner_id { + stop_runner(ctx, runner_id).await?; + } + + // Continue waiting on req while draining + while let Some(event) = es.next().await { + match event { + Ok(sse::Event::Open) => {} + Ok(sse::Event::Message(msg)) => { + tracing::debug!(%msg.data, "received outbound req message"); + + // If runner_id is none at this point it means we did not send the stopping signal yet, so + // send it now + if runner_id.is_none() { + stop_runner(ctx, Id::parse(&msg.data)?).await?; + } + } + Err(sse::Error::StreamEnded) => break, + Err(err) => return Err(err.into()), + } + } + + tracing::info!("outbound req stopped"); + + Ok(()) +} + +async fn stop_runner(ctx: &StandaloneCtx, runner_id: Id) -> Result<()> { + let res = ctx + .signal(protocol::ToServer::Stopping) + .to_workflow::() + .tag("runner_id", runner_id) + .send() + .await; + + if let Some(WorkflowError::WorkflowNotFound) = res + .as_ref() + .err() + .and_then(|x| x.chain().find_map(|x| x.downcast_ref::())) + { + tracing::warn!( + ?runner_id, + "runner workflow not found, likely already stopped" + ); + } else { + res?; + } + + Ok(()) +} diff --git a/packages/infra/engine/Cargo.toml b/packages/infra/engine/Cargo.toml index 2556c213f3..975ea3d788 100644 --- a/packages/infra/engine/Cargo.toml +++ b/packages/infra/engine/Cargo.toml @@ -11,16 +11,15 @@ path = "src/main.rs" [dependencies] anyhow.workspace = true -gas.workspace = true chrono.workspace = true clap.workspace = true colored_json.workspace = true -udb-util.workspace = true -universaldb.workspace = true futures-util.workspace = true +gas.workspace = true hex.workspace = true include_dir.workspace = true lz4_flex.workspace = true +pegboard-outbound.workspace = true pegboard-runner-ws.workspace = true reqwest.workspace = true rivet-api-peer.workspace = true @@ -38,15 +37,17 @@ rivet-term.workspace = true rivet-util.workspace = true rivet-workflow-worker.workspace = true rustyline.workspace = true -serde.workspace = true serde_json.workspace = true serde_yaml.workspace = true +serde.workspace = true strum.workspace = true tabled.workspace = true tempfile.workspace = true thiserror.workspace = true tokio.workspace = true tracing.workspace = true +udb-util.workspace = true +universaldb.workspace = true url.workspace = true uuid.workspace = true diff --git a/packages/infra/engine/src/run_config.rs b/packages/infra/engine/src/run_config.rs index ee683efc15..890b12842b 100644 --- a/packages/infra/engine/src/run_config.rs +++ b/packages/infra/engine/src/run_config.rs @@ -25,6 +25,11 @@ pub fn config(_rivet_config: rivet_config::Config) -> Result { Service::new("bootstrap", ServiceKind::Oneshot, |config, pools| { Box::pin(rivet_bootstrap::start(config, pools)) }), + Service::new( + "pegboard_outbound", + ServiceKind::Standalone, + |config, pools| Box::pin(pegboard_outbound::start(config, pools)), + ), ]; Ok(RunConfigData { services }) diff --git a/packages/services/namespace/Cargo.toml b/packages/services/namespace/Cargo.toml index fef88ecd1c..32d46a6735 100644 --- a/packages/services/namespace/Cargo.toml +++ b/packages/services/namespace/Cargo.toml @@ -8,12 +8,14 @@ edition.workspace = true [dependencies] anyhow.workspace = true gas.workspace = true -udb-util.workspace = true -universaldb.workspace = true rivet-api-builder.workspace = true rivet-api-util.workspace = true +rivet-data.workspace = true rivet-error.workspace = true rivet-util.workspace = true serde.workspace = true tracing.workspace = true +udb-util.workspace = true +universaldb.workspace = true utoipa.workspace = true +versioned-data-util.workspace = true diff --git a/packages/services/namespace/src/keys.rs b/packages/services/namespace/src/keys.rs index 3803feb6f6..c4e13afde4 100644 --- a/packages/services/namespace/src/keys.rs +++ b/packages/services/namespace/src/keys.rs @@ -3,6 +3,7 @@ use std::result::Result::Ok; use anyhow::*; use gas::prelude::*; use udb_util::prelude::*; +use versioned_data_util::OwnedVersionedData; pub fn subspace() -> udb_util::Subspace { udb_util::Subspace::new(&(RIVET, NAMESPACE)) @@ -144,6 +145,55 @@ impl<'de> TupleUnpack<'de> for CreateTsKey { } } +#[derive(Debug)] +pub struct RunnerKindKey { + namespace_id: Id, +} + +impl RunnerKindKey { + pub fn new(namespace_id: Id) -> Self { + RunnerKindKey { namespace_id } + } +} + +impl FormalKey for RunnerKindKey { + type Value = crate::types::RunnerKind; + + fn deserialize(&self, raw: &[u8]) -> Result { + Ok( + rivet_data::versioned::NamespaceRunnerKind::deserialize_with_embedded_version(raw)? + .into(), + ) + } + + fn serialize(&self, value: Self::Value) -> Result> { + rivet_data::versioned::NamespaceRunnerKind::latest(value.into()) + .serialize_with_embedded_version( + rivet_data::PEGBOARD_NAMESPACE_RUNNER_ALLOC_IDX_VERSION, + ) + } +} + +impl TuplePack for RunnerKindKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = (DATA, self.namespace_id, CREATE_TS); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for RunnerKindKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, namespace_id, _)) = <(usize, Id, usize)>::unpack(input, tuple_depth)?; + let v = RunnerKindKey { namespace_id }; + + Ok((input, v)) + } +} + #[derive(Debug)] pub struct ByNameKey { name: String, diff --git a/packages/services/namespace/src/ops/get_global.rs b/packages/services/namespace/src/ops/get_global.rs index 5dc5a58b08..a62eeda288 100644 --- a/packages/services/namespace/src/ops/get_global.rs +++ b/packages/services/namespace/src/ops/get_global.rs @@ -4,20 +4,19 @@ use crate::types::Namespace; #[derive(Debug)] pub struct Input { - // TODO: Accept vec - pub namespace_id: Id, + pub namespace_ids: Vec, } #[operation] -pub async fn namespace_get_global(ctx: &OperationCtx, input: &Input) -> Result> { +pub async fn namespace_get_global(ctx: &OperationCtx, input: &Input) -> Result> { if ctx.config().is_leader() { let namespaces_res = ctx .op(crate::ops::get_local::Input { - namespace_ids: vec![input.namespace_id], + namespace_ids: input.namespace_ids.clone(), }) .await?; - Ok(namespaces_res.namespaces.into_iter().next()) + Ok(namespaces_res.namespaces) } else { let leader_dc = ctx.config().leader_dc()?; let client = rivet_pools::reqwest::client().await?; @@ -25,51 +24,42 @@ pub async fn namespace_get_global(ctx: &OperationCtx, input: &Input) -> Result>(), + ) + .send() + .await?; - let res = rivet_api_util::parse_response::(res).await; + let res = rivet_api_util::parse_response::(res).await?; - let res = match res { - Ok(res) => Ok(Some(res.namespace)), - Err(err) => { - // Explicitly handle namespace not found error - if let Some(error) = err.chain().find_map(|x| { - x.downcast_ref::() - }) { - if error.1.group == "namespace" && error.1.code == "not_found" { - Ok(None) - } else { - Err(err) - } - } else { - Err(err) - } - } - }; - - cache.resolve(&key, res?); + for ns in res.namespaces { + let namespace_id = ns.namespace_id; + cache.resolve(&&namespace_id, ns); + } Ok(cache) } } }) .await - .map(|x| x.flatten()) } } // TODO: Cyclical dependency with api_peer #[derive(Deserialize)] -struct GetResponse { - namespace: Namespace, +struct ListResponse { + namespaces: Vec, } diff --git a/packages/services/namespace/src/ops/get_local.rs b/packages/services/namespace/src/ops/get_local.rs index 913579092f..ed6663d589 100644 --- a/packages/services/namespace/src/ops/get_local.rs +++ b/packages/services/namespace/src/ops/get_local.rs @@ -1,6 +1,6 @@ use futures_util::{StreamExt, TryStreamExt}; use gas::prelude::*; -use udb_util::{FormalKey, SERIALIZABLE}; +use udb_util::{SERIALIZABLE, TxnExt}; use universaldb as udb; use crate::{errors, keys, types::Namespace}; @@ -45,39 +45,40 @@ pub(crate) async fn get_inner( namespace_id: Id, tx: &udb::RetryableTransaction, ) -> std::result::Result, udb::FdbBindingError> { + let txs = tx.subspace(keys::subspace()); + let name_key = keys::NameKey::new(namespace_id); let display_name_key = keys::DisplayNameKey::new(namespace_id); let create_ts_key = keys::CreateTsKey::new(namespace_id); + let runner_kind_key = keys::RunnerKindKey::new(namespace_id); - let (name_entry, display_name_entry, create_ts_entry) = tokio::try_join!( - tx.get(&keys::subspace().pack(&name_key), SERIALIZABLE), - tx.get(&keys::subspace().pack(&display_name_key), SERIALIZABLE), - tx.get(&keys::subspace().pack(&create_ts_key), SERIALIZABLE), + let (name, display_name, create_ts, runner_kind) = tokio::try_join!( + txs.read_opt(&name_key, SERIALIZABLE), + txs.read_opt(&display_name_key, SERIALIZABLE), + txs.read_opt(&create_ts_key, SERIALIZABLE), + txs.read_opt(&runner_kind_key, SERIALIZABLE), )?; // Namespace not found - let Some(name_entry) = name_entry else { + let Some(name) = name else { return Ok(None); }; - let name = name_key - .deserialize(&name_entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; - let display_name = display_name_key - .deserialize(&display_name_entry.ok_or(udb::FdbBindingError::CustomError( - format!("key should exist: {display_name_key:?}").into(), - ))?) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; - let create_ts = create_ts_key - .deserialize(&create_ts_entry.ok_or(udb::FdbBindingError::CustomError( - format!("key should exist: {create_ts_key:?}").into(), - ))?) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let display_name = display_name.ok_or(udb::FdbBindingError::CustomError( + format!("key should exist: {display_name_key:?}").into(), + ))?; + let create_ts = create_ts.ok_or(udb::FdbBindingError::CustomError( + format!("key should exist: {create_ts_key:?}").into(), + ))?; + let runner_kind = runner_kind.ok_or(udb::FdbBindingError::CustomError( + format!("key should exist: {runner_kind_key:?}").into(), + ))?; Ok(Some(Namespace { namespace_id, name, display_name, create_ts, + runner_kind, })) } diff --git a/packages/services/namespace/src/types.rs b/packages/services/namespace/src/types.rs index c7e6f71d59..05e924ad34 100644 --- a/packages/services/namespace/src/types.rs +++ b/packages/services/namespace/src/types.rs @@ -7,4 +7,58 @@ pub struct Namespace { pub name: String, pub display_name: String, pub create_ts: i64, + pub runner_kind: RunnerKind, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Hash, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum RunnerKind { + Outbound { + url: String, + slots_per_runner: u32, + min_runners: u32, + max_runners: u32, + runners_margin: u32, + }, + Custom, +} + +impl From for rivet_data::generated::namespace_runner_kind_v1::Data { + fn from(value: RunnerKind) -> Self { + match value { + RunnerKind::Outbound { + url, + slots_per_runner, + min_runners, + max_runners, + runners_margin, + } => rivet_data::generated::namespace_runner_kind_v1::Data::Outbound( + rivet_data::generated::namespace_runner_kind_v1::Outbound { + url, + slots_per_runner, + min_runners, + max_runners, + runners_margin, + }, + ), + RunnerKind::Custom => rivet_data::generated::namespace_runner_kind_v1::Data::Custom, + } + } +} + +impl From for RunnerKind { + fn from(value: rivet_data::generated::namespace_runner_kind_v1::Data) -> Self { + match value { + rivet_data::generated::namespace_runner_kind_v1::Data::Outbound(o) => { + RunnerKind::Outbound { + url: o.url, + slots_per_runner: o.slots_per_runner, + min_runners: o.min_runners, + max_runners: o.max_runners, + runners_margin: o.runners_margin, + } + } + rivet_data::generated::namespace_runner_kind_v1::Data::Custom => RunnerKind::Custom, + } + } } diff --git a/packages/services/namespace/src/workflows/namespace.rs b/packages/services/namespace/src/workflows/namespace.rs index 16575347fe..90078b23e6 100644 --- a/packages/services/namespace/src/workflows/namespace.rs +++ b/packages/services/namespace/src/workflows/namespace.rs @@ -1,10 +1,9 @@ use futures_util::FutureExt; use gas::prelude::*; use serde::{Deserialize, Serialize}; -use udb_util::{FormalKey, SERIALIZABLE}; -use universaldb as udb; +use udb_util::{SERIALIZABLE, TxnExt}; -use crate::{errors, keys}; +use crate::{errors, keys, types::RunnerKind}; #[derive(Debug, Deserialize, Serialize)] pub struct Input { @@ -59,7 +58,7 @@ pub async fn namespace(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> { // Does nothing yet ctx.repeat(|ctx| { async move { - ctx.listen::().await?; + ctx.listen::().await?; Ok(Loop::<()>::Continue) } @@ -79,7 +78,7 @@ pub struct Failed { } #[signal("namespace_update")] -pub struct NamespaceUpdate {} +pub struct Update {} #[derive(Debug, Clone, Serialize, Deserialize, Hash)] pub struct ValidateInput { @@ -156,45 +155,29 @@ async fn insert_fdb( let display_name = input.display_name.clone(); async move { - let name_key = keys::NameKey::new(namespace_id); - let name_idx_key = keys::ByNameKey::new(name.clone()); - let display_name_key = keys::DisplayNameKey::new(namespace_id); - let create_ts_key = keys::CreateTsKey::new(namespace_id); + let txs = tx.subspace(keys::subspace()); - let name_idx_entry = tx - .get(&keys::subspace().pack(&name_idx_key), SERIALIZABLE) - .await?; + let name_idx_key = keys::ByNameKey::new(name.clone()); - if name_idx_entry.is_some() { + if txs.exists(&name_idx_key, SERIALIZABLE).await? { return Ok(Err(errors::Namespace::NameNotUnique)); } - tx.set( - &keys::subspace().pack(&name_key), - &name_key - .serialize(name) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ); - tx.set( - &keys::subspace().pack(&display_name_key), - &display_name_key - .serialize(display_name) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ); - tx.set( - &keys::subspace().pack(&create_ts_key), - &create_ts_key - .serialize(input.create_ts) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ); + txs.write(&keys::NameKey::new(namespace_id), name)?; + txs.write(&keys::DisplayNameKey::new(namespace_id), display_name)?; + txs.write(&keys::CreateTsKey::new(namespace_id), input.create_ts)?; + txs.write(&keys::RunnerKindKey::new(namespace_id), RunnerKind::Custom)?; + + // RunnerKind::Outbound { + // url: "http://runner:5051/start".to_string(), + // slots_per_runner: 10, + // min_runners: 1, + // max_runners: 1, + // runners_margin: 0, + // } // Insert idx - tx.set( - &keys::subspace().pack(&name_idx_key), - &name_idx_key - .serialize(namespace_id) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ); + txs.write(&name_idx_key, namespace_id)?; Ok(Ok(())) } diff --git a/packages/services/pegboard/Cargo.toml b/packages/services/pegboard/Cargo.toml index 02e85d00e6..945f36bdfa 100644 --- a/packages/services/pegboard/Cargo.toml +++ b/packages/services/pegboard/Cargo.toml @@ -16,7 +16,7 @@ rivet-api-client.workspace = true rivet-api-types.workspace = true rivet-api-util.workspace = true rivet-error.workspace = true -rivet-key-data.workspace = true +rivet-data.workspace = true rivet-metrics.workspace = true rivet-runner-protocol.workspace = true rivet-types.workspace = true diff --git a/packages/services/pegboard/src/keys/datacenter.rs b/packages/services/pegboard/src/keys/datacenter.rs deleted file mode 100644 index bd89e1db22..0000000000 --- a/packages/services/pegboard/src/keys/datacenter.rs +++ /dev/null @@ -1,249 +0,0 @@ -use std::result::Result::Ok; - -use anyhow::*; -use gas::prelude::*; -use udb_util::prelude::*; -use versioned_data_util::OwnedVersionedData; - -#[derive(Debug)] -pub struct RunnerAllocIdxKey { - pub namespace_id: Id, - pub name: String, - pub version: u32, - pub remaining_millislots: u32, - pub last_ping_ts: i64, - pub runner_id: Id, -} - -impl RunnerAllocIdxKey { - pub fn new( - namespace_id: Id, - name: String, - version: u32, - remaining_millislots: u32, - last_ping_ts: i64, - runner_id: Id, - ) -> Self { - RunnerAllocIdxKey { - namespace_id, - name, - version, - remaining_millislots, - last_ping_ts, - runner_id, - } - } - - pub fn subspace(namespace_id: Id, name: String) -> RunnerAllocIdxSubspaceKey { - RunnerAllocIdxSubspaceKey::new(namespace_id, name) - } - - pub fn entire_subspace() -> RunnerAllocIdxSubspaceKey { - RunnerAllocIdxSubspaceKey::entire() - } -} - -impl FormalKey for RunnerAllocIdxKey { - type Value = rivet_key_data::converted::RunnerAllocIdxKeyData; - - fn deserialize(&self, raw: &[u8]) -> Result { - rivet_key_data::versioned::RunnerAllocIdxKeyData::deserialize_with_embedded_version(raw)? - .try_into() - } - - fn serialize(&self, value: Self::Value) -> Result> { - rivet_key_data::versioned::RunnerAllocIdxKeyData::latest(value.try_into()?) - .serialize_with_embedded_version( - rivet_key_data::PEGBOARD_DATACENTER_RUNNER_ALLOC_IDX_VERSION, - ) - } -} - -impl TuplePack for RunnerAllocIdxKey { - fn pack( - &self, - w: &mut W, - tuple_depth: TupleDepth, - ) -> std::io::Result { - let t = ( - DATACENTER, - RUNNER_ALLOC_IDX, - self.namespace_id, - &self.name, - // Stored in reverse order (higher versions are first) - -(self.version as i32), - // Stored in reverse order (higher remaining slots are first) - -(self.remaining_millislots as i32), - self.last_ping_ts, - self.runner_id, - ); - t.pack(w, tuple_depth) - } -} - -impl<'de> TupleUnpack<'de> for RunnerAllocIdxKey { - fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let ( - input, - (_, _, namespace_id, name, version, remaining_millislots, last_ping_ts, runner_id), - ) = <(usize, usize, Id, String, i32, i32, i64, Id)>::unpack(input, tuple_depth)?; - - let v = RunnerAllocIdxKey { - namespace_id, - name, - version: -version as u32, - remaining_millislots: -remaining_millislots as u32, - last_ping_ts, - runner_id, - }; - - Ok((input, v)) - } -} - -pub struct RunnerAllocIdxSubspaceKey { - pub namespace_id: Option, - pub name: Option, -} - -impl RunnerAllocIdxSubspaceKey { - pub fn new(namespace_id: Id, name: String) -> Self { - RunnerAllocIdxSubspaceKey { - namespace_id: Some(namespace_id), - name: Some(name), - } - } - - pub fn entire() -> Self { - RunnerAllocIdxSubspaceKey { - namespace_id: None, - name: None, - } - } -} - -impl TuplePack for RunnerAllocIdxSubspaceKey { - fn pack( - &self, - w: &mut W, - tuple_depth: TupleDepth, - ) -> std::io::Result { - let mut offset = VersionstampOffset::None { size: 0 }; - - let t = (DATACENTER, RUNNER_ALLOC_IDX); - offset += t.pack(w, tuple_depth)?; - - if let Some(namespace_id) = &self.namespace_id { - offset += namespace_id.pack(w, tuple_depth)?; - - if let Some(name) = &self.name { - offset += name.pack(w, tuple_depth)?; - } - } - - Ok(offset) - } -} - -#[derive(Debug)] -pub struct PendingActorByRunnerNameSelectorKey { - pub namespace_id: Id, - pub runner_name_selector: String, - pub ts: i64, - pub actor_id: Id, -} - -impl PendingActorByRunnerNameSelectorKey { - pub fn new(namespace_id: Id, runner_name_selector: String, ts: i64, actor_id: Id) -> Self { - PendingActorByRunnerNameSelectorKey { - namespace_id, - runner_name_selector, - ts, - actor_id, - } - } - - pub fn subspace( - namespace_id: Id, - runner_name_selector: String, - ) -> PendingActorByRunnerNameSelectorSubspaceKey { - PendingActorByRunnerNameSelectorSubspaceKey::new(namespace_id, runner_name_selector) - } -} - -impl FormalKey for PendingActorByRunnerNameSelectorKey { - /// Generation. - type Value = u32; - - fn deserialize(&self, raw: &[u8]) -> Result { - Ok(u32::from_be_bytes(raw.try_into()?)) - } - - fn serialize(&self, value: Self::Value) -> Result> { - Ok(value.to_be_bytes().to_vec()) - } -} - -impl TuplePack for PendingActorByRunnerNameSelectorKey { - fn pack( - &self, - w: &mut W, - tuple_depth: TupleDepth, - ) -> std::io::Result { - let t = ( - DATACENTER, - PENDING_ACTOR_BY_RUNNER_NAME_SELECTOR, - self.namespace_id, - &self.runner_name_selector, - self.ts, - self.actor_id, - ); - t.pack(w, tuple_depth) - } -} - -impl<'de> TupleUnpack<'de> for PendingActorByRunnerNameSelectorKey { - fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let (input, (_, _, namespace_id, runner_name_selector, ts, actor_id)) = - <(usize, usize, Id, String, i64, Id)>::unpack(input, tuple_depth)?; - - let v = PendingActorByRunnerNameSelectorKey { - namespace_id, - runner_name_selector, - ts, - actor_id, - }; - - Ok((input, v)) - } -} - -pub struct PendingActorByRunnerNameSelectorSubspaceKey { - pub namespace_id: Id, - pub runner_name_selector: String, -} - -impl PendingActorByRunnerNameSelectorSubspaceKey { - pub fn new(namespace_id: Id, runner_name_selector: String) -> Self { - PendingActorByRunnerNameSelectorSubspaceKey { - namespace_id, - runner_name_selector, - } - } -} - -impl TuplePack for PendingActorByRunnerNameSelectorSubspaceKey { - fn pack( - &self, - w: &mut W, - tuple_depth: TupleDepth, - ) -> std::io::Result { - let t = ( - DATACENTER, - PENDING_ACTOR_BY_RUNNER_NAME_SELECTOR, - self.namespace_id, - &self.runner_name_selector, - ); - t.pack(w, tuple_depth) - } -} diff --git a/packages/services/pegboard/src/keys/mod.rs b/packages/services/pegboard/src/keys/mod.rs index 9e93b8983a..402214f8a0 100644 --- a/packages/services/pegboard/src/keys/mod.rs +++ b/packages/services/pegboard/src/keys/mod.rs @@ -1,7 +1,6 @@ use udb_util::prelude::*; pub mod actor; -pub mod datacenter; pub mod epoxy; pub mod ns; pub mod runner; diff --git a/packages/services/pegboard/src/keys/ns.rs b/packages/services/pegboard/src/keys/ns.rs index dd1aea42a9..236c61bc9c 100644 --- a/packages/services/pegboard/src/keys/ns.rs +++ b/packages/services/pegboard/src/keys/ns.rs @@ -5,6 +5,249 @@ use gas::prelude::*; use udb_util::prelude::*; use versioned_data_util::OwnedVersionedData; +#[derive(Debug)] +pub struct RunnerAllocIdxKey { + pub namespace_id: Id, + pub name: String, + pub version: u32, + pub remaining_millislots: u32, + pub last_ping_ts: i64, + pub runner_id: Id, +} + +impl RunnerAllocIdxKey { + pub fn new( + namespace_id: Id, + name: String, + version: u32, + remaining_millislots: u32, + last_ping_ts: i64, + runner_id: Id, + ) -> Self { + RunnerAllocIdxKey { + namespace_id, + name, + version, + remaining_millislots, + last_ping_ts, + runner_id, + } + } + + pub fn subspace(namespace_id: Id, name: String) -> RunnerAllocIdxSubspaceKey { + RunnerAllocIdxSubspaceKey::new(namespace_id, name) + } + + pub fn entire_subspace() -> RunnerAllocIdxSubspaceKey { + RunnerAllocIdxSubspaceKey::entire() + } +} + +impl FormalKey for RunnerAllocIdxKey { + type Value = rivet_data::converted::RunnerAllocIdxKeyData; + + fn deserialize(&self, raw: &[u8]) -> Result { + rivet_data::versioned::RunnerAllocIdxKeyData::deserialize_with_embedded_version(raw)? + .try_into() + } + + fn serialize(&self, value: Self::Value) -> Result> { + rivet_data::versioned::RunnerAllocIdxKeyData::latest(value.try_into()?) + .serialize_with_embedded_version( + rivet_data::PEGBOARD_NAMESPACE_RUNNER_ALLOC_IDX_VERSION, + ) + } +} + +impl TuplePack for RunnerAllocIdxKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + NAMESPACE, + RUNNER_ALLOC_IDX, + self.namespace_id, + &self.name, + // Stored in reverse order (higher versions are first) + -(self.version as i32), + // Stored in reverse order (higher remaining slots are first) + -(self.remaining_millislots as i32), + self.last_ping_ts, + self.runner_id, + ); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for RunnerAllocIdxKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let ( + input, + (_, _, namespace_id, name, version, remaining_millislots, last_ping_ts, runner_id), + ) = <(usize, usize, Id, String, i32, i32, i64, Id)>::unpack(input, tuple_depth)?; + + let v = RunnerAllocIdxKey { + namespace_id, + name, + version: -version as u32, + remaining_millislots: -remaining_millislots as u32, + last_ping_ts, + runner_id, + }; + + Ok((input, v)) + } +} + +pub struct RunnerAllocIdxSubspaceKey { + pub namespace_id: Option, + pub name: Option, +} + +impl RunnerAllocIdxSubspaceKey { + pub fn new(namespace_id: Id, name: String) -> Self { + RunnerAllocIdxSubspaceKey { + namespace_id: Some(namespace_id), + name: Some(name), + } + } + + pub fn entire() -> Self { + RunnerAllocIdxSubspaceKey { + namespace_id: None, + name: None, + } + } +} + +impl TuplePack for RunnerAllocIdxSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let mut offset = VersionstampOffset::None { size: 0 }; + + let t = (NAMESPACE, RUNNER_ALLOC_IDX); + offset += t.pack(w, tuple_depth)?; + + if let Some(namespace_id) = &self.namespace_id { + offset += namespace_id.pack(w, tuple_depth)?; + + if let Some(name) = &self.name { + offset += name.pack(w, tuple_depth)?; + } + } + + Ok(offset) + } +} + +#[derive(Debug)] +pub struct PendingActorByRunnerNameSelectorKey { + pub namespace_id: Id, + pub runner_name_selector: String, + pub ts: i64, + pub actor_id: Id, +} + +impl PendingActorByRunnerNameSelectorKey { + pub fn new(namespace_id: Id, runner_name_selector: String, ts: i64, actor_id: Id) -> Self { + PendingActorByRunnerNameSelectorKey { + namespace_id, + runner_name_selector, + ts, + actor_id, + } + } + + pub fn subspace( + namespace_id: Id, + runner_name_selector: String, + ) -> PendingActorByRunnerNameSelectorSubspaceKey { + PendingActorByRunnerNameSelectorSubspaceKey::new(namespace_id, runner_name_selector) + } +} + +impl FormalKey for PendingActorByRunnerNameSelectorKey { + /// Generation. + type Value = u32; + + fn deserialize(&self, raw: &[u8]) -> Result { + Ok(u32::from_be_bytes(raw.try_into()?)) + } + + fn serialize(&self, value: Self::Value) -> Result> { + Ok(value.to_be_bytes().to_vec()) + } +} + +impl TuplePack for PendingActorByRunnerNameSelectorKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + NAMESPACE, + PENDING_ACTOR_BY_RUNNER_NAME_SELECTOR, + self.namespace_id, + &self.runner_name_selector, + self.ts, + self.actor_id, + ); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for PendingActorByRunnerNameSelectorKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, namespace_id, runner_name_selector, ts, actor_id)) = + <(usize, usize, Id, String, i64, Id)>::unpack(input, tuple_depth)?; + + let v = PendingActorByRunnerNameSelectorKey { + namespace_id, + runner_name_selector, + ts, + actor_id, + }; + + Ok((input, v)) + } +} + +pub struct PendingActorByRunnerNameSelectorSubspaceKey { + pub namespace_id: Id, + pub runner_name_selector: String, +} + +impl PendingActorByRunnerNameSelectorSubspaceKey { + pub fn new(namespace_id: Id, runner_name_selector: String) -> Self { + PendingActorByRunnerNameSelectorSubspaceKey { + namespace_id, + runner_name_selector, + } + } +} + +impl TuplePack for PendingActorByRunnerNameSelectorSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + NAMESPACE, + PENDING_ACTOR_BY_RUNNER_NAME_SELECTOR, + self.namespace_id, + &self.runner_name_selector, + ); + t.pack(w, tuple_depth) + } +} + #[derive(Debug)] pub struct ActiveActorKey { namespace_id: Id, @@ -320,18 +563,15 @@ impl ActorByKeyKey { } impl FormalKey for ActorByKeyKey { - type Value = rivet_key_data::converted::ActorByKeyKeyData; + type Value = rivet_data::converted::ActorByKeyKeyData; fn deserialize(&self, raw: &[u8]) -> Result { - rivet_key_data::versioned::ActorByKeyKeyData::deserialize_with_embedded_version(raw)? - .try_into() + rivet_data::versioned::ActorByKeyKeyData::deserialize_with_embedded_version(raw)?.try_into() } fn serialize(&self, value: Self::Value) -> Result> { - rivet_key_data::versioned::ActorByKeyKeyData::latest(value.try_into()?) - .serialize_with_embedded_version( - rivet_key_data::PEGBOARD_NAMESPACE_ACTOR_BY_KEY_VERSION, - ) + rivet_data::versioned::ActorByKeyKeyData::latest(value.try_into()?) + .serialize_with_embedded_version(rivet_data::PEGBOARD_NAMESPACE_ACTOR_BY_KEY_VERSION) } } @@ -938,18 +1178,16 @@ impl RunnerByKeyKey { } impl FormalKey for RunnerByKeyKey { - type Value = rivet_key_data::converted::RunnerByKeyKeyData; + type Value = rivet_data::converted::RunnerByKeyKeyData; fn deserialize(&self, raw: &[u8]) -> Result { - rivet_key_data::versioned::RunnerByKeyKeyData::deserialize_with_embedded_version(raw)? + rivet_data::versioned::RunnerByKeyKeyData::deserialize_with_embedded_version(raw)? .try_into() } fn serialize(&self, value: Self::Value) -> Result> { - rivet_key_data::versioned::RunnerByKeyKeyData::latest(value.try_into()?) - .serialize_with_embedded_version( - rivet_key_data::PEGBOARD_NAMESPACE_RUNNER_BY_KEY_VERSION, - ) + rivet_data::versioned::RunnerByKeyKeyData::latest(value.try_into()?) + .serialize_with_embedded_version(rivet_data::PEGBOARD_NAMESPACE_RUNNER_BY_KEY_VERSION) } } @@ -1002,16 +1240,15 @@ impl ActorNameKey { } impl FormalKey for ActorNameKey { - type Value = rivet_key_data::converted::ActorNameKeyData; + type Value = rivet_data::converted::ActorNameKeyData; fn deserialize(&self, raw: &[u8]) -> Result { - rivet_key_data::versioned::ActorNameKeyData::deserialize_with_embedded_version(raw)? - .try_into() + rivet_data::versioned::ActorNameKeyData::deserialize_with_embedded_version(raw)?.try_into() } fn serialize(&self, value: Self::Value) -> Result> { - rivet_key_data::versioned::ActorNameKeyData::latest(value.try_into()?) - .serialize_with_embedded_version(rivet_key_data::PEGBOARD_NAMESPACE_ACTOR_NAME_VERSION) + rivet_data::versioned::ActorNameKeyData::latest(value.try_into()?) + .serialize_with_embedded_version(rivet_data::PEGBOARD_NAMESPACE_ACTOR_NAME_VERSION) } } @@ -1128,3 +1365,87 @@ impl TuplePack for RunnerNameSubspaceKey { t.pack(w, tuple_depth) } } + +#[derive(Debug)] +pub struct OutboundDesiredSlotsKey { + pub namespace_id: Id, + pub runner_name_selector: String, +} + +impl OutboundDesiredSlotsKey { + pub fn new(namespace_id: Id, runner_name_selector: String) -> Self { + OutboundDesiredSlotsKey { + namespace_id, + runner_name_selector, + } + } + + pub fn subspace() -> OutboundDesiredSlotsSubspaceKey { + OutboundDesiredSlotsSubspaceKey::new() + } +} + +impl FormalKey for OutboundDesiredSlotsKey { + /// Count. + type Value = u32; + + fn deserialize(&self, raw: &[u8]) -> Result { + // NOTE: Atomic ops use little endian + Ok(u32::from_le_bytes(raw.try_into()?)) + } + + fn serialize(&self, value: Self::Value) -> Result> { + // NOTE: Atomic ops use little endian + Ok(value.to_le_bytes().to_vec()) + } +} + +impl TuplePack for OutboundDesiredSlotsKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + NAMESPACE, + OUTBOUND, + DESIRED_SLOTS, + self.namespace_id, + &self.runner_name_selector, + ); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for OutboundDesiredSlotsKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, namespace_id, runner_name_selector)) = + <(usize, usize, Id, String)>::unpack(input, tuple_depth)?; + + let v = OutboundDesiredSlotsKey { + namespace_id, + runner_name_selector, + }; + + Ok((input, v)) + } +} + +pub struct OutboundDesiredSlotsSubspaceKey {} + +impl OutboundDesiredSlotsSubspaceKey { + pub fn new() -> Self { + OutboundDesiredSlotsSubspaceKey {} + } +} + +impl TuplePack for OutboundDesiredSlotsSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = (NAMESPACE, OUTBOUND, DESIRED_SLOTS); + t.pack(w, tuple_depth) + } +} diff --git a/packages/services/pegboard/src/keys/runner.rs b/packages/services/pegboard/src/keys/runner.rs index 9e83aac404..528ba2f6b4 100644 --- a/packages/services/pegboard/src/keys/runner.rs +++ b/packages/services/pegboard/src/keys/runner.rs @@ -524,15 +524,15 @@ impl AddressKey { } impl FormalKey for AddressKey { - type Value = ::Latest; + type Value = ::Latest; fn deserialize(&self, raw: &[u8]) -> Result { - rivet_key_data::versioned::AddressKeyData::deserialize_with_embedded_version(raw) + rivet_data::versioned::AddressKeyData::deserialize_with_embedded_version(raw) } fn serialize(&self, value: Self::Value) -> Result> { - rivet_key_data::versioned::AddressKeyData::latest(value) - .serialize_with_embedded_version(rivet_key_data::PEGBOARD_RUNNER_ADDRESS_VERSION) + rivet_data::versioned::AddressKeyData::latest(value) + .serialize_with_embedded_version(rivet_data::PEGBOARD_RUNNER_ADDRESS_VERSION) } } @@ -816,7 +816,7 @@ impl MetadataKey { impl FormalChunkedKey for MetadataKey { type ChunkKey = MetadataChunkKey; - type Value = rivet_key_data::converted::MetadataKeyData; + type Value = rivet_data::converted::MetadataKeyData; fn chunk(&self, chunk: usize) -> Self::ChunkKey { MetadataChunkKey { @@ -826,7 +826,7 @@ impl FormalChunkedKey for MetadataKey { } fn combine(&self, chunks: Vec) -> Result { - rivet_key_data::versioned::MetadataKeyData::deserialize_with_embedded_version( + rivet_data::versioned::MetadataKeyData::deserialize_with_embedded_version( &chunks .iter() .map(|x| x.value().iter().map(|x| *x)) @@ -838,8 +838,8 @@ impl FormalChunkedKey for MetadataKey { fn split(&self, value: Self::Value) -> Result>> { Ok( - rivet_key_data::versioned::MetadataKeyData::latest(value.try_into()?) - .serialize_with_embedded_version(rivet_key_data::PEGBOARD_RUNNER_METADATA_VERSION)? + rivet_data::versioned::MetadataKeyData::latest(value.try_into()?) + .serialize_with_embedded_version(rivet_data::PEGBOARD_RUNNER_METADATA_VERSION)? .chunks(udb_util::CHUNK_SIZE) .map(|x| x.to_vec()) .collect(), diff --git a/packages/services/pegboard/src/lib.rs b/packages/services/pegboard/src/lib.rs index 8a08a5b9a9..b5dd33dd0a 100644 --- a/packages/services/pegboard/src/lib.rs +++ b/packages/services/pegboard/src/lib.rs @@ -2,6 +2,7 @@ use gas::prelude::*; pub mod errors; pub mod keys; +pub mod messages; mod metrics; pub mod ops; pub mod pubsub_subjects; diff --git a/packages/services/pegboard/src/messages.rs b/packages/services/pegboard/src/messages.rs new file mode 100644 index 0000000000..e3ad78680d --- /dev/null +++ b/packages/services/pegboard/src/messages.rs @@ -0,0 +1,4 @@ +use gas::prelude::*; + +#[message("pegboard_bump_outbound_autoscaler")] +pub struct BumpOutboundAutoscaler {} diff --git a/packages/services/pegboard/src/ops/actor/create.rs b/packages/services/pegboard/src/ops/actor/create.rs index 5dd72c7c75..66a7a1b42c 100644 --- a/packages/services/pegboard/src/ops/actor/create.rs +++ b/packages/services/pegboard/src/ops/actor/create.rs @@ -125,8 +125,12 @@ async fn forward_to_datacenter( // Get namespace name for the remote call let namespace = ctx - .op(namespace::ops::get_global::Input { namespace_id }) + .op(namespace::ops::get_global::Input { + namespace_ids: vec![namespace_id], + }) .await? + .into_iter() + .next() .ok_or_else(|| namespace::errors::Namespace::NotFound.build())?; // Generate a new actor ID with the correct datacenter label diff --git a/packages/services/pegboard/src/ops/actor/get_for_key.rs b/packages/services/pegboard/src/ops/actor/get_for_key.rs index d6960e5689..f850f88c69 100644 --- a/packages/services/pegboard/src/ops/actor/get_for_key.rs +++ b/packages/services/pegboard/src/ops/actor/get_for_key.rs @@ -60,9 +60,11 @@ pub async fn pegboard_actor_get_for_key(ctx: &OperationCtx, input: &Input) -> Re // Get namespace name for the remote call let namespace = ctx .op(namespace::ops::get_global::Input { - namespace_id: input.namespace_id, + namespace_ids: vec![input.namespace_id], }) .await? + .into_iter() + .next() .ok_or_else(|| namespace::errors::Namespace::NotFound.build())?; // Make request to remote datacenter diff --git a/packages/services/pegboard/src/ops/actor/list_names.rs b/packages/services/pegboard/src/ops/actor/list_names.rs index 727a0bed6e..8fe72d2100 100644 --- a/packages/services/pegboard/src/ops/actor/list_names.rs +++ b/packages/services/pegboard/src/ops/actor/list_names.rs @@ -1,6 +1,6 @@ use futures_util::{StreamExt, TryStreamExt}; use gas::prelude::*; -use rivet_key_data::converted::ActorNameKeyData; +use rivet_data::converted::ActorNameKeyData; use udb_util::{SNAPSHOT, TxnExt}; use universaldb::{self as udb, options::StreamingMode}; diff --git a/packages/services/pegboard/src/ops/runner/get.rs b/packages/services/pegboard/src/ops/runner/get.rs index d967eccee3..22cbbcc6e4 100644 --- a/packages/services/pegboard/src/ops/runner/get.rs +++ b/packages/services/pegboard/src/ops/runner/get.rs @@ -1,7 +1,7 @@ use anyhow::Result; use futures_util::TryStreamExt; use gas::prelude::*; -use rivet_key_data::generated::pegboard_runner_address_v1::Data as AddressKeyData; +use rivet_data::generated::pegboard_runner_address_v1::Data as AddressKeyData; use rivet_types::runners::Runner; use udb_util::{FormalChunkedKey, SERIALIZABLE, SNAPSHOT, TxnExt}; use universaldb::{self as udb, options::StreamingMode}; diff --git a/packages/services/pegboard/src/ops/runner/update_alloc_idx.rs b/packages/services/pegboard/src/ops/runner/update_alloc_idx.rs index 7b94ce7f8f..39519590df 100644 --- a/packages/services/pegboard/src/ops/runner/update_alloc_idx.rs +++ b/packages/services/pegboard/src/ops/runner/update_alloc_idx.rs @@ -121,7 +121,7 @@ pub async fn pegboard_runner_update_alloc_idx(ctx: &OperationCtx, input: &Input) let remaining_millislots = (remaining_slots * 1000) / total_slots; - let old_alloc_key = keys::datacenter::RunnerAllocIdxKey::new( + let old_alloc_key = keys::ns::RunnerAllocIdxKey::new( namespace_id, name.clone(), version, @@ -140,7 +140,7 @@ pub async fn pegboard_runner_update_alloc_idx(ctx: &OperationCtx, input: &Input) Action::AddIdx => { txs.write( &old_alloc_key, - rivet_key_data::converted::RunnerAllocIdxKeyData { + rivet_data::converted::RunnerAllocIdxKeyData { workflow_id, remaining_slots, total_slots, @@ -162,7 +162,7 @@ pub async fn pegboard_runner_update_alloc_idx(ctx: &OperationCtx, input: &Input) txs.delete(&old_alloc_key); txs.write( - &keys::datacenter::RunnerAllocIdxKey::new( + &keys::ns::RunnerAllocIdxKey::new( namespace_id, name.clone(), version, @@ -170,7 +170,7 @@ pub async fn pegboard_runner_update_alloc_idx(ctx: &OperationCtx, input: &Input) last_ping_ts, runner.runner_id, ), - rivet_key_data::converted::RunnerAllocIdxKeyData { + rivet_data::converted::RunnerAllocIdxKeyData { workflow_id, remaining_slots, total_slots, diff --git a/packages/services/pegboard/src/workflows/actor/actor_keys.rs b/packages/services/pegboard/src/workflows/actor/actor_keys.rs index 1ff028160a..e5cc89a10c 100644 --- a/packages/services/pegboard/src/workflows/actor/actor_keys.rs +++ b/packages/services/pegboard/src/workflows/actor/actor_keys.rs @@ -4,7 +4,7 @@ use epoxy::{ }; use futures_util::TryStreamExt; use gas::prelude::*; -use rivet_key_data::converted::ActorByKeyKeyData; +use rivet_data::converted::ActorByKeyKeyData; use udb_util::prelude::*; use universaldb::{self as udb, FdbBindingError, options::StreamingMode}; diff --git a/packages/services/pegboard/src/workflows/actor/destroy.rs b/packages/services/pegboard/src/workflows/actor/destroy.rs index c408ef6500..44862d219d 100644 --- a/packages/services/pegboard/src/workflows/actor/destroy.rs +++ b/packages/services/pegboard/src/workflows/actor/destroy.rs @@ -1,8 +1,9 @@ use gas::prelude::*; -use rivet_key_data::converted::ActorByKeyKeyData; +use namespace::types::RunnerKind; +use rivet_data::converted::ActorByKeyKeyData; use rivet_runner_protocol::protocol; use udb_util::{SERIALIZABLE, TxnExt}; -use universaldb as udb; +use universaldb::{self as udb, options::MutationType}; use super::{DestroyComplete, DestroyStarted, State}; @@ -85,6 +86,7 @@ async fn update_state_and_fdb( state.namespace_id, &state.runner_name_selector, runner_id, + &state.ns_runner_kind, &tx, ) .await?; @@ -162,6 +164,7 @@ pub(crate) async fn clear_slot( namespace_id: Id, runner_name_selector: &str, runner_id: Id, + ns_runner_kind: &RunnerKind, tx: &udb::RetryableTransaction, ) -> Result<(), udb::FdbBindingError> { let txs = tx.subspace(keys::subspace()); @@ -198,7 +201,7 @@ pub(crate) async fn clear_slot( // Write new remaining slots txs.write(&runner_remaining_slots_key, new_runner_remaining_slots)?; - let old_runner_alloc_key = keys::datacenter::RunnerAllocIdxKey::new( + let old_runner_alloc_key = keys::ns::RunnerAllocIdxKey::new( namespace_id, runner_name_selector.to_string(), runner_version, @@ -213,7 +216,7 @@ pub(crate) async fn clear_slot( txs.delete(&old_runner_alloc_key); let new_remaining_millislots = (new_runner_remaining_slots * 1000) / runner_total_slots; - let new_runner_alloc_key = keys::datacenter::RunnerAllocIdxKey::new( + let new_runner_alloc_key = keys::ns::RunnerAllocIdxKey::new( namespace_id, runner_name_selector.to_string(), runner_version, @@ -224,7 +227,7 @@ pub(crate) async fn clear_slot( txs.write( &new_runner_alloc_key, - rivet_key_data::converted::RunnerAllocIdxKeyData { + rivet_data::converted::RunnerAllocIdxKeyData { workflow_id: runner_workflow_id, remaining_slots: new_runner_remaining_slots, total_slots: runner_total_slots, @@ -232,6 +235,14 @@ pub(crate) async fn clear_slot( )?; } + if let RunnerKind::Outbound { .. } = ns_runner_kind { + txs.atomic_op( + &keys::ns::OutboundDesiredSlotsKey::new(namespace_id, runner_name_selector.to_string()), + &(-1i32).to_le_bytes(), + MutationType::Add, + ); + } + Ok(()) } diff --git a/packages/services/pegboard/src/workflows/actor/mod.rs b/packages/services/pegboard/src/workflows/actor/mod.rs index a512312465..3bff9c38fb 100644 --- a/packages/services/pegboard/src/workflows/actor/mod.rs +++ b/packages/services/pegboard/src/workflows/actor/mod.rs @@ -1,5 +1,6 @@ use futures_util::FutureExt; use gas::prelude::*; +use namespace::types::RunnerKind; use rivet_runner_protocol::protocol; use rivet_types::actors::CrashPolicy; @@ -45,6 +46,9 @@ pub struct State { pub create_ts: i64, pub create_complete_ts: Option, + + pub ns_runner_kind: RunnerKind, + pub start_ts: Option, // NOTE: This is not the alarm ts, this is when the actor started sleeping. See `LifecycleState` for alarm pub sleep_ts: Option, @@ -66,6 +70,7 @@ impl State { runner_name_selector: String, crash_policy: CrashPolicy, create_ts: i64, + ns_runner_kind: RunnerKind, ) -> Self { State { name, @@ -78,6 +83,8 @@ impl State { create_ts, create_complete_ts: None, + ns_runner_kind, + start_ts: None, pending_allocation_ts: None, sleep_ts: None, @@ -115,15 +122,18 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> }) .await?; - if let Err(error) = validation_res { - ctx.msg(Failed { error }) - .tag("actor_id", input.actor_id) - .send() - .await?; + let metadata = match validation_res { + Ok(metadata) => metadata, + Err(error) => { + ctx.msg(Failed { error }) + .tag("actor_id", input.actor_id) + .send() + .await?; - // TODO(RVT-3928): return Ok(Err); - return Ok(()); - } + // TODO(RVT-3928): return Ok(Err); + return Ok(()); + } + }; ctx.activity(setup::InitStateAndUdbInput { actor_id: input.actor_id, @@ -133,6 +143,7 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> runner_name_selector: input.runner_name_selector.clone(), crash_policy: input.crash_policy, create_ts: ctx.create_ts(), + ns_runner_kind: metadata.ns_runner_kind, }) .await?; @@ -156,6 +167,19 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> .tag("actor_id", input.actor_id) .send() .await?; + + // Destroyed early + ctx.workflow(destroy::Input { + namespace_id: input.namespace_id, + actor_id: input.actor_id, + name: input.name.clone(), + key: input.key.clone(), + generation: 0, + kill: false, + }) + .output() + .await?; + return Ok(()); } actor_keys::ReserveKeyOutput::KeyExists { existing_actor_id } => { @@ -170,8 +194,6 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> .await?; // Destroyed early - // - // This will also deallocate any key that was already allocated to Epoxy ctx.workflow(destroy::Input { namespace_id: input.namespace_id, actor_id: input.actor_id, @@ -335,7 +357,7 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> state.alarm_ts = None; state.sleeping = false; - if runtime::reschedule_actor(ctx, &input, state, true).await? { + if runtime::reschedule_actor(ctx, &input, state).await? { // Destroyed early return Ok(Loop::Break(runtime::LifecycleRes { generation: state.generation, @@ -466,7 +488,7 @@ async fn handle_stopped( .await?; } - if runtime::reschedule_actor(ctx, &input, state, false).await? { + if runtime::reschedule_actor(ctx, &input, state).await? { // Destroyed early return Ok(Some(runtime::LifecycleRes { generation: state.generation, diff --git a/packages/services/pegboard/src/workflows/actor/runtime.rs b/packages/services/pegboard/src/workflows/actor/runtime.rs index 6a6a93a14e..b1cb7887e8 100644 --- a/packages/services/pegboard/src/workflows/actor/runtime.rs +++ b/packages/services/pegboard/src/workflows/actor/runtime.rs @@ -3,18 +3,16 @@ use std::time::Instant; use futures_util::StreamExt; use futures_util::{FutureExt, TryStreamExt}; use gas::prelude::*; +use namespace::types::RunnerKind; use rivet_metrics::KeyValue; use rivet_runner_protocol::protocol; use udb_util::{FormalKey, SERIALIZABLE, SNAPSHOT, TxnExt}; use universaldb::{ self as udb, - options::{ConflictRangeType, StreamingMode}, + options::{ConflictRangeType, MutationType, StreamingMode}, }; -use crate::{ - keys, metrics, - workflows::runner::{AllocatePendingActorsInput, RUNNER_ELIGIBLE_THRESHOLD_MS}, -}; +use crate::{keys, metrics, workflows::runner::RUNNER_ELIGIBLE_THRESHOLD_MS}; use super::{ ACTOR_START_THRESHOLD_MS, Allocate, BASE_RETRY_TIMEOUT_MS, Destroy, Input, PendingAllocation, @@ -105,6 +103,7 @@ async fn allocate_actor( let start_instant = Instant::now(); let mut state = ctx.state::()?; let namespace_id = state.namespace_id; + let ns_runner_kind = &state.ns_runner_kind; // NOTE: This txn should closely resemble the one found in the allocate_pending_actors activity of the // client wf @@ -114,13 +113,24 @@ async fn allocate_actor( let ping_threshold_ts = util::timestamp::now() - RUNNER_ELIGIBLE_THRESHOLD_MS; let txs = tx.subspace(keys::subspace()); + // Increment desired slots if namespace has an outbound runner kind + if let RunnerKind::Outbound { .. } = ns_runner_kind { + txs.atomic_op( + &keys::ns::OutboundDesiredSlotsKey::new( + namespace_id, + input.runner_name_selector.clone(), + ), + &1u32.to_le_bytes(), + MutationType::Add, + ); + } + // Check if a queue exists - let pending_actor_subspace = txs.subspace( - &keys::datacenter::PendingActorByRunnerNameSelectorKey::subspace( + let pending_actor_subspace = + txs.subspace(&keys::ns::PendingActorByRunnerNameSelectorKey::subspace( namespace_id, input.runner_name_selector.clone(), - ), - ); + )); let queue_exists = txs .get_ranges_keyvalues( udb::RangeOption { @@ -137,11 +147,10 @@ async fn allocate_actor( .is_some(); if !queue_exists { - let runner_alloc_subspace = - txs.subspace(&keys::datacenter::RunnerAllocIdxKey::subspace( - namespace_id, - input.runner_name_selector.clone(), - )); + let runner_alloc_subspace = txs.subspace(&keys::ns::RunnerAllocIdxKey::subspace( + namespace_id, + input.runner_name_selector.clone(), + )); let mut stream = txs.get_ranges_keyvalues( udb::RangeOption { @@ -161,7 +170,7 @@ async fn allocate_actor( }; let (old_runner_alloc_key, old_runner_alloc_key_data) = - txs.read_entry::(&entry)?; + txs.read_entry::(&entry)?; if let Some(highest_version) = highest_version { // We have passed all of the runners with the highest version. This is reachable if @@ -196,7 +205,7 @@ async fn allocate_actor( // Write new allocation key with 1 less slot txs.write( - &keys::datacenter::RunnerAllocIdxKey::new( + &keys::ns::RunnerAllocIdxKey::new( namespace_id, input.runner_name_selector.clone(), old_runner_alloc_key.version, @@ -204,7 +213,7 @@ async fn allocate_actor( old_runner_alloc_key.last_ping_ts, old_runner_alloc_key.runner_id, ), - rivet_key_data::converted::RunnerAllocIdxKeyData { + rivet_data::converted::RunnerAllocIdxKeyData { workflow_id: old_runner_alloc_key_data.workflow_id, remaining_slots: new_remaining_slots, total_slots: old_runner_alloc_key_data.total_slots, @@ -250,7 +259,7 @@ async fn allocate_actor( // want. If a runner reads from the queue while this is being inserted, one of the two txns will // retry and we ensure the actor does not end up in queue limbo. txs.write( - &keys::datacenter::PendingActorByRunnerNameSelectorKey::new( + &keys::ns::PendingActorByRunnerNameSelectorKey::new( namespace_id, input.runner_name_selector.clone(), pending_ts, @@ -299,7 +308,7 @@ pub async fn set_not_connectable(ctx: &ActivityCtx, input: &SetNotConnectableInp Ok(()) }) - .custom_instrument(tracing::info_span!("actor_deallocate_tx")) + .custom_instrument(tracing::info_span!("actor_set_not_connectable_tx")) .await?; state.connectable_ts = None; @@ -318,11 +327,13 @@ pub async fn deallocate(ctx: &ActivityCtx, input: &DeallocateInput) -> Result<() let runner_name_selector = &state.runner_name_selector; let namespace_id = state.namespace_id; let runner_id = state.runner_id; + let ns_runner_kind = &state.ns_runner_kind; ctx.udb()? .run(|tx, _mc| async move { - let connectable_key = keys::actor::ConnectableKey::new(input.actor_id); - tx.clear(&keys::subspace().pack(&connectable_key)); + let txs = tx.subspace(keys::subspace()); + + txs.delete(&keys::actor::ConnectableKey::new(input.actor_id)); if let Some(runner_id) = runner_id { destroy::clear_slot( @@ -330,9 +341,19 @@ pub async fn deallocate(ctx: &ActivityCtx, input: &DeallocateInput) -> Result<() namespace_id, runner_name_selector, runner_id, + ns_runner_kind, &tx, ) .await?; + } else if let RunnerKind::Outbound { .. } = ns_runner_kind { + txs.atomic_op( + &keys::ns::OutboundDesiredSlotsKey::new( + namespace_id, + runner_name_selector.clone(), + ), + &(-1i32).to_le_bytes(), + MutationType::Add, + ); } Ok(()) @@ -370,6 +391,10 @@ pub async fn spawn_actor( "failed to allocate (no availability), waiting for allocation", ); + ctx.msg(crate::messages::BumpOutboundAutoscaler {}) + .send() + .await?; + // If allocation fails, the allocate txn already inserted this actor into the queue. Now we wait for // an `Allocate` signal match ctx.listen::().await? { @@ -441,35 +466,9 @@ pub async fn reschedule_actor( ctx: &mut WorkflowCtx, input: &Input, state: &mut LifecycleState, - sleeping: bool, ) -> Result { tracing::debug!(actor_id=?input.actor_id, "rescheduling actor"); - // There shouldn't be an allocation if the actor is sleeping - if !sleeping { - ctx.activity(DeallocateInput { - actor_id: input.actor_id, - }) - .await?; - - // Allocate other pending actors from queue - let res = ctx - .activity(AllocatePendingActorsInput { - namespace_id: input.namespace_id, - name: input.runner_name_selector.clone(), - }) - .await?; - - // Dispatch pending allocs - for alloc in res.allocations { - ctx.signal(alloc.signal) - .to_workflow::() - .tag("actor_id", alloc.actor_id) - .send() - .await?; - } - } - let next_generation = state.generation + 1; // Waits for the actor to be ready (or destroyed) and automatically retries if failed to allocate. @@ -563,7 +562,7 @@ pub async fn clear_pending_allocation( .udb()? .run(|tx, _mc| async move { let pending_alloc_key = - keys::subspace().pack(&keys::datacenter::PendingActorByRunnerNameSelectorKey::new( + keys::subspace().pack(&keys::ns::PendingActorByRunnerNameSelectorKey::new( input.namespace_id, input.runner_name_selector.clone(), input.pending_allocation_ts, diff --git a/packages/services/pegboard/src/workflows/actor/setup.rs b/packages/services/pegboard/src/workflows/actor/setup.rs index 9d55ff596f..421b7228fc 100644 --- a/packages/services/pegboard/src/workflows/actor/setup.rs +++ b/packages/services/pegboard/src/workflows/actor/setup.rs @@ -1,5 +1,6 @@ use gas::prelude::*; -use rivet_key_data::converted::ActorNameKeyData; +use namespace::types::RunnerKind; +use rivet_data::converted::ActorNameKeyData; use rivet_types::actors::CrashPolicy; use udb_util::{SERIALIZABLE, TxnExt}; @@ -17,20 +18,25 @@ pub struct ValidateInput { pub input: Option, } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ValidateOutput { + pub ns_runner_kind: RunnerKind, +} + #[activity(Validate)] pub async fn validate( ctx: &ActivityCtx, input: &ValidateInput, -) -> Result> { +) -> Result> { let ns_res = ctx .op(namespace::ops::get_global::Input { - namespace_id: input.namespace_id, + namespace_ids: vec![input.namespace_id], }) .await?; - if ns_res.is_none() { + let Some(ns) = ns_res.into_iter().next() else { return Ok(Err(errors::Actor::NamespaceNotFound)); - } + }; if input .input @@ -55,7 +61,9 @@ pub async fn validate( } } - Ok(Ok(())) + Ok(Ok(ValidateOutput { + ns_runner_kind: ns.runner_kind, + })) } #[derive(Debug, Clone, Serialize, Deserialize, Hash)] @@ -67,6 +75,7 @@ pub struct InitStateAndUdbInput { pub runner_name_selector: String, pub crash_policy: CrashPolicy, pub create_ts: i64, + pub ns_runner_kind: RunnerKind, } #[activity(InitStateAndFdb)] @@ -80,6 +89,7 @@ pub async fn insert_state_and_fdb(ctx: &ActivityCtx, input: &InitStateAndUdbInpu input.runner_name_selector.clone(), input.crash_policy, input.create_ts, + input.ns_runner_kind.clone(), )); ctx.udb()? diff --git a/packages/services/pegboard/src/workflows/runner.rs b/packages/services/pegboard/src/workflows/runner.rs index ff7b5a64cc..e5d64f17d6 100644 --- a/packages/services/pegboard/src/workflows/runner.rs +++ b/packages/services/pegboard/src/workflows/runner.rs @@ -1,6 +1,6 @@ use futures_util::{FutureExt, StreamExt, TryStreamExt}; use gas::prelude::*; -use rivet_key_data::{ +use rivet_data::{ converted::{ActorNameKeyData, MetadataKeyData, RunnerByKeyKeyData}, generated::pegboard_runner_address_v1::Data as AddressKeyData, }; @@ -639,7 +639,7 @@ async fn insert_fdb(ctx: &ActivityCtx, input: &InsertFdbInput) -> Result<()> { // Insert into index (same as the `update_alloc_idx` op with `AddIdx`) txs.write( - &keys::datacenter::RunnerAllocIdxKey::new( + &keys::ns::RunnerAllocIdxKey::new( input.namespace_id, input.name.clone(), input.version, @@ -647,7 +647,7 @@ async fn insert_fdb(ctx: &ActivityCtx, input: &InsertFdbInput) -> Result<()> { last_ping_ts, input.runner_id, ), - rivet_key_data::converted::RunnerAllocIdxKeyData { + rivet_data::converted::RunnerAllocIdxKeyData { workflow_id: ctx.workflow_id(), remaining_slots, total_slots: input.total_slots, @@ -998,12 +998,11 @@ pub(crate) async fn allocate_pending_actors( let txs = tx.subspace(keys::subspace()); let mut results = Vec::new(); - let pending_actor_subspace = txs.subspace( - &keys::datacenter::PendingActorByRunnerNameSelectorKey::subspace( + let pending_actor_subspace = + txs.subspace(&keys::ns::PendingActorByRunnerNameSelectorKey::subspace( input.namespace_id, input.name.clone(), - ), - ); + )); let mut queue_stream = txs.get_ranges_keyvalues( udb::RangeOption { mode: StreamingMode::Iterator, @@ -1021,15 +1020,12 @@ pub(crate) async fn allocate_pending_actors( }; let (queue_key, generation) = - txs.read_entry::( - &queue_entry, - )?; + txs.read_entry::(&queue_entry)?; - let runner_alloc_subspace = - txs.subspace(&keys::datacenter::RunnerAllocIdxKey::subspace( - input.namespace_id, - input.name.clone(), - )); + let runner_alloc_subspace = txs.subspace(&keys::ns::RunnerAllocIdxKey::subspace( + input.namespace_id, + input.name.clone(), + )); let mut stream = txs.get_ranges_keyvalues( udb::RangeOption { @@ -1051,7 +1047,7 @@ pub(crate) async fn allocate_pending_actors( }; let (old_runner_alloc_key, old_runner_alloc_key_data) = - txs.read_entry::(&entry)?; + txs.read_entry::(&entry)?; if let Some(highest_version) = highest_version { // We have passed all of the runners with the highest version. This is reachable if @@ -1088,7 +1084,7 @@ pub(crate) async fn allocate_pending_actors( // Write new allocation key with 1 less slot txs.write( - &keys::datacenter::RunnerAllocIdxKey::new( + &keys::ns::RunnerAllocIdxKey::new( input.namespace_id, input.name.clone(), old_runner_alloc_key.version, @@ -1096,7 +1092,7 @@ pub(crate) async fn allocate_pending_actors( old_runner_alloc_key.last_ping_ts, old_runner_alloc_key.runner_id, ), - rivet_key_data::converted::RunnerAllocIdxKeyData { + rivet_data::converted::RunnerAllocIdxKeyData { workflow_id: old_runner_alloc_key_data.workflow_id, remaining_slots: new_remaining_slots, total_slots: old_runner_alloc_key_data.total_slots, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 91185e953a..d2358c0f78 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -741,6 +741,9 @@ importers: '@rivetkit/engine-runner-protocol': specifier: workspace:* version: link:../runner-protocol + hono: + specifier: ^4.0.0 + version: 4.8.12 ws: specifier: ^8.18.3 version: 8.18.3 diff --git a/sdks/rust/key-data/Cargo.toml b/sdks/rust/data/Cargo.toml similarity index 95% rename from sdks/rust/key-data/Cargo.toml rename to sdks/rust/data/Cargo.toml index 339d60c362..3282e73918 100644 --- a/sdks/rust/key-data/Cargo.toml +++ b/sdks/rust/data/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "rivet-key-data" +name = "rivet-data" version.workspace = true authors.workspace = true license.workspace = true diff --git a/sdks/rust/key-data/build.rs b/sdks/rust/data/build.rs similarity index 99% rename from sdks/rust/key-data/build.rs rename to sdks/rust/data/build.rs index 898e7ea645..e6b18ab845 100644 --- a/sdks/rust/key-data/build.rs +++ b/sdks/rust/data/build.rs @@ -61,7 +61,7 @@ fn main() { .and_then(|p| p.parent()) .expect("Failed to find workspace root"); - let schema_dir = workspace_root.join("sdks").join("schemas").join("key-data"); + let schema_dir = workspace_root.join("sdks").join("schemas").join("data"); println!("cargo:rerun-if-changed={}", schema_dir.display()); diff --git a/sdks/rust/key-data/src/converted.rs b/sdks/rust/data/src/converted.rs similarity index 94% rename from sdks/rust/key-data/src/converted.rs rename to sdks/rust/data/src/converted.rs index 44e954bc84..859682e546 100644 --- a/sdks/rust/key-data/src/converted.rs +++ b/sdks/rust/data/src/converted.rs @@ -9,10 +9,10 @@ pub struct RunnerAllocIdxKeyData { pub total_slots: u32, } -impl TryFrom for RunnerAllocIdxKeyData { +impl TryFrom for RunnerAllocIdxKeyData { type Error = anyhow::Error; - fn try_from(value: pegboard_datacenter_runner_alloc_idx_v1::Data) -> Result { + fn try_from(value: pegboard_namespace_runner_alloc_idx_v1::Data) -> Result { Ok(RunnerAllocIdxKeyData { workflow_id: Id::from_slice(&value.workflow_id)?, remaining_slots: value.remaining_slots, @@ -21,11 +21,11 @@ impl TryFrom for RunnerAllocIdxKe } } -impl TryFrom for pegboard_datacenter_runner_alloc_idx_v1::Data { +impl TryFrom for pegboard_namespace_runner_alloc_idx_v1::Data { type Error = anyhow::Error; fn try_from(value: RunnerAllocIdxKeyData) -> Result { - Ok(pegboard_datacenter_runner_alloc_idx_v1::Data { + Ok(pegboard_namespace_runner_alloc_idx_v1::Data { workflow_id: value.workflow_id.as_bytes(), remaining_slots: value.remaining_slots, total_slots: value.total_slots, diff --git a/sdks/rust/key-data/src/generated.rs b/sdks/rust/data/src/generated.rs similarity index 100% rename from sdks/rust/key-data/src/generated.rs rename to sdks/rust/data/src/generated.rs diff --git a/sdks/rust/key-data/src/lib.rs b/sdks/rust/data/src/lib.rs similarity index 84% rename from sdks/rust/key-data/src/lib.rs rename to sdks/rust/data/src/lib.rs index 3292e22e83..a93ceaea1d 100644 --- a/sdks/rust/key-data/src/lib.rs +++ b/sdks/rust/data/src/lib.rs @@ -2,9 +2,9 @@ pub mod converted; pub mod generated; pub mod versioned; -pub const PEGBOARD_DATACENTER_RUNNER_ALLOC_IDX_VERSION: u16 = 1; pub const PEGBOARD_RUNNER_ADDRESS_VERSION: u16 = 1; pub const PEGBOARD_RUNNER_METADATA_VERSION: u16 = 1; pub const PEGBOARD_NAMESPACE_ACTOR_BY_KEY_VERSION: u16 = 1; +pub const PEGBOARD_NAMESPACE_RUNNER_ALLOC_IDX_VERSION: u16 = 1; pub const PEGBOARD_NAMESPACE_RUNNER_BY_KEY_VERSION: u16 = 1; pub const PEGBOARD_NAMESPACE_ACTOR_NAME_VERSION: u16 = 1; diff --git a/sdks/rust/key-data/src/versioned.rs b/sdks/rust/data/src/versioned.rs similarity index 83% rename from sdks/rust/key-data/src/versioned.rs rename to sdks/rust/data/src/versioned.rs index af709553f3..002363d361 100644 --- a/sdks/rust/key-data/src/versioned.rs +++ b/sdks/rust/data/src/versioned.rs @@ -4,13 +4,13 @@ use versioned_data_util::OwnedVersionedData; use crate::generated::*; pub enum RunnerAllocIdxKeyData { - V1(pegboard_datacenter_runner_alloc_idx_v1::Data), + V1(pegboard_namespace_runner_alloc_idx_v1::Data), } impl OwnedVersionedData for RunnerAllocIdxKeyData { - type Latest = pegboard_datacenter_runner_alloc_idx_v1::Data; + type Latest = pegboard_namespace_runner_alloc_idx_v1::Data; - fn latest(latest: pegboard_datacenter_runner_alloc_idx_v1::Data) -> Self { + fn latest(latest: pegboard_namespace_runner_alloc_idx_v1::Data) -> Self { RunnerAllocIdxKeyData::V1(latest) } @@ -206,3 +206,37 @@ impl OwnedVersionedData for ActorNameKeyData { } } } + +pub enum NamespaceRunnerKind { + V1(namespace_runner_kind_v1::Data), +} + +impl OwnedVersionedData for NamespaceRunnerKind { + type Latest = namespace_runner_kind_v1::Data; + + fn latest(latest: namespace_runner_kind_v1::Data) -> Self { + NamespaceRunnerKind::V1(latest) + } + + fn into_latest(self) -> Result { + #[allow(irrefutable_let_patterns)] + if let NamespaceRunnerKind::V1(data) = self { + Ok(data) + } else { + bail!("version not latest"); + } + } + + fn deserialize_version(payload: &[u8], version: u16) -> Result { + match version { + 1 => Ok(NamespaceRunnerKind::V1(serde_bare::from_slice(payload)?)), + _ => bail!("invalid version: {version}"), + } + } + + fn serialize_version(self, _version: u16) -> Result> { + match self { + NamespaceRunnerKind::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + } + } +} diff --git a/sdks/schemas/data/namespace.runner_kind.v1.bare b/sdks/schemas/data/namespace.runner_kind.v1.bare new file mode 100644 index 0000000000..1cc263b54e --- /dev/null +++ b/sdks/schemas/data/namespace.runner_kind.v1.bare @@ -0,0 +1,14 @@ +type Outbound struct { + url: str + slots_per_runner: u32 + min_runners: u32 + max_runners: u32 + runners_margin: u32 +} + +type Custom void + +type Data union { + Outbound | + Custom +} diff --git a/sdks/schemas/key-data/pegboard.namespace.actor_by_key.v1.bare b/sdks/schemas/data/pegboard.namespace.actor_by_key.v1.bare similarity index 100% rename from sdks/schemas/key-data/pegboard.namespace.actor_by_key.v1.bare rename to sdks/schemas/data/pegboard.namespace.actor_by_key.v1.bare diff --git a/sdks/schemas/key-data/pegboard.namespace.actor_name.v1.bare b/sdks/schemas/data/pegboard.namespace.actor_name.v1.bare similarity index 100% rename from sdks/schemas/key-data/pegboard.namespace.actor_name.v1.bare rename to sdks/schemas/data/pegboard.namespace.actor_name.v1.bare diff --git a/sdks/schemas/key-data/pegboard.datacenter.runner_alloc_idx.v1.bare b/sdks/schemas/data/pegboard.namespace.runner_alloc_idx.v1.bare similarity index 100% rename from sdks/schemas/key-data/pegboard.datacenter.runner_alloc_idx.v1.bare rename to sdks/schemas/data/pegboard.namespace.runner_alloc_idx.v1.bare diff --git a/sdks/schemas/key-data/pegboard.namespace.runner_by_key.v1.bare b/sdks/schemas/data/pegboard.namespace.runner_by_key.v1.bare similarity index 100% rename from sdks/schemas/key-data/pegboard.namespace.runner_by_key.v1.bare rename to sdks/schemas/data/pegboard.namespace.runner_by_key.v1.bare diff --git a/sdks/schemas/key-data/pegboard.runner.address.v1.bare b/sdks/schemas/data/pegboard.runner.address.v1.bare similarity index 100% rename from sdks/schemas/key-data/pegboard.runner.address.v1.bare rename to sdks/schemas/data/pegboard.runner.address.v1.bare diff --git a/sdks/schemas/key-data/pegboard.runner.metadata.v1.bare b/sdks/schemas/data/pegboard.runner.metadata.v1.bare similarity index 100% rename from sdks/schemas/key-data/pegboard.runner.metadata.v1.bare rename to sdks/schemas/data/pegboard.runner.metadata.v1.bare diff --git a/sdks/typescript/runner/src/mod.ts b/sdks/typescript/runner/src/mod.ts index 2623d63848..611f2bbcf0 100644 --- a/sdks/typescript/runner/src/mod.ts +++ b/sdks/typescript/runner/src/mod.ts @@ -35,6 +35,7 @@ export interface RunnerConfig { metadata?: Record; onConnected: () => void; onDisconnected: () => void; + onShutdown: () => void; fetch: (actorId: string, request: Request) => Promise; websocket?: (actorId: string, ws: any, request: Request) => Promise; onActorStart: ( @@ -362,9 +363,9 @@ export class Runner { //console.log("Tunnel shutdown completed"); } - if (exit) { - process.exit(0); - } + if (exit) process.exit(0); + + this.#config.onShutdown(); } // MARK: Networking diff --git a/sdks/typescript/test-runner/package.json b/sdks/typescript/test-runner/package.json index 7092ba4acf..ed441804ed 100644 --- a/sdks/typescript/test-runner/package.json +++ b/sdks/typescript/test-runner/package.json @@ -11,6 +11,7 @@ "@rivetkit/engine-runner": "workspace:*", "@hono/node-server": "^1.18.2", "@rivetkit/engine-runner-protocol": "workspace:*", + "hono": "^4.0.0", "ws": "^8.18.3" }, "devDependencies": { @@ -22,4 +23,4 @@ "typescript": "^5.3.3", "vitest": "^1.6.0" } -} +} \ No newline at end of file diff --git a/sdks/typescript/test-runner/src/main.ts b/sdks/typescript/test-runner/src/main.ts index 596dda5009..fbe681326d 100644 --- a/sdks/typescript/test-runner/src/main.ts +++ b/sdks/typescript/test-runner/src/main.ts @@ -2,6 +2,8 @@ import { Runner } from "@rivetkit/engine-runner"; import type { RunnerConfig, ActorConfig } from "@rivetkit/engine-runner"; import WebSocket from "ws"; import { serve } from "@hono/node-server"; +import { streamSSE } from "hono/streaming"; +import { Hono } from 'hono' const INTERNAL_SERVER_PORT = process.env.INTERNAL_SERVER_PORT ? Number(process.env.INTERNAL_SERVER_PORT) @@ -16,120 +18,150 @@ const RIVET_RUNNER_TOTAL_SLOTS = process.env.RIVET_RUNNER_TOTAL_SLOTS ? Number(process.env.RIVET_RUNNER_TOTAL_SLOTS) : 100; const RIVET_ENDPOINT = process.env.RIVET_ENDPOINT ?? "http://localhost:6420"; +const AUTOSTART = process.env.NO_AUTOSTART == undefined; let runnerStarted = Promise.withResolvers(); +let runnerStopped = Promise.withResolvers(); let websocketOpen = Promise.withResolvers(); let websocketClosed = Promise.withResolvers(); let runner: Runner | null = null; const actorWebSockets = new Map(); -// Start internal server -serve({ - fetch: async (request: Request) => { - const url = new URL(request.url); - if (url.pathname == "/wait-ready") { - await runnerStarted.promise; - return new Response(JSON.stringify(runner?.runnerId), { - status: 200, - }); - } else if (url.pathname == "/has-actor") { - let actorIdQuery = url.searchParams.get("actor"); - let generationQuery = url.searchParams.get("generation"); - let generation = generationQuery - ? Number(generationQuery) - : undefined; - - if (!actorIdQuery || !runner?.hasActor(actorIdQuery, generation)) { - return new Response(undefined, { status: 404 }); - } - } else if (url.pathname == "/shutdown") { - await runner?.shutdown(true); - } +// Create internal server +const app = new Hono(); + +app.get('/wait-ready', async (c) => { + await runnerStarted.promise; + return c.json(runner?.runnerId); +}); + +app.get('/has-actor', async (c) => { + let actorIdQuery = c.req.query('actor'); + let generationQuery = c.req.query('generation'); + let generation = generationQuery ? Number(generationQuery) : undefined; + + if (!actorIdQuery || !runner?.hasActor(actorIdQuery, generation)) { + return c.text('', 404); + } + return c.text('ok'); +}); + +app.get('/shutdown', async (c) => { + await runner?.shutdown(true); + return c.text('ok'); +}); + +app.get('/start', async (c) => { + return streamSSE(c, async (stream) => { + if (runner == null) runner = await startRunner(); + + stream.writeSSE({ data: runner.runnerId! }); + + await runnerStopped.promise; + }); +}); - return new Response("ok", { status: 200 }); - }, +app.get('*', (c) => c.text('ok')); + +serve({ + fetch: app.fetch, port: INTERNAL_SERVER_PORT, }); console.log(`Internal HTTP server listening on port ${INTERNAL_SERVER_PORT}`); -// Use objects to hold the current promise resolvers so callbacks always get the latest -const startedRef = { current: Promise.withResolvers() }; -const stoppedRef = { current: Promise.withResolvers() }; - -const config: RunnerConfig = { - version: RIVET_RUNNER_VERSION, - endpoint: RIVET_ENDPOINT, - namespace: RIVET_NAMESPACE, - runnerName: "test-runner", - runnerKey: RIVET_RUNNER_KEY, - totalSlots: RIVET_RUNNER_TOTAL_SLOTS, - prepopulateActorNames: {}, - onConnected: () => { - runnerStarted.resolve(undefined); - }, - onDisconnected: () => {}, - fetch: async (actorId: string, request: Request) => { - console.log( - `[TEST-RUNNER] Fetch called for actor ${actorId}, URL: ${request.url}`, - ); - const url = new URL(request.url); - if (url.pathname === "/ping") { - // Return the actor ID in response - const responseData = { - actorId, - status: "ok", - timestamp: Date.now(), - }; - console.log(`[TEST-RUNNER] Returning ping response:`, responseData); - return new Response(JSON.stringify(responseData), { - status: 200, - headers: { "Content-Type": "application/json" }, +if (AUTOSTART) runner = await startRunner(); + +async function startRunner(): Promise { + const config: RunnerConfig = { + version: RIVET_RUNNER_VERSION, + endpoint: RIVET_ENDPOINT, + namespace: RIVET_NAMESPACE, + runnerName: "test-runner", + runnerKey: RIVET_RUNNER_KEY, + totalSlots: RIVET_RUNNER_TOTAL_SLOTS, + prepopulateActorNames: {}, + onConnected: () => { + runnerStarted.resolve(undefined); + }, + onDisconnected: () => { }, + onShutdown: () => { + runnerStopped.resolve(undefined); + }, + fetch: async (actorId: string, request: Request) => { + console.log(`[TEST-RUNNER] Fetch called for actor ${actorId}, URL: ${request.url}`); + const url = new URL(request.url); + if (url.pathname === "/ping") { + // Return the actor ID in response + const responseData = { + actorId, + status: "ok", + timestamp: Date.now(), + }; + console.log(`[TEST-RUNNER] Returning ping response:`, responseData); + return new Response( + JSON.stringify(responseData), + { + status: 200, + headers: { "Content-Type": "application/json" }, + }, + ); + } + + return new Response("ok", { status: 200 }); + }, + onActorStart: async ( + _actorId: string, + _generation: number, + _config: ActorConfig, + ) => { + console.log( + `Actor ${_actorId} started (generation ${_generation})`, + ); + }, + onActorStop: async (_actorId: string, _generation: number) => { + console.log( + `Actor ${_actorId} stopped (generation ${_generation})`, + ); + }, + websocket: async ( + actorId: string, + ws: WebSocket, + request: Request, + ) => { + console.log(`WebSocket connected for actor ${actorId}`); + websocketOpen.resolve(undefined); + actorWebSockets.set(actorId, ws); + + // Echo server - send back any messages received + ws.addEventListener("message", (event) => { + const data = event.data; + console.log( + `WebSocket message from actor ${actorId}:`, + data, + ); + ws.send(`Echo: ${data}`); + }); + + ws.addEventListener("close", () => { + console.log(`WebSocket closed for actor ${actorId}`); + actorWebSockets.delete(actorId); + websocketClosed.resolve(undefined); + }); + + ws.addEventListener("error", (error) => { + console.error(`WebSocket error for actor ${actorId}:`, error); }); - } - - return new Response("ok", { status: 200 }); - }, - onActorStart: async ( - _actorId: string, - _generation: number, - _config: ActorConfig, - ) => { - console.log(`Actor ${_actorId} started (generation ${_generation})`); - startedRef.current.resolve(undefined); - }, - onActorStop: async (_actorId: string, _generation: number) => { - console.log(`Actor ${_actorId} stopped (generation ${_generation})`); - stoppedRef.current.resolve(undefined); - }, - websocket: async (actorId: string, ws: WebSocket, request: Request) => { - console.log(`WebSocket connected for actor ${actorId}`); - websocketOpen.resolve(undefined); - actorWebSockets.set(actorId, ws); - - // Echo server - send back any messages received - ws.addEventListener("message", (event) => { - const data = event.data; - console.log(`WebSocket message from actor ${actorId}:`, data); - ws.send(`Echo: ${data}`); - }); - - ws.addEventListener("close", () => { - console.log(`WebSocket closed for actor ${actorId}`); - actorWebSockets.delete(actorId); - websocketClosed.resolve(undefined); - }); - - ws.addEventListener("error", (error) => { - console.error(`WebSocket error for actor ${actorId}:`, error); - }); - }, -}; - -runner = new Runner(config); - -// Start runner -await runner.start(); - -// Wait for runner to be ready -console.log("Waiting runner start..."); -await runnerStarted.promise; + }, + }; + + runner = new Runner(config); + + // Start runner + await runner.start(); + + // Wait for runner to be ready + console.log("Waiting runner start..."); + await runnerStarted.promise; + + return runner; +} \ No newline at end of file From 3df4e46516c5e6ae4f53b6e23c4ed91514fa8127 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Mon, 8 Sep 2025 14:03:56 -0700 Subject: [PATCH 09/17] feat(ns): implement namespace updates, global cache purge --- Cargo.lock | 16 ++ Cargo.toml | 6 +- docker/dev/docker-compose.yml | 2 +- docker/dev/rivet-engine/config.jsonc | 7 - out/errors/namespace.invalid_update.json | 5 + out/openapi.json | 184 ++++++++++++------ packages/common/cache/build/src/key.rs | 16 ++ .../gasoline/core/src/db/kv/keys/history.rs | 1 + .../gasoline/core/src/db/kv/keys/workflow.rs | 2 + packages/common/gasoline/macros/src/lib.rs | 3 +- packages/common/types/Cargo.toml | 4 +- packages/common/types/README.md | 3 + packages/common/types/src/keys/mod.rs | 1 + .../common/types/src/keys/pegboard/mod.rs | 1 + packages/common/types/src/keys/pegboard/ns.rs | 109 +++++++++++ packages/common/types/src/lib.rs | 2 + packages/common/types/src/msgs/mod.rs | 1 + .../types/src/msgs/pegboard.rs} | 0 packages/common/udb-util/src/ext.rs | 48 ++++- packages/common/udb-util/src/keys.rs | 1 + packages/common/universaldb/src/atomic.rs | 9 +- packages/core/actor-kv/src/entry.rs | 1 + packages/core/actor-kv/src/key.rs | 2 +- packages/core/api-peer/src/internal.rs | 29 +++ packages/core/api-peer/src/lib.rs | 1 + packages/core/api-peer/src/namespaces.rs | 79 ++++++-- packages/core/api-peer/src/router.rs | 5 +- packages/core/api-public/src/namespaces.rs | 51 +++++ packages/core/api-public/src/router.rs | 5 + packages/core/bootstrap/src/lib.rs | 3 +- packages/core/pegboard-outbound/Cargo.toml | 1 + packages/core/pegboard-outbound/src/lib.rs | 95 +++++---- packages/services/internal/Cargo.toml | 12 ++ packages/services/internal/README.md | 1 + packages/services/internal/src/lib.rs | 1 + .../services/internal/src/ops/cache/mod.rs | 1 + .../internal/src/ops/cache/purge_global.rs | 78 ++++++++ packages/services/internal/src/ops/mod.rs | 1 + packages/services/namespace/Cargo.toml | 5 +- packages/services/namespace/src/errors.rs | 7 + packages/services/namespace/src/keys.rs | 2 +- .../src/ops/resolve_for_name_local.rs | 20 +- packages/services/namespace/src/types.rs | 6 + .../namespace/src/workflows/namespace.rs | 148 ++++++++++++-- packages/services/pegboard/Cargo.toml | 4 +- packages/services/pegboard/src/keys/ns.rs | 84 -------- packages/services/pegboard/src/lib.rs | 1 - .../pegboard/src/workflows/actor/destroy.rs | 5 +- .../pegboard/src/workflows/actor/runtime.rs | 6 +- .../data/namespace.runner_kind.v1.bare | 1 + 50 files changed, 814 insertions(+), 262 deletions(-) create mode 100644 out/errors/namespace.invalid_update.json create mode 100644 packages/common/types/README.md create mode 100644 packages/common/types/src/keys/mod.rs create mode 100644 packages/common/types/src/keys/pegboard/mod.rs create mode 100644 packages/common/types/src/keys/pegboard/ns.rs create mode 100644 packages/common/types/src/msgs/mod.rs rename packages/{services/pegboard/src/messages.rs => common/types/src/msgs/pegboard.rs} (100%) create mode 100644 packages/core/api-peer/src/internal.rs create mode 100644 packages/services/internal/Cargo.toml create mode 100644 packages/services/internal/README.md create mode 100644 packages/services/internal/src/lib.rs create mode 100644 packages/services/internal/src/ops/cache/mod.rs create mode 100644 packages/services/internal/src/ops/cache/purge_global.rs create mode 100644 packages/services/internal/src/ops/mod.rs diff --git a/Cargo.lock b/Cargo.lock index f1f2a7a52d..0758fb75e0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2378,6 +2378,16 @@ dependencies = [ "generic-array", ] +[[package]] +name = "internal" +version = "0.0.1" +dependencies = [ + "anyhow", + "gasoline", + "rivet-api-client", + "serde", +] + [[package]] name = "io-uring" version = "0.7.9" @@ -2771,15 +2781,18 @@ version = "0.0.1" dependencies = [ "anyhow", "gasoline", + "internal", "rivet-api-builder", "rivet-api-util", "rivet-data", "rivet-error", + "rivet-types", "rivet-util", "serde", "tracing", "udb-util", "universaldb", + "url", "utoipa", "versioned-data-util", ] @@ -3327,6 +3340,7 @@ dependencies = [ "reqwest-eventsource", "rivet-config", "rivet-runner-protocol", + "rivet-types", "tracing", "udb-util", "universaldb", @@ -4694,7 +4708,9 @@ dependencies = [ "rivet-runner-protocol", "rivet-util", "serde", + "udb-util", "utoipa", + "versioned-data-util", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index cc32db3c24..5a96596a8e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [workspace] resolver = "2" -members = ["packages/common/api-builder","packages/common/api-client","packages/common/api-types","packages/common/api-util","packages/common/cache/build","packages/common/cache/result","packages/common/clickhouse-inserter","packages/common/clickhouse-user-query","packages/common/config","packages/common/env","packages/common/error/core","packages/common/error/macros","packages/common/gasoline/core","packages/common/gasoline/macros","packages/common/logs","packages/common/metrics","packages/common/pools","packages/common/runtime","packages/common/service-manager","packages/common/telemetry","packages/common/test-deps","packages/common/test-deps-docker","packages/common/types","packages/common/udb-util","packages/common/universaldb","packages/common/universalpubsub","packages/common/util/core","packages/common/util/id","packages/common/versioned-data-util","packages/core/actor-kv","packages/core/api-peer","packages/core/api-public","packages/core/bootstrap","packages/core/dump-openapi","packages/core/guard/core","packages/core/guard/server","packages/core/pegboard-gateway","packages/core/pegboard-outbound","packages/core/pegboard-runner-ws","packages/core/pegboard-tunnel","packages/core/workflow-worker","packages/infra/engine","packages/services/epoxy","packages/services/namespace","packages/services/pegboard","sdks/rust/api-full","sdks/rust/bare_gen","sdks/rust/data","sdks/rust/epoxy-protocol","sdks/rust/runner-protocol","sdks/rust/tunnel-protocol","sdks/rust/ups-protocol"] +members = ["packages/common/api-builder","packages/common/api-client","packages/common/api-types","packages/common/api-util","packages/common/cache/build","packages/common/cache/result","packages/common/clickhouse-inserter","packages/common/clickhouse-user-query","packages/common/config","packages/common/env","packages/common/error/core","packages/common/error/macros","packages/common/gasoline/core","packages/common/gasoline/macros","packages/common/logs","packages/common/metrics","packages/common/pools","packages/common/runtime","packages/common/service-manager","packages/common/telemetry","packages/common/test-deps","packages/common/test-deps-docker","packages/common/types","packages/common/udb-util","packages/common/universaldb","packages/common/universalpubsub","packages/common/util/core","packages/common/util/id","packages/common/versioned-data-util","packages/core/actor-kv","packages/core/api-peer","packages/core/api-public","packages/core/bootstrap","packages/core/dump-openapi","packages/core/guard/core","packages/core/guard/server","packages/core/pegboard-gateway","packages/core/pegboard-outbound","packages/core/pegboard-runner-ws","packages/core/pegboard-tunnel","packages/core/workflow-worker","packages/infra/engine","packages/services/epoxy","packages/services/internal","packages/services/namespace","packages/services/pegboard","sdks/rust/api-full","sdks/rust/bare_gen","sdks/rust/data","sdks/rust/epoxy-protocol","sdks/rust/runner-protocol","sdks/rust/tunnel-protocol","sdks/rust/ups-protocol"] [workspace.package] version = "0.0.1" @@ -280,7 +280,6 @@ path = "packages/common/error/core" [workspace.dependencies.rivet-error-macros] path = "packages/common/error/macros" - [workspace.dependencies.gas] package = "gasoline" path = "packages/common/gasoline/core" @@ -382,6 +381,9 @@ path = "packages/infra/engine" [workspace.dependencies.epoxy] path = "packages/services/epoxy" +[workspace.dependencies.internal] +path = "packages/services/internal" + [workspace.dependencies.namespace] path = "packages/services/namespace" diff --git a/docker/dev/docker-compose.yml b/docker/dev/docker-compose.yml index 64de5f178d..d1b9fd8cfb 100644 --- a/docker/dev/docker-compose.yml +++ b/docker/dev/docker-compose.yml @@ -187,7 +187,7 @@ services: environment: - RIVET_ENDPOINT=http://rivet-engine:6420 - RUNNER_HOST=runner - # - NO_AUTOSTART=1 + - NO_AUTOSTART=1 stop_grace_period: 4s ports: - '5050:5050' diff --git a/docker/dev/rivet-engine/config.jsonc b/docker/dev/rivet-engine/config.jsonc index 810d5be6e5..ba8f2b4dea 100644 --- a/docker/dev/rivet-engine/config.jsonc +++ b/docker/dev/rivet-engine/config.jsonc @@ -32,13 +32,6 @@ "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" }, - "postgres_notify": { - "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine", - "memory_optimization": false - }, - // "memory": { - // "channel": "default" - // }, "cache": { "driver": "in_memory" }, diff --git a/out/errors/namespace.invalid_update.json b/out/errors/namespace.invalid_update.json new file mode 100644 index 0000000000..2dda6ea2dc --- /dev/null +++ b/out/errors/namespace.invalid_update.json @@ -0,0 +1,5 @@ +{ + "code": "invalid_update", + "group": "namespace", + "message": "Failed to update namespace." +} \ No newline at end of file diff --git a/out/openapi.json b/out/openapi.json index ff97997a91..f4278895ed 100644 --- a/out/openapi.json +++ b/out/openapi.json @@ -469,7 +469,7 @@ { "name": "namespace_id", "in": "query", - "required": true, + "required": false, "schema": { "type": "array", "items": { @@ -548,6 +548,44 @@ } } } + }, + "put": { + "tags": [ + "namespaces" + ], + "operationId": "namespaces_update", + "parameters": [ + { + "name": "namespace_id", + "in": "path", + "required": true, + "schema": { + "$ref": "#/components/schemas/RivetId" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/NamespacesUpdateRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "description": "", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/NamespacesUpdateResponse" + } + } + } + } + } } }, "/runners": { @@ -1071,7 +1109,7 @@ "$ref": "#/components/schemas/RivetId" }, "runner_kind": { - "$ref": "#/components/schemas/RunnerKind" + "$ref": "#/components/schemas/NamespacesRunnerKind" } } }, @@ -1133,6 +1171,95 @@ }, "additionalProperties": false }, + "NamespacesRunnerKind": { + "oneOf": [ + { + "type": "object", + "required": [ + "outbound" + ], + "properties": { + "outbound": { + "type": "object", + "required": [ + "url", + "request_lifespan", + "slots_per_runner", + "min_runners", + "max_runners", + "runners_margin" + ], + "properties": { + "max_runners": { + "type": "integer", + "format": "int32", + "minimum": 0 + }, + "min_runners": { + "type": "integer", + "format": "int32", + "minimum": 0 + }, + "request_lifespan": { + "type": "integer", + "format": "int32", + "description": "Seconds.", + "minimum": 0 + }, + "runners_margin": { + "type": "integer", + "format": "int32", + "minimum": 0 + }, + "slots_per_runner": { + "type": "integer", + "format": "int32", + "minimum": 0 + }, + "url": { + "type": "string" + } + } + } + } + }, + { + "type": "string", + "enum": [ + "custom" + ] + } + ] + }, + "NamespacesUpdate": { + "oneOf": [ + { + "type": "object", + "required": [ + "update_runner_kind" + ], + "properties": { + "update_runner_kind": { + "type": "object", + "required": [ + "runner_kind" + ], + "properties": { + "runner_kind": { + "$ref": "#/components/schemas/NamespacesRunnerKind" + } + } + } + } + } + ] + }, + "NamespacesUpdateRequest": { + "$ref": "#/components/schemas/NamespacesUpdate" + }, + "NamespacesUpdateResponse": { + "type": "object" + }, "Pagination": { "type": "object", "properties": { @@ -1253,59 +1380,6 @@ }, "additionalProperties": false }, - "RunnerKind": { - "oneOf": [ - { - "type": "object", - "required": [ - "outbound" - ], - "properties": { - "outbound": { - "type": "object", - "required": [ - "url", - "slots_per_runner", - "min_runners", - "max_runners", - "runners_margin" - ], - "properties": { - "max_runners": { - "type": "integer", - "format": "int32", - "minimum": 0 - }, - "min_runners": { - "type": "integer", - "format": "int32", - "minimum": 0 - }, - "runners_margin": { - "type": "integer", - "format": "int32", - "minimum": 0 - }, - "slots_per_runner": { - "type": "integer", - "format": "int32", - "minimum": 0 - }, - "url": { - "type": "string" - } - } - } - } - }, - { - "type": "string", - "enum": [ - "custom" - ] - } - ] - }, "RunnersGetResponse": { "type": "object", "required": [ diff --git a/packages/common/cache/build/src/key.rs b/packages/common/cache/build/src/key.rs index 22944adf49..0290799b4a 100644 --- a/packages/common/cache/build/src/key.rs +++ b/packages/common/cache/build/src/key.rs @@ -1,3 +1,4 @@ +use serde::{Deserialize, Serialize}; use std::fmt::Debug; /// A type that can be serialized in to a key that can be used in the cache. @@ -78,3 +79,18 @@ impl_to_string!(i32); impl_to_string!(i64); impl_to_string!(i128); impl_to_string!(isize); + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] +pub struct RawCacheKey(String); + +impl CacheKey for RawCacheKey { + fn cache_key(&self) -> String { + self.0.clone() + } +} + +impl From for RawCacheKey { + fn from(value: String) -> Self { + RawCacheKey(value) + } +} diff --git a/packages/common/gasoline/core/src/db/kv/keys/history.rs b/packages/common/gasoline/core/src/db/kv/keys/history.rs index 8e6c1d95bd..9b8a838fd1 100644 --- a/packages/common/gasoline/core/src/db/kv/keys/history.rs +++ b/packages/common/gasoline/core/src/db/kv/keys/history.rs @@ -1151,6 +1151,7 @@ impl<'de> TupleUnpack<'de> for InnerEventTypeKey { } } +#[derive(Debug)] pub struct TagKey { workflow_id: Id, location: Location, diff --git a/packages/common/gasoline/core/src/db/kv/keys/workflow.rs b/packages/common/gasoline/core/src/db/kv/keys/workflow.rs index eea0f0e3b4..4a5e01d8ef 100644 --- a/packages/common/gasoline/core/src/db/kv/keys/workflow.rs +++ b/packages/common/gasoline/core/src/db/kv/keys/workflow.rs @@ -71,6 +71,7 @@ impl TuplePack for LeaseSubspaceKey { } } +#[derive(Debug)] pub struct TagKey { workflow_id: Id, pub k: String, @@ -882,6 +883,7 @@ impl TuplePack for PendingSignalSubspaceKey { } } +#[derive(Debug)] pub struct ByNameAndTagKey { workflow_name: String, k: String, diff --git a/packages/common/gasoline/macros/src/lib.rs b/packages/common/gasoline/macros/src/lib.rs index 0afef2ea64..178aec078e 100644 --- a/packages/common/gasoline/macros/src/lib.rs +++ b/packages/common/gasoline/macros/src/lib.rs @@ -166,6 +166,7 @@ pub fn operation(attr: TokenStream, item: TokenStream) -> TokenStream { let struct_ident = Ident::new(&name, proc_macro2::Span::call_site()); let fn_name = item_fn.sig.ident.to_string(); + let generics = &item_fn.sig.generics; let fn_body = item_fn.block; let vis = item_fn.vis; @@ -186,7 +187,7 @@ pub fn operation(attr: TokenStream, item: TokenStream) -> TokenStream { const NAME: &'static str = #fn_name; const TIMEOUT: std::time::Duration = std::time::Duration::from_secs(#timeout); - async fn run(#ctx_ident: #ctx_ty, #input_ident: &Self::Input) -> Result { + async fn run #generics (#ctx_ident: #ctx_ty, #input_ident: &Self::Input) -> Result { #fn_body } } diff --git a/packages/common/types/Cargo.toml b/packages/common/types/Cargo.toml index 9ae35c64aa..cad21dfa05 100644 --- a/packages/common/types/Cargo.toml +++ b/packages/common/types/Cargo.toml @@ -9,8 +9,10 @@ license.workspace = true anyhow.workspace = true gas.workspace = true rivet-api-builder.workspace = true -rivet-runner-protocol.workspace = true rivet-data.workspace = true +rivet-runner-protocol.workspace = true rivet-util.workspace = true serde.workspace = true +udb-util.workspace = true utoipa.workspace = true +versioned-data-util.workspace = true diff --git a/packages/common/types/README.md b/packages/common/types/README.md new file mode 100644 index 0000000000..1f22f90271 --- /dev/null +++ b/packages/common/types/README.md @@ -0,0 +1,3 @@ +# Common + +This pkg exists to get around cargo cyclical deps. diff --git a/packages/common/types/src/keys/mod.rs b/packages/common/types/src/keys/mod.rs new file mode 100644 index 0000000000..38311a6f72 --- /dev/null +++ b/packages/common/types/src/keys/mod.rs @@ -0,0 +1 @@ +pub mod pegboard; diff --git a/packages/common/types/src/keys/pegboard/mod.rs b/packages/common/types/src/keys/pegboard/mod.rs new file mode 100644 index 0000000000..7e0a481030 --- /dev/null +++ b/packages/common/types/src/keys/pegboard/mod.rs @@ -0,0 +1 @@ +pub mod ns; diff --git a/packages/common/types/src/keys/pegboard/ns.rs b/packages/common/types/src/keys/pegboard/ns.rs new file mode 100644 index 0000000000..a768f14f61 --- /dev/null +++ b/packages/common/types/src/keys/pegboard/ns.rs @@ -0,0 +1,109 @@ +use std::result::Result::Ok; + +use anyhow::*; +use gas::prelude::*; +use udb_util::prelude::*; + +#[derive(Debug)] +pub struct OutboundDesiredSlotsKey { + pub namespace_id: Id, + pub runner_name_selector: String, +} + +impl OutboundDesiredSlotsKey { + pub fn new(namespace_id: Id, runner_name_selector: String) -> Self { + OutboundDesiredSlotsKey { + namespace_id, + runner_name_selector, + } + } + + pub fn subspace(namespace_id: Id) -> OutboundDesiredSlotsSubspaceKey { + OutboundDesiredSlotsSubspaceKey::new(namespace_id) + } + + pub fn entire_subspace() -> OutboundDesiredSlotsSubspaceKey { + OutboundDesiredSlotsSubspaceKey::entire() + } +} + +impl FormalKey for OutboundDesiredSlotsKey { + /// Count. + type Value = u32; + + fn deserialize(&self, raw: &[u8]) -> Result { + // NOTE: Atomic ops use little endian + Ok(u32::from_le_bytes(raw.try_into()?)) + } + + fn serialize(&self, value: Self::Value) -> Result> { + // NOTE: Atomic ops use little endian + Ok(value.to_le_bytes().to_vec()) + } +} + +impl TuplePack for OutboundDesiredSlotsKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + NAMESPACE, + OUTBOUND, + DESIRED_SLOTS, + self.namespace_id, + &self.runner_name_selector, + ); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for OutboundDesiredSlotsKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, _, namespace_id, runner_name_selector)) = + <(usize, usize, usize, Id, String)>::unpack(input, tuple_depth)?; + + let v = OutboundDesiredSlotsKey { + namespace_id, + runner_name_selector, + }; + + Ok((input, v)) + } +} + +pub struct OutboundDesiredSlotsSubspaceKey { + namespace_id: Option, +} + +impl OutboundDesiredSlotsSubspaceKey { + pub fn new(namespace_id: Id) -> Self { + OutboundDesiredSlotsSubspaceKey { + namespace_id: Some(namespace_id), + } + } + + pub fn entire() -> Self { + OutboundDesiredSlotsSubspaceKey { namespace_id: None } + } +} + +impl TuplePack for OutboundDesiredSlotsSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let mut offset = VersionstampOffset::None { size: 0 }; + + let t = (NAMESPACE, OUTBOUND, DESIRED_SLOTS); + offset += t.pack(w, tuple_depth)?; + + if let Some(namespace_id) = self.namespace_id { + offset += namespace_id.pack(w, tuple_depth)?; + } + + Ok(offset) + } +} diff --git a/packages/common/types/src/lib.rs b/packages/common/types/src/lib.rs index fc3902d998..19ca637eaf 100644 --- a/packages/common/types/src/lib.rs +++ b/packages/common/types/src/lib.rs @@ -1,3 +1,5 @@ pub mod actors; pub mod datacenters; +pub mod keys; +pub mod msgs; pub mod runners; diff --git a/packages/common/types/src/msgs/mod.rs b/packages/common/types/src/msgs/mod.rs new file mode 100644 index 0000000000..38311a6f72 --- /dev/null +++ b/packages/common/types/src/msgs/mod.rs @@ -0,0 +1 @@ +pub mod pegboard; diff --git a/packages/services/pegboard/src/messages.rs b/packages/common/types/src/msgs/pegboard.rs similarity index 100% rename from packages/services/pegboard/src/messages.rs rename to packages/common/types/src/msgs/pegboard.rs diff --git a/packages/common/udb-util/src/ext.rs b/packages/common/udb-util/src/ext.rs index 064e8a8fef..82ce1e159d 100644 --- a/packages/common/udb-util/src/ext.rs +++ b/packages/common/udb-util/src/ext.rs @@ -1,4 +1,4 @@ -use std::{fmt::Debug, ops::Deref, result::Result::Ok}; +use std::{ops::Deref, result::Result::Ok}; use anyhow::*; use futures_util::TryStreamExt; @@ -44,6 +44,7 @@ impl<'a> TxnSubspace<'a> { ) -> Result { self.subspace .unpack(key) + .context("failed unpacking key") .map_err(|x| udb::FdbBindingError::CustomError(x.into())) } @@ -55,13 +56,19 @@ impl<'a> TxnSubspace<'a> { self.tx.set( &self.subspace.pack(key), &key.serialize(value) + .with_context(|| { + format!( + "failed serializing key value of {}", + std::any::type_name::() + ) + }) .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, ); Ok(()) } - pub async fn read<'de, T: Debug + FormalKey + TuplePack + TupleUnpack<'de>>( + pub async fn read<'de, T: FormalKey + TuplePack + TupleUnpack<'de>>( &self, key: &'de T, snapshot: bool, @@ -99,6 +106,11 @@ impl<'a> TxnSubspace<'a> { self.tx.clear(&self.subspace.pack(key)); } + pub fn delete_key_subspace(&self, key: &T) { + self.tx + .clear_subspace_range(&self.subspace(&self.subspace.pack(key))); + } + pub fn read_entry TupleUnpack<'de>>( &self, entry: &udb::future::FdbValue, @@ -106,6 +118,12 @@ impl<'a> TxnSubspace<'a> { let key = self.unpack::(entry.key())?; let value = key .deserialize(entry.value()) + .with_context(|| { + format!( + "failed deserializing key value of {}", + std::any::type_name::() + ) + }) .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; Ok((key, value)) @@ -131,7 +149,7 @@ impl<'a> TxnSubspace<'a> { .map_err(Into::into) } - pub fn atomic_op<'de, T: Debug + FormalKey + TuplePack + TupleUnpack<'de>>( + pub fn atomic_op<'de, T: FormalKey + TuplePack + TupleUnpack<'de>>( &self, key: &'de T, param: &[u8], @@ -157,7 +175,7 @@ pub trait SliceExt { } pub trait OptSliceExt { - fn read<'de, T: Debug + FormalKey + TupleUnpack<'de>>( + fn read<'de, T: FormalKey + TupleUnpack<'de>>( &self, key: &'de T, ) -> Result; @@ -173,18 +191,30 @@ impl SliceExt for udb::future::FdbSlice { key: &'de T, ) -> Result { key.deserialize(self) + .with_context(|| { + format!( + "failed deserializing key value of {}", + std::any::type_name::() + ) + }) .map_err(|x| udb::FdbBindingError::CustomError(x.into())) } } impl OptSliceExt for Option { - fn read<'de, T: Debug + FormalKey + TupleUnpack<'de>>( + fn read<'de, T: FormalKey + TupleUnpack<'de>>( &self, key: &'de T, ) -> Result { key.deserialize(&self.as_ref().ok_or(udb::FdbBindingError::CustomError( - format!("key should exist: {key:?}").into(), + format!("key should exist: {}", std::any::type_name::()).into(), ))?) + .with_context(|| { + format!( + "failed deserializing key value of {}", + std::any::type_name::() + ) + }) .map_err(|x| udb::FdbBindingError::CustomError(x.into())) } @@ -195,6 +225,12 @@ impl OptSliceExt for Option { if let Some(data) = self { key.deserialize(data) .map(Some) + .with_context(|| { + format!( + "failed deserializing key value of {}", + std::any::type_name::() + ) + }) .map_err(|x| udb::FdbBindingError::CustomError(x.into())) } else { Ok(None) diff --git a/packages/common/udb-util/src/keys.rs b/packages/common/udb-util/src/keys.rs index c0406fe67f..f873ca7066 100644 --- a/packages/common/udb-util/src/keys.rs +++ b/packages/common/udb-util/src/keys.rs @@ -121,4 +121,5 @@ define_keys! { (93, INSTANCE_BALLOT, "instance_ballot"), (94, OUTBOUND, "outbound"), (95, DESIRED_SLOTS, "desired_slots"), + (96, RUNNER_KIND, "runner_kind"), } diff --git a/packages/common/universaldb/src/atomic.rs b/packages/common/universaldb/src/atomic.rs index 3cbbfa5ec9..09e6eb52d6 100644 --- a/packages/common/universaldb/src/atomic.rs +++ b/packages/common/universaldb/src/atomic.rs @@ -36,7 +36,7 @@ fn apply_add(current: Option<&[u8]>, param: &[u8]) -> Vec { let param_int = bytes_to_i64_le(param); let result = current_int.wrapping_add(param_int); - i64_to_bytes_le(result, param.len().max(current.len()).max(8)) + i64_to_bytes_le(result, param.len().max(current.len())) } fn apply_bit_and(current: Option<&[u8]>, param: &[u8]) -> Vec { @@ -174,7 +174,7 @@ fn bytes_to_i64_le(bytes: &[u8]) -> i64 { } let mut padded = [0u8; 8]; - let len = bytes.len().min(8); + let len = bytes.len(); padded[..len].copy_from_slice(&bytes[..len]); i64::from_le_bytes(padded) @@ -182,10 +182,9 @@ fn bytes_to_i64_le(bytes: &[u8]) -> i64 { fn i64_to_bytes_le(value: i64, min_len: usize) -> Vec { let bytes = value.to_le_bytes(); - let len = min_len.max(8); - let mut result = vec![0u8; len]; - result[..8].copy_from_slice(&bytes); + let mut result = vec![0u8; min_len]; + result.copy_from_slice(&bytes[..min_len]); result } diff --git a/packages/core/actor-kv/src/entry.rs b/packages/core/actor-kv/src/entry.rs index d9a1f89792..143348be2f 100644 --- a/packages/core/actor-kv/src/entry.rs +++ b/packages/core/actor-kv/src/entry.rs @@ -99,6 +99,7 @@ impl<'de> TupleUnpack<'de> for EntryValueChunkKey { } } +#[derive(Debug)] pub struct EntryMetadataKey { pub key: KeyWrapper, } diff --git a/packages/core/actor-kv/src/key.rs b/packages/core/actor-kv/src/key.rs index 800aacfda7..a186caed6f 100644 --- a/packages/core/actor-kv/src/key.rs +++ b/packages/core/actor-kv/src/key.rs @@ -3,7 +3,7 @@ use universaldb::tuple::{ Bytes, PackResult, TupleDepth, TuplePack, TupleUnpack, VersionstampOffset, }; -#[derive(Clone, PartialEq)] +#[derive(Debug, Clone, PartialEq)] pub struct KeyWrapper(pub rp::KvKey); impl KeyWrapper { diff --git a/packages/core/api-peer/src/internal.rs b/packages/core/api-peer/src/internal.rs new file mode 100644 index 0000000000..47c84cd57b --- /dev/null +++ b/packages/core/api-peer/src/internal.rs @@ -0,0 +1,29 @@ +use anyhow::Result; +use gas::prelude::*; +use rivet_api_builder::ApiCtx; +use rivet_util::Id; +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize)] +pub struct CachePurgeRequest { + pub base_key: String, + pub keys: Vec, +} + +#[derive(Serialize)] +pub struct CachePurgeResponse {} + +pub async fn cache_purge( + ctx: ApiCtx, + _path: (), + _query: (), + body: CachePurgeRequest, +) -> Result { + ctx.cache() + .clone() + .request() + .purge(&body.base_key, body.keys) + .await?; + + Ok(CachePurgeResponse {}) +} diff --git a/packages/core/api-peer/src/lib.rs b/packages/core/api-peer/src/lib.rs index c994ada92c..c951c31208 100644 --- a/packages/core/api-peer/src/lib.rs +++ b/packages/core/api-peer/src/lib.rs @@ -3,6 +3,7 @@ use std::net::SocketAddr; use anyhow::*; pub mod actors; +pub mod internal; pub mod namespaces; pub mod router; pub mod runners; diff --git a/packages/core/api-peer/src/namespaces.rs b/packages/core/api-peer/src/namespaces.rs index 36b0a3c9fc..6fae8668ae 100644 --- a/packages/core/api-peer/src/namespaces.rs +++ b/packages/core/api-peer/src/namespaces.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use gas::prelude::*; use rivet_api_builder::ApiCtx; use rivet_api_types::pagination::Pagination; use rivet_util::Id; @@ -73,6 +74,7 @@ pub struct ListQuery { pub limit: Option, pub cursor: Option, pub name: Option, + #[serde(default)] pub namespace_id: Vec, } @@ -84,15 +86,6 @@ pub struct ListResponse { pub pagination: Pagination, } -#[utoipa::path( - get, - operation_id = "namespaces_list", - path = "/namespaces", - params(ListQuery), - responses( - (status = 200, body = ListResponse), - ), -)] pub async fn list(ctx: ApiCtx, _path: (), query: ListQuery) -> Result { // If name filter is provided, resolve and return only that namespace if let Some(name) = query.name { @@ -155,15 +148,6 @@ pub struct CreateResponse { pub namespace: namespace::types::Namespace, } -#[utoipa::path( - post, - operation_id = "namespaces_create", - path = "/namespaces", - request_body(content = CreateRequest, content_type = "application/json"), - responses( - (status = 200, body = CreateResponse), - ), -)] pub async fn create( ctx: ApiCtx, _path: (), @@ -211,3 +195,62 @@ pub async fn create( Ok(CreateResponse { namespace }) } + +#[derive(Debug, Serialize, Deserialize, IntoParams)] +#[serde(deny_unknown_fields)] +#[into_params(parameter_in = Query)] +pub struct UpdateQuery {} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UpdatePath { + pub namespace_id: Id, +} + +#[derive(Deserialize, Serialize, ToSchema)] +#[serde(deny_unknown_fields)] +#[schema(as = NamespacesUpdateRequest)] +pub struct UpdateRequest(namespace::workflows::namespace::Update); + +#[derive(Serialize, ToSchema)] +#[schema(as = NamespacesUpdateResponse)] +pub struct UpdateResponse {} + +pub async fn update( + ctx: ApiCtx, + path: UpdatePath, + _query: UpdateQuery, + body: UpdateRequest, +) -> Result { + let mut sub = ctx + .subscribe::(( + "namespace_id", + path.namespace_id, + )) + .await?; + + let res = ctx + .signal(body.0) + .to_workflow::() + .tag("namespace_id", path.namespace_id) + .send() + .await; + + if let Some(WorkflowError::WorkflowNotFound) = res + .as_ref() + .err() + .and_then(|x| x.chain().find_map(|x| x.downcast_ref::())) + { + return Err(namespace::errors::Namespace::NotFound.build()); + } else { + res?; + } + + sub.next() + .await? + .into_body() + .res + .map_err(|err| err.build())?; + + Ok(UpdateResponse {}) +} diff --git a/packages/core/api-peer/src/router.rs b/packages/core/api-peer/src/router.rs index bf279263f1..14c4ce1f83 100644 --- a/packages/core/api-peer/src/router.rs +++ b/packages/core/api-peer/src/router.rs @@ -1,6 +1,6 @@ use rivet_api_builder::{create_router, prelude::*}; -use crate::{actors, namespaces, runners}; +use crate::{actors, internal, namespaces, runners}; pub async fn router( name: &'static str, @@ -14,6 +14,7 @@ pub async fn router( .route("/namespaces", get(namespaces::list)) .route("/namespaces", post(namespaces::create)) .route("/namespaces/{namespace_id}", get(namespaces::get)) + .route("/namespaces/{namespace_id}", put(namespaces::update)) .route( "/namespaces/resolve/{name}", get(namespaces::resolve_for_name), @@ -28,6 +29,8 @@ pub async fn router( .route("/runners", get(runners::list)) .route("/runners/{runner_id}", get(runners::get)) .route("/runners/names", get(runners::list_names)) + // MARK: Internal + .route("/cache/purge", post(internal::cache_purge)) }) .await } diff --git a/packages/core/api-public/src/namespaces.rs b/packages/core/api-public/src/namespaces.rs index 61f3c20f2c..f1c4173124 100644 --- a/packages/core/api-public/src/namespaces.rs +++ b/packages/core/api-public/src/namespaces.rs @@ -137,3 +137,54 @@ async fn create_inner( .await } } + +#[utoipa::path( + put, + operation_id = "namespaces_update", + path = "/namespaces/{namespace_id}", + params( + ("namespace_id" = Id, Path), + UpdateQuery, + ), + request_body(content = UpdateRequest, content_type = "application/json"), + responses( + (status = 200, body = UpdateResponse), + ), +)] +pub async fn update( + Extension(ctx): Extension, + headers: HeaderMap, + Path(path): Path, + Query(query): Query, + Json(body): Json, +) -> Response { + match update_inner(ctx, headers, path, query, body).await { + Ok(response) => response, + Err(err) => ApiError::from(err).into_response(), + } +} + +async fn update_inner( + ctx: ApiCtx, + headers: HeaderMap, + path: UpdatePath, + query: UpdateQuery, + body: UpdateRequest, +) -> Result { + if ctx.config().is_leader() { + let res = rivet_api_peer::namespaces::update(ctx, path, query, body).await?; + Ok(Json(res).into_response()) + } else { + let leader_dc = ctx.config().leader_dc()?; + request_remote_datacenter_raw( + &ctx, + leader_dc.datacenter_label, + &format!("/namespaces/{}", path.namespace_id), + axum::http::Method::PUT, + headers, + Some(&query), + Some(&body), + ) + .await + } +} diff --git a/packages/core/api-public/src/router.rs b/packages/core/api-public/src/router.rs index c0a0599bc5..a3075ffaa4 100644 --- a/packages/core/api-public/src/router.rs +++ b/packages/core/api-public/src/router.rs @@ -22,6 +22,7 @@ use crate::{actors, datacenters, namespaces, runners, ui}; runners::list_names, namespaces::list, namespaces::get, + namespaces::update, namespaces::create, datacenters::list, ))] @@ -46,6 +47,10 @@ pub async fn router( "/namespaces/{namespace_id}", axum::routing::get(namespaces::get), ) + .route( + "/namespaces/{namespace_id}", + axum::routing::put(namespaces::update), + ) // MARK: Actors .route("/actors", axum::routing::get(actors::list::list)) .route("/actors", post(actors::create::create)) diff --git a/packages/core/bootstrap/src/lib.rs b/packages/core/bootstrap/src/lib.rs index 3db5b704d4..304ccafbee 100644 --- a/packages/core/bootstrap/src/lib.rs +++ b/packages/core/bootstrap/src/lib.rs @@ -75,7 +75,8 @@ async fn create_default_namespace(ctx: &StandaloneCtx) -> Result<()> { .op(namespace::ops::resolve_for_name_local::Input { name: "default".to_string(), }) - .await?; + .await + .context("failed resolving default name")?; if existing_namespace.is_none() { // Create namespace diff --git a/packages/core/pegboard-outbound/Cargo.toml b/packages/core/pegboard-outbound/Cargo.toml index 13d4b5c334..039d612d4c 100644 --- a/packages/core/pegboard-outbound/Cargo.toml +++ b/packages/core/pegboard-outbound/Cargo.toml @@ -12,6 +12,7 @@ gas.workspace = true reqwest-eventsource.workspace = true rivet-config.workspace = true rivet-runner-protocol.workspace = true +rivet-types.workspace = true tracing.workspace = true udb-util.workspace = true universaldb.workspace = true diff --git a/packages/core/pegboard-outbound/src/lib.rs b/packages/core/pegboard-outbound/src/lib.rs index 8cb0e1044f..aea4ebd81d 100644 --- a/packages/core/pegboard-outbound/src/lib.rs +++ b/packages/core/pegboard-outbound/src/lib.rs @@ -17,8 +17,6 @@ use tokio::{sync::oneshot, task::JoinHandle, time::Duration}; use udb_util::{SNAPSHOT, TxnExt}; use universaldb::{self as udb, options::StreamingMode}; -const OUTBOUND_REQUEST_LIFESPAN: Duration = Duration::from_secs(14 * 60 + 30); - struct OutboundConnection { handle: JoinHandle<()>, shutdown_tx: oneshot::Sender<()>, @@ -39,7 +37,7 @@ pub async fn start(config: rivet_config::Config, pools: rivet_pools::Pools) -> R )?; let mut sub = ctx - .subscribe::(()) + .subscribe::(()) .await?; let mut outbound_connections = HashMap::new(); @@ -54,36 +52,35 @@ async fn tick( ctx: &StandaloneCtx, outbound_connections: &mut HashMap<(Id, String), Vec>, ) -> Result<()> { - let outbound_data = ctx - .udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); - let outbound_desired_subspace = - txs.subspace(&keys::ns::OutboundDesiredSlotsKey::subspace()); - - txs.get_ranges_keyvalues( - udb::RangeOption { - mode: StreamingMode::WantAll, - ..(&outbound_desired_subspace).into() - }, - // NOTE: This is a snapshot to prevent conflict with updates to this subspace - SNAPSHOT, - ) - .map(|res| match res { - Ok(entry) => { - let (key, desired_slots) = - txs.read_entry::(&entry)?; - - Ok((key.namespace_id, key.runner_name_selector, desired_slots)) - } - Err(err) => Err(err.into()), + let outbound_data = + ctx.udb()? + .run(|tx, _mc| async move { + let txs = tx.subspace(keys::subspace()); + let outbound_desired_subspace = txs.subspace( + &rivet_types::keys::pegboard::ns::OutboundDesiredSlotsKey::entire_subspace(), + ); + + txs.get_ranges_keyvalues( + udb::RangeOption { + mode: StreamingMode::WantAll, + ..(&outbound_desired_subspace).into() + }, + // NOTE: This is a snapshot to prevent conflict with updates to this subspace + SNAPSHOT, + ) + .map(|res| match res { + Ok(entry) => { + let (key, desired_slots) = + txs.read_entry::(&entry)?; + + Ok((key.namespace_id, key.runner_name_selector, desired_slots)) + } + Err(err) => Err(err.into()), + }) + .try_collect::>() + .await }) - .try_collect::>() - .await - - // outbound/{ns_id}/{runner_name_selector}/desired_slots - }) - .await?; + .await?; let mut namespace_ids = outbound_data .iter() @@ -103,6 +100,7 @@ async fn tick( let RunnerKind::Outbound { url, + request_lifespan, slots_per_runner, min_runners, max_runners, @@ -123,11 +121,9 @@ async fn tick( // Remove finished and draining connections from list curr.retain(|conn| !conn.handle.is_finished() && !conn.draining.load(Ordering::SeqCst)); - let desired_count = (desired_slots - .div_ceil(*slots_per_runner) - .max(*min_runners) - .min(*max_runners) + let desired_count = (desired_slots.div_ceil(*slots_per_runner).max(*min_runners) + runners_margin) + .min(*max_runners) .try_into()?; // Calculate diff @@ -147,8 +143,14 @@ async fn tick( } } - let starting_connections = - std::iter::repeat_with(|| spawn_connection(ctx.clone(), url.clone())).take(start_count); + let starting_connections = std::iter::repeat_with(|| { + spawn_connection( + ctx.clone(), + url.clone(), + Duration::from_secs(*request_lifespan as u64), + ) + }) + .take(start_count); curr.extend(starting_connections); } @@ -164,13 +166,19 @@ async fn tick( Ok(()) } -fn spawn_connection(ctx: StandaloneCtx, url: String) -> OutboundConnection { +fn spawn_connection( + ctx: StandaloneCtx, + url: String, + request_lifespan: Duration, +) -> OutboundConnection { let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); let draining = Arc::new(AtomicBool::new(false)); let draining2 = draining.clone(); let handle = tokio::spawn(async move { - if let Err(err) = outbound_handler(&ctx, url, shutdown_rx, draining2).await { + if let Err(err) = + outbound_handler(&ctx, url, request_lifespan, shutdown_rx, draining2).await + { tracing::error!(?err, "outbound req failed"); // TODO: Add backoff @@ -178,7 +186,7 @@ fn spawn_connection(ctx: StandaloneCtx, url: String) -> OutboundConnection { // On error, bump the autoscaler loop again let _ = ctx - .msg(pegboard::messages::BumpOutboundAutoscaler {}) + .msg(rivet_types::msgs::pegboard::BumpOutboundAutoscaler {}) .send() .await; } @@ -194,6 +202,7 @@ fn spawn_connection(ctx: StandaloneCtx, url: String) -> OutboundConnection { async fn outbound_handler( ctx: &StandaloneCtx, url: String, + request_lifespan: Duration, shutdown_rx: oneshot::Receiver<()>, draining: Arc, ) -> Result<()> { @@ -226,13 +235,13 @@ async fn outbound_handler( tokio::select! { res = stream_handler => return res.map_err(Into::into), - _ = tokio::time::sleep(OUTBOUND_REQUEST_LIFESPAN) => {} + _ = tokio::time::sleep(request_lifespan) => {} _ = shutdown_rx => {} } draining.store(true, Ordering::SeqCst); - ctx.msg(pegboard::messages::BumpOutboundAutoscaler {}) + ctx.msg(rivet_types::msgs::pegboard::BumpOutboundAutoscaler {}) .send() .await?; diff --git a/packages/services/internal/Cargo.toml b/packages/services/internal/Cargo.toml new file mode 100644 index 0000000000..e9f9ec8a88 --- /dev/null +++ b/packages/services/internal/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "internal" +version.workspace = true +authors.workspace = true +license.workspace = true +edition.workspace = true + +[dependencies] +anyhow.workspace = true +gas.workspace = true +rivet-api-client.workspace = true +serde.workspace = true diff --git a/packages/services/internal/README.md b/packages/services/internal/README.md new file mode 100644 index 0000000000..0c60920680 --- /dev/null +++ b/packages/services/internal/README.md @@ -0,0 +1 @@ +TODO: move somewhere else diff --git a/packages/services/internal/src/lib.rs b/packages/services/internal/src/lib.rs new file mode 100644 index 0000000000..01eafd2ecc --- /dev/null +++ b/packages/services/internal/src/lib.rs @@ -0,0 +1 @@ +pub mod ops; diff --git a/packages/services/internal/src/ops/cache/mod.rs b/packages/services/internal/src/ops/cache/mod.rs new file mode 100644 index 0000000000..65ba4904bd --- /dev/null +++ b/packages/services/internal/src/ops/cache/mod.rs @@ -0,0 +1 @@ +pub mod purge_global; diff --git a/packages/services/internal/src/ops/cache/purge_global.rs b/packages/services/internal/src/ops/cache/purge_global.rs new file mode 100644 index 0000000000..b86edac77a --- /dev/null +++ b/packages/services/internal/src/ops/cache/purge_global.rs @@ -0,0 +1,78 @@ +use std::fmt::Debug; + +use futures_util::StreamExt; +use gas::prelude::*; +use rivet_api_client::{HeaderMap, Method, request_remote_datacenter}; +use rivet_cache::RawCacheKey; +use serde::Serialize; + +#[derive(Clone, Debug, Default)] +pub struct Input { + pub base_key: String, + pub keys: Vec, +} + +#[operation] +pub async fn cache_purge_global(ctx: &OperationCtx, input: &Input) -> Result<()> { + let dcs = &ctx.config().topology().datacenters; + + let results = futures_util::stream::iter(dcs.clone().into_iter().map(|dc| { + let ctx = ctx.clone(); + let input = input.clone(); + + async move { + if dc.datacenter_label == ctx.config().dc_label() { + // Local datacenter + ctx.cache() + .clone() + .request() + .purge(input.base_key, input.keys) + .await + } else { + // Remote datacenter - HTTP request + request_remote_datacenter( + ctx.config(), + dc.datacenter_label, + "/cache/purge", + Method::POST, + HeaderMap::new(), + Option::<&()>::None, + Some(&CachePurgeRequest { + base_key: input.base_key, + keys: input.keys, + }), + ) + .await + } + } + })) + .buffer_unordered(16) + .collect::>() + .await; + + // Aggregate results + let result_count = results.len(); + let mut errors = Vec::new(); + for res in results { + if let Err(err) = res { + tracing::error!(?err, "failed to request edge dc"); + errors.push(err); + } + } + + // Error only if all requests failed + if result_count == errors.len() { + if let Some(res) = errors.into_iter().next() { + return Err(res).context("all datacenter requests failed"); + } + } + + Ok(()) +} + +// TODO: This is cloned from api-peer because of a cyclical dependency +#[derive(Serialize)] +pub struct CachePurgeRequest { + pub base_key: String, + pub keys: Vec, +} diff --git a/packages/services/internal/src/ops/mod.rs b/packages/services/internal/src/ops/mod.rs new file mode 100644 index 0000000000..a5c08fdb0d --- /dev/null +++ b/packages/services/internal/src/ops/mod.rs @@ -0,0 +1 @@ +pub mod cache; diff --git a/packages/services/namespace/Cargo.toml b/packages/services/namespace/Cargo.toml index 32d46a6735..823fb28bf2 100644 --- a/packages/services/namespace/Cargo.toml +++ b/packages/services/namespace/Cargo.toml @@ -8,14 +8,17 @@ edition.workspace = true [dependencies] anyhow.workspace = true gas.workspace = true +internal.workspace = true rivet-api-builder.workspace = true rivet-api-util.workspace = true rivet-data.workspace = true rivet-error.workspace = true +rivet-types.workspace = true rivet-util.workspace = true serde.workspace = true tracing.workspace = true udb-util.workspace = true universaldb.workspace = true +url.workspace = true utoipa.workspace = true -versioned-data-util.workspace = true +versioned-data-util.workspace = true \ No newline at end of file diff --git a/packages/services/namespace/src/errors.rs b/packages/services/namespace/src/errors.rs index 2bf4acfcf3..a76a0419ac 100644 --- a/packages/services/namespace/src/errors.rs +++ b/packages/services/namespace/src/errors.rs @@ -19,4 +19,11 @@ pub enum Namespace { #[error("not_leader", "Attempting to run operation in non-leader datacenter.")] NotLeader, + + #[error( + "invalid_update", + "Failed to update namespace.", + "Failed to update namespace: {reason}" + )] + InvalidUpdate { reason: String }, } diff --git a/packages/services/namespace/src/keys.rs b/packages/services/namespace/src/keys.rs index c4e13afde4..0ed14ec8a2 100644 --- a/packages/services/namespace/src/keys.rs +++ b/packages/services/namespace/src/keys.rs @@ -180,7 +180,7 @@ impl TuplePack for RunnerKindKey { w: &mut W, tuple_depth: TupleDepth, ) -> std::io::Result { - let t = (DATA, self.namespace_id, CREATE_TS); + let t = (DATA, self.namespace_id, RUNNER_KIND); t.pack(w, tuple_depth) } } diff --git a/packages/services/namespace/src/ops/resolve_for_name_local.rs b/packages/services/namespace/src/ops/resolve_for_name_local.rs index baac221ddc..d70e382000 100644 --- a/packages/services/namespace/src/ops/resolve_for_name_local.rs +++ b/packages/services/namespace/src/ops/resolve_for_name_local.rs @@ -1,6 +1,5 @@ use gas::prelude::*; -use udb_util::{FormalKey, SERIALIZABLE}; -use universaldb as udb; +use udb_util::{SERIALIZABLE, TxnExt}; use crate::{errors, keys, ops::get_local::get_inner, types::Namespace}; @@ -22,21 +21,16 @@ pub async fn namespace_resolve_for_name_local( .run(|tx, _mc| { let name = input.name.clone(); async move { - let name_idx_key = keys::ByNameKey::new(name.clone()); + let txs = tx.subspace(keys::subspace()); - let name_idx_entry = tx - .get(&keys::subspace().pack(&name_idx_key), SERIALIZABLE) - .await?; - - // Namespace not found - let Some(name_idx_entry) = name_idx_entry else { + let Some(namespace_id) = txs + .read_opt(&keys::ByNameKey::new(name.clone()), SERIALIZABLE) + .await? + else { + // Namespace not found return Ok(None); }; - let namespace_id = name_idx_key - .deserialize(&name_idx_entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; - get_inner(namespace_id, &tx).await } }) diff --git a/packages/services/namespace/src/types.rs b/packages/services/namespace/src/types.rs index 05e924ad34..a9709e4efb 100644 --- a/packages/services/namespace/src/types.rs +++ b/packages/services/namespace/src/types.rs @@ -12,9 +12,12 @@ pub struct Namespace { #[derive(Debug, Clone, Serialize, Deserialize, Hash, ToSchema)] #[serde(rename_all = "snake_case")] +#[schema(as = NamespacesRunnerKind)] pub enum RunnerKind { Outbound { url: String, + /// Seconds. + request_lifespan: u32, slots_per_runner: u32, min_runners: u32, max_runners: u32, @@ -28,6 +31,7 @@ impl From for rivet_data::generated::namespace_runner_kind_v1::Data match value { RunnerKind::Outbound { url, + request_lifespan, slots_per_runner, min_runners, max_runners, @@ -35,6 +39,7 @@ impl From for rivet_data::generated::namespace_runner_kind_v1::Data } => rivet_data::generated::namespace_runner_kind_v1::Data::Outbound( rivet_data::generated::namespace_runner_kind_v1::Outbound { url, + request_lifespan, slots_per_runner, min_runners, max_runners, @@ -52,6 +57,7 @@ impl From for RunnerKind rivet_data::generated::namespace_runner_kind_v1::Data::Outbound(o) => { RunnerKind::Outbound { url: o.url, + request_lifespan: o.request_lifespan, slots_per_runner: o.slots_per_runner, min_runners: o.min_runners, max_runners: o.max_runners, diff --git a/packages/services/namespace/src/workflows/namespace.rs b/packages/services/namespace/src/workflows/namespace.rs index 90078b23e6..af5d5b4855 100644 --- a/packages/services/namespace/src/workflows/namespace.rs +++ b/packages/services/namespace/src/workflows/namespace.rs @@ -1,7 +1,9 @@ use futures_util::FutureExt; use gas::prelude::*; +use rivet_cache::CacheKey; use serde::{Deserialize, Serialize}; use udb_util::{SERIALIZABLE, TxnExt}; +use utoipa::ToSchema; use crate::{errors, keys, types::RunnerKind}; @@ -57,8 +59,34 @@ pub async fn namespace(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> { // Does nothing yet ctx.repeat(|ctx| { + let namespace_id = input.namespace_id; + async move { - ctx.listen::().await?; + let update = ctx.listen::().await?; + + let res = ctx + .activity(UpdateInput { + namespace_id, + update, + }) + .await?; + + if let Ok(update_res) = &res { + ctx.activity(PurgeCacheInput { namespace_id }).await?; + + if update_res.bump_autoscaler { + ctx.msg(rivet_types::msgs::pegboard::BumpOutboundAutoscaler {}) + .send() + .await?; + } + } + + ctx.msg(UpdateResult { + res: res.map(|_| ()), + }) + .tag("namespace_id", namespace_id) + .send() + .await?; Ok(Loop::<()>::Continue) } @@ -78,7 +106,17 @@ pub struct Failed { } #[signal("namespace_update")] -pub struct Update {} +#[derive(Debug, Clone, Hash, ToSchema)] +#[schema(as = NamespacesUpdate)] +#[serde(rename_all = "snake_case")] +pub enum Update { + UpdateRunnerKind { runner_kind: RunnerKind }, +} + +#[message("namespace_update_result")] +pub struct UpdateResult { + pub res: Result<(), errors::Namespace>, +} #[derive(Debug, Clone, Serialize, Deserialize, Hash)] pub struct ValidateInput { @@ -147,8 +185,7 @@ async fn insert_fdb( ctx: &ActivityCtx, input: &InsertFdbInput, ) -> Result> { - let res = ctx - .udb()? + ctx.udb()? .run(|tx, _mc| { let namespace_id = input.namespace_id; let name = input.name.clone(); @@ -168,14 +205,6 @@ async fn insert_fdb( txs.write(&keys::CreateTsKey::new(namespace_id), input.create_ts)?; txs.write(&keys::RunnerKindKey::new(namespace_id), RunnerKind::Custom)?; - // RunnerKind::Outbound { - // url: "http://runner:5051/start".to_string(), - // slots_per_runner: 10, - // min_runners: 1, - // max_runners: 1, - // runners_margin: 0, - // } - // Insert idx txs.write(&name_idx_key, namespace_id)?; @@ -183,7 +212,98 @@ async fn insert_fdb( } }) .custom_instrument(tracing::info_span!("namespace_create_tx")) - .await?; + .await + .map_err(Into::into) +} + +#[derive(Debug, Clone, Serialize, Deserialize, Hash)] +struct UpdateInput { + namespace_id: Id, + update: Update, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Hash)] +struct UpdateOutput { + bump_autoscaler: bool, +} + +#[activity(UpdateActivity)] +async fn update( + ctx: &ActivityCtx, + input: &UpdateInput, +) -> Result> { + ctx + .udb()? + .run(|tx, _mc| { + let namespace_id = input.namespace_id; + let update = input.update.clone(); + + async move { + let txs = tx.subspace(keys::subspace()); + + let bump_autoscaler = match update { + Update::UpdateRunnerKind { runner_kind } => { + let bump_autoscaler = match &runner_kind { + RunnerKind::Outbound { + url, + slots_per_runner, + .. + } => { + // Validate url + if let Err(err) = url::Url::parse(url) { + return Ok(Err(errors::Namespace::InvalidUpdate { + reason: format!("invalid outbound url: {err}"), + })); + } + + // Validate slots per runner + if *slots_per_runner == 0 { + return Ok(Err(errors::Namespace::InvalidUpdate { + reason: "`slots_per_runner` cannot be 0".to_string(), + })); + } + + true + } + RunnerKind::Custom => { + // Clear outbound data + txs.delete_key_subspace(&rivet_types::keys::pegboard::ns::OutboundDesiredSlotsKey::subspace(namespace_id)); + + false + } + }; + + txs.write(&keys::RunnerKindKey::new(namespace_id), runner_kind)?; + + bump_autoscaler + } + }; + + Ok(Ok(UpdateOutput { bump_autoscaler })) + } + }) + .custom_instrument(tracing::info_span!("namespace_create_tx")) + .await + .map_err(Into::into) +} + +#[derive(Debug, Clone, Serialize, Deserialize, Hash)] +struct PurgeCacheInput { + namespace_id: Id, +} - Ok(res) +#[activity(PurgeCache)] +async fn purge_cache(ctx: &ActivityCtx, input: &PurgeCacheInput) -> Result<()> { + let res = ctx + .op(internal::ops::cache::purge_global::Input { + base_key: "namespace.get_global".to_string(), + keys: vec![input.namespace_id.cache_key().into()], + }) + .await; + + if let Err(err) = res { + tracing::error!(?err, "failed to purge global namespace cache"); + } + + Ok(()) } diff --git a/packages/services/pegboard/Cargo.toml b/packages/services/pegboard/Cargo.toml index 945f36bdfa..f9569fa571 100644 --- a/packages/services/pegboard/Cargo.toml +++ b/packages/services/pegboard/Cargo.toml @@ -7,16 +7,16 @@ edition.workspace = true [dependencies] anyhow.workspace = true -gas.workspace = true epoxy.workspace = true +gas.workspace = true lazy_static.workspace = true namespace.workspace = true nix.workspace = true rivet-api-client.workspace = true rivet-api-types.workspace = true rivet-api-util.workspace = true -rivet-error.workspace = true rivet-data.workspace = true +rivet-error.workspace = true rivet-metrics.workspace = true rivet-runner-protocol.workspace = true rivet-types.workspace = true diff --git a/packages/services/pegboard/src/keys/ns.rs b/packages/services/pegboard/src/keys/ns.rs index 236c61bc9c..5ccf65fcb5 100644 --- a/packages/services/pegboard/src/keys/ns.rs +++ b/packages/services/pegboard/src/keys/ns.rs @@ -1365,87 +1365,3 @@ impl TuplePack for RunnerNameSubspaceKey { t.pack(w, tuple_depth) } } - -#[derive(Debug)] -pub struct OutboundDesiredSlotsKey { - pub namespace_id: Id, - pub runner_name_selector: String, -} - -impl OutboundDesiredSlotsKey { - pub fn new(namespace_id: Id, runner_name_selector: String) -> Self { - OutboundDesiredSlotsKey { - namespace_id, - runner_name_selector, - } - } - - pub fn subspace() -> OutboundDesiredSlotsSubspaceKey { - OutboundDesiredSlotsSubspaceKey::new() - } -} - -impl FormalKey for OutboundDesiredSlotsKey { - /// Count. - type Value = u32; - - fn deserialize(&self, raw: &[u8]) -> Result { - // NOTE: Atomic ops use little endian - Ok(u32::from_le_bytes(raw.try_into()?)) - } - - fn serialize(&self, value: Self::Value) -> Result> { - // NOTE: Atomic ops use little endian - Ok(value.to_le_bytes().to_vec()) - } -} - -impl TuplePack for OutboundDesiredSlotsKey { - fn pack( - &self, - w: &mut W, - tuple_depth: TupleDepth, - ) -> std::io::Result { - let t = ( - NAMESPACE, - OUTBOUND, - DESIRED_SLOTS, - self.namespace_id, - &self.runner_name_selector, - ); - t.pack(w, tuple_depth) - } -} - -impl<'de> TupleUnpack<'de> for OutboundDesiredSlotsKey { - fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let (input, (_, _, namespace_id, runner_name_selector)) = - <(usize, usize, Id, String)>::unpack(input, tuple_depth)?; - - let v = OutboundDesiredSlotsKey { - namespace_id, - runner_name_selector, - }; - - Ok((input, v)) - } -} - -pub struct OutboundDesiredSlotsSubspaceKey {} - -impl OutboundDesiredSlotsSubspaceKey { - pub fn new() -> Self { - OutboundDesiredSlotsSubspaceKey {} - } -} - -impl TuplePack for OutboundDesiredSlotsSubspaceKey { - fn pack( - &self, - w: &mut W, - tuple_depth: TupleDepth, - ) -> std::io::Result { - let t = (NAMESPACE, OUTBOUND, DESIRED_SLOTS); - t.pack(w, tuple_depth) - } -} diff --git a/packages/services/pegboard/src/lib.rs b/packages/services/pegboard/src/lib.rs index b5dd33dd0a..8a08a5b9a9 100644 --- a/packages/services/pegboard/src/lib.rs +++ b/packages/services/pegboard/src/lib.rs @@ -2,7 +2,6 @@ use gas::prelude::*; pub mod errors; pub mod keys; -pub mod messages; mod metrics; pub mod ops; pub mod pubsub_subjects; diff --git a/packages/services/pegboard/src/workflows/actor/destroy.rs b/packages/services/pegboard/src/workflows/actor/destroy.rs index 44862d219d..3faf1f171c 100644 --- a/packages/services/pegboard/src/workflows/actor/destroy.rs +++ b/packages/services/pegboard/src/workflows/actor/destroy.rs @@ -237,7 +237,10 @@ pub(crate) async fn clear_slot( if let RunnerKind::Outbound { .. } = ns_runner_kind { txs.atomic_op( - &keys::ns::OutboundDesiredSlotsKey::new(namespace_id, runner_name_selector.to_string()), + &rivet_types::keys::pegboard::ns::OutboundDesiredSlotsKey::new( + namespace_id, + runner_name_selector.to_string(), + ), &(-1i32).to_le_bytes(), MutationType::Add, ); diff --git a/packages/services/pegboard/src/workflows/actor/runtime.rs b/packages/services/pegboard/src/workflows/actor/runtime.rs index b1cb7887e8..3371834f83 100644 --- a/packages/services/pegboard/src/workflows/actor/runtime.rs +++ b/packages/services/pegboard/src/workflows/actor/runtime.rs @@ -116,7 +116,7 @@ async fn allocate_actor( // Increment desired slots if namespace has an outbound runner kind if let RunnerKind::Outbound { .. } = ns_runner_kind { txs.atomic_op( - &keys::ns::OutboundDesiredSlotsKey::new( + &rivet_types::keys::pegboard::ns::OutboundDesiredSlotsKey::new( namespace_id, input.runner_name_selector.clone(), ), @@ -347,7 +347,7 @@ pub async fn deallocate(ctx: &ActivityCtx, input: &DeallocateInput) -> Result<() .await?; } else if let RunnerKind::Outbound { .. } = ns_runner_kind { txs.atomic_op( - &keys::ns::OutboundDesiredSlotsKey::new( + &rivet_types::keys::pegboard::ns::OutboundDesiredSlotsKey::new( namespace_id, runner_name_selector.clone(), ), @@ -391,7 +391,7 @@ pub async fn spawn_actor( "failed to allocate (no availability), waiting for allocation", ); - ctx.msg(crate::messages::BumpOutboundAutoscaler {}) + ctx.msg(rivet_types::msgs::pegboard::BumpOutboundAutoscaler {}) .send() .await?; diff --git a/sdks/schemas/data/namespace.runner_kind.v1.bare b/sdks/schemas/data/namespace.runner_kind.v1.bare index 1cc263b54e..05c8eaa739 100644 --- a/sdks/schemas/data/namespace.runner_kind.v1.bare +++ b/sdks/schemas/data/namespace.runner_kind.v1.bare @@ -1,5 +1,6 @@ type Outbound struct { url: str + request_lifespan: u32 slots_per_runner: u32 min_runners: u32 max_runners: u32 From 7d7900b83d031adc9fd98a69324a1a47f408948f Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Mon, 8 Sep 2025 17:42:39 -0700 Subject: [PATCH 10/17] fix: configure runner config per runner name --- Cargo.lock | 38 +-- Cargo.toml | 9 +- out/errors/runner_config.invalid.json | 5 + out/errors/runner_config.not_found.json | 5 + out/openapi.json | 220 +++++++++++++--- packages/common/pools/src/reqwest.rs | 3 +- .../common/types/src/keys/pegboard/mod.rs | 6 + packages/common/types/src/keys/pegboard/ns.rs | 48 ++-- packages/common/types/src/msgs/pegboard.rs | 4 +- packages/common/udb-util/src/ext.rs | 10 +- packages/common/udb-util/src/keys.rs | 4 +- packages/core/api-peer/src/internal.rs | 17 +- .../src/{namespaces.rs => namespaces/mod.rs} | 63 +---- .../api-peer/src/namespaces/runner_configs.rs | 192 ++++++++++++++ packages/core/api-peer/src/router.rs | 22 +- .../core/api-public/src/actors/list_names.rs | 34 ++- .../src/{namespaces.rs => namespaces/mod.rs} | 53 +--- .../src/namespaces/runner_configs.rs | 214 ++++++++++++++++ packages/core/api-public/src/router.rs | 21 +- packages/core/guard/server/src/cache/actor.rs | 1 - .../Cargo.toml | 2 +- .../src/lib.rs | 112 ++++---- packages/infra/engine/Cargo.toml | 2 +- packages/infra/engine/src/run_config.rs | 4 +- packages/services/internal/Cargo.toml | 1 + .../ops/bump_serverless_autoscaler_global.rs | 60 +++++ packages/services/internal/src/ops/mod.rs | 1 + packages/services/namespace/Cargo.toml | 1 + packages/services/namespace/src/errors.rs | 10 + packages/services/namespace/src/keys.rs | 239 +++++++++++++++--- .../services/namespace/src/ops/get_global.rs | 16 +- .../services/namespace/src/ops/get_local.rs | 17 +- packages/services/namespace/src/ops/mod.rs | 1 + .../namespace/src/ops/runner_config/delete.rs | 55 ++++ .../src/ops/runner_config/get_global.rs | 96 +++++++ .../src/ops/runner_config/get_local.rs | 65 +++++ .../namespace/src/ops/runner_config/list.rs | 94 +++++++ .../namespace/src/ops/runner_config/mod.rs | 5 + .../namespace/src/ops/runner_config/upsert.rs | 93 +++++++ packages/services/namespace/src/types.rs | 38 +-- .../namespace/src/workflows/namespace.rs | 137 +--------- packages/services/pegboard/src/keys/mod.rs | 2 +- .../pegboard/src/ops/actor/list_names.rs | 4 +- .../pegboard/src/workflows/actor/destroy.rs | 9 +- .../pegboard/src/workflows/actor/mod.rs | 26 +- .../pegboard/src/workflows/actor/runtime.rs | 46 ++-- .../pegboard/src/workflows/actor/setup.rs | 16 +- sdks/rust/data/src/versioned.rs | 18 +- ...1.bare => namespace.runner_config.v1.bare} | 7 +- 49 files changed, 1582 insertions(+), 564 deletions(-) create mode 100644 out/errors/runner_config.invalid.json create mode 100644 out/errors/runner_config.not_found.json rename packages/core/api-peer/src/{namespaces.rs => namespaces/mod.rs} (78%) create mode 100644 packages/core/api-peer/src/namespaces/runner_configs.rs rename packages/core/api-public/src/{namespaces.rs => namespaces/mod.rs} (73%) create mode 100644 packages/core/api-public/src/namespaces/runner_configs.rs rename packages/core/{pegboard-outbound => pegboard-serverless}/Cargo.toml (94%) rename packages/core/{pegboard-outbound => pegboard-serverless}/src/lib.rs (73%) create mode 100644 packages/services/internal/src/ops/bump_serverless_autoscaler_global.rs create mode 100644 packages/services/namespace/src/ops/runner_config/delete.rs create mode 100644 packages/services/namespace/src/ops/runner_config/get_global.rs create mode 100644 packages/services/namespace/src/ops/runner_config/get_local.rs create mode 100644 packages/services/namespace/src/ops/runner_config/list.rs create mode 100644 packages/services/namespace/src/ops/runner_config/mod.rs create mode 100644 packages/services/namespace/src/ops/runner_config/upsert.rs rename sdks/schemas/data/{namespace.runner_kind.v1.bare => namespace.runner_config.v1.bare} (69%) diff --git a/Cargo.lock b/Cargo.lock index 0758fb75e0..770f168860 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2385,6 +2385,7 @@ dependencies = [ "anyhow", "gasoline", "rivet-api-client", + "rivet-types", "serde", ] @@ -2789,6 +2790,7 @@ dependencies = [ "rivet-types", "rivet-util", "serde", + "strum", "tracing", "udb-util", "universaldb", @@ -3329,44 +3331,44 @@ dependencies = [ ] [[package]] -name = "pegboard-outbound" +name = "pegboard-runner-ws" version = "0.0.1" dependencies = [ "anyhow", - "epoxy", "gasoline", + "hyper 1.6.0", "namespace", "pegboard", - "reqwest-eventsource", + "pegboard-actor-kv", "rivet-config", + "rivet-error", + "rivet-metrics", "rivet-runner-protocol", - "rivet-types", + "rivet-runtime", + "serde", + "serde_json", + "tokio-tungstenite", "tracing", - "udb-util", - "universaldb", + "url", + "versioned-data-util", ] [[package]] -name = "pegboard-runner-ws" +name = "pegboard-serverless" version = "0.0.1" dependencies = [ "anyhow", + "epoxy", "gasoline", - "hyper 1.6.0", "namespace", "pegboard", - "pegboard-actor-kv", + "reqwest-eventsource", "rivet-config", - "rivet-error", - "rivet-metrics", "rivet-runner-protocol", - "rivet-runtime", - "serde", - "serde_json", - "tokio-tungstenite", + "rivet-types", "tracing", - "url", - "versioned-data-util", + "udb-util", + "universaldb", ] [[package]] @@ -4337,8 +4339,8 @@ dependencies = [ "lz4_flex", "namespace", "pegboard", - "pegboard-outbound", "pegboard-runner-ws", + "pegboard-serverless", "portpicker", "rand 0.8.5", "reqwest", diff --git a/Cargo.toml b/Cargo.toml index 5a96596a8e..4ed65e5644 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [workspace] resolver = "2" -members = ["packages/common/api-builder","packages/common/api-client","packages/common/api-types","packages/common/api-util","packages/common/cache/build","packages/common/cache/result","packages/common/clickhouse-inserter","packages/common/clickhouse-user-query","packages/common/config","packages/common/env","packages/common/error/core","packages/common/error/macros","packages/common/gasoline/core","packages/common/gasoline/macros","packages/common/logs","packages/common/metrics","packages/common/pools","packages/common/runtime","packages/common/service-manager","packages/common/telemetry","packages/common/test-deps","packages/common/test-deps-docker","packages/common/types","packages/common/udb-util","packages/common/universaldb","packages/common/universalpubsub","packages/common/util/core","packages/common/util/id","packages/common/versioned-data-util","packages/core/actor-kv","packages/core/api-peer","packages/core/api-public","packages/core/bootstrap","packages/core/dump-openapi","packages/core/guard/core","packages/core/guard/server","packages/core/pegboard-gateway","packages/core/pegboard-outbound","packages/core/pegboard-runner-ws","packages/core/pegboard-tunnel","packages/core/workflow-worker","packages/infra/engine","packages/services/epoxy","packages/services/internal","packages/services/namespace","packages/services/pegboard","sdks/rust/api-full","sdks/rust/bare_gen","sdks/rust/data","sdks/rust/epoxy-protocol","sdks/rust/runner-protocol","sdks/rust/tunnel-protocol","sdks/rust/ups-protocol"] +members = ["packages/common/api-builder","packages/common/api-client","packages/common/api-types","packages/common/api-util","packages/common/cache/build","packages/common/cache/result","packages/common/clickhouse-inserter","packages/common/clickhouse-user-query","packages/common/config","packages/common/env","packages/common/error/core","packages/common/error/macros","packages/common/gasoline/core","packages/common/gasoline/macros","packages/common/logs","packages/common/metrics","packages/common/pools","packages/common/runtime","packages/common/service-manager","packages/common/telemetry","packages/common/test-deps","packages/common/test-deps-docker","packages/common/types","packages/common/udb-util","packages/common/universaldb","packages/common/universalpubsub","packages/common/util/core","packages/common/util/id","packages/common/versioned-data-util","packages/core/actor-kv","packages/core/api-peer","packages/core/api-public","packages/core/bootstrap","packages/core/dump-openapi","packages/core/guard/core","packages/core/guard/server","packages/core/pegboard-gateway","packages/core/pegboard-runner-ws","packages/core/pegboard-serverless","packages/core/pegboard-tunnel","packages/core/workflow-worker","packages/infra/engine","packages/services/epoxy","packages/services/internal","packages/services/namespace","packages/services/pegboard","sdks/rust/api-full","sdks/rust/bare_gen","sdks/rust/data","sdks/rust/epoxy-protocol","sdks/rust/runner-protocol","sdks/rust/tunnel-protocol","sdks/rust/ups-protocol"] [workspace.package] version = "0.0.1" @@ -280,6 +280,7 @@ path = "packages/common/error/core" [workspace.dependencies.rivet-error-macros] path = "packages/common/error/macros" + [workspace.dependencies.gas] package = "gasoline" path = "packages/common/gasoline/core" @@ -363,12 +364,12 @@ path = "packages/core/guard/server" [workspace.dependencies.pegboard-gateway] path = "packages/core/pegboard-gateway" -[workspace.dependencies.pegboard-outbound] -path = "packages/core/pegboard-outbound" - [workspace.dependencies.pegboard-runner-ws] path = "packages/core/pegboard-runner-ws" +[workspace.dependencies.pegboard-serverless] +path = "packages/core/pegboard-serverless" + [workspace.dependencies.pegboard-tunnel] path = "packages/core/pegboard-tunnel" diff --git a/out/errors/runner_config.invalid.json b/out/errors/runner_config.invalid.json new file mode 100644 index 0000000000..314730a13e --- /dev/null +++ b/out/errors/runner_config.invalid.json @@ -0,0 +1,5 @@ +{ + "code": "invalid", + "group": "runner_config", + "message": "Invalid runner config." +} \ No newline at end of file diff --git a/out/errors/runner_config.not_found.json b/out/errors/runner_config.not_found.json new file mode 100644 index 0000000000..a78fba06a6 --- /dev/null +++ b/out/errors/runner_config.not_found.json @@ -0,0 +1,5 @@ +{ + "code": "not_found", + "group": "runner_config", + "message": "No config for this runner exists." +} \ No newline at end of file diff --git a/out/openapi.json b/out/openapi.json index f4278895ed..8a3df3d7c5 100644 --- a/out/openapi.json +++ b/out/openapi.json @@ -548,12 +548,105 @@ } } } + } + }, + "/namespaces/{namespace_id}/runner-configs": { + "get": { + "tags": [ + "namespaces::runner_configs" + ], + "operationId": "namespaces_runner_configs_list", + "parameters": [ + { + "name": "namespace_id", + "in": "path", + "required": true, + "schema": { + "$ref": "#/components/schemas/RivetId" + } + }, + { + "name": "limit", + "in": "query", + "required": false, + "schema": { + "type": "integer", + "minimum": 0 + } + }, + { + "name": "cursor", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "variant", + "in": "query", + "required": false, + "schema": { + "$ref": "#/components/schemas/RunnerConfigVariant" + } + } + ], + "responses": { + "200": { + "description": "", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/NamespacesRunnerConfigsListResponse" + } + } + } + } + } + } + }, + "/namespaces/{namespace_id}/runner-configs/{runner_name}": { + "get": { + "tags": [ + "namespaces::runner_configs" + ], + "operationId": "namespaces_runner_configs_get", + "parameters": [ + { + "name": "namespace_id", + "in": "path", + "required": true, + "schema": { + "$ref": "#/components/schemas/RivetId" + } + }, + { + "name": "runner_name", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/NamespacesRunnerConfigsGetResponse" + } + } + } + } + } }, "put": { "tags": [ - "namespaces" + "namespaces::runner_configs" ], - "operationId": "namespaces_update", + "operationId": "namespaces_runner_configs_upsert", "parameters": [ { "name": "namespace_id", @@ -562,13 +655,21 @@ "schema": { "$ref": "#/components/schemas/RivetId" } + }, + { + "name": "runner_name", + "in": "path", + "required": true, + "schema": { + "type": "string" + } } ], "requestBody": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/NamespacesUpdateRequest" + "$ref": "#/components/schemas/NamespacesRunnerConfigsUpsertRequest" } } }, @@ -580,7 +681,43 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/NamespacesUpdateResponse" + "$ref": "#/components/schemas/NamespacesRunnerConfigsUpsertResponse" + } + } + } + } + } + }, + "delete": { + "tags": [ + "namespaces::runner_configs" + ], + "operationId": "namespaces_runner_configs_delete", + "parameters": [ + { + "name": "namespace_id", + "in": "path", + "required": true, + "schema": { + "$ref": "#/components/schemas/RivetId" + } + }, + { + "name": "runner_name", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/NamespacesRunnerConfigsDeleteResponse" } } } @@ -1091,8 +1228,7 @@ "namespace_id", "name", "display_name", - "create_ts", - "runner_kind" + "create_ts" ], "properties": { "create_ts": { @@ -1107,9 +1243,6 @@ }, "namespace_id": { "$ref": "#/components/schemas/RivetId" - }, - "runner_kind": { - "$ref": "#/components/schemas/NamespacesRunnerKind" } } }, @@ -1171,15 +1304,15 @@ }, "additionalProperties": false }, - "NamespacesRunnerKind": { + "NamespacesRunnerConfig": { "oneOf": [ { "type": "object", "required": [ - "outbound" + "serverless" ], "properties": { - "outbound": { + "serverless": { "type": "object", "required": [ "url", @@ -1222,42 +1355,49 @@ } } } - }, - { - "type": "string", - "enum": [ - "custom" - ] } ] }, - "NamespacesUpdate": { - "oneOf": [ - { + "NamespacesRunnerConfigsDeleteResponse": { + "type": "object" + }, + "NamespacesRunnerConfigsGetResponse": { + "type": "object", + "required": [ + "runner_config" + ], + "properties": { + "runner_config": { + "$ref": "#/components/schemas/NamespacesRunnerConfig" + } + } + }, + "NamespacesRunnerConfigsListResponse": { + "type": "object", + "required": [ + "runner_configs", + "pagination" + ], + "properties": { + "pagination": { + "$ref": "#/components/schemas/Pagination" + }, + "runner_configs": { "type": "object", - "required": [ - "update_runner_kind" - ], - "properties": { - "update_runner_kind": { - "type": "object", - "required": [ - "runner_kind" - ], - "properties": { - "runner_kind": { - "$ref": "#/components/schemas/NamespacesRunnerKind" - } - } - } + "additionalProperties": { + "$ref": "#/components/schemas/NamespacesRunnerConfig" + }, + "propertyNames": { + "type": "string" } } - ] + }, + "additionalProperties": false }, - "NamespacesUpdateRequest": { - "$ref": "#/components/schemas/NamespacesUpdate" + "NamespacesRunnerConfigsUpsertRequest": { + "$ref": "#/components/schemas/NamespacesRunnerConfig" }, - "NamespacesUpdateResponse": { + "NamespacesRunnerConfigsUpsertResponse": { "type": "object" }, "Pagination": { diff --git a/packages/common/pools/src/reqwest.rs b/packages/common/pools/src/reqwest.rs index 78f1f2e7cb..23476203dd 100644 --- a/packages/common/pools/src/reqwest.rs +++ b/packages/common/pools/src/reqwest.rs @@ -2,6 +2,7 @@ use reqwest::Client; use tokio::sync::OnceCell; static CLIENT: OnceCell = OnceCell::const_new(); +static CLIENT_NO_TIMEOUT: OnceCell = OnceCell::const_new(); pub async fn client() -> Result { CLIENT @@ -15,7 +16,7 @@ pub async fn client() -> Result { } pub async fn client_no_timeout() -> Result { - CLIENT + CLIENT_NO_TIMEOUT .get_or_try_init(|| async { Client::builder().build() }) .await .cloned() diff --git a/packages/common/types/src/keys/pegboard/mod.rs b/packages/common/types/src/keys/pegboard/mod.rs index 7e0a481030..1e3ff30358 100644 --- a/packages/common/types/src/keys/pegboard/mod.rs +++ b/packages/common/types/src/keys/pegboard/mod.rs @@ -1 +1,7 @@ +use udb_util::prelude::*; + pub mod ns; + +pub fn subspace() -> udb_util::Subspace { + udb_util::Subspace::new(&(RIVET, PEGBOARD)) +} diff --git a/packages/common/types/src/keys/pegboard/ns.rs b/packages/common/types/src/keys/pegboard/ns.rs index a768f14f61..fa04e89a6e 100644 --- a/packages/common/types/src/keys/pegboard/ns.rs +++ b/packages/common/types/src/keys/pegboard/ns.rs @@ -5,29 +5,29 @@ use gas::prelude::*; use udb_util::prelude::*; #[derive(Debug)] -pub struct OutboundDesiredSlotsKey { +pub struct ServerlessDesiredSlotsKey { pub namespace_id: Id, - pub runner_name_selector: String, + pub runner_name: String, } -impl OutboundDesiredSlotsKey { - pub fn new(namespace_id: Id, runner_name_selector: String) -> Self { - OutboundDesiredSlotsKey { +impl ServerlessDesiredSlotsKey { + pub fn new(namespace_id: Id, runner_name: String) -> Self { + ServerlessDesiredSlotsKey { namespace_id, - runner_name_selector, + runner_name, } } - pub fn subspace(namespace_id: Id) -> OutboundDesiredSlotsSubspaceKey { - OutboundDesiredSlotsSubspaceKey::new(namespace_id) + pub fn subspace(namespace_id: Id) -> ServerlessDesiredSlotsSubspaceKey { + ServerlessDesiredSlotsSubspaceKey::new(namespace_id) } - pub fn entire_subspace() -> OutboundDesiredSlotsSubspaceKey { - OutboundDesiredSlotsSubspaceKey::entire() + pub fn entire_subspace() -> ServerlessDesiredSlotsSubspaceKey { + ServerlessDesiredSlotsSubspaceKey::entire() } } -impl FormalKey for OutboundDesiredSlotsKey { +impl FormalKey for ServerlessDesiredSlotsKey { /// Count. type Value = u32; @@ -42,7 +42,7 @@ impl FormalKey for OutboundDesiredSlotsKey { } } -impl TuplePack for OutboundDesiredSlotsKey { +impl TuplePack for ServerlessDesiredSlotsKey { fn pack( &self, w: &mut W, @@ -50,46 +50,46 @@ impl TuplePack for OutboundDesiredSlotsKey { ) -> std::io::Result { let t = ( NAMESPACE, - OUTBOUND, + SERVERLESS, DESIRED_SLOTS, self.namespace_id, - &self.runner_name_selector, + &self.runner_name, ); t.pack(w, tuple_depth) } } -impl<'de> TupleUnpack<'de> for OutboundDesiredSlotsKey { +impl<'de> TupleUnpack<'de> for ServerlessDesiredSlotsKey { fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let (input, (_, _, _, namespace_id, runner_name_selector)) = + let (input, (_, _, _, namespace_id, runner_name)) = <(usize, usize, usize, Id, String)>::unpack(input, tuple_depth)?; - let v = OutboundDesiredSlotsKey { + let v = ServerlessDesiredSlotsKey { namespace_id, - runner_name_selector, + runner_name, }; Ok((input, v)) } } -pub struct OutboundDesiredSlotsSubspaceKey { +pub struct ServerlessDesiredSlotsSubspaceKey { namespace_id: Option, } -impl OutboundDesiredSlotsSubspaceKey { +impl ServerlessDesiredSlotsSubspaceKey { pub fn new(namespace_id: Id) -> Self { - OutboundDesiredSlotsSubspaceKey { + ServerlessDesiredSlotsSubspaceKey { namespace_id: Some(namespace_id), } } pub fn entire() -> Self { - OutboundDesiredSlotsSubspaceKey { namespace_id: None } + ServerlessDesiredSlotsSubspaceKey { namespace_id: None } } } -impl TuplePack for OutboundDesiredSlotsSubspaceKey { +impl TuplePack for ServerlessDesiredSlotsSubspaceKey { fn pack( &self, w: &mut W, @@ -97,7 +97,7 @@ impl TuplePack for OutboundDesiredSlotsSubspaceKey { ) -> std::io::Result { let mut offset = VersionstampOffset::None { size: 0 }; - let t = (NAMESPACE, OUTBOUND, DESIRED_SLOTS); + let t = (NAMESPACE, SERVERLESS, DESIRED_SLOTS); offset += t.pack(w, tuple_depth)?; if let Some(namespace_id) = self.namespace_id { diff --git a/packages/common/types/src/msgs/pegboard.rs b/packages/common/types/src/msgs/pegboard.rs index e3ad78680d..ece441706b 100644 --- a/packages/common/types/src/msgs/pegboard.rs +++ b/packages/common/types/src/msgs/pegboard.rs @@ -1,4 +1,4 @@ use gas::prelude::*; -#[message("pegboard_bump_outbound_autoscaler")] -pub struct BumpOutboundAutoscaler {} +#[message("pegboard_bump_serverless_autoscaler")] +pub struct BumpServerlessAutoscaler {} diff --git a/packages/common/udb-util/src/ext.rs b/packages/common/udb-util/src/ext.rs index 82ce1e159d..0d5bcc9c0f 100644 --- a/packages/common/udb-util/src/ext.rs +++ b/packages/common/udb-util/src/ext.rs @@ -44,7 +44,7 @@ impl<'a> TxnSubspace<'a> { ) -> Result { self.subspace .unpack(key) - .context("failed unpacking key") + .with_context(|| format!("failed unpacking key of {}", std::any::type_name::())) .map_err(|x| udb::FdbBindingError::CustomError(x.into())) } @@ -59,7 +59,7 @@ impl<'a> TxnSubspace<'a> { .with_context(|| { format!( "failed serializing key value of {}", - std::any::type_name::() + std::any::type_name::(), ) }) .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, @@ -194,7 +194,7 @@ impl SliceExt for udb::future::FdbSlice { .with_context(|| { format!( "failed deserializing key value of {}", - std::any::type_name::() + std::any::type_name::(), ) }) .map_err(|x| udb::FdbBindingError::CustomError(x.into())) @@ -212,7 +212,7 @@ impl OptSliceExt for Option { .with_context(|| { format!( "failed deserializing key value of {}", - std::any::type_name::() + std::any::type_name::(), ) }) .map_err(|x| udb::FdbBindingError::CustomError(x.into())) @@ -228,7 +228,7 @@ impl OptSliceExt for Option { .with_context(|| { format!( "failed deserializing key value of {}", - std::any::type_name::() + std::any::type_name::(), ) }) .map_err(|x| udb::FdbBindingError::CustomError(x.into())) diff --git a/packages/common/udb-util/src/keys.rs b/packages/common/udb-util/src/keys.rs index f873ca7066..177a4eb685 100644 --- a/packages/common/udb-util/src/keys.rs +++ b/packages/common/udb-util/src/keys.rs @@ -119,7 +119,7 @@ define_keys! { (91, METRIC, "metric"), (92, CURRENT_BALLOT, "current_ballot"), (93, INSTANCE_BALLOT, "instance_ballot"), - (94, OUTBOUND, "outbound"), + (94, SERVERLESS, "serverless"), (95, DESIRED_SLOTS, "desired_slots"), - (96, RUNNER_KIND, "runner_kind"), + (96, BY_VARIANT, "by_variant"), } diff --git a/packages/core/api-peer/src/internal.rs b/packages/core/api-peer/src/internal.rs index 47c84cd57b..3aa04b075e 100644 --- a/packages/core/api-peer/src/internal.rs +++ b/packages/core/api-peer/src/internal.rs @@ -1,7 +1,6 @@ use anyhow::Result; use gas::prelude::*; use rivet_api_builder::ApiCtx; -use rivet_util::Id; use serde::{Deserialize, Serialize}; #[derive(Serialize, Deserialize)] @@ -27,3 +26,19 @@ pub async fn cache_purge( Ok(CachePurgeResponse {}) } + +#[derive(Serialize)] +pub struct BumpServerlessAutoscalerResponse {} + +pub async fn bump_serverless_autoscaler( + ctx: ApiCtx, + _path: (), + _query: (), + _body: (), +) -> Result { + ctx.msg(rivet_types::msgs::pegboard::BumpServerlessAutoscaler {}) + .send() + .await?; + + Ok(BumpServerlessAutoscalerResponse {}) +} diff --git a/packages/core/api-peer/src/namespaces.rs b/packages/core/api-peer/src/namespaces/mod.rs similarity index 78% rename from packages/core/api-peer/src/namespaces.rs rename to packages/core/api-peer/src/namespaces/mod.rs index 6fae8668ae..c0c8dcffd3 100644 --- a/packages/core/api-peer/src/namespaces.rs +++ b/packages/core/api-peer/src/namespaces/mod.rs @@ -6,6 +6,8 @@ use rivet_util::Id; use serde::{Deserialize, Serialize}; use utoipa::{IntoParams, ToSchema}; +pub mod runner_configs; + #[derive(Debug, Serialize, Deserialize, IntoParams)] #[serde(deny_unknown_fields)] #[into_params(parameter_in = Query)] @@ -29,7 +31,6 @@ pub async fn get(ctx: ApiCtx, path: GetPath, _query: GetQuery) -> Result Result { - let mut sub = ctx - .subscribe::(( - "namespace_id", - path.namespace_id, - )) - .await?; - - let res = ctx - .signal(body.0) - .to_workflow::() - .tag("namespace_id", path.namespace_id) - .send() - .await; - - if let Some(WorkflowError::WorkflowNotFound) = res - .as_ref() - .err() - .and_then(|x| x.chain().find_map(|x| x.downcast_ref::())) - { - return Err(namespace::errors::Namespace::NotFound.build()); - } else { - res?; - } - - sub.next() - .await? - .into_body() - .res - .map_err(|err| err.build())?; - - Ok(UpdateResponse {}) -} diff --git a/packages/core/api-peer/src/namespaces/runner_configs.rs b/packages/core/api-peer/src/namespaces/runner_configs.rs new file mode 100644 index 0000000000..39dec5a601 --- /dev/null +++ b/packages/core/api-peer/src/namespaces/runner_configs.rs @@ -0,0 +1,192 @@ +use std::collections::HashMap; + +use anyhow::Result; +use rivet_api_builder::ApiCtx; +use rivet_api_types::pagination::Pagination; +use rivet_util::Id; +use serde::{Deserialize, Serialize}; +use utoipa::{IntoParams, ToSchema}; + +#[derive(Debug, Serialize, Deserialize, IntoParams)] +#[serde(deny_unknown_fields)] +#[into_params(parameter_in = Query)] +pub struct GetQuery {} + +#[derive(Deserialize, Serialize, ToSchema)] +#[schema(as = NamespacesRunnerConfigsGetResponse)] +pub struct GetResponse { + pub runner_config: namespace::types::RunnerConfig, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub struct GetPath { + pub namespace_id: Id, + pub runner_name: String, +} + +pub async fn get(ctx: ApiCtx, path: GetPath, _query: GetQuery) -> Result { + let runner_config = ctx + .op(namespace::ops::runner_config::get_local::Input { + runners: vec![(path.namespace_id, path.runner_name)], + }) + .await? + .into_iter() + .next() + .ok_or_else(|| namespace::errors::RunnerConfig::NotFound.build())?; + + Ok(GetResponse { + runner_config: runner_config.config, + }) +} + +#[derive(Debug, Serialize, Deserialize, Clone, IntoParams)] +#[serde(deny_unknown_fields)] +#[into_params(parameter_in = Query)] +pub struct ListQuery { + pub limit: Option, + pub cursor: Option, + pub variant: Option, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ListPath { + pub namespace_id: Id, +} + +#[derive(Deserialize, Serialize, ToSchema)] +#[serde(deny_unknown_fields)] +#[schema(as = NamespacesRunnerConfigsListResponse)] +pub struct ListResponse { + pub runner_configs: HashMap, + pub pagination: Pagination, +} + +pub async fn list(ctx: ApiCtx, path: ListPath, query: ListQuery) -> Result { + ctx.op(namespace::ops::get_local::Input { + namespace_ids: vec![path.namespace_id], + }) + .await? + .first() + .ok_or_else(|| namespace::errors::Namespace::NotFound.build())?; + + // Parse variant from cursor if needed + let (variant, after_name) = if let Some(cursor) = query.cursor { + if let Some((variant, after_name)) = cursor.split_once(":") { + if query.variant.is_some() { + (query.variant, Some(after_name.to_string())) + } else { + ( + namespace::keys::RunnerConfigVariant::parse(variant), + Some(after_name.to_string()), + ) + } + } else { + (query.variant, None) + } + } else { + (query.variant, None) + }; + + let runner_configs_res = ctx + .op(namespace::ops::runner_config::list::Input { + namespace_id: path.namespace_id, + variant, + after_name, + limit: query.limit.unwrap_or(100), + }) + .await?; + + let cursor = runner_configs_res + .last() + .map(|(name, config)| format!("{}:{}", config.variant(), name)); + + Ok(ListResponse { + // TODO: Implement ComposeSchema for FakeMap so we don't have to reallocate + runner_configs: runner_configs_res.into_iter().collect(), + pagination: Pagination { cursor }, + }) +} + +#[derive(Debug, Serialize, Deserialize, IntoParams)] +#[serde(deny_unknown_fields)] +#[into_params(parameter_in = Query)] +pub struct UpsertQuery {} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UpsertPath { + pub namespace_id: Id, + pub runner_name: String, +} + +#[derive(Deserialize, Serialize, ToSchema)] +#[serde(deny_unknown_fields)] +#[schema(as = NamespacesRunnerConfigsUpsertRequest)] +pub struct UpsertRequest(namespace::types::RunnerConfig); + +#[derive(Deserialize, Serialize, ToSchema)] +#[schema(as = NamespacesRunnerConfigsUpsertResponse)] +pub struct UpsertResponse {} + +pub async fn upsert( + ctx: ApiCtx, + path: UpsertPath, + _query: UpsertQuery, + body: UpsertRequest, +) -> Result { + ctx.op(namespace::ops::get_local::Input { + namespace_ids: vec![path.namespace_id], + }) + .await? + .first() + .ok_or_else(|| namespace::errors::Namespace::NotFound.build())?; + + ctx.op(namespace::ops::runner_config::upsert::Input { + namespace_id: path.namespace_id, + name: path.runner_name, + config: body.0, + }) + .await?; + + Ok(UpsertResponse {}) +} + +#[derive(Debug, Serialize, Deserialize, IntoParams)] +#[serde(deny_unknown_fields)] +#[into_params(parameter_in = Query)] +pub struct DeleteQuery {} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub struct DeletePath { + pub namespace_id: Id, + pub runner_name: String, +} + +#[derive(Deserialize, Serialize, ToSchema)] +#[serde(deny_unknown_fields)] +#[schema(as = NamespacesRunnerConfigsDeleteRequest)] +pub struct DeleteRequest(namespace::types::RunnerConfig); + +#[derive(Deserialize, Serialize, ToSchema)] +#[schema(as = NamespacesRunnerConfigsDeleteResponse)] +pub struct DeleteResponse {} + +pub async fn delete(ctx: ApiCtx, path: DeletePath, _query: DeleteQuery) -> Result { + ctx.op(namespace::ops::get_local::Input { + namespace_ids: vec![path.namespace_id], + }) + .await? + .first() + .ok_or_else(|| namespace::errors::Namespace::NotFound.build())?; + + ctx.op(namespace::ops::runner_config::delete::Input { + namespace_id: path.namespace_id, + name: path.runner_name, + }) + .await?; + + Ok(DeleteResponse {}) +} diff --git a/packages/core/api-peer/src/router.rs b/packages/core/api-peer/src/router.rs index 14c4ce1f83..fc8dd3dfce 100644 --- a/packages/core/api-peer/src/router.rs +++ b/packages/core/api-peer/src/router.rs @@ -14,11 +14,27 @@ pub async fn router( .route("/namespaces", get(namespaces::list)) .route("/namespaces", post(namespaces::create)) .route("/namespaces/{namespace_id}", get(namespaces::get)) - .route("/namespaces/{namespace_id}", put(namespaces::update)) .route( "/namespaces/resolve/{name}", get(namespaces::resolve_for_name), ) + // MARK: Runner configs + .route( + "/namespaces/{namespace_id}/runner-configs", + get(namespaces::runner_configs::list), + ) + .route( + "/namespaces/{namespace_id}/runner-configs/{runner_name}", + put(namespaces::runner_configs::upsert), + ) + .route( + "/namespaces/{namespace_id}/runner-configs/{runner_name}", + get(namespaces::runner_configs::get), + ) + .route( + "/namespaces/{namespace_id}/runner-configs/{runner_name}", + delete(namespaces::runner_configs::delete), + ) // MARK: Actors .route("/actors", get(actors::list::list)) .route("/actors", post(actors::create::create)) @@ -31,6 +47,10 @@ pub async fn router( .route("/runners/names", get(runners::list_names)) // MARK: Internal .route("/cache/purge", post(internal::cache_purge)) + .route( + "/bump-serverless-autoscaler", + post(internal::bump_serverless_autoscaler), + ) }) .await } diff --git a/packages/core/api-public/src/actors/list_names.rs b/packages/core/api-public/src/actors/list_names.rs index 40522cbebc..aadbd9b6ec 100644 --- a/packages/core/api-public/src/actors/list_names.rs +++ b/packages/core/api-public/src/actors/list_names.rs @@ -47,25 +47,21 @@ async fn list_names_inner( }; // Fanout to all datacenters - let mut all_names = fanout_to_datacenters::< - ListNamesResponse, - _, - _, - _, - _, - rivet_util::serde::FakeMap, - >( - ctx, - headers, - "/actors/names", - peer_query, - |ctx, query| async move { rivet_api_peer::actors::list_names::list_names(ctx, (), query).await }, - |res, agg| agg.extend(res.names), - ) - .await?; + let mut all_names = + fanout_to_datacenters::>( + ctx, + headers, + "/actors/names", + peer_query, + |ctx, query| async move { + rivet_api_peer::actors::list_names::list_names(ctx, (), query).await + }, + |res, agg| agg.extend(res.names), + ) + .await?; // Sort by name for consistency - all_names.sort(); + all_names.sort_by(|a, b| a.0.cmp(&b.0)); // Truncate to the requested limit all_names.truncate(query.limit.unwrap_or(100)); @@ -73,8 +69,8 @@ async fn list_names_inner( let cursor = all_names.last().map(|(name, _)| name.to_string()); Ok(ListNamesResponse { - // TODO: Implement ComposeSchema for FakeMap so we don't have to use .into() - names: all_names.into(), + // TODO: Implement ComposeSchema for FakeMap so we don't have to reallocate + names: all_names.into_iter().collect(), pagination: Pagination { cursor }, }) } diff --git a/packages/core/api-public/src/namespaces.rs b/packages/core/api-public/src/namespaces/mod.rs similarity index 73% rename from packages/core/api-public/src/namespaces.rs rename to packages/core/api-public/src/namespaces/mod.rs index f1c4173124..7c2ef444bc 100644 --- a/packages/core/api-public/src/namespaces.rs +++ b/packages/core/api-public/src/namespaces/mod.rs @@ -10,6 +10,8 @@ use rivet_util::Id; use rivet_api_client::{request_remote_datacenter, request_remote_datacenter_raw}; use rivet_api_peer::namespaces::*; +pub mod runner_configs; + #[utoipa::path( get, operation_id = "namespaces_list", @@ -137,54 +139,3 @@ async fn create_inner( .await } } - -#[utoipa::path( - put, - operation_id = "namespaces_update", - path = "/namespaces/{namespace_id}", - params( - ("namespace_id" = Id, Path), - UpdateQuery, - ), - request_body(content = UpdateRequest, content_type = "application/json"), - responses( - (status = 200, body = UpdateResponse), - ), -)] -pub async fn update( - Extension(ctx): Extension, - headers: HeaderMap, - Path(path): Path, - Query(query): Query, - Json(body): Json, -) -> Response { - match update_inner(ctx, headers, path, query, body).await { - Ok(response) => response, - Err(err) => ApiError::from(err).into_response(), - } -} - -async fn update_inner( - ctx: ApiCtx, - headers: HeaderMap, - path: UpdatePath, - query: UpdateQuery, - body: UpdateRequest, -) -> Result { - if ctx.config().is_leader() { - let res = rivet_api_peer::namespaces::update(ctx, path, query, body).await?; - Ok(Json(res).into_response()) - } else { - let leader_dc = ctx.config().leader_dc()?; - request_remote_datacenter_raw( - &ctx, - leader_dc.datacenter_label, - &format!("/namespaces/{}", path.namespace_id), - axum::http::Method::PUT, - headers, - Some(&query), - Some(&body), - ) - .await - } -} diff --git a/packages/core/api-public/src/namespaces/runner_configs.rs b/packages/core/api-public/src/namespaces/runner_configs.rs new file mode 100644 index 0000000000..250bf7ee3e --- /dev/null +++ b/packages/core/api-public/src/namespaces/runner_configs.rs @@ -0,0 +1,214 @@ +use anyhow::Result; +use axum::{ + extract::{Extension, Path, Query}, + http::HeaderMap, + response::{IntoResponse, Json, Response}, +}; +use rivet_api_builder::{ApiCtx, ApiError}; +use rivet_util::Id; + +use rivet_api_client::request_remote_datacenter; +use rivet_api_peer::namespaces::runner_configs::*; + +#[utoipa::path( + get, + operation_id = "namespaces_runner_configs_get", + path = "/namespaces/{namespace_id}/runner-configs/{runner_name}", + params( + ("namespace_id" = Id, Path), + ("runner_name" = String, Path), + GetQuery, + ), + responses( + (status = 200, body = GetResponse), + ), +)] +pub async fn get( + Extension(ctx): Extension, + headers: HeaderMap, + Path(path): Path, + Query(query): Query, +) -> Response { + match get_inner(ctx, headers, path, query).await { + Ok(response) => Json(response).into_response(), + Err(err) => ApiError::from(err).into_response(), + } +} + +async fn get_inner( + ctx: ApiCtx, + headers: HeaderMap, + path: GetPath, + query: GetQuery, +) -> Result { + if ctx.config().is_leader() { + rivet_api_peer::namespaces::runner_configs::get(ctx, path, query).await + } else { + let leader_dc = ctx.config().leader_dc()?; + request_remote_datacenter::( + ctx.config(), + leader_dc.datacenter_label, + &format!( + "/namespaces/{}/runner-configs/{}", + path.namespace_id, path.runner_name + ), + axum::http::Method::GET, + headers, + Some(&query), + Option::<&()>::None, + ) + .await + } +} + +#[utoipa::path( + get, + operation_id = "namespaces_runner_configs_list", + path = "/namespaces/{namespace_id}/runner-configs", + params( + ("namespace_id" = Id, Path), + ListQuery, + ), + responses( + (status = 200, body = ListResponse), + ), +)] +pub async fn list( + Extension(ctx): Extension, + headers: HeaderMap, + Path(path): Path, + Query(query): Query, +) -> Response { + match list_inner(ctx, headers, path, query).await { + Ok(response) => Json(response).into_response(), + Err(err) => ApiError::from(err).into_response(), + } +} + +async fn list_inner( + ctx: ApiCtx, + headers: HeaderMap, + path: ListPath, + query: ListQuery, +) -> Result { + if ctx.config().is_leader() { + rivet_api_peer::namespaces::runner_configs::list(ctx, path, query).await + } else { + let leader_dc = ctx.config().leader_dc()?; + request_remote_datacenter::( + ctx.config(), + leader_dc.datacenter_label, + &format!("/namespaces/{}/runner-configs", path.namespace_id), + axum::http::Method::GET, + headers, + Some(&query), + Option::<&()>::None, + ) + .await + } +} + +#[utoipa::path( + put, + operation_id = "namespaces_runner_configs_upsert", + path = "/namespaces/{namespace_id}/runner-configs/{runner_name}", + params( + ("namespace_id" = Id, Path), + ("runner_name" = String, Path), + UpsertQuery, + ), + request_body(content = UpsertRequest, content_type = "application/json"), + responses( + (status = 200, body = UpsertResponse), + ), +)] +pub async fn upsert( + Extension(ctx): Extension, + headers: HeaderMap, + Path(path): Path, + Query(query): Query, + Json(body): Json, +) -> Response { + match upsert_inner(ctx, headers, path, query, body).await { + Ok(response) => Json(response).into_response(), + Err(err) => ApiError::from(err).into_response(), + } +} + +async fn upsert_inner( + ctx: ApiCtx, + headers: HeaderMap, + path: UpsertPath, + query: UpsertQuery, + body: UpsertRequest, +) -> Result { + if ctx.config().is_leader() { + rivet_api_peer::namespaces::runner_configs::upsert(ctx, path, query, body).await + } else { + let leader_dc = ctx.config().leader_dc()?; + request_remote_datacenter::( + ctx.config(), + leader_dc.datacenter_label, + &format!( + "/namespaces/{}/runner-configs/{}", + path.namespace_id, path.runner_name + ), + axum::http::Method::PUT, + headers, + Option::<&()>::None, + Some(&body), + ) + .await + } +} + +#[utoipa::path( + delete, + operation_id = "namespaces_runner_configs_delete", + path = "/namespaces/{namespace_id}/runner-configs/{runner_name}", + params( + ("namespace_id" = Id, Path), + ("runner_name" = String, Path), + DeleteQuery, + ), + responses( + (status = 200, body = DeleteResponse), + ), +)] +pub async fn delete( + Extension(ctx): Extension, + headers: HeaderMap, + Path(path): Path, + Query(query): Query, +) -> Response { + match delete_inner(ctx, headers, path, query).await { + Ok(response) => Json(response).into_response(), + Err(err) => ApiError::from(err).into_response(), + } +} + +async fn delete_inner( + ctx: ApiCtx, + headers: HeaderMap, + path: DeletePath, + query: DeleteQuery, +) -> Result { + if ctx.config().is_leader() { + rivet_api_peer::namespaces::runner_configs::delete(ctx, path, query).await + } else { + let leader_dc = ctx.config().leader_dc()?; + request_remote_datacenter::( + ctx.config(), + leader_dc.datacenter_label, + &format!( + "/namespaces/{}/runner-configs/{}", + path.namespace_id, path.runner_name + ), + axum::http::Method::DELETE, + headers, + Some(&query), + Option::<&()>::None, + ) + .await + } +} diff --git a/packages/core/api-public/src/router.rs b/packages/core/api-public/src/router.rs index a3075ffaa4..9c06fd15b8 100644 --- a/packages/core/api-public/src/router.rs +++ b/packages/core/api-public/src/router.rs @@ -22,8 +22,11 @@ use crate::{actors, datacenters, namespaces, runners, ui}; runners::list_names, namespaces::list, namespaces::get, - namespaces::update, namespaces::create, + namespaces::runner_configs::list, + namespaces::runner_configs::get, + namespaces::runner_configs::upsert, + namespaces::runner_configs::delete, datacenters::list, ))] pub struct ApiDoc; @@ -48,8 +51,20 @@ pub async fn router( axum::routing::get(namespaces::get), ) .route( - "/namespaces/{namespace_id}", - axum::routing::put(namespaces::update), + "/namespaces/{namespace_id}/runner-configs", + axum::routing::get(namespaces::runner_configs::list), + ) + .route( + "/namespaces/{namespace_id}/runner-configs/{runner_name}", + axum::routing::get(namespaces::runner_configs::get), + ) + .route( + "/namespaces/{namespace_id}/runner-configs/{runner_name}", + axum::routing::put(namespaces::runner_configs::upsert), + ) + .route( + "/namespaces/{namespace_id}/runner-configs/{runner_name}", + axum::routing::delete(namespaces::runner_configs::delete), ) // MARK: Actors .route("/actors", axum::routing::get(actors::list::list)) diff --git a/packages/core/guard/server/src/cache/actor.rs b/packages/core/guard/server/src/cache/actor.rs index 4cd073ed2d..1a954cb6fb 100644 --- a/packages/core/guard/server/src/cache/actor.rs +++ b/packages/core/guard/server/src/cache/actor.rs @@ -7,7 +7,6 @@ use anyhow::Result; use gas::prelude::*; use crate::routing::pegboard_gateway::{X_RIVET_ACTOR, X_RIVET_PORT}; -use hyper::header::HeaderName; #[tracing::instrument(skip_all)] pub fn build_cache_key(target: &str, path: &str, headers: &hyper::HeaderMap) -> Result { diff --git a/packages/core/pegboard-outbound/Cargo.toml b/packages/core/pegboard-serverless/Cargo.toml similarity index 94% rename from packages/core/pegboard-outbound/Cargo.toml rename to packages/core/pegboard-serverless/Cargo.toml index 039d612d4c..78eaca5978 100644 --- a/packages/core/pegboard-outbound/Cargo.toml +++ b/packages/core/pegboard-serverless/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "pegboard-outbound" +name = "pegboard-serverless" version.workspace = true authors.workspace = true license.workspace = true diff --git a/packages/core/pegboard-outbound/src/lib.rs b/packages/core/pegboard-serverless/src/lib.rs similarity index 73% rename from packages/core/pegboard-outbound/src/lib.rs rename to packages/core/pegboard-serverless/src/lib.rs index aea4ebd81d..9d6ff472bf 100644 --- a/packages/core/pegboard-outbound/src/lib.rs +++ b/packages/core/pegboard-serverless/src/lib.rs @@ -9,7 +9,7 @@ use std::{ use anyhow::Result; use futures_util::{StreamExt, TryStreamExt}; use gas::prelude::*; -use namespace::types::RunnerKind; +use namespace::types::RunnerConfig; use pegboard::keys; use reqwest_eventsource as sse; use rivet_runner_protocol::protocol; @@ -31,13 +31,13 @@ pub async fn start(config: rivet_config::Config, pools: rivet_pools::Pools) -> R config.clone(), pools, cache, - "pegboard-outbound", + "pegboard-serverless", Id::new_v1(config.dc_label()), Id::new_v1(config.dc_label()), )?; let mut sub = ctx - .subscribe::(()) + .subscribe::(()) .await?; let mut outbound_connections = HashMap::new(); @@ -52,77 +52,77 @@ async fn tick( ctx: &StandaloneCtx, outbound_connections: &mut HashMap<(Id, String), Vec>, ) -> Result<()> { - let outbound_data = - ctx.udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); - let outbound_desired_subspace = txs.subspace( - &rivet_types::keys::pegboard::ns::OutboundDesiredSlotsKey::entire_subspace(), - ); - - txs.get_ranges_keyvalues( - udb::RangeOption { - mode: StreamingMode::WantAll, - ..(&outbound_desired_subspace).into() - }, - // NOTE: This is a snapshot to prevent conflict with updates to this subspace - SNAPSHOT, - ) - .map(|res| match res { - Ok(entry) => { - let (key, desired_slots) = - txs.read_entry::(&entry)?; - - Ok((key.namespace_id, key.runner_name_selector, desired_slots)) - } - Err(err) => Err(err.into()), - }) - .try_collect::>() - .await - }) - .await?; + let serverless_data = ctx + .udb()? + .run(|tx, _mc| async move { + let txs = tx.subspace(keys::subspace()); + + let serverless_desired_subspace = txs.subspace( + &rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::entire_subspace(), + ); + + txs.get_ranges_keyvalues( + udb::RangeOption { + mode: StreamingMode::WantAll, + ..(&serverless_desired_subspace).into() + }, + // NOTE: This is a snapshot to prevent conflict with updates to this subspace + SNAPSHOT, + ) + .map(|res| match res { + Ok(entry) => { + let (key, desired_slots) = + txs.read_entry::(&entry)?; - let mut namespace_ids = outbound_data - .iter() - .map(|(ns_id, _, _)| *ns_id) - .collect::>(); - namespace_ids.dedup(); + Ok((key.namespace_id, key.runner_name, desired_slots)) + } + Err(err) => Err(err.into()), + }) + .try_collect::>() + .await + }) + .await?; - let namespaces = ctx - .op(namespace::ops::get_global::Input { namespace_ids }) + let runner_configs = ctx + .op(namespace::ops::runner_config::get_global::Input { + runners: serverless_data + .iter() + .map(|(ns_id, runner_name, _)| (*ns_id, runner_name.clone())) + .collect(), + }) .await?; - for (ns_id, runner_name_selector, desired_slots) in &outbound_data { - let namespace = namespaces + for (ns_id, runner_name, desired_slots) in &serverless_data { + let runner_config = runner_configs .iter() - .find(|ns| ns.namespace_id == *ns_id) - .context("ns not found")?; + .find(|rc| rc.namespace_id == *ns_id) + .context("runner config not found")?; - let RunnerKind::Outbound { + let RunnerConfig::Serverless { url, request_lifespan, slots_per_runner, min_runners, max_runners, runners_margin, - } = &namespace.runner_kind + } = &runner_config.config else { tracing::warn!( ?ns_id, - "this namespace should not be in the outbound subspace (wrong runner kind)" + "this runner config should not be in the serverless subspace (wrong config kind)" ); continue; }; let curr = outbound_connections - .entry((*ns_id, runner_name_selector.clone())) + .entry((*ns_id, runner_name.clone())) .or_insert_with(Vec::new); // Remove finished and draining connections from list curr.retain(|conn| !conn.handle.is_finished() && !conn.draining.load(Ordering::SeqCst)); let desired_count = (desired_slots.div_ceil(*slots_per_runner).max(*min_runners) - + runners_margin) + + *runners_margin) .min(*max_runners) .try_into()?; @@ -137,7 +137,7 @@ async fn tick( for conn in draining_connections { if conn.shutdown_tx.send(()).is_err() { tracing::warn!( - "outbound connection shutdown channel dropped, likely already stopped" + "serverless connection shutdown channel dropped, likely already stopped" ); } } @@ -146,7 +146,7 @@ async fn tick( let starting_connections = std::iter::repeat_with(|| { spawn_connection( ctx.clone(), - url.clone(), + url.to_string(), Duration::from_secs(*request_lifespan as u64), ) }) @@ -155,12 +155,10 @@ async fn tick( } // Remove entries that aren't returned from udb - outbound_connections.retain(|(ns_id, runner_name_selector), _| { - outbound_data + outbound_connections.retain(|(ns_id, runner_name), _| { + serverless_data .iter() - .any(|(ns_id2, runner_name_selector2, _)| { - ns_id == ns_id2 && runner_name_selector == runner_name_selector2 - }) + .any(|(ns_id2, runner_name2, _)| ns_id == ns_id2 && runner_name == runner_name2) }); Ok(()) @@ -186,7 +184,7 @@ fn spawn_connection( // On error, bump the autoscaler loop again let _ = ctx - .msg(rivet_types::msgs::pegboard::BumpOutboundAutoscaler {}) + .msg(rivet_types::msgs::pegboard::BumpServerlessAutoscaler {}) .send() .await; } @@ -241,7 +239,7 @@ async fn outbound_handler( draining.store(true, Ordering::SeqCst); - ctx.msg(rivet_types::msgs::pegboard::BumpOutboundAutoscaler {}) + ctx.msg(rivet_types::msgs::pegboard::BumpServerlessAutoscaler {}) .send() .await?; diff --git a/packages/infra/engine/Cargo.toml b/packages/infra/engine/Cargo.toml index 975ea3d788..f141806f6c 100644 --- a/packages/infra/engine/Cargo.toml +++ b/packages/infra/engine/Cargo.toml @@ -19,7 +19,7 @@ gas.workspace = true hex.workspace = true include_dir.workspace = true lz4_flex.workspace = true -pegboard-outbound.workspace = true +pegboard-serverless.workspace = true pegboard-runner-ws.workspace = true reqwest.workspace = true rivet-api-peer.workspace = true diff --git a/packages/infra/engine/src/run_config.rs b/packages/infra/engine/src/run_config.rs index 890b12842b..7cfb989853 100644 --- a/packages/infra/engine/src/run_config.rs +++ b/packages/infra/engine/src/run_config.rs @@ -26,9 +26,9 @@ pub fn config(_rivet_config: rivet_config::Config) -> Result { Box::pin(rivet_bootstrap::start(config, pools)) }), Service::new( - "pegboard_outbound", + "pegboard_serverless", ServiceKind::Standalone, - |config, pools| Box::pin(pegboard_outbound::start(config, pools)), + |config, pools| Box::pin(pegboard_serverless::start(config, pools)), ), ]; diff --git a/packages/services/internal/Cargo.toml b/packages/services/internal/Cargo.toml index e9f9ec8a88..fd65a9323d 100644 --- a/packages/services/internal/Cargo.toml +++ b/packages/services/internal/Cargo.toml @@ -9,4 +9,5 @@ edition.workspace = true anyhow.workspace = true gas.workspace = true rivet-api-client.workspace = true +rivet-types.workspace = true serde.workspace = true diff --git a/packages/services/internal/src/ops/bump_serverless_autoscaler_global.rs b/packages/services/internal/src/ops/bump_serverless_autoscaler_global.rs new file mode 100644 index 0000000000..2c76300b27 --- /dev/null +++ b/packages/services/internal/src/ops/bump_serverless_autoscaler_global.rs @@ -0,0 +1,60 @@ +use std::fmt::Debug; + +use futures_util::StreamExt; +use gas::prelude::*; +use rivet_api_client::{HeaderMap, Method, request_remote_datacenter}; + +#[derive(Clone, Debug, Default)] +pub struct Input {} + +#[operation] +pub async fn bump_serverless_autoscaler_global(ctx: &OperationCtx, input: &Input) -> Result<()> { + let dcs = &ctx.config().topology().datacenters; + + let results = futures_util::stream::iter(dcs.clone().into_iter().map(|dc| { + let ctx = ctx.clone(); + + async move { + if dc.datacenter_label == ctx.config().dc_label() { + // Local datacenter + ctx.msg(rivet_types::msgs::pegboard::BumpServerlessAutoscaler {}) + .send() + .await + } else { + // Remote datacenter - HTTP request + request_remote_datacenter( + ctx.config(), + dc.datacenter_label, + "/bump-serverless-autoscaler", + Method::POST, + HeaderMap::new(), + Option::<&()>::None, + Option::<&()>::None, + ) + .await + } + } + })) + .buffer_unordered(16) + .collect::>() + .await; + + // Aggregate results + let result_count = results.len(); + let mut errors = Vec::new(); + for res in results { + if let Err(err) = res { + tracing::error!(?err, "failed to request edge dc"); + errors.push(err); + } + } + + // Error only if all requests failed + if result_count == errors.len() { + if let Some(res) = errors.into_iter().next() { + return Err(res).context("all datacenter requests failed"); + } + } + + Ok(()) +} diff --git a/packages/services/internal/src/ops/mod.rs b/packages/services/internal/src/ops/mod.rs index a5c08fdb0d..ab866e20d5 100644 --- a/packages/services/internal/src/ops/mod.rs +++ b/packages/services/internal/src/ops/mod.rs @@ -1 +1,2 @@ +pub mod bump_serverless_autoscaler_global; pub mod cache; diff --git a/packages/services/namespace/Cargo.toml b/packages/services/namespace/Cargo.toml index 823fb28bf2..46ec60f082 100644 --- a/packages/services/namespace/Cargo.toml +++ b/packages/services/namespace/Cargo.toml @@ -16,6 +16,7 @@ rivet-error.workspace = true rivet-types.workspace = true rivet-util.workspace = true serde.workspace = true +strum.workspace = true tracing.workspace = true udb-util.workspace = true universaldb.workspace = true diff --git a/packages/services/namespace/src/errors.rs b/packages/services/namespace/src/errors.rs index a76a0419ac..02e297e6a7 100644 --- a/packages/services/namespace/src/errors.rs +++ b/packages/services/namespace/src/errors.rs @@ -27,3 +27,13 @@ pub enum Namespace { )] InvalidUpdate { reason: String }, } + +#[derive(RivetError, Debug, Deserialize, Serialize)] +#[error("runner_config")] +pub enum RunnerConfig { + #[error("invalid", "Invalid runner config.", "Invalid runner config: {reason}")] + Invalid { reason: String }, + + #[error("not_found", "No config for this runner exists.")] + NotFound, +} diff --git a/packages/services/namespace/src/keys.rs b/packages/services/namespace/src/keys.rs index 0ed14ec8a2..93e133de25 100644 --- a/packages/services/namespace/src/keys.rs +++ b/packages/services/namespace/src/keys.rs @@ -2,7 +2,9 @@ use std::result::Result::Ok; use anyhow::*; use gas::prelude::*; +use serde::Serialize; use udb_util::prelude::*; +use utoipa::ToSchema; use versioned_data_util::OwnedVersionedData; pub fn subspace() -> udb_util::Subspace { @@ -146,95 +148,272 @@ impl<'de> TupleUnpack<'de> for CreateTsKey { } #[derive(Debug)] -pub struct RunnerKindKey { - namespace_id: Id, +pub struct ByNameKey { + name: String, } -impl RunnerKindKey { - pub fn new(namespace_id: Id) -> Self { - RunnerKindKey { namespace_id } +impl ByNameKey { + pub fn new(name: String) -> Self { + ByNameKey { name } + } +} + +impl FormalKey for ByNameKey { + /// Namespace id. + type Value = Id; + + fn deserialize(&self, raw: &[u8]) -> Result { + Ok(Id::from_slice(raw)?) + } + + fn serialize(&self, value: Self::Value) -> Result> { + Ok(value.as_bytes()) + } +} + +impl TuplePack for ByNameKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = (BY_NAME, &self.name); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for ByNameKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, name)) = <(usize, String)>::unpack(input, tuple_depth)?; + + let v = ByNameKey { name }; + + Ok((input, v)) + } +} + +#[derive(Clone, Copy, Debug, Serialize, Deserialize, strum::FromRepr, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum RunnerConfigVariant { + Serverless = 0, +} + +impl RunnerConfigVariant { + pub fn parse(v: &str) -> Option { + match v { + "serverless" => Some(RunnerConfigVariant::Serverless), + _ => None, + } + } +} + +impl std::fmt::Display for RunnerConfigVariant { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + RunnerConfigVariant::Serverless => write!(f, "serverless"), + } } } -impl FormalKey for RunnerKindKey { - type Value = crate::types::RunnerKind; +#[derive(Debug)] +pub struct RunnerConfigKey { + pub namespace_id: Id, + pub name: String, +} + +impl RunnerConfigKey { + pub fn new(namespace_id: Id, name: String) -> Self { + RunnerConfigKey { namespace_id, name } + } + + pub fn subspace(namespace_id: Id) -> RunnerConfigSubspaceKey { + RunnerConfigSubspaceKey::new(namespace_id) + } +} + +impl FormalKey for RunnerConfigKey { + type Value = crate::types::RunnerConfig; fn deserialize(&self, raw: &[u8]) -> Result { Ok( - rivet_data::versioned::NamespaceRunnerKind::deserialize_with_embedded_version(raw)? + rivet_data::versioned::NamespaceRunnerConfig::deserialize_with_embedded_version(raw)? .into(), ) } fn serialize(&self, value: Self::Value) -> Result> { - rivet_data::versioned::NamespaceRunnerKind::latest(value.into()) + rivet_data::versioned::NamespaceRunnerConfig::latest(value.into()) .serialize_with_embedded_version( rivet_data::PEGBOARD_NAMESPACE_RUNNER_ALLOC_IDX_VERSION, ) } } -impl TuplePack for RunnerKindKey { +impl TuplePack for RunnerConfigKey { fn pack( &self, w: &mut W, tuple_depth: TupleDepth, ) -> std::io::Result { - let t = (DATA, self.namespace_id, RUNNER_KIND); + let t = (RUNNER, CONFIG, DATA, self.namespace_id, &self.name); t.pack(w, tuple_depth) } } -impl<'de> TupleUnpack<'de> for RunnerKindKey { +impl<'de> TupleUnpack<'de> for RunnerConfigKey { fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let (input, (_, namespace_id, _)) = <(usize, Id, usize)>::unpack(input, tuple_depth)?; - let v = RunnerKindKey { namespace_id }; + let (input, (_, _, _, namespace_id, name)) = + <(usize, usize, usize, Id, String)>::unpack(input, tuple_depth)?; + + let v = RunnerConfigKey { namespace_id, name }; Ok((input, v)) } } +pub struct RunnerConfigSubspaceKey { + pub namespace_id: Id, +} + +impl RunnerConfigSubspaceKey { + pub fn new(namespace_id: Id) -> Self { + RunnerConfigSubspaceKey { namespace_id } + } +} + +impl TuplePack for RunnerConfigSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let mut offset = VersionstampOffset::None { size: 0 }; + + let t = (RUNNER, CONFIG, DATA, self.namespace_id); + offset += t.pack(w, tuple_depth)?; + + Ok(offset) + } +} + #[derive(Debug)] -pub struct ByNameKey { - name: String, +pub struct RunnerConfigByVariantKey { + pub namespace_id: Id, + pub variant: RunnerConfigVariant, + pub name: String, } -impl ByNameKey { - pub fn new(name: String) -> Self { - ByNameKey { name } +impl RunnerConfigByVariantKey { + pub fn new(namespace_id: Id, variant: RunnerConfigVariant, name: String) -> Self { + RunnerConfigByVariantKey { + namespace_id, + name, + variant, + } + } + + pub fn subspace(namespace_id: Id) -> RunnerConfigByVariantSubspaceKey { + RunnerConfigByVariantSubspaceKey::new(namespace_id) + } + + pub fn subspace_with_variant( + namespace_id: Id, + variant: RunnerConfigVariant, + ) -> RunnerConfigByVariantSubspaceKey { + RunnerConfigByVariantSubspaceKey::new_with_variant(namespace_id, variant) } } -impl FormalKey for ByNameKey { - /// Namespace id. - type Value = Id; +impl FormalKey for RunnerConfigByVariantKey { + type Value = crate::types::RunnerConfig; fn deserialize(&self, raw: &[u8]) -> Result { - Ok(Id::from_slice(raw)?) + Ok( + rivet_data::versioned::NamespaceRunnerConfig::deserialize_with_embedded_version(raw)? + .into(), + ) } fn serialize(&self, value: Self::Value) -> Result> { - Ok(value.as_bytes()) + rivet_data::versioned::NamespaceRunnerConfig::latest(value.into()) + .serialize_with_embedded_version( + rivet_data::PEGBOARD_NAMESPACE_RUNNER_ALLOC_IDX_VERSION, + ) } } -impl TuplePack for ByNameKey { +impl TuplePack for RunnerConfigByVariantKey { fn pack( &self, w: &mut W, tuple_depth: TupleDepth, ) -> std::io::Result { - let t = (BY_NAME, &self.name); + let t = ( + RUNNER, + CONFIG, + BY_VARIANT, + self.namespace_id, + self.variant as usize, + &self.name, + ); t.pack(w, tuple_depth) } } -impl<'de> TupleUnpack<'de> for ByNameKey { +impl<'de> TupleUnpack<'de> for RunnerConfigByVariantKey { fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let (input, (_, name)) = <(usize, String)>::unpack(input, tuple_depth)?; - - let v = ByNameKey { name }; + let (input, (_, _, _, namespace_id, variant, name)) = + <(usize, usize, usize, Id, usize, String)>::unpack(input, tuple_depth)?; + let variant = RunnerConfigVariant::from_repr(variant).ok_or_else(|| { + PackError::Message(format!("invalid runner config variant `{variant}` in key").into()) + })?; + + let v = RunnerConfigByVariantKey { + namespace_id, + variant, + name, + }; Ok((input, v)) } } + +pub struct RunnerConfigByVariantSubspaceKey { + pub namespace_id: Id, + pub variant: Option, +} + +impl RunnerConfigByVariantSubspaceKey { + pub fn new(namespace_id: Id) -> Self { + RunnerConfigByVariantSubspaceKey { + namespace_id, + variant: None, + } + } + + pub fn new_with_variant(namespace_id: Id, variant: RunnerConfigVariant) -> Self { + RunnerConfigByVariantSubspaceKey { + namespace_id, + variant: Some(variant), + } + } +} + +impl TuplePack for RunnerConfigByVariantSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let mut offset = VersionstampOffset::None { size: 0 }; + + let t = (RUNNER, CONFIG, BY_VARIANT, self.namespace_id); + offset += t.pack(w, tuple_depth)?; + + if let Some(variant) = self.variant { + offset += (variant as usize).pack(w, tuple_depth)?; + } + + Ok(offset) + } +} diff --git a/packages/services/namespace/src/ops/get_global.rs b/packages/services/namespace/src/ops/get_global.rs index a62eeda288..613f8030ee 100644 --- a/packages/services/namespace/src/ops/get_global.rs +++ b/packages/services/namespace/src/ops/get_global.rs @@ -10,13 +10,10 @@ pub struct Input { #[operation] pub async fn namespace_get_global(ctx: &OperationCtx, input: &Input) -> Result> { if ctx.config().is_leader() { - let namespaces_res = ctx - .op(crate::ops::get_local::Input { - namespace_ids: input.namespace_ids.clone(), - }) - .await?; - - Ok(namespaces_res.namespaces) + ctx.op(super::get_local::Input { + namespace_ids: input.namespace_ids.clone(), + }) + .await } else { let leader_dc = ctx.config().leader_dc()?; let client = rivet_pools::reqwest::client().await?; @@ -43,7 +40,8 @@ pub async fn namespace_get_global(ctx: &OperationCtx, input: &Input) -> Result(res).await?; + let res = + rivet_api_util::parse_response::(res).await?; for ns in res.namespaces { let namespace_id = ns.namespace_id; @@ -60,6 +58,6 @@ pub async fn namespace_get_global(ctx: &OperationCtx, input: &Input) -> Result, } diff --git a/packages/services/namespace/src/ops/get_local.rs b/packages/services/namespace/src/ops/get_local.rs index ed6663d589..3f1b5f4474 100644 --- a/packages/services/namespace/src/ops/get_local.rs +++ b/packages/services/namespace/src/ops/get_local.rs @@ -10,13 +10,8 @@ pub struct Input { pub namespace_ids: Vec, } -#[derive(Debug)] -pub struct Output { - pub namespaces: Vec, -} - #[operation] -pub async fn namespace_get(ctx: &OperationCtx, input: &Input) -> Result { +pub async fn namespace_get_local(ctx: &OperationCtx, input: &Input) -> Result> { if !ctx.config().is_leader() { return Err(errors::Namespace::NotLeader.build()); } @@ -38,7 +33,7 @@ pub async fn namespace_get(ctx: &OperationCtx, input: &Input) -> Result .custom_instrument(tracing::info_span!("namespace_get_tx")) .await?; - Ok(Output { namespaces }) + Ok(namespaces) } pub(crate) async fn get_inner( @@ -50,13 +45,11 @@ pub(crate) async fn get_inner( let name_key = keys::NameKey::new(namespace_id); let display_name_key = keys::DisplayNameKey::new(namespace_id); let create_ts_key = keys::CreateTsKey::new(namespace_id); - let runner_kind_key = keys::RunnerKindKey::new(namespace_id); - let (name, display_name, create_ts, runner_kind) = tokio::try_join!( + let (name, display_name, create_ts) = tokio::try_join!( txs.read_opt(&name_key, SERIALIZABLE), txs.read_opt(&display_name_key, SERIALIZABLE), txs.read_opt(&create_ts_key, SERIALIZABLE), - txs.read_opt(&runner_kind_key, SERIALIZABLE), )?; // Namespace not found @@ -70,15 +63,11 @@ pub(crate) async fn get_inner( let create_ts = create_ts.ok_or(udb::FdbBindingError::CustomError( format!("key should exist: {create_ts_key:?}").into(), ))?; - let runner_kind = runner_kind.ok_or(udb::FdbBindingError::CustomError( - format!("key should exist: {runner_kind_key:?}").into(), - ))?; Ok(Some(Namespace { namespace_id, name, display_name, create_ts, - runner_kind, })) } diff --git a/packages/services/namespace/src/ops/mod.rs b/packages/services/namespace/src/ops/mod.rs index 74fc79b4a9..f08fd1f5b5 100644 --- a/packages/services/namespace/src/ops/mod.rs +++ b/packages/services/namespace/src/ops/mod.rs @@ -3,3 +3,4 @@ pub mod get_local; pub mod list; pub mod resolve_for_name_global; pub mod resolve_for_name_local; +pub mod runner_config; diff --git a/packages/services/namespace/src/ops/runner_config/delete.rs b/packages/services/namespace/src/ops/runner_config/delete.rs new file mode 100644 index 0000000000..09ba689150 --- /dev/null +++ b/packages/services/namespace/src/ops/runner_config/delete.rs @@ -0,0 +1,55 @@ +use gas::prelude::*; +use rivet_cache::CacheKey; +use udb_util::{SERIALIZABLE, TxnExt}; + +use crate::{errors, keys}; + +#[derive(Debug)] +pub struct Input { + pub namespace_id: Id, + pub name: String, +} + +#[operation] +pub async fn namespace_runner_config_delete(ctx: &OperationCtx, input: &Input) -> Result<()> { + if !ctx.config().is_leader() { + return Err(errors::Namespace::NotLeader.build()); + } + + ctx.udb()? + .run(|tx, _mc| async move { + let txs = tx.subspace(keys::subspace()); + + // Read existing config to determine variant + let runner_config_key = + keys::RunnerConfigKey::new(input.namespace_id, input.name.clone()); + + if let Some(config) = txs.read_opt(&runner_config_key, SERIALIZABLE).await? { + txs.delete(&runner_config_key); + + // Clear secondary idx + txs.delete(&keys::RunnerConfigByVariantKey::new( + input.namespace_id, + config.variant(), + input.name.clone(), + )); + } + + Ok(()) + }) + .custom_instrument(tracing::info_span!("runner_config_upsert_tx")) + .await?; + + // Purge cache in all dcs + ctx.op(internal::ops::cache::purge_global::Input { + base_key: "namespace.runner_config.{}.get_global".to_string(), + keys: vec![(input.namespace_id, input.name.as_str()).cache_key().into()], + }) + .await?; + + // Bump autoscaler in all dcs + ctx.op(internal::ops::bump_serverless_autoscaler_global::Input {}) + .await?; + + Ok(()) +} diff --git a/packages/services/namespace/src/ops/runner_config/get_global.rs b/packages/services/namespace/src/ops/runner_config/get_global.rs new file mode 100644 index 0000000000..64a2ade433 --- /dev/null +++ b/packages/services/namespace/src/ops/runner_config/get_global.rs @@ -0,0 +1,96 @@ +use std::collections::HashMap; + +use gas::prelude::*; + +use crate::types::RunnerConfig; + +#[derive(Debug)] +pub struct Input { + pub runners: Vec<(Id, String)>, +} + +#[operation] +pub async fn namespace_runner_config_get_global( + ctx: &OperationCtx, + input: &Input, +) -> Result> { + if ctx.config().is_leader() { + ctx.op(super::get_local::Input { + runners: input.runners.clone(), + }) + .await + } else { + let leader_dc = ctx.config().leader_dc()?; + let client = rivet_pools::reqwest::client().await?; + + ctx.cache() + .clone() + .request() + .fetch_all_json( + &format!("namespace.runner_config.get_global"), + input.runners.clone(), + { + let leader_dc = leader_dc.clone(); + let client = client.clone(); + + move |mut cache, runners| { + let leader_dc = leader_dc.clone(); + let client = client.clone(); + + async move { + let mut runner_names_by_namespace_id = + HashMap::with_capacity(runners.len()); + + for (namespace_id, runner_name) in runners { + let runner_names = runner_names_by_namespace_id + .entry(namespace_id) + .or_insert_with(Vec::new); + runner_names.push(runner_name); + } + + // TODO: Parallelize + for (namespace_id, runner_names) in runner_names_by_namespace_id { + let url = leader_dc + .api_peer_url + .join(&format!("/namespaces/{namespace_id}/runner-configs"))?; + let res = client + .get(url) + .query( + &runner_names + .iter() + .map(|runner_name| ("runner", runner_name)) + .collect::>(), + ) + .send() + .await?; + + let res = + rivet_api_util::parse_response::(res) + .await?; + + for (runner_name, runner_config) in res.runner_configs { + cache.resolve( + &(namespace_id, runner_name.clone()), + super::get_local::RunnerConfig { + namespace_id, + name: runner_name, + config: runner_config, + }, + ); + } + } + + Ok(cache) + } + } + }, + ) + .await + } +} + +// TODO: Cyclical dependency with api_peer +#[derive(Deserialize)] +struct RunnerConfigListResponse { + runner_configs: HashMap, +} diff --git a/packages/services/namespace/src/ops/runner_config/get_local.rs b/packages/services/namespace/src/ops/runner_config/get_local.rs new file mode 100644 index 0000000000..fd23ad9562 --- /dev/null +++ b/packages/services/namespace/src/ops/runner_config/get_local.rs @@ -0,0 +1,65 @@ +use futures_util::{StreamExt, TryStreamExt}; +use gas::prelude::*; +use serde::{Deserialize, Serialize}; +use udb_util::{SERIALIZABLE, TxnExt}; + +use crate::{errors, keys}; + +#[derive(Debug)] +pub struct Input { + pub runners: Vec<(Id, String)>, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct RunnerConfig { + pub namespace_id: Id, + pub name: String, + pub config: crate::types::RunnerConfig, +} + +#[operation] +pub async fn namespace_runner_config_get_local( + ctx: &OperationCtx, + input: &Input, +) -> Result> { + if !ctx.config().is_leader() { + return Err(errors::Namespace::NotLeader.build()); + } + + let runner_configs = ctx + .udb()? + .run(|tx, _mc| async move { + futures_util::stream::iter(input.runners.clone()) + .map(|(namespace_id, runner_name)| { + let tx = tx.clone(); + + async move { + let txs = tx.subspace(keys::subspace()); + + let runner_config_key = + keys::RunnerConfigKey::new(namespace_id, runner_name.clone()); + + // Runner config not found + let Some(runner_config) = + txs.read_opt(&runner_config_key, SERIALIZABLE).await? + else { + return Ok(None); + }; + + Ok(Some(RunnerConfig { + namespace_id, + name: runner_name, + config: runner_config, + })) + } + }) + .buffer_unordered(1024) + .try_filter_map(|x| std::future::ready(Ok(x))) + .try_collect::>() + .await + }) + .custom_instrument(tracing::info_span!("runner_config_get_local_tx")) + .await?; + + Ok(runner_configs) +} diff --git a/packages/services/namespace/src/ops/runner_config/list.rs b/packages/services/namespace/src/ops/runner_config/list.rs new file mode 100644 index 0000000000..1b15b21fb2 --- /dev/null +++ b/packages/services/namespace/src/ops/runner_config/list.rs @@ -0,0 +1,94 @@ +use futures_util::{StreamExt, TryStreamExt}; +use gas::prelude::*; +use udb_util::{SERIALIZABLE, TxnExt}; +use universaldb::{self as udb, options::StreamingMode}; + +use crate::{errors, keys, types::RunnerConfig}; + +#[derive(Debug)] +pub struct Input { + pub namespace_id: Id, + pub variant: Option, + pub after_name: Option, + pub limit: usize, +} + +#[operation] +pub async fn namespace_runner_config_list( + ctx: &OperationCtx, + input: &Input, +) -> Result> { + if !ctx.config().is_leader() { + return Err(errors::Namespace::NotLeader.build()); + } + + let runner_configs = ctx + .udb()? + .run(|tx, _mc| async move { + let txs = tx.subspace(keys::subspace()); + + let (start, end) = if let Some(variant) = input.variant { + let (start, end) = txs + .subspace(&keys::RunnerConfigByVariantKey::subspace_with_variant( + input.namespace_id, + variant, + )) + .range(); + + let start = if let Some(name) = &input.after_name { + txs.pack(&keys::RunnerConfigByVariantKey::new( + input.namespace_id, + variant, + name.clone(), + )) + } else { + start + }; + + (start, end) + } else { + let (start, end) = txs + .subspace(&keys::RunnerConfigKey::subspace(input.namespace_id)) + .range(); + + let start = if let Some(name) = &input.after_name { + txs.pack(&keys::RunnerConfigKey::new( + input.namespace_id, + name.clone(), + )) + } else { + start + }; + + (start, end) + }; + + txs.get_ranges_keyvalues( + udb::RangeOption { + mode: StreamingMode::WantAll, + limit: Some(input.limit), + ..(start, end).into() + }, + SERIALIZABLE, + ) + .map(|res| match res { + Ok(entry) => { + if input.variant.is_some() { + let (key, config) = + txs.read_entry::(&entry)?; + Ok((key.name, config)) + } else { + let (key, config) = txs.read_entry::(&entry)?; + Ok((key.name, config)) + } + } + Err(err) => Err(err.into()), + }) + .try_collect() + .await + }) + .custom_instrument(tracing::info_span!("runner_config_get_local_tx")) + .await?; + + Ok(runner_configs) +} diff --git a/packages/services/namespace/src/ops/runner_config/mod.rs b/packages/services/namespace/src/ops/runner_config/mod.rs new file mode 100644 index 0000000000..3c44a67d92 --- /dev/null +++ b/packages/services/namespace/src/ops/runner_config/mod.rs @@ -0,0 +1,5 @@ +pub mod delete; +pub mod get_global; +pub mod get_local; +pub mod list; +pub mod upsert; diff --git a/packages/services/namespace/src/ops/runner_config/upsert.rs b/packages/services/namespace/src/ops/runner_config/upsert.rs new file mode 100644 index 0000000000..e530f8e3ce --- /dev/null +++ b/packages/services/namespace/src/ops/runner_config/upsert.rs @@ -0,0 +1,93 @@ +use gas::prelude::*; +use rivet_cache::CacheKey; +use udb_util::TxnExt; +use universaldb::options::MutationType; + +use crate::{errors, keys, types::RunnerConfig}; + +#[derive(Debug)] +pub struct Input { + pub namespace_id: Id, + pub name: String, + pub config: RunnerConfig, +} + +#[operation] +pub async fn namespace_runner_config_upsert(ctx: &OperationCtx, input: &Input) -> Result<()> { + if !ctx.config().is_leader() { + return Err(errors::Namespace::NotLeader.build()); + } + + ctx.udb()? + .run(|tx, _mc| async move { + let txs = tx.subspace(keys::subspace()); + + // TODO: Once other types of configs get added, delete previous config before writing + txs.write( + &keys::RunnerConfigKey::new(input.namespace_id, input.name.clone()), + input.config.clone(), + )?; + + // Write to secondary idx + txs.write( + &keys::RunnerConfigByVariantKey::new( + input.namespace_id, + input.config.variant(), + input.name.clone(), + ), + input.config.clone(), + )?; + + match &input.config { + RunnerConfig::Serverless { + url, + slots_per_runner, + .. + } => { + // Validate url + if let Err(err) = url::Url::parse(url) { + return Ok(Err(errors::RunnerConfig::Invalid { + reason: format!("invalid serverless url: {err}"), + })); + } + + // Validate slots per runner + if *slots_per_runner == 0 { + return Ok(Err(errors::RunnerConfig::Invalid { + reason: "`slots_per_runner` cannot be 0".to_string(), + })); + } + + // Sets desired count to 0 if it doesn't exist + let txs = tx.subspace(rivet_types::keys::pegboard::subspace()); + txs.atomic_op( + &rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new( + input.namespace_id, + input.name.clone(), + ), + &0u32.to_le_bytes(), + MutationType::Add, + ); + } + } + + Ok(Ok(())) + }) + .custom_instrument(tracing::info_span!("runner_config_upsert_tx")) + .await? + .map_err(|err| err.build())?; + + // Purge cache in all dcs + let variant_str = serde_json::to_string(&input.config.variant())?; + ctx.op(internal::ops::cache::purge_global::Input { + base_key: format!("namespace.runner_config.{variant_str}.get_global"), + keys: vec![(input.namespace_id, input.name.as_str()).cache_key().into()], + }) + .await?; + + // Bump autoscaler in all dcs + ctx.op(internal::ops::bump_serverless_autoscaler_global::Input {}) + .await?; + + Ok(()) +} diff --git a/packages/services/namespace/src/types.rs b/packages/services/namespace/src/types.rs index a9709e4efb..2dc477ce11 100644 --- a/packages/services/namespace/src/types.rs +++ b/packages/services/namespace/src/types.rs @@ -1,20 +1,21 @@ use gas::prelude::*; use utoipa::ToSchema; +use crate::keys; + #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] pub struct Namespace { pub namespace_id: Id, pub name: String, pub display_name: String, pub create_ts: i64, - pub runner_kind: RunnerKind, } #[derive(Debug, Clone, Serialize, Deserialize, Hash, ToSchema)] #[serde(rename_all = "snake_case")] -#[schema(as = NamespacesRunnerKind)] -pub enum RunnerKind { - Outbound { +#[schema(as = NamespacesRunnerConfig)] +pub enum RunnerConfig { + Serverless { url: String, /// Seconds. request_lifespan: u32, @@ -23,21 +24,28 @@ pub enum RunnerKind { max_runners: u32, runners_margin: u32, }, - Custom, } -impl From for rivet_data::generated::namespace_runner_kind_v1::Data { - fn from(value: RunnerKind) -> Self { +impl RunnerConfig { + pub fn variant(&self) -> keys::RunnerConfigVariant { + match self { + RunnerConfig::Serverless { .. } => keys::RunnerConfigVariant::Serverless, + } + } +} + +impl From for rivet_data::generated::namespace_runner_config_v1::Data { + fn from(value: RunnerConfig) -> Self { match value { - RunnerKind::Outbound { + RunnerConfig::Serverless { url, request_lifespan, slots_per_runner, min_runners, max_runners, runners_margin, - } => rivet_data::generated::namespace_runner_kind_v1::Data::Outbound( - rivet_data::generated::namespace_runner_kind_v1::Outbound { + } => rivet_data::generated::namespace_runner_config_v1::Data::Serverless( + rivet_data::generated::namespace_runner_config_v1::Serverless { url, request_lifespan, slots_per_runner, @@ -46,16 +54,15 @@ impl From for rivet_data::generated::namespace_runner_kind_v1::Data runners_margin, }, ), - RunnerKind::Custom => rivet_data::generated::namespace_runner_kind_v1::Data::Custom, } } } -impl From for RunnerKind { - fn from(value: rivet_data::generated::namespace_runner_kind_v1::Data) -> Self { +impl From for RunnerConfig { + fn from(value: rivet_data::generated::namespace_runner_config_v1::Data) -> Self { match value { - rivet_data::generated::namespace_runner_kind_v1::Data::Outbound(o) => { - RunnerKind::Outbound { + rivet_data::generated::namespace_runner_config_v1::Data::Serverless(o) => { + RunnerConfig::Serverless { url: o.url, request_lifespan: o.request_lifespan, slots_per_runner: o.slots_per_runner, @@ -64,7 +71,6 @@ impl From for RunnerKind runners_margin: o.runners_margin, } } - rivet_data::generated::namespace_runner_kind_v1::Data::Custom => RunnerKind::Custom, } } } diff --git a/packages/services/namespace/src/workflows/namespace.rs b/packages/services/namespace/src/workflows/namespace.rs index af5d5b4855..4feac7a0d2 100644 --- a/packages/services/namespace/src/workflows/namespace.rs +++ b/packages/services/namespace/src/workflows/namespace.rs @@ -1,11 +1,9 @@ use futures_util::FutureExt; use gas::prelude::*; -use rivet_cache::CacheKey; use serde::{Deserialize, Serialize}; use udb_util::{SERIALIZABLE, TxnExt}; -use utoipa::ToSchema; -use crate::{errors, keys, types::RunnerKind}; +use crate::{errors, keys}; #[derive(Debug, Deserialize, Serialize)] pub struct Input { @@ -59,34 +57,8 @@ pub async fn namespace(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> { // Does nothing yet ctx.repeat(|ctx| { - let namespace_id = input.namespace_id; - async move { - let update = ctx.listen::().await?; - - let res = ctx - .activity(UpdateInput { - namespace_id, - update, - }) - .await?; - - if let Ok(update_res) = &res { - ctx.activity(PurgeCacheInput { namespace_id }).await?; - - if update_res.bump_autoscaler { - ctx.msg(rivet_types::msgs::pegboard::BumpOutboundAutoscaler {}) - .send() - .await?; - } - } - - ctx.msg(UpdateResult { - res: res.map(|_| ()), - }) - .tag("namespace_id", namespace_id) - .send() - .await?; + ctx.listen::().await?; Ok(Loop::<()>::Continue) } @@ -106,17 +78,7 @@ pub struct Failed { } #[signal("namespace_update")] -#[derive(Debug, Clone, Hash, ToSchema)] -#[schema(as = NamespacesUpdate)] -#[serde(rename_all = "snake_case")] -pub enum Update { - UpdateRunnerKind { runner_kind: RunnerKind }, -} - -#[message("namespace_update_result")] -pub struct UpdateResult { - pub res: Result<(), errors::Namespace>, -} +pub struct Update {} #[derive(Debug, Clone, Serialize, Deserialize, Hash)] pub struct ValidateInput { @@ -203,7 +165,6 @@ async fn insert_fdb( txs.write(&keys::NameKey::new(namespace_id), name)?; txs.write(&keys::DisplayNameKey::new(namespace_id), display_name)?; txs.write(&keys::CreateTsKey::new(namespace_id), input.create_ts)?; - txs.write(&keys::RunnerKindKey::new(namespace_id), RunnerKind::Custom)?; // Insert idx txs.write(&name_idx_key, namespace_id)?; @@ -215,95 +176,3 @@ async fn insert_fdb( .await .map_err(Into::into) } - -#[derive(Debug, Clone, Serialize, Deserialize, Hash)] -struct UpdateInput { - namespace_id: Id, - update: Update, -} - -#[derive(Debug, Clone, Serialize, Deserialize, Hash)] -struct UpdateOutput { - bump_autoscaler: bool, -} - -#[activity(UpdateActivity)] -async fn update( - ctx: &ActivityCtx, - input: &UpdateInput, -) -> Result> { - ctx - .udb()? - .run(|tx, _mc| { - let namespace_id = input.namespace_id; - let update = input.update.clone(); - - async move { - let txs = tx.subspace(keys::subspace()); - - let bump_autoscaler = match update { - Update::UpdateRunnerKind { runner_kind } => { - let bump_autoscaler = match &runner_kind { - RunnerKind::Outbound { - url, - slots_per_runner, - .. - } => { - // Validate url - if let Err(err) = url::Url::parse(url) { - return Ok(Err(errors::Namespace::InvalidUpdate { - reason: format!("invalid outbound url: {err}"), - })); - } - - // Validate slots per runner - if *slots_per_runner == 0 { - return Ok(Err(errors::Namespace::InvalidUpdate { - reason: "`slots_per_runner` cannot be 0".to_string(), - })); - } - - true - } - RunnerKind::Custom => { - // Clear outbound data - txs.delete_key_subspace(&rivet_types::keys::pegboard::ns::OutboundDesiredSlotsKey::subspace(namespace_id)); - - false - } - }; - - txs.write(&keys::RunnerKindKey::new(namespace_id), runner_kind)?; - - bump_autoscaler - } - }; - - Ok(Ok(UpdateOutput { bump_autoscaler })) - } - }) - .custom_instrument(tracing::info_span!("namespace_create_tx")) - .await - .map_err(Into::into) -} - -#[derive(Debug, Clone, Serialize, Deserialize, Hash)] -struct PurgeCacheInput { - namespace_id: Id, -} - -#[activity(PurgeCache)] -async fn purge_cache(ctx: &ActivityCtx, input: &PurgeCacheInput) -> Result<()> { - let res = ctx - .op(internal::ops::cache::purge_global::Input { - base_key: "namespace.get_global".to_string(), - keys: vec![input.namespace_id.cache_key().into()], - }) - .await; - - if let Err(err) = res { - tracing::error!(?err, "failed to purge global namespace cache"); - } - - Ok(()) -} diff --git a/packages/services/pegboard/src/keys/mod.rs b/packages/services/pegboard/src/keys/mod.rs index 402214f8a0..3cb17c5bbb 100644 --- a/packages/services/pegboard/src/keys/mod.rs +++ b/packages/services/pegboard/src/keys/mod.rs @@ -6,7 +6,7 @@ pub mod ns; pub mod runner; pub fn subspace() -> udb_util::Subspace { - udb_util::Subspace::new(&(RIVET, PEGBOARD)) + rivet_types::keys::pegboard::subspace() } pub fn actor_kv_subspace() -> udb_util::Subspace { diff --git a/packages/services/pegboard/src/ops/actor/list_names.rs b/packages/services/pegboard/src/ops/actor/list_names.rs index 8fe72d2100..0cd3752ce3 100644 --- a/packages/services/pegboard/src/ops/actor/list_names.rs +++ b/packages/services/pegboard/src/ops/actor/list_names.rs @@ -15,7 +15,7 @@ pub struct Input { #[derive(Debug)] pub struct Output { - pub names: util::serde::FakeMap, + pub names: Vec<(String, ActorNameKeyData)>, } #[operation] @@ -55,7 +55,7 @@ pub async fn pegboard_actor_list_names(ctx: &OperationCtx, input: &Input) -> Res } Err(err) => Err(Into::::into(err)), }) - .try_collect::>() + .try_collect::>() .await }) .custom_instrument(tracing::info_span!("actor_list_names_tx")) diff --git a/packages/services/pegboard/src/workflows/actor/destroy.rs b/packages/services/pegboard/src/workflows/actor/destroy.rs index 3faf1f171c..7ba691c760 100644 --- a/packages/services/pegboard/src/workflows/actor/destroy.rs +++ b/packages/services/pegboard/src/workflows/actor/destroy.rs @@ -1,5 +1,4 @@ use gas::prelude::*; -use namespace::types::RunnerKind; use rivet_data::converted::ActorByKeyKeyData; use rivet_runner_protocol::protocol; use udb_util::{SERIALIZABLE, TxnExt}; @@ -86,7 +85,7 @@ async fn update_state_and_fdb( state.namespace_id, &state.runner_name_selector, runner_id, - &state.ns_runner_kind, + state.for_serverless, &tx, ) .await?; @@ -164,7 +163,7 @@ pub(crate) async fn clear_slot( namespace_id: Id, runner_name_selector: &str, runner_id: Id, - ns_runner_kind: &RunnerKind, + for_serverless: bool, tx: &udb::RetryableTransaction, ) -> Result<(), udb::FdbBindingError> { let txs = tx.subspace(keys::subspace()); @@ -235,9 +234,9 @@ pub(crate) async fn clear_slot( )?; } - if let RunnerKind::Outbound { .. } = ns_runner_kind { + if for_serverless { txs.atomic_op( - &rivet_types::keys::pegboard::ns::OutboundDesiredSlotsKey::new( + &rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new( namespace_id, runner_name_selector.to_string(), ), diff --git a/packages/services/pegboard/src/workflows/actor/mod.rs b/packages/services/pegboard/src/workflows/actor/mod.rs index 3bff9c38fb..117cb85b6e 100644 --- a/packages/services/pegboard/src/workflows/actor/mod.rs +++ b/packages/services/pegboard/src/workflows/actor/mod.rs @@ -1,6 +1,5 @@ use futures_util::FutureExt; use gas::prelude::*; -use namespace::types::RunnerKind; use rivet_runner_protocol::protocol; use rivet_types::actors::CrashPolicy; @@ -47,7 +46,7 @@ pub struct State { pub create_ts: i64, pub create_complete_ts: Option, - pub ns_runner_kind: RunnerKind, + pub for_serverless: bool, pub start_ts: Option, // NOTE: This is not the alarm ts, this is when the actor started sleeping. See `LifecycleState` for alarm @@ -70,7 +69,6 @@ impl State { runner_name_selector: String, crash_policy: CrashPolicy, create_ts: i64, - ns_runner_kind: RunnerKind, ) -> Self { State { name, @@ -83,7 +81,7 @@ impl State { create_ts, create_complete_ts: None, - ns_runner_kind, + for_serverless: false, start_ts: None, pending_allocation_ts: None, @@ -122,18 +120,15 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> }) .await?; - let metadata = match validation_res { - Ok(metadata) => metadata, - Err(error) => { - ctx.msg(Failed { error }) - .tag("actor_id", input.actor_id) - .send() - .await?; + if let Err(error) = validation_res { + ctx.msg(Failed { error }) + .tag("actor_id", input.actor_id) + .send() + .await?; - // TODO(RVT-3928): return Ok(Err); - return Ok(()); - } - }; + // TODO(RVT-3928): return Ok(Err); + return Ok(()); + } ctx.activity(setup::InitStateAndUdbInput { actor_id: input.actor_id, @@ -143,7 +138,6 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> runner_name_selector: input.runner_name_selector.clone(), crash_policy: input.crash_policy, create_ts: ctx.create_ts(), - ns_runner_kind: metadata.ns_runner_kind, }) .await?; diff --git a/packages/services/pegboard/src/workflows/actor/runtime.rs b/packages/services/pegboard/src/workflows/actor/runtime.rs index 3371834f83..9172583612 100644 --- a/packages/services/pegboard/src/workflows/actor/runtime.rs +++ b/packages/services/pegboard/src/workflows/actor/runtime.rs @@ -3,7 +3,6 @@ use std::time::Instant; use futures_util::StreamExt; use futures_util::{FutureExt, TryStreamExt}; use gas::prelude::*; -use namespace::types::RunnerKind; use rivet_metrics::KeyValue; use rivet_runner_protocol::protocol; use udb_util::{FormalKey, SERIALIZABLE, SNAPSHOT, TxnExt}; @@ -103,20 +102,30 @@ async fn allocate_actor( let start_instant = Instant::now(); let mut state = ctx.state::()?; let namespace_id = state.namespace_id; - let ns_runner_kind = &state.ns_runner_kind; // NOTE: This txn should closely resemble the one found in the allocate_pending_actors activity of the // client wf - let res = ctx + let (for_serverless, res) = ctx .udb()? .run(|tx, _mc| async move { let ping_threshold_ts = util::timestamp::now() - RUNNER_ELIGIBLE_THRESHOLD_MS; let txs = tx.subspace(keys::subspace()); - // Increment desired slots if namespace has an outbound runner kind - if let RunnerKind::Outbound { .. } = ns_runner_kind { + // Check if runner is an serverless runner + let for_serverless = txs + .exists( + &namespace::keys::RunnerConfigByVariantKey::new( + namespace_id, + namespace::keys::RunnerConfigVariant::Serverless, + input.runner_name_selector.clone(), + ), + SERIALIZABLE, + ) + .await?; + + if for_serverless { txs.atomic_op( - &rivet_types::keys::pegboard::ns::OutboundDesiredSlotsKey::new( + &rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new( namespace_id, input.runner_name_selector.clone(), ), @@ -244,10 +253,13 @@ async fn allocate_actor( // Set actor as not sleeping txs.delete(&keys::actor::SleepTsKey::new(input.actor_id)); - return Ok(Ok(AllocateActorOutput { - runner_id: old_runner_alloc_key.runner_id, - runner_workflow_id: old_runner_alloc_key_data.workflow_id, - })); + return Ok(( + for_serverless, + Ok(AllocateActorOutput { + runner_id: old_runner_alloc_key.runner_id, + runner_workflow_id: old_runner_alloc_key_data.workflow_id, + }), + )); } } @@ -268,7 +280,7 @@ async fn allocate_actor( input.generation, )?; - return Ok(Err(pending_ts)); + return Ok((for_serverless, Err(pending_ts))); }) .custom_instrument(tracing::info_span!("actor_allocate_tx")) .await?; @@ -277,6 +289,8 @@ async fn allocate_actor( metrics::ACTOR_ALLOCATE_DURATION .record(dt, &[KeyValue::new("did_reserve", res.is_ok().to_string())]); + state.for_serverless = for_serverless; + match &res { Ok(res) => { state.sleep_ts = None; @@ -327,7 +341,7 @@ pub async fn deallocate(ctx: &ActivityCtx, input: &DeallocateInput) -> Result<() let runner_name_selector = &state.runner_name_selector; let namespace_id = state.namespace_id; let runner_id = state.runner_id; - let ns_runner_kind = &state.ns_runner_kind; + let for_serverless = state.for_serverless; ctx.udb()? .run(|tx, _mc| async move { @@ -341,13 +355,13 @@ pub async fn deallocate(ctx: &ActivityCtx, input: &DeallocateInput) -> Result<() namespace_id, runner_name_selector, runner_id, - ns_runner_kind, + for_serverless, &tx, ) .await?; - } else if let RunnerKind::Outbound { .. } = ns_runner_kind { + } else if for_serverless { txs.atomic_op( - &rivet_types::keys::pegboard::ns::OutboundDesiredSlotsKey::new( + &rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new( namespace_id, runner_name_selector.clone(), ), @@ -391,7 +405,7 @@ pub async fn spawn_actor( "failed to allocate (no availability), waiting for allocation", ); - ctx.msg(rivet_types::msgs::pegboard::BumpOutboundAutoscaler {}) + ctx.msg(rivet_types::msgs::pegboard::BumpServerlessAutoscaler {}) .send() .await?; diff --git a/packages/services/pegboard/src/workflows/actor/setup.rs b/packages/services/pegboard/src/workflows/actor/setup.rs index 421b7228fc..068cc562f0 100644 --- a/packages/services/pegboard/src/workflows/actor/setup.rs +++ b/packages/services/pegboard/src/workflows/actor/setup.rs @@ -1,5 +1,4 @@ use gas::prelude::*; -use namespace::types::RunnerKind; use rivet_data::converted::ActorNameKeyData; use rivet_types::actors::CrashPolicy; use udb_util::{SERIALIZABLE, TxnExt}; @@ -18,23 +17,18 @@ pub struct ValidateInput { pub input: Option, } -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ValidateOutput { - pub ns_runner_kind: RunnerKind, -} - #[activity(Validate)] pub async fn validate( ctx: &ActivityCtx, input: &ValidateInput, -) -> Result> { +) -> Result> { let ns_res = ctx .op(namespace::ops::get_global::Input { namespace_ids: vec![input.namespace_id], }) .await?; - let Some(ns) = ns_res.into_iter().next() else { + if ns_res.is_empty() { return Ok(Err(errors::Actor::NamespaceNotFound)); }; @@ -61,9 +55,7 @@ pub async fn validate( } } - Ok(Ok(ValidateOutput { - ns_runner_kind: ns.runner_kind, - })) + Ok(Ok(())) } #[derive(Debug, Clone, Serialize, Deserialize, Hash)] @@ -75,7 +67,6 @@ pub struct InitStateAndUdbInput { pub runner_name_selector: String, pub crash_policy: CrashPolicy, pub create_ts: i64, - pub ns_runner_kind: RunnerKind, } #[activity(InitStateAndFdb)] @@ -89,7 +80,6 @@ pub async fn insert_state_and_fdb(ctx: &ActivityCtx, input: &InitStateAndUdbInpu input.runner_name_selector.clone(), input.crash_policy, input.create_ts, - input.ns_runner_kind.clone(), )); ctx.udb()? diff --git a/sdks/rust/data/src/versioned.rs b/sdks/rust/data/src/versioned.rs index 002363d361..551e518158 100644 --- a/sdks/rust/data/src/versioned.rs +++ b/sdks/rust/data/src/versioned.rs @@ -207,20 +207,20 @@ impl OwnedVersionedData for ActorNameKeyData { } } -pub enum NamespaceRunnerKind { - V1(namespace_runner_kind_v1::Data), +pub enum NamespaceRunnerConfig { + V1(namespace_runner_config_v1::Data), } -impl OwnedVersionedData for NamespaceRunnerKind { - type Latest = namespace_runner_kind_v1::Data; +impl OwnedVersionedData for NamespaceRunnerConfig { + type Latest = namespace_runner_config_v1::Data; - fn latest(latest: namespace_runner_kind_v1::Data) -> Self { - NamespaceRunnerKind::V1(latest) + fn latest(latest: namespace_runner_config_v1::Data) -> Self { + NamespaceRunnerConfig::V1(latest) } fn into_latest(self) -> Result { #[allow(irrefutable_let_patterns)] - if let NamespaceRunnerKind::V1(data) = self { + if let NamespaceRunnerConfig::V1(data) = self { Ok(data) } else { bail!("version not latest"); @@ -229,14 +229,14 @@ impl OwnedVersionedData for NamespaceRunnerKind { fn deserialize_version(payload: &[u8], version: u16) -> Result { match version { - 1 => Ok(NamespaceRunnerKind::V1(serde_bare::from_slice(payload)?)), + 1 => Ok(NamespaceRunnerConfig::V1(serde_bare::from_slice(payload)?)), _ => bail!("invalid version: {version}"), } } fn serialize_version(self, _version: u16) -> Result> { match self { - NamespaceRunnerKind::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), + NamespaceRunnerConfig::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), } } } diff --git a/sdks/schemas/data/namespace.runner_kind.v1.bare b/sdks/schemas/data/namespace.runner_config.v1.bare similarity index 69% rename from sdks/schemas/data/namespace.runner_kind.v1.bare rename to sdks/schemas/data/namespace.runner_config.v1.bare index 05c8eaa739..a25e630013 100644 --- a/sdks/schemas/data/namespace.runner_kind.v1.bare +++ b/sdks/schemas/data/namespace.runner_config.v1.bare @@ -1,4 +1,4 @@ -type Outbound struct { +type Serverless struct { url: str request_lifespan: u32 slots_per_runner: u32 @@ -7,9 +7,6 @@ type Outbound struct { runners_margin: u32 } -type Custom void - type Data union { - Outbound | - Custom + Serverless } From 36bef7ee2b4005ed467dd98908cdce960e2e3c3c Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Tue, 9 Sep 2025 14:07:46 -0700 Subject: [PATCH 11/17] chore: remove port name --- frontend/src/queries/actor-engine.ts | 1 - out/openapi.json | 72 ----- packages/common/config/src/config/db.rs | 1 + packages/common/types/src/runners.rs | 100 ------- packages/common/udb-util/src/keys.rs | 2 +- packages/core/guard/server/src/cache/actor.rs | 13 +- packages/core/guard/server/src/errors.rs | 3 +- .../server/src/routing/pegboard_gateway.rs | 24 +- packages/core/pegboard-gateway/src/lib.rs | 19 +- packages/core/pegboard-runner-ws/src/lib.rs | 7 - packages/core/pegboard-tunnel/src/lib.rs | 6 +- .../core/pegboard-tunnel/tests/integration.rs | 6 +- .../infra/engine/tests/actors_lifecycle.rs | 16 +- packages/infra/engine/tests/common/actors.rs | 24 +- packages/services/pegboard/src/keys/runner.rs | 72 ----- .../services/pegboard/src/ops/runner/get.rs | 40 --- .../services/pegboard/src/pubsub_subjects.rs | 14 +- .../services/pegboard/src/workflows/runner.rs | 37 +-- scripts/tests/actor_e2e.ts | 2 - scripts/tests/actor_e2e_multidc.ts | 5 +- sdks/rust/data/src/converted.rs | 54 ---- sdks/rust/data/src/versioned.rs | 34 --- sdks/rust/runner-protocol/src/protocol.rs | 22 -- sdks/rust/runner-protocol/src/versioned.rs | 39 --- .../data/pegboard.runner.address.v1.bare | 20 -- sdks/schemas/runner-protocol/v1.bare | 18 -- sdks/typescript/runner-protocol/src/index.ts | 255 ++++-------------- .../runner/benches/actor-lifecycle.bench.ts | 2 - sdks/typescript/runner/src/mod.ts | 3 - .../typescript/runner/tests/lifecycle.test.ts | 7 - tests/load/actor-lifecycle/actor.ts | 1 - tests/load/actor-lifecycle/types.ts | 6 - 32 files changed, 80 insertions(+), 845 deletions(-) delete mode 100644 sdks/schemas/data/pegboard.runner.address.v1.bare diff --git a/frontend/src/queries/actor-engine.ts b/frontend/src/queries/actor-engine.ts index 271635d8d6..020bf572eb 100644 --- a/frontend/src/queries/actor-engine.ts +++ b/frontend/src/queries/actor-engine.ts @@ -18,7 +18,6 @@ export const createEngineActorContext = ({ headers: { "x-rivet-actor": actorId, "x-rivet-target": "actor", - "x-rivet-port": "main", ...(token ? { authorization: `Bearer ${token}` } : {}), }, }; diff --git a/out/openapi.json b/out/openapi.json index 8a3df3d7c5..637d75cf88 100644 --- a/out/openapi.json +++ b/out/openapi.json @@ -1426,23 +1426,11 @@ "version", "total_slots", "remaining_slots", - "addresses_http", - "addresses_tcp", - "addresses_udp", "create_ts", "last_ping_ts", "last_rtt" ], "properties": { - "addresses_http": { - "$ref": "#/components/schemas/StringHttpAddressHashableMap" - }, - "addresses_tcp": { - "$ref": "#/components/schemas/StringTcpAddressHashableMap" - }, - "addresses_udp": { - "$ref": "#/components/schemas/StringUdpAddressHashableMap" - }, "create_ts": { "type": "integer", "format": "int64" @@ -1569,66 +1557,6 @@ } }, "additionalProperties": false - }, - "StringHttpAddressHashableMap": { - "type": "object", - "additionalProperties": { - "type": "object", - "required": [ - "hostname", - "port" - ], - "properties": { - "hostname": { - "type": "string" - }, - "port": { - "type": "integer", - "format": "int32", - "minimum": 0 - } - } - } - }, - "StringTcpAddressHashableMap": { - "type": "object", - "additionalProperties": { - "type": "object", - "required": [ - "hostname", - "port" - ], - "properties": { - "hostname": { - "type": "string" - }, - "port": { - "type": "integer", - "format": "int32", - "minimum": 0 - } - } - } - }, - "StringUdpAddressHashableMap": { - "type": "object", - "additionalProperties": { - "type": "object", - "required": [ - "hostname", - "port" - ], - "properties": { - "hostname": { - "type": "string" - }, - "port": { - "type": "integer", - "format": "int32", - "minimum": 0 - } - } - } } } } diff --git a/packages/common/config/src/config/db.rs b/packages/common/config/src/config/db.rs index 5303ec5c05..44179363f0 100644 --- a/packages/common/config/src/config/db.rs +++ b/packages/common/config/src/config/db.rs @@ -19,6 +19,7 @@ impl Default for Database { } #[derive(Debug, Serialize, Deserialize, Clone, JsonSchema)] +#[serde(deny_unknown_fields)] pub struct Postgres { /// URL to connect to Postgres with /// diff --git a/packages/common/types/src/runners.rs b/packages/common/types/src/runners.rs index d70ed49eda..84db5e0b13 100644 --- a/packages/common/types/src/runners.rs +++ b/packages/common/types/src/runners.rs @@ -1,9 +1,6 @@ use gas::prelude::*; -use rivet_data::generated::pegboard_runner_address_v1; -use rivet_runner_protocol::protocol; use rivet_util::Id; use serde::{Deserialize, Serialize}; -use std::ops::Deref; use utoipa::ToSchema; #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] @@ -17,9 +14,6 @@ pub struct Runner { pub version: u32, pub total_slots: u32, pub remaining_slots: u32, - pub addresses_http: StringHttpAddressHashableMap, - pub addresses_tcp: StringTcpAddressHashableMap, - pub addresses_udp: StringUdpAddressHashableMap, pub create_ts: i64, pub drain_ts: Option, pub stop_ts: Option, @@ -28,97 +22,3 @@ pub struct Runner { pub last_rtt: u32, pub metadata: Option>, } - -// HACK: We can't define ToSchema on HashableMap directly, so we have to define concrete types -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct StringHttpAddressHashableMap( - util::serde::HashableMap, -); - -impl From> - for StringHttpAddressHashableMap -{ - fn from(value: util::serde::HashableMap) -> Self { - Self(value) - } -} - -impl Deref for StringHttpAddressHashableMap { - type Target = util::serde::HashableMap; - - fn deref(&self) -> &Self::Target { - &self.0 - } -} - -impl utoipa::ToSchema for StringHttpAddressHashableMap {} - -impl utoipa::PartialSchema for StringHttpAddressHashableMap { - fn schema() -> utoipa::openapi::RefOr { - utoipa::openapi::ObjectBuilder::new() - .additional_properties(Some(protocol::RunnerAddressHttp::schema())) - .into() - } -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct StringTcpAddressHashableMap( - util::serde::HashableMap, -); - -impl From> - for StringTcpAddressHashableMap -{ - fn from(value: util::serde::HashableMap) -> Self { - Self(value) - } -} - -impl Deref for StringTcpAddressHashableMap { - type Target = util::serde::HashableMap; - - fn deref(&self) -> &Self::Target { - &self.0 - } -} - -impl utoipa::ToSchema for StringTcpAddressHashableMap {} - -impl utoipa::PartialSchema for StringTcpAddressHashableMap { - fn schema() -> utoipa::openapi::RefOr { - utoipa::openapi::ObjectBuilder::new() - .additional_properties(Some(protocol::RunnerAddressTcp::schema())) - .into() - } -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct StringUdpAddressHashableMap( - util::serde::HashableMap, -); - -impl From> - for StringUdpAddressHashableMap -{ - fn from(value: util::serde::HashableMap) -> Self { - Self(value) - } -} - -impl Deref for StringUdpAddressHashableMap { - type Target = util::serde::HashableMap; - - fn deref(&self) -> &Self::Target { - &self.0 - } -} - -impl utoipa::ToSchema for StringUdpAddressHashableMap {} - -impl utoipa::PartialSchema for StringUdpAddressHashableMap { - fn schema() -> utoipa::openapi::RefOr { - utoipa::openapi::ObjectBuilder::new() - .additional_properties(Some(protocol::RunnerAddressUdp::schema())) - .into() - } -} diff --git a/packages/common/udb-util/src/keys.rs b/packages/common/udb-util/src/keys.rs index 177a4eb685..e267e80a88 100644 --- a/packages/common/udb-util/src/keys.rs +++ b/packages/common/udb-util/src/keys.rs @@ -65,7 +65,7 @@ define_keys! { (37, TOTAL_MEMORY, "total_memory"), (38, TOTAL_CPU, "total_cpu"), (39, NAMESPACE, "namespace"), - (40, ADDRESS, "address"), + // 40 (41, DISPLAY_NAME, "display_name"), (42, CONNECTABLE, "connectable"), (43, SLEEP_TS, "sleep_ts"), diff --git a/packages/core/guard/server/src/cache/actor.rs b/packages/core/guard/server/src/cache/actor.rs index 1a954cb6fb..7b30025890 100644 --- a/packages/core/guard/server/src/cache/actor.rs +++ b/packages/core/guard/server/src/cache/actor.rs @@ -6,7 +6,7 @@ use std::{ use anyhow::Result; use gas::prelude::*; -use crate::routing::pegboard_gateway::{X_RIVET_ACTOR, X_RIVET_PORT}; +use crate::routing::pegboard_gateway::X_RIVET_ACTOR; #[tracing::instrument(skip_all)] pub fn build_cache_key(target: &str, path: &str, headers: &hyper::HeaderMap) -> Result { @@ -22,19 +22,10 @@ pub fn build_cache_key(target: &str, path: &str, headers: &hyper::HeaderMap) -> })?; let actor_id = Id::parse(actor_id_str.to_str()?)?; - let port_name = headers.get(X_RIVET_PORT).ok_or_else(|| { - crate::errors::MissingHeader { - header: X_RIVET_PORT.to_string(), - } - .build() - })?; - let port_name = port_name.to_str()?; - - // Create a hash using target, actor_id and port_name + // Create a hash using target, actor_id, and path let mut hasher = DefaultHasher::new(); target.hash(&mut hasher); actor_id.hash(&mut hasher); - port_name.hash(&mut hasher); path.hash(&mut hasher); let hash = hasher.finish(); diff --git a/packages/core/guard/server/src/errors.rs b/packages/core/guard/server/src/errors.rs index abc94bcbdd..1e9f5270b0 100644 --- a/packages/core/guard/server/src/errors.rs +++ b/packages/core/guard/server/src/errors.rs @@ -43,11 +43,10 @@ pub struct WrongAddrProtocol { "guard", "actor_not_found", "Actor not found.", - "Actor with ID {actor_id} and port {port_name} not found." + "Actor with ID {actor_id} not found." )] pub struct ActorNotFound { pub actor_id: Id, - pub port_name: String, } #[derive(RivetError, Serialize)] diff --git a/packages/core/guard/server/src/routing/pegboard_gateway.rs b/packages/core/guard/server/src/routing/pegboard_gateway.rs index 58088a1ac4..7015a8a686 100644 --- a/packages/core/guard/server/src/routing/pegboard_gateway.rs +++ b/packages/core/guard/server/src/routing/pegboard_gateway.rs @@ -10,7 +10,6 @@ use crate::{errors, shared_state::SharedState}; const ACTOR_READY_TIMEOUT: Duration = Duration::from_secs(10); pub const X_RIVET_ACTOR: HeaderName = HeaderName::from_static("x-rivet-actor"); -pub const X_RIVET_PORT: HeaderName = HeaderName::from_static("x-rivet-port"); /// Route requests to actor services based on hostname and path #[tracing::instrument(skip_all)] @@ -65,16 +64,8 @@ pub async fn route_request( }))); } - let port_name = headers.get(X_RIVET_PORT).ok_or_else(|| { - crate::errors::MissingHeader { - header: X_RIVET_PORT.to_string(), - } - .build() - })?; - let port_name = port_name.to_str()?; - // Lookup actor - find_actor(ctx, shared_state, actor_id, port_name, path).await + find_actor(ctx, shared_state, actor_id, path).await } struct FoundActor { @@ -83,13 +74,12 @@ struct FoundActor { destroyed: bool, } -/// Find an actor by actor_id and port_name -#[tracing::instrument(skip_all, fields(%actor_id, %port_name, %path))] +/// Find an actor by actor_id +#[tracing::instrument(skip_all, fields(%actor_id, %path))] async fn find_actor( ctx: &StandaloneCtx, shared_state: &SharedState, actor_id: Id, - port_name: &str, path: &str, ) -> Result> { // TODO: Optimize this down to a single FDB call @@ -136,11 +126,7 @@ async fn find_actor( .await??; let Some(actor) = actor_res else { - return Err(errors::ActorNotFound { - actor_id, - port_name: port_name.to_string(), - } - .build()); + return Err(errors::ActorNotFound { actor_id }.build()); }; if actor.destroyed { @@ -187,6 +173,7 @@ async fn find_actor( tracing::debug!(?actor_id, ?runner_id, "actor ready"); + // TODO: Remove round trip, return key from get_runner op above // Get runner key from runner_id let runner_key = ctx .udb()? @@ -204,7 +191,6 @@ async fn find_actor( shared_state.pegboard_gateway.clone(), actor_id, runner_key, - port_name.to_string(), ); Ok(Some(RoutingOutput::CustomServe(std::sync::Arc::new( gateway, diff --git a/packages/core/pegboard-gateway/src/lib.rs b/packages/core/pegboard-gateway/src/lib.rs index 8d752eba42..67a23eff2a 100644 --- a/packages/core/pegboard-gateway/src/lib.rs +++ b/packages/core/pegboard-gateway/src/lib.rs @@ -30,7 +30,6 @@ pub struct PegboardGateway { shared_state: SharedState, actor_id: Id, runner_key: String, - port_name: String, } impl PegboardGateway { @@ -39,14 +38,12 @@ impl PegboardGateway { shared_state: SharedState, actor_id: Id, runner_key: String, - port_name: String, ) -> Self { Self { ctx, shared_state, actor_id, runner_key, - port_name, } } } @@ -139,11 +136,9 @@ impl PegboardGateway { .to_bytes(); // Build subject to publish to - let tunnel_subject = pegboard::pubsub_subjects::TunnelRunnerReceiverSubject::new( - &self.runner_key, - &self.port_name, - ) - .to_string(); + let tunnel_subject = + pegboard::pubsub_subjects::TunnelRunnerReceiverSubject::new(&self.runner_key) + .to_string(); // Start listening for request responses let (request_id, mut msg_rx) = self @@ -235,11 +230,9 @@ impl PegboardGateway { } // Build subject to publish to - let tunnel_subject = pegboard::pubsub_subjects::TunnelRunnerReceiverSubject::new( - &self.runner_key, - &self.port_name, - ) - .to_string(); + let tunnel_subject = + pegboard::pubsub_subjects::TunnelRunnerReceiverSubject::new(&self.runner_key) + .to_string(); // Start listening for WebSocket messages let (request_id, mut msg_rx) = self diff --git a/packages/core/pegboard-runner-ws/src/lib.rs b/packages/core/pegboard-runner-ws/src/lib.rs index 576c7d567e..334c140ff0 100644 --- a/packages/core/pegboard-runner-ws/src/lib.rs +++ b/packages/core/pegboard-runner-ws/src/lib.rs @@ -304,9 +304,6 @@ async fn build_connection( name, version, total_slots, - addresses_http, - addresses_tcp, - addresses_udp, .. } = &packet { @@ -360,10 +357,6 @@ async fn build_connection( key: runner_key.clone(), version: version.clone(), total_slots: *total_slots, - - addresses_http: addresses_http.clone().unwrap_or_default(), - addresses_tcp: addresses_tcp.clone().unwrap_or_default(), - addresses_udp: addresses_udp.clone().unwrap_or_default(), }) .tag("runner_id", runner_id) .unique() diff --git a/packages/core/pegboard-tunnel/src/lib.rs b/packages/core/pegboard-tunnel/src/lib.rs index a71d3dc62f..ea07758b75 100644 --- a/packages/core/pegboard-tunnel/src/lib.rs +++ b/packages/core/pegboard-tunnel/src/lib.rs @@ -91,11 +91,8 @@ impl CustomServeTrait for PegboardTunnelCustomServe { } }; - let port_name = "main".to_string(); // Use "main" as default port name - tracing::info!( ?runner_key, - ?port_name, ?protocol_version, ?path, "tunnel WebSocket connection established" @@ -104,8 +101,7 @@ impl CustomServeTrait for PegboardTunnelCustomServe { // Subscribe to pubsub topic for this runner before accepting the client websocket so // that failures can be retried by the proxy. let topic = - pegboard::pubsub_subjects::TunnelRunnerReceiverSubject::new(&runner_key, &port_name) - .to_string(); + pegboard::pubsub_subjects::TunnelRunnerReceiverSubject::new(&runner_key).to_string(); tracing::info!(%topic, ?runner_key, "subscribing to runner receiver topic"); let mut sub = match ups.subscribe(&topic).await { Result::Ok(s) => s, diff --git a/packages/core/pegboard-tunnel/tests/integration.rs b/packages/core/pegboard-tunnel/tests/integration.rs index 70051af60a..df5d1ecf8f 100644 --- a/packages/core/pegboard-tunnel/tests/integration.rs +++ b/packages/core/pegboard-tunnel/tests/integration.rs @@ -48,13 +48,12 @@ async fn test_tunnel_bidirectional_forwarding() -> Result<()> { // Use the same placeholders as the tunnel implementation // TODO: Update when tunnel properly extracts these from connection let runner_id = Id::nil(); - let port_name = "default"; // Give tunnel time to set up pubsub subscription after WebSocket connection sleep(Duration::from_secs(1)).await; // Test 1: pubsub to WebSocket forwarding - test_pubsub_to_websocket(&ups, &mut ws_stream, runner_id, port_name).await?; + test_pubsub_to_websocket(&ups, &mut ws_stream, runner_id).await?; // Test 2: WebSocket to pubsub forwarding test_websocket_to_pubsub(&ups, &mut ws_stream, runner_id).await?; @@ -69,7 +68,6 @@ async fn test_pubsub_to_websocket( ups: &PubSub, ws_stream: &mut WebSocketStream>, runner_id: Id, - port_name: &str, ) -> Result<()> { // Create a test request message let request_id = rand::random::(); @@ -96,7 +94,7 @@ async fn test_pubsub_to_websocket( )?; // Publish to pubsub topic using proper subject - let topic = TunnelHttpRunnerSubject::new(&runner_id.to_string(), port_name).to_string(); + let topic = TunnelHttpRunnerSubject::new(&runner_id.to_string()).to_string(); ups.request(&topic, &serialized).await?; // Wait for message on WebSocket diff --git a/packages/infra/engine/tests/actors_lifecycle.rs b/packages/infra/engine/tests/actors_lifecycle.rs index 46aeff7e9d..d98d831bfd 100644 --- a/packages/infra/engine/tests/actors_lifecycle.rs +++ b/packages/infra/engine/tests/actors_lifecycle.rs @@ -30,14 +30,12 @@ async fn actor_lifecycle_inner(ctx: &common::TestCtx, multi_dc: bool) { let actor_id = common::create_actor(&namespace, target_dc.guard_port()).await; // Test ping via guard - let ping_response = - common::ping_actor_via_guard(ctx.leader_dc().guard_port(), &actor_id, "main").await; + let ping_response = common::ping_actor_via_guard(ctx.leader_dc().guard_port(), &actor_id).await; assert_eq!(ping_response["status"], "ok"); // Test websocket via guard let ws_response = - common::ping_actor_websocket_via_guard(ctx.leader_dc().guard_port(), &actor_id, "main") - .await; + common::ping_actor_websocket_via_guard(ctx.leader_dc().guard_port(), &actor_id).await; assert_eq!(ws_response["status"], "ok"); // Validate runner state @@ -115,13 +113,12 @@ async fn actor_lifecycle_with_same_key_inner(ctx: &common::TestCtx, dc_choice: D // Test ping via guard let ping_response = - common::ping_actor_via_guard(ctx.leader_dc().guard_port(), &actor_id1, "main").await; + common::ping_actor_via_guard(ctx.leader_dc().guard_port(), &actor_id1).await; assert_eq!(ping_response["status"], "ok"); // Test websocket via guard let ws_response = - common::ping_actor_websocket_via_guard(ctx.leader_dc().guard_port(), &actor_id1, "main") - .await; + common::ping_actor_websocket_via_guard(ctx.leader_dc().guard_port(), &actor_id1).await; assert_eq!(ws_response["status"], "ok"); // Destroy @@ -150,13 +147,12 @@ async fn actor_lifecycle_with_same_key_inner(ctx: &common::TestCtx, dc_choice: D // Test ping via guard let ping_response = - common::ping_actor_via_guard(ctx.leader_dc().guard_port(), &actor_id2, "main").await; + common::ping_actor_via_guard(ctx.leader_dc().guard_port(), &actor_id2).await; assert_eq!(ping_response["status"], "ok"); // Test websocket via guard let ws_response = - common::ping_actor_websocket_via_guard(ctx.leader_dc().guard_port(), &actor_id2, "main") - .await; + common::ping_actor_websocket_via_guard(ctx.leader_dc().guard_port(), &actor_id2).await; assert_eq!(ws_response["status"], "ok"); // Destroy diff --git a/packages/infra/engine/tests/common/actors.rs b/packages/infra/engine/tests/common/actors.rs index 9108a1ef23..d3495c6dc5 100644 --- a/packages/infra/engine/tests/common/actors.rs +++ b/packages/infra/engine/tests/common/actors.rs @@ -94,24 +94,14 @@ pub async fn create_actor(namespace_name: &str, guard_port: u16) -> String { } /// Pings actor via Guard. -pub async fn ping_actor_via_guard( - guard_port: u16, - actor_id: &str, - addr_name: &str, -) -> serde_json::Value { - tracing::info!( - ?guard_port, - ?actor_id, - ?addr_name, - "sending request to actor via guard" - ); +pub async fn ping_actor_via_guard(guard_port: u16, actor_id: &str) -> serde_json::Value { + tracing::info!(?guard_port, ?actor_id, "sending request to actor via guard"); let client = reqwest::Client::new(); let response = client .get(format!("http://127.0.0.1:{}/ping", guard_port)) .header("X-Rivet-Target", "actor") .header("X-Rivet-Actor", actor_id) - .header("X-Rivet-Port", addr_name) .send() .await .expect("Failed to send ping request through guard"); @@ -425,11 +415,7 @@ pub async fn bulk_create_actors( } /// Tests WebSocket connection to actor via Guard using a simple ping pong. -pub async fn ping_actor_websocket_via_guard( - guard_port: u16, - actor_id: &str, - addr_name: &str, -) -> serde_json::Value { +pub async fn ping_actor_websocket_via_guard(guard_port: u16, actor_id: &str) -> serde_json::Value { use tokio_tungstenite::{ connect_async, tungstenite::{Message, client::IntoClientRequest}, @@ -438,7 +424,6 @@ pub async fn ping_actor_websocket_via_guard( tracing::info!( ?guard_port, ?actor_id, - ?addr_name, "testing websocket connection to actor via guard" ); @@ -456,9 +441,6 @@ pub async fn ping_actor_websocket_via_guard( request .headers_mut() .insert("X-Rivet-Actor", actor_id.parse().unwrap()); - request - .headers_mut() - .insert("X-Rivet-Port", addr_name.parse().unwrap()); // Connect to WebSocket let (ws_stream, response) = connect_async(request) diff --git a/packages/services/pegboard/src/keys/runner.rs b/packages/services/pegboard/src/keys/runner.rs index 528ba2f6b4..9c20248557 100644 --- a/packages/services/pegboard/src/keys/runner.rs +++ b/packages/services/pegboard/src/keys/runner.rs @@ -507,78 +507,6 @@ impl<'de> TupleUnpack<'de> for VersionKey { } } -#[derive(Debug)] -pub struct AddressKey { - pub runner_id: Id, - pub name: String, -} - -impl AddressKey { - pub fn new(runner_id: Id, name: String) -> Self { - AddressKey { runner_id, name } - } - - pub fn subspace(runner_id: Id) -> AddressSubspaceKey { - AddressSubspaceKey::new(runner_id) - } -} - -impl FormalKey for AddressKey { - type Value = ::Latest; - - fn deserialize(&self, raw: &[u8]) -> Result { - rivet_data::versioned::AddressKeyData::deserialize_with_embedded_version(raw) - } - - fn serialize(&self, value: Self::Value) -> Result> { - rivet_data::versioned::AddressKeyData::latest(value) - .serialize_with_embedded_version(rivet_data::PEGBOARD_RUNNER_ADDRESS_VERSION) - } -} - -impl TuplePack for AddressKey { - fn pack( - &self, - w: &mut W, - tuple_depth: TupleDepth, - ) -> std::io::Result { - let t = (RUNNER, DATA, self.runner_id, ADDRESS, &self.name); - t.pack(w, tuple_depth) - } -} - -impl<'de> TupleUnpack<'de> for AddressKey { - fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let (input, (_, _, runner_id, _, name)) = - <(usize, usize, Id, usize, String)>::unpack(input, tuple_depth)?; - - let v = AddressKey { runner_id, name }; - - Ok((input, v)) - } -} - -pub struct AddressSubspaceKey { - runner_id: Id, -} - -impl AddressSubspaceKey { - pub fn new(runner_id: Id) -> Self { - AddressSubspaceKey { runner_id } - } -} - -impl TuplePack for AddressSubspaceKey { - fn pack( - &self, - w: &mut W, - tuple_depth: TupleDepth, - ) -> std::io::Result { - let t = (RUNNER, DATA, self.runner_id, ADDRESS); - t.pack(w, tuple_depth) - } -} - #[derive(Debug)] pub struct StopTsKey { runner_id: Id, diff --git a/packages/services/pegboard/src/ops/runner/get.rs b/packages/services/pegboard/src/ops/runner/get.rs index 22cbbcc6e4..484eb8e0de 100644 --- a/packages/services/pegboard/src/ops/runner/get.rs +++ b/packages/services/pegboard/src/ops/runner/get.rs @@ -1,7 +1,6 @@ use anyhow::Result; use futures_util::TryStreamExt; use gas::prelude::*; -use rivet_data::generated::pegboard_runner_address_v1::Data as AddressKeyData; use rivet_types::runners::Runner; use udb_util::{FormalChunkedKey, SERIALIZABLE, SNAPSHOT, TxnExt}; use universaldb::{self as udb, options::StreamingMode}; @@ -87,7 +86,6 @@ pub(crate) async fn get_inner( stop_ts, last_ping_ts, last_rtt, - (addresses_http, addresses_tcp, addresses_udp), metadata_chunks, ) = tokio::try_join!( // NOTE: These are not SERIALIZABLE because this op is meant for basic information (i.e. data for the @@ -104,41 +102,6 @@ pub(crate) async fn get_inner( txs.read_opt(&stop_ts_key, SNAPSHOT), txs.read_opt(&last_ping_ts_key, SNAPSHOT), txs.read_opt(&last_rtt_key, SNAPSHOT), - async { - // Get addresses by scanning all address keys for this runner - let mut addresses_http = util::serde::HashableMap::new(); - let mut addresses_tcp = util::serde::HashableMap::new(); - let mut addresses_udp = util::serde::HashableMap::new(); - - let address_subspace = txs.subspace(&keys::runner::AddressKey::subspace(runner_id)); - - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { - mode: StreamingMode::Iterator, - ..(&address_subspace).into() - }, - SNAPSHOT, - ); - - while let Some(entry) = stream.try_next().await? { - let (address_key, address_data) = - txs.read_entry::(&entry)?; - - match address_data { - AddressKeyData::Http(addr) => { - addresses_http.insert(address_key.name.clone(), addr); - } - AddressKeyData::Tcp(addr) => { - addresses_tcp.insert(address_key.name.clone(), addr); - } - AddressKeyData::Udp(addr) => { - addresses_udp.insert(address_key.name.clone(), addr); - } - } - } - - Ok((addresses_http, addresses_tcp, addresses_udp)) - }, async { txs.get_ranges_keyvalues( udb::RangeOption { @@ -173,9 +136,6 @@ pub(crate) async fn get_inner( version, total_slots, remaining_slots, - addresses_http: addresses_http.into(), - addresses_tcp: addresses_tcp.into(), - addresses_udp: addresses_udp.into(), create_ts, last_connected_ts: connected_ts, drain_ts, diff --git a/packages/services/pegboard/src/pubsub_subjects.rs b/packages/services/pegboard/src/pubsub_subjects.rs index ae9e9438b6..d51f2c47c9 100644 --- a/packages/services/pegboard/src/pubsub_subjects.rs +++ b/packages/services/pegboard/src/pubsub_subjects.rs @@ -2,25 +2,17 @@ use gas::prelude::*; pub struct TunnelRunnerReceiverSubject<'a> { runner_key: &'a str, - port_name: &'a str, } impl<'a> TunnelRunnerReceiverSubject<'a> { - pub fn new(runner_key: &'a str, port_name: &'a str) -> Self { - Self { - runner_key, - port_name, - } + pub fn new(runner_key: &'a str) -> Self { + Self { runner_key } } } impl std::fmt::Display for TunnelRunnerReceiverSubject<'_> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "pegboard.tunnel.runner_receiver.{}.{}", - self.runner_key, self.port_name - ) + write!(f, "pegboard.tunnel.runner_receiver.{}", self.runner_key) } } diff --git a/packages/services/pegboard/src/workflows/runner.rs b/packages/services/pegboard/src/workflows/runner.rs index e5d64f17d6..b88cc019c2 100644 --- a/packages/services/pegboard/src/workflows/runner.rs +++ b/packages/services/pegboard/src/workflows/runner.rs @@ -1,9 +1,6 @@ use futures_util::{FutureExt, StreamExt, TryStreamExt}; use gas::prelude::*; -use rivet_data::{ - converted::{ActorNameKeyData, MetadataKeyData, RunnerByKeyKeyData}, - generated::pegboard_runner_address_v1::Data as AddressKeyData, -}; +use rivet_data::converted::{ActorNameKeyData, MetadataKeyData, RunnerByKeyKeyData}; use rivet_runner_protocol::protocol; use udb_util::{FormalChunkedKey, SERIALIZABLE, SNAPSHOT, TxnExt}; use universaldb::{ @@ -27,10 +24,6 @@ pub struct Input { pub key: String, pub version: u32, pub total_slots: u32, - - pub addresses_http: util::serde::HashableMap, - pub addresses_tcp: util::serde::HashableMap, - pub addresses_udp: util::serde::HashableMap, } #[derive(Debug, Serialize, Deserialize)] @@ -72,9 +65,6 @@ pub async fn pegboard_runner(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> key: input.key.clone(), namespace_id: input.namespace_id, create_ts: ctx.create_ts(), - addresses_http: input.addresses_http.clone(), - addresses_tcp: input.addresses_tcp.clone(), - addresses_udp: input.addresses_udp.clone(), }) .await?; @@ -437,10 +427,6 @@ struct InitInput { key: String, namespace_id: Id, create_ts: i64, - - addresses_http: util::serde::HashableMap, - addresses_tcp: util::serde::HashableMap, - addresses_udp: util::serde::HashableMap, } #[derive(Debug, Serialize, Deserialize)] @@ -460,27 +446,6 @@ async fn init(ctx: &ActivityCtx, input: &InitInput) -> Result { .run(|tx, _mc| async move { let txs = tx.subspace(keys::subspace()); - for (name, port_http) in &input.addresses_http { - txs.write( - &keys::runner::AddressKey::new(input.runner_id, name.into()), - AddressKeyData::Http(port_http.clone().into()), - )?; - } - - for (name, port_tcp) in &input.addresses_tcp { - txs.write( - &keys::runner::AddressKey::new(input.runner_id, name.into()), - AddressKeyData::Tcp(port_tcp.clone().into()), - )?; - } - - for (name, port_udp) in &input.addresses_udp { - txs.write( - &keys::runner::AddressKey::new(input.runner_id, name.into()), - AddressKeyData::Udp(port_udp.clone().into()), - )?; - } - let runner_by_key_key = keys::ns::RunnerByKeyKey::new( input.namespace_id, input.name.clone(), diff --git a/scripts/tests/actor_e2e.ts b/scripts/tests/actor_e2e.ts index 070ac21fab..0290460e2f 100755 --- a/scripts/tests/actor_e2e.ts +++ b/scripts/tests/actor_e2e.ts @@ -19,7 +19,6 @@ async function main() { headers: { "X-Rivet-Target": "actor", "X-Rivet-Actor": actorResponse.actor.actor_id, - "X-Rivet-Port": "main", }, }); @@ -65,7 +64,6 @@ function testWebSocket(actorId: string): Promise { headers: { "X-Rivet-Target": "actor", "X-Rivet-Actor": actorId, - "X-Rivet-Port": "main", }, }); diff --git a/scripts/tests/actor_e2e_multidc.ts b/scripts/tests/actor_e2e_multidc.ts index f2caca24ec..6e1f4c9802 100755 --- a/scripts/tests/actor_e2e_multidc.ts +++ b/scripts/tests/actor_e2e_multidc.ts @@ -37,7 +37,7 @@ async function createActorInDc( async function testActorInDc(dc: string) { console.log(`\n=== Testing actor in ${dc} ===`); - + try { // Create an actor in the specified datacenter console.log(`Creating actor in ${dc}...`); @@ -55,7 +55,6 @@ async function testActorInDc(dc: string) { headers: { "x-rivet-target": "actor", "x-rivet-actor": actorResponse.actor.actor_id, - "x-rivet-addr": "main", }, }); @@ -83,7 +82,7 @@ async function testActorInDc(dc: string) { async function main() { const datacenters = ["dc-a", "dc-b", "dc-c"]; const results: Record = {}; - + console.log("Starting multi-datacenter actor E2E test..."); console.log(`Testing datacenters: ${datacenters.join(", ")}`); diff --git a/sdks/rust/data/src/converted.rs b/sdks/rust/data/src/converted.rs index 859682e546..4c956ffda9 100644 --- a/sdks/rust/data/src/converted.rs +++ b/sdks/rust/data/src/converted.rs @@ -33,60 +33,6 @@ impl TryFrom for pegboard_namespace_runner_alloc_idx_v1:: } } -impl From for rivet_runner_protocol::protocol::RunnerAddressHttp { - fn from(value: pegboard_runner_address_v1::Http) -> Self { - rivet_runner_protocol::protocol::RunnerAddressHttp { - hostname: value.hostname, - port: value.port, - } - } -} - -impl From for pegboard_runner_address_v1::Http { - fn from(value: rivet_runner_protocol::protocol::RunnerAddressHttp) -> Self { - pegboard_runner_address_v1::Http { - hostname: value.hostname, - port: value.port, - } - } -} - -impl From for rivet_runner_protocol::protocol::RunnerAddressTcp { - fn from(value: pegboard_runner_address_v1::Tcp) -> Self { - rivet_runner_protocol::protocol::RunnerAddressTcp { - hostname: value.hostname, - port: value.port, - } - } -} - -impl From for pegboard_runner_address_v1::Tcp { - fn from(value: rivet_runner_protocol::protocol::RunnerAddressTcp) -> Self { - pegboard_runner_address_v1::Tcp { - hostname: value.hostname, - port: value.port, - } - } -} - -impl From for rivet_runner_protocol::protocol::RunnerAddressUdp { - fn from(value: pegboard_runner_address_v1::Udp) -> Self { - rivet_runner_protocol::protocol::RunnerAddressUdp { - hostname: value.hostname, - port: value.port, - } - } -} - -impl From for pegboard_runner_address_v1::Udp { - fn from(value: rivet_runner_protocol::protocol::RunnerAddressUdp) -> Self { - pegboard_runner_address_v1::Udp { - hostname: value.hostname, - port: value.port, - } - } -} - pub struct MetadataKeyData { pub metadata: serde_json::Map, } diff --git a/sdks/rust/data/src/versioned.rs b/sdks/rust/data/src/versioned.rs index 551e518158..b898f38096 100644 --- a/sdks/rust/data/src/versioned.rs +++ b/sdks/rust/data/src/versioned.rs @@ -37,40 +37,6 @@ impl OwnedVersionedData for RunnerAllocIdxKeyData { } } -pub enum AddressKeyData { - V1(pegboard_runner_address_v1::Data), -} - -impl OwnedVersionedData for AddressKeyData { - type Latest = pegboard_runner_address_v1::Data; - - fn latest(latest: pegboard_runner_address_v1::Data) -> Self { - AddressKeyData::V1(latest) - } - - fn into_latest(self) -> Result { - #[allow(irrefutable_let_patterns)] - if let AddressKeyData::V1(data) = self { - Ok(data) - } else { - bail!("version not latest"); - } - } - - fn deserialize_version(payload: &[u8], version: u16) -> Result { - match version { - 1 => Ok(AddressKeyData::V1(serde_bare::from_slice(payload)?)), - _ => bail!("invalid version: {version}"), - } - } - - fn serialize_version(self, _version: u16) -> Result> { - match self { - AddressKeyData::V1(data) => serde_bare::to_vec(&data).map_err(Into::into), - } - } -} - pub enum MetadataKeyData { V1(pegboard_runner_metadata_v1::Data), } diff --git a/sdks/rust/runner-protocol/src/protocol.rs b/sdks/rust/runner-protocol/src/protocol.rs index a037941e2a..67d7644a18 100644 --- a/sdks/rust/runner-protocol/src/protocol.rs +++ b/sdks/rust/runner-protocol/src/protocol.rs @@ -24,10 +24,6 @@ pub enum ToServer { version: u32, total_slots: u32, - addresses_http: Option>, - addresses_tcp: Option>, - addresses_udp: Option>, - last_command_idx: Option, prepopulate_actor_names: Option>, metadata: Option, @@ -156,24 +152,6 @@ pub enum StopCode { Error, } -#[derive(Debug, Clone, Serialize, Deserialize, Hash, utoipa::ToSchema)] -pub struct RunnerAddressHttp { - pub hostname: String, - pub port: u16, -} - -#[derive(Debug, Clone, Serialize, Deserialize, Hash, utoipa::ToSchema)] -pub struct RunnerAddressTcp { - pub hostname: String, - pub port: u16, -} - -#[derive(Debug, Clone, Serialize, Deserialize, Hash, utoipa::ToSchema)] -pub struct RunnerAddressUdp { - pub hostname: String, - pub port: u16, -} - #[derive(Debug, Clone, Copy, Serialize, Deserialize, Hash)] #[serde(rename_all = "snake_case")] pub enum WebsocketCloseReason { diff --git a/sdks/rust/runner-protocol/src/versioned.rs b/sdks/rust/runner-protocol/src/versioned.rs index 34c54168b6..954338905c 100644 --- a/sdks/rust/runner-protocol/src/versioned.rs +++ b/sdks/rust/runner-protocol/src/versioned.rs @@ -252,33 +252,6 @@ impl TryFrom for protocol::StopCode { } } -impl From for protocol::RunnerAddressHttp { - fn from(value: v1::RunnerAddressHttp) -> Self { - protocol::RunnerAddressHttp { - hostname: value.hostname, - port: value.port, - } - } -} - -impl From for protocol::RunnerAddressTcp { - fn from(value: v1::RunnerAddressTcp) -> Self { - protocol::RunnerAddressTcp { - hostname: value.hostname, - port: value.port, - } - } -} - -impl From for protocol::RunnerAddressUdp { - fn from(value: v1::RunnerAddressUdp) -> Self { - protocol::RunnerAddressUdp { - hostname: value.hostname, - port: value.port, - } - } -} - impl TryFrom for protocol::ToServer { type Error = anyhow::Error; @@ -288,18 +261,6 @@ impl TryFrom for protocol::ToServer { name: init.name, version: init.version, total_slots: init.total_slots, - addresses_http: init - .addresses_http - .map(|addrs| Ok(addrs.into_iter().map(|(k, v)| (k, v.into())).collect())) - .transpose()?, - addresses_tcp: init - .addresses_tcp - .map(|addrs| Ok(addrs.into_iter().map(|(k, v)| (k, v.into())).collect())) - .transpose()?, - addresses_udp: init - .addresses_udp - .map(|addrs| Ok(addrs.into_iter().map(|(k, v)| (k, v.into())).collect())) - .transpose()?, last_command_idx: init.last_command_idx, prepopulate_actor_names: init .prepopulate_actor_names diff --git a/sdks/schemas/data/pegboard.runner.address.v1.bare b/sdks/schemas/data/pegboard.runner.address.v1.bare deleted file mode 100644 index 620ee155aa..0000000000 --- a/sdks/schemas/data/pegboard.runner.address.v1.bare +++ /dev/null @@ -1,20 +0,0 @@ -type Http struct { - hostname: str - port: u16 -} - -type Tcp struct { - hostname: str - port: u16 -} - -type Udp struct { - hostname: str - port: u16 -} - -type Data union { - Http | - Tcp | - Udp -} diff --git a/sdks/schemas/runner-protocol/v1.bare b/sdks/schemas/runner-protocol/v1.bare index 5db5afcea4..e04ca50e24 100644 --- a/sdks/schemas/runner-protocol/v1.bare +++ b/sdks/schemas/runner-protocol/v1.bare @@ -34,21 +34,6 @@ type ActorName struct { metadata: Json } -type RunnerAddressHttp struct { - hostname: str - port: u16 -} - -type RunnerAddressTcp struct { - hostname: str - port: u16 -} - -type RunnerAddressUdp struct { - hostname: str - port: u16 -} - type StopCode enum { OK ERROR @@ -136,9 +121,6 @@ type ToServerInit struct { name: str version: u32 totalSlots: u32 - addressesHttp: optional> - addressesTcp: optional> - addressesUdp: optional> lastCommandIdx: optional prepopulateActorNames: optional> metadata: optional diff --git a/sdks/typescript/runner-protocol/src/index.ts b/sdks/typescript/runner-protocol/src/index.ts index 96c00212c6..84a6c9463d 100644 --- a/sdks/typescript/runner-protocol/src/index.ts +++ b/sdks/typescript/runner-protocol/src/index.ts @@ -3,7 +3,6 @@ import * as bare from "@bare-ts/lib" const DEFAULT_CONFIG = /* @__PURE__ */ bare.Config({}) export type i64 = bigint -export type u16 = number export type u32 = number export type u64 = bigint @@ -155,57 +154,6 @@ export function writeActorName(bc: bare.ByteCursor, x: ActorName): void { writeJson(bc, x.metadata) } -export type RunnerAddressHttp = { - readonly hostname: string - readonly port: u16 -} - -export function readRunnerAddressHttp(bc: bare.ByteCursor): RunnerAddressHttp { - return { - hostname: bare.readString(bc), - port: bare.readU16(bc), - } -} - -export function writeRunnerAddressHttp(bc: bare.ByteCursor, x: RunnerAddressHttp): void { - bare.writeString(bc, x.hostname) - bare.writeU16(bc, x.port) -} - -export type RunnerAddressTcp = { - readonly hostname: string - readonly port: u16 -} - -export function readRunnerAddressTcp(bc: bare.ByteCursor): RunnerAddressTcp { - return { - hostname: bare.readString(bc), - port: bare.readU16(bc), - } -} - -export function writeRunnerAddressTcp(bc: bare.ByteCursor, x: RunnerAddressTcp): void { - bare.writeString(bc, x.hostname) - bare.writeU16(bc, x.port) -} - -export type RunnerAddressUdp = { - readonly hostname: string - readonly port: u16 -} - -export function readRunnerAddressUdp(bc: bare.ByteCursor): RunnerAddressUdp { - return { - hostname: bare.readString(bc), - port: bare.readU16(bc), - } -} - -export function writeRunnerAddressUdp(bc: bare.ByteCursor, x: RunnerAddressUdp): void { - bare.writeString(bc, x.hostname) - bare.writeU16(bc, x.port) -} - export enum StopCode { Ok = "Ok", Error = "Error", @@ -590,109 +538,7 @@ export function writeCommandWrapper(bc: bare.ByteCursor, x: CommandWrapper): voi writeCommand(bc, x.inner) } -function read3(bc: bare.ByteCursor): ReadonlyMap { - const len = bare.readUintSafe(bc) - const result = new Map() - for (let i = 0; i < len; i++) { - const offset = bc.offset - const key = bare.readString(bc) - if (result.has(key)) { - bc.offset = offset - throw new bare.BareError(offset, "duplicated key") - } - result.set(key, readRunnerAddressHttp(bc)) - } - return result -} - -function write3(bc: bare.ByteCursor, x: ReadonlyMap): void { - bare.writeUintSafe(bc, x.size) - for (const kv of x) { - bare.writeString(bc, kv[0]) - writeRunnerAddressHttp(bc, kv[1]) - } -} - -function read4(bc: bare.ByteCursor): ReadonlyMap | null { - return bare.readBool(bc) ? read3(bc) : null -} - -function write4(bc: bare.ByteCursor, x: ReadonlyMap | null): void { - bare.writeBool(bc, x != null) - if (x != null) { - write3(bc, x) - } -} - -function read5(bc: bare.ByteCursor): ReadonlyMap { - const len = bare.readUintSafe(bc) - const result = new Map() - for (let i = 0; i < len; i++) { - const offset = bc.offset - const key = bare.readString(bc) - if (result.has(key)) { - bc.offset = offset - throw new bare.BareError(offset, "duplicated key") - } - result.set(key, readRunnerAddressTcp(bc)) - } - return result -} - -function write5(bc: bare.ByteCursor, x: ReadonlyMap): void { - bare.writeUintSafe(bc, x.size) - for (const kv of x) { - bare.writeString(bc, kv[0]) - writeRunnerAddressTcp(bc, kv[1]) - } -} - -function read6(bc: bare.ByteCursor): ReadonlyMap | null { - return bare.readBool(bc) ? read5(bc) : null -} - -function write6(bc: bare.ByteCursor, x: ReadonlyMap | null): void { - bare.writeBool(bc, x != null) - if (x != null) { - write5(bc, x) - } -} - -function read7(bc: bare.ByteCursor): ReadonlyMap { - const len = bare.readUintSafe(bc) - const result = new Map() - for (let i = 0; i < len; i++) { - const offset = bc.offset - const key = bare.readString(bc) - if (result.has(key)) { - bc.offset = offset - throw new bare.BareError(offset, "duplicated key") - } - result.set(key, readRunnerAddressUdp(bc)) - } - return result -} - -function write7(bc: bare.ByteCursor, x: ReadonlyMap): void { - bare.writeUintSafe(bc, x.size) - for (const kv of x) { - bare.writeString(bc, kv[0]) - writeRunnerAddressUdp(bc, kv[1]) - } -} - -function read8(bc: bare.ByteCursor): ReadonlyMap | null { - return bare.readBool(bc) ? read7(bc) : null -} - -function write8(bc: bare.ByteCursor, x: ReadonlyMap | null): void { - bare.writeBool(bc, x != null) - if (x != null) { - write7(bc, x) - } -} - -function read9(bc: bare.ByteCursor): ReadonlyMap { +function read3(bc: bare.ByteCursor): ReadonlyMap { const len = bare.readUintSafe(bc) const result = new Map() for (let i = 0; i < len; i++) { @@ -707,7 +553,7 @@ function read9(bc: bare.ByteCursor): ReadonlyMap { return result } -function write9(bc: bare.ByteCursor, x: ReadonlyMap): void { +function write3(bc: bare.ByteCursor, x: ReadonlyMap): void { bare.writeUintSafe(bc, x.size) for (const kv of x) { bare.writeString(bc, kv[0]) @@ -715,22 +561,22 @@ function write9(bc: bare.ByteCursor, x: ReadonlyMap): void { } } -function read10(bc: bare.ByteCursor): ReadonlyMap | null { - return bare.readBool(bc) ? read9(bc) : null +function read4(bc: bare.ByteCursor): ReadonlyMap | null { + return bare.readBool(bc) ? read3(bc) : null } -function write10(bc: bare.ByteCursor, x: ReadonlyMap | null): void { +function write4(bc: bare.ByteCursor, x: ReadonlyMap | null): void { bare.writeBool(bc, x != null) if (x != null) { - write9(bc, x) + write3(bc, x) } } -function read11(bc: bare.ByteCursor): Json | null { +function read5(bc: bare.ByteCursor): Json | null { return bare.readBool(bc) ? readJson(bc) : null } -function write11(bc: bare.ByteCursor, x: Json | null): void { +function write5(bc: bare.ByteCursor, x: Json | null): void { bare.writeBool(bc, x != null) if (x != null) { writeJson(bc, x) @@ -741,9 +587,6 @@ export type ToServerInit = { readonly name: string readonly version: u32 readonly totalSlots: u32 - readonly addressesHttp: ReadonlyMap | null - readonly addressesTcp: ReadonlyMap | null - readonly addressesUdp: ReadonlyMap | null readonly lastCommandIdx: i64 | null readonly prepopulateActorNames: ReadonlyMap | null readonly metadata: Json | null @@ -754,12 +597,9 @@ export function readToServerInit(bc: bare.ByteCursor): ToServerInit { name: bare.readString(bc), version: bare.readU32(bc), totalSlots: bare.readU32(bc), - addressesHttp: read4(bc), - addressesTcp: read6(bc), - addressesUdp: read8(bc), lastCommandIdx: read1(bc), - prepopulateActorNames: read10(bc), - metadata: read11(bc), + prepopulateActorNames: read4(bc), + metadata: read5(bc), } } @@ -767,12 +607,9 @@ export function writeToServerInit(bc: bare.ByteCursor, x: ToServerInit): void { bare.writeString(bc, x.name) bare.writeU32(bc, x.version) bare.writeU32(bc, x.totalSlots) - write4(bc, x.addressesHttp) - write6(bc, x.addressesTcp) - write8(bc, x.addressesUdp) write1(bc, x.lastCommandIdx) - write10(bc, x.prepopulateActorNames) - write11(bc, x.metadata) + write4(bc, x.prepopulateActorNames) + write5(bc, x.metadata) } export type ToServerEvents = readonly EventWrapper[] @@ -826,7 +663,7 @@ export function writeToServerPing(bc: bare.ByteCursor, x: ToServerPing): void { bare.writeI64(bc, x.ts) } -function read12(bc: bare.ByteCursor): readonly KvKey[] { +function read6(bc: bare.ByteCursor): readonly KvKey[] { const len = bare.readUintSafe(bc) if (len === 0) { return [] @@ -838,7 +675,7 @@ function read12(bc: bare.ByteCursor): readonly KvKey[] { return result } -function write12(bc: bare.ByteCursor, x: readonly KvKey[]): void { +function write6(bc: bare.ByteCursor, x: readonly KvKey[]): void { bare.writeUintSafe(bc, x.length) for (let i = 0; i < x.length; i++) { writeKvKey(bc, x[i]) @@ -851,30 +688,30 @@ export type KvGetRequest = { export function readKvGetRequest(bc: bare.ByteCursor): KvGetRequest { return { - keys: read12(bc), + keys: read6(bc), } } export function writeKvGetRequest(bc: bare.ByteCursor, x: KvGetRequest): void { - write12(bc, x.keys) + write6(bc, x.keys) } -function read13(bc: bare.ByteCursor): boolean | null { +function read7(bc: bare.ByteCursor): boolean | null { return bare.readBool(bc) ? bare.readBool(bc) : null } -function write13(bc: bare.ByteCursor, x: boolean | null): void { +function write7(bc: bare.ByteCursor, x: boolean | null): void { bare.writeBool(bc, x != null) if (x != null) { bare.writeBool(bc, x) } } -function read14(bc: bare.ByteCursor): u64 | null { +function read8(bc: bare.ByteCursor): u64 | null { return bare.readBool(bc) ? bare.readU64(bc) : null } -function write14(bc: bare.ByteCursor, x: u64 | null): void { +function write8(bc: bare.ByteCursor, x: u64 | null): void { bare.writeBool(bc, x != null) if (x != null) { bare.writeU64(bc, x) @@ -890,18 +727,18 @@ export type KvListRequest = { export function readKvListRequest(bc: bare.ByteCursor): KvListRequest { return { query: readKvListQuery(bc), - reverse: read13(bc), - limit: read14(bc), + reverse: read7(bc), + limit: read8(bc), } } export function writeKvListRequest(bc: bare.ByteCursor, x: KvListRequest): void { writeKvListQuery(bc, x.query) - write13(bc, x.reverse) - write14(bc, x.limit) + write7(bc, x.reverse) + write8(bc, x.limit) } -function read15(bc: bare.ByteCursor): readonly KvValue[] { +function read9(bc: bare.ByteCursor): readonly KvValue[] { const len = bare.readUintSafe(bc) if (len === 0) { return [] @@ -913,7 +750,7 @@ function read15(bc: bare.ByteCursor): readonly KvValue[] { return result } -function write15(bc: bare.ByteCursor, x: readonly KvValue[]): void { +function write9(bc: bare.ByteCursor, x: readonly KvValue[]): void { bare.writeUintSafe(bc, x.length) for (let i = 0; i < x.length; i++) { writeKvValue(bc, x[i]) @@ -927,14 +764,14 @@ export type KvPutRequest = { export function readKvPutRequest(bc: bare.ByteCursor): KvPutRequest { return { - keys: read12(bc), - values: read15(bc), + keys: read6(bc), + values: read9(bc), } } export function writeKvPutRequest(bc: bare.ByteCursor, x: KvPutRequest): void { - write12(bc, x.keys) - write15(bc, x.values) + write6(bc, x.keys) + write9(bc, x.values) } export type KvDeleteRequest = { @@ -943,12 +780,12 @@ export type KvDeleteRequest = { export function readKvDeleteRequest(bc: bare.ByteCursor): KvDeleteRequest { return { - keys: read12(bc), + keys: read6(bc), } } export function writeKvDeleteRequest(bc: bare.ByteCursor, x: KvDeleteRequest): void { - write12(bc, x.keys) + write6(bc, x.keys) } export type KvDropRequest = null @@ -1197,7 +1034,7 @@ export function writeKvErrorResponse(bc: bare.ByteCursor, x: KvErrorResponse): v bare.writeString(bc, x.message) } -function read16(bc: bare.ByteCursor): readonly KvMetadata[] { +function read10(bc: bare.ByteCursor): readonly KvMetadata[] { const len = bare.readUintSafe(bc) if (len === 0) { return [] @@ -1209,7 +1046,7 @@ function read16(bc: bare.ByteCursor): readonly KvMetadata[] { return result } -function write16(bc: bare.ByteCursor, x: readonly KvMetadata[]): void { +function write10(bc: bare.ByteCursor, x: readonly KvMetadata[]): void { bare.writeUintSafe(bc, x.length) for (let i = 0; i < x.length; i++) { writeKvMetadata(bc, x[i]) @@ -1224,16 +1061,16 @@ export type KvGetResponse = { export function readKvGetResponse(bc: bare.ByteCursor): KvGetResponse { return { - keys: read12(bc), - values: read15(bc), - metadata: read16(bc), + keys: read6(bc), + values: read9(bc), + metadata: read10(bc), } } export function writeKvGetResponse(bc: bare.ByteCursor, x: KvGetResponse): void { - write12(bc, x.keys) - write15(bc, x.values) - write16(bc, x.metadata) + write6(bc, x.keys) + write9(bc, x.values) + write10(bc, x.metadata) } export type KvListResponse = { @@ -1244,16 +1081,16 @@ export type KvListResponse = { export function readKvListResponse(bc: bare.ByteCursor): KvListResponse { return { - keys: read12(bc), - values: read15(bc), - metadata: read16(bc), + keys: read6(bc), + values: read9(bc), + metadata: read10(bc), } } export function writeKvListResponse(bc: bare.ByteCursor, x: KvListResponse): void { - write12(bc, x.keys) - write15(bc, x.values) - write16(bc, x.metadata) + write6(bc, x.keys) + write9(bc, x.values) + write10(bc, x.metadata) } export type KvPutResponse = null diff --git a/sdks/typescript/runner/benches/actor-lifecycle.bench.ts b/sdks/typescript/runner/benches/actor-lifecycle.bench.ts index e9a73a8f4a..e4853feb0d 100644 --- a/sdks/typescript/runner/benches/actor-lifecycle.bench.ts +++ b/sdks/typescript/runner/benches/actor-lifecycle.bench.ts @@ -93,7 +93,6 @@ // headers: { // "x-rivet-target": "actor", // "x-rivet-actor": actorId, -// "x-rivet-addr": "main", // }, // }); // if (!pingResponse.ok) throw "Request failed"; @@ -110,7 +109,6 @@ // // headers: { // // "x-rivet-target": "actor", // // "x-rivet-actor": wakeActorId, -// // "x-rivet-addr": "main", // // }, // // }); // // diff --git a/sdks/typescript/runner/src/mod.ts b/sdks/typescript/runner/src/mod.ts index 611f2bbcf0..c7a7ce1a8c 100644 --- a/sdks/typescript/runner/src/mod.ts +++ b/sdks/typescript/runner/src/mod.ts @@ -450,9 +450,6 @@ export class Runner { name: this.#config.runnerName, version: this.#config.version, totalSlots: this.#config.totalSlots, - addressesHttp: new Map(), // No addresses needed with tunnel - addressesTcp: null, - addressesUdp: null, lastCommandIdx: this.#lastCommandIdx >= 0 ? BigInt(this.#lastCommandIdx) diff --git a/sdks/typescript/runner/tests/lifecycle.test.ts b/sdks/typescript/runner/tests/lifecycle.test.ts index 347c8d6a49..c84976890b 100644 --- a/sdks/typescript/runner/tests/lifecycle.test.ts +++ b/sdks/typescript/runner/tests/lifecycle.test.ts @@ -242,7 +242,6 @@ // headers: { // "x-rivet-target": "actor", // "x-rivet-actor": actorId, -// "x-rivet-addr": "main", // }, // }); // expect(actorPingResponse.ok).toBe(true); @@ -255,7 +254,6 @@ // headers: { // "x-rivet-target": "actor", // "x-rivet-actor": actorId, -// "x-rivet-addr": "main", // }, // }); // @@ -305,7 +303,6 @@ // headers: { // "x-rivet-target": "actor", // "x-rivet-actor": actorId, -// "x-rivet-addr": "main", // }, // }); // console.log(`Wake response status: ${wakeResponse.status}`); @@ -340,7 +337,6 @@ // headers: { // "x-rivet-target": "actor", // "x-rivet-actor": actorId, -// "x-rivet-addr": "main", // }, // }); // @@ -367,7 +363,6 @@ // headers: { // "x-rivet-target": "actor", // "x-rivet-actor": actorId, -// "x-rivet-addr": "main", // }, // }); // @@ -413,7 +408,6 @@ // headers: { // "x-rivet-target": "actor", // "x-rivet-actor": actorId, -// "x-rivet-addr": "main", // }, // }); // expect(destroyedPingResponse.status).toBe(404); @@ -424,7 +418,6 @@ // headers: { // "x-rivet-target": "actor", // "x-rivet-actor": actorId, -// "x-rivet-addr": "main", // }, // }); // diff --git a/tests/load/actor-lifecycle/actor.ts b/tests/load/actor-lifecycle/actor.ts index 7132978d5c..8181df1cfa 100644 --- a/tests/load/actor-lifecycle/actor.ts +++ b/tests/load/actor-lifecycle/actor.ts @@ -28,7 +28,6 @@ export function waitForHealth(url: string, actorId: string): boolean { headers: { "x-rivet-target": "actor", "x-rivet-actor": actorId, - "x-rivet-addr": "main", } }); if (response.status === 200) { diff --git a/tests/load/actor-lifecycle/types.ts b/tests/load/actor-lifecycle/types.ts index 9160e758c7..dafde22396 100644 --- a/tests/load/actor-lifecycle/types.ts +++ b/tests/load/actor-lifecycle/types.ts @@ -10,12 +10,6 @@ export interface Config { export interface Actor { actor_id: string; - addresses_http: { - main: { - hostname: string; - port: number; - }; - }; } export interface CreateActorResponse { From 50dfdaee1ab2835dbedfca7c02100aaeae2714e6 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Tue, 9 Sep 2025 18:44:58 -0700 Subject: [PATCH 12/17] fix(udb): revamp udb api --- Cargo.lock | 30 +- Cargo.toml | 5 +- packages/common/gasoline/core/Cargo.toml | 1 - .../common/gasoline/core/src/ctx/workflow.rs | 57 +- .../common/gasoline/core/src/db/kv/debug.rs | 509 +++---- .../gasoline/core/src/db/kv/keys/history.rs | 163 +- .../gasoline/core/src/db/kv/keys/metric.rs | 2 +- .../gasoline/core/src/db/kv/keys/signal.rs | 6 +- .../gasoline/core/src/db/kv/keys/wake.rs | 2 +- .../core/src/db/kv/keys/worker_instance.rs | 2 +- .../gasoline/core/src/db/kv/keys/workflow.rs | 14 +- .../common/gasoline/core/src/db/kv/mod.rs | 796 ++++------ packages/common/gasoline/core/src/error.rs | 7 +- .../gasoline/core/src/history/location.rs | 4 +- packages/common/gasoline/core/src/worker.rs | 3 +- packages/common/pools/Cargo.toml | 7 +- packages/common/pools/src/db/udb.rs | 15 +- packages/common/types/Cargo.toml | 2 +- .../common/types/src/keys/pegboard/mod.rs | 6 +- packages/common/types/src/keys/pegboard/ns.rs | 2 +- packages/common/udb-util/Cargo.toml | 16 - packages/common/udb-util/src/ext.rs | 322 ---- packages/common/universaldb/Cargo.toml | 9 +- packages/common/universaldb/src/database.rs | 26 +- packages/common/universaldb/src/driver/mod.rs | 56 +- .../src/driver/postgres/database.rs | 100 +- .../src/driver/postgres/transaction.rs | 263 ++-- .../src/driver/postgres/transaction_task.rs | 117 +- .../driver/rocksdb/conflict_range_tracker.rs | 31 +- .../src/driver/rocksdb/database.rs | 83 +- .../src/driver/rocksdb/transaction.rs | 131 +- .../src/driver/rocksdb/transaction_task.rs | 200 +-- packages/common/universaldb/src/error.rs | 32 + packages/common/universaldb/src/future.rs | 129 -- .../universaldb/src/inherited/README.md | 7 - .../common/universaldb/src/inherited/error.rs | 1313 ----------------- .../common/universaldb/src/inherited/mod.rs | 4 - .../keyselector.rs => key_selector.rs} | 10 - packages/common/universaldb/src/lib.rs | 25 +- .../{udb-util => universaldb}/src/metrics.rs | 2 +- .../src/{inherited => }/options.rs | 157 -- packages/common/universaldb/src/prelude.rs | 8 + .../rangeoption.rs => range_option.rs} | 6 +- .../common/universaldb/src/transaction.rs | 331 ++++- packages/common/universaldb/src/tx_ops.rs | 51 +- packages/common/universaldb/src/types.rs | 49 - packages/common/universaldb/src/utils.rs | 36 - .../universaldb/src/utils/cherry_pick.rs | 88 ++ .../src => universaldb/src/utils}/codes.rs | 0 packages/common/universaldb/src/utils/ext.rs | 58 + .../src/utils}/formal_key.rs | 12 +- .../src => universaldb/src/utils}/keys.rs | 0 .../lib.rs => universaldb/src/utils/mod.rs} | 36 +- .../src => universaldb/src/utils}/subspace.rs | 19 +- packages/common/universaldb/src/value.rs | 159 ++ .../common/universaldb/tests/integration.rs | 641 ++++---- .../universaldb/tests/integration_gas.rs | 32 +- packages/common/util/id/Cargo.toml | 4 +- packages/common/util/id/src/lib.rs | 12 +- packages/core/actor-kv/Cargo.toml | 3 +- packages/core/actor-kv/src/entry.rs | 2 +- packages/core/actor-kv/src/key.rs | 10 +- packages/core/actor-kv/src/lib.rs | 111 +- packages/core/guard/server/Cargo.toml | 9 +- .../server/src/routing/pegboard_gateway.rs | 23 +- .../core/guard/server/src/shared_state.rs | 1 - packages/core/pegboard-gateway/Cargo.toml | 2 +- packages/core/pegboard-serverless/Cargo.toml | 1 - packages/core/pegboard-serverless/src/lib.rs | 18 +- packages/infra/engine/Cargo.toml | 1 - packages/infra/engine/src/commands/udb/cli.rs | 49 +- packages/infra/engine/src/util/udb.rs | 8 +- packages/services/epoxy/Cargo.toml | 1 - packages/services/epoxy/src/keys/keys.rs | 2 +- packages/services/epoxy/src/keys/mod.rs | 6 +- packages/services/epoxy/src/keys/replica.rs | 4 +- .../epoxy/src/ops/explicit_prepare.rs | 13 +- .../services/epoxy/src/ops/kv/get_local.rs | 7 +- .../epoxy/src/ops/kv/get_optimistic.rs | 14 +- packages/services/epoxy/src/ops/propose.rs | 36 +- .../epoxy/src/ops/read_cluster_config.rs | 8 +- packages/services/epoxy/src/replica/ballot.rs | 31 +- .../services/epoxy/src/replica/commit_kv.rs | 18 +- .../services/epoxy/src/replica/decide_path.rs | 5 +- .../epoxy/src/replica/lead_consensus.rs | 20 +- packages/services/epoxy/src/replica/log.rs | 39 +- .../epoxy/src/replica/message_request.rs | 48 +- .../epoxy/src/replica/messages/accept.rs | 11 +- .../epoxy/src/replica/messages/accepted.rs | 5 +- .../epoxy/src/replica/messages/commit.rs | 5 +- .../epoxy/src/replica/messages/committed.rs | 5 +- .../replica/messages/download_instances.rs | 23 +- .../epoxy/src/replica/messages/pre_accept.rs | 11 +- .../epoxy/src/replica/messages/prepare.rs | 21 +- .../epoxy/src/replica/update_config.rs | 11 +- packages/services/epoxy/src/replica/utils.rs | 21 +- packages/services/epoxy/src/utils.rs | 6 +- .../epoxy/src/workflows/replica/setup.rs | 66 +- packages/services/epoxy/tests/reconfigure.rs | 36 +- packages/services/namespace/Cargo.toml | 1 - packages/services/namespace/src/keys.rs | 6 +- .../services/namespace/src/ops/get_local.rs | 25 +- packages/services/namespace/src/ops/list.rs | 12 +- .../src/ops/resolve_for_name_local.rs | 10 +- .../namespace/src/ops/runner_config/delete.rs | 12 +- .../src/ops/runner_config/get_local.rs | 8 +- .../namespace/src/ops/runner_config/list.rs | 26 +- .../namespace/src/ops/runner_config/upsert.rs | 13 +- .../namespace/src/workflows/namespace.rs | 26 +- packages/services/pegboard/Cargo.toml | 1 - packages/services/pegboard/src/keys/actor.rs | 2 +- .../services/pegboard/src/keys/epoxy/ns.rs | 2 +- packages/services/pegboard/src/keys/mod.rs | 8 +- packages/services/pegboard/src/keys/ns.rs | 2 +- packages/services/pegboard/src/keys/runner.rs | 6 +- .../services/pegboard/src/ops/actor/get.rs | 11 +- .../src/ops/actor/get_reservation_for_key.rs | 2 +- .../pegboard/src/ops/actor/get_runner.rs | 13 +- .../pegboard/src/ops/actor/list_for_ns.rs | 49 +- .../pegboard/src/ops/actor/list_names.rs | 29 +- .../services/pegboard/src/ops/runner/get.rs | 61 +- .../pegboard/src/ops/runner/get_by_key.rs | 15 +- .../pegboard/src/ops/runner/list_for_ns.rs | 61 +- .../pegboard/src/ops/runner/list_names.rs | 30 +- .../src/ops/runner/update_alloc_idx.rs | 38 +- .../src/workflows/actor/actor_keys.rs | 22 +- .../pegboard/src/workflows/actor/destroy.rs | 62 +- .../pegboard/src/workflows/actor/runtime.rs | 95 +- .../pegboard/src/workflows/actor/setup.rs | 26 +- .../services/pegboard/src/workflows/runner.rs | 193 ++- 130 files changed, 2856 insertions(+), 4930 deletions(-) delete mode 100644 packages/common/udb-util/Cargo.toml delete mode 100644 packages/common/udb-util/src/ext.rs create mode 100644 packages/common/universaldb/src/error.rs delete mode 100644 packages/common/universaldb/src/future.rs delete mode 100644 packages/common/universaldb/src/inherited/README.md delete mode 100644 packages/common/universaldb/src/inherited/error.rs delete mode 100644 packages/common/universaldb/src/inherited/mod.rs rename packages/common/universaldb/src/{inherited/keyselector.rs => key_selector.rs} (82%) rename packages/common/{udb-util => universaldb}/src/metrics.rs (94%) rename packages/common/universaldb/src/{inherited => }/options.rs (77%) create mode 100644 packages/common/universaldb/src/prelude.rs rename packages/common/universaldb/src/{inherited/rangeoption.rs => range_option.rs} (96%) delete mode 100644 packages/common/universaldb/src/types.rs delete mode 100644 packages/common/universaldb/src/utils.rs create mode 100644 packages/common/universaldb/src/utils/cherry_pick.rs rename packages/common/{udb-util/src => universaldb/src/utils}/codes.rs (100%) create mode 100644 packages/common/universaldb/src/utils/ext.rs rename packages/common/{udb-util/src => universaldb/src/utils}/formal_key.rs (54%) rename packages/common/{udb-util/src => universaldb/src/utils}/keys.rs (100%) rename packages/common/{udb-util/src/lib.rs => universaldb/src/utils/mod.rs} (69%) rename packages/common/{udb-util/src => universaldb/src/utils}/subspace.rs (88%) create mode 100644 packages/common/universaldb/src/value.rs diff --git a/Cargo.lock b/Cargo.lock index 770f168860..cd762d64fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1413,7 +1413,6 @@ dependencies = [ "tokio", "tracing", "tracing-slog", - "udb-util", "universaldb", "url", "uuid", @@ -1723,7 +1722,6 @@ dependencies = [ "tracing-logfmt", "tracing-opentelemetry", "tracing-subscriber", - "udb-util", "universaldb", "universalpubsub", "url", @@ -2792,7 +2790,6 @@ dependencies = [ "serde", "strum", "tracing", - "udb-util", "universaldb", "url", "utoipa", @@ -3280,7 +3277,6 @@ dependencies = [ "serde_json", "strum", "tracing", - "udb-util", "universaldb", "utoipa", "versioned-data-util", @@ -3301,7 +3297,6 @@ dependencies = [ "tracing", "tracing-logfmt", "tracing-subscriber", - "udb-util", "universaldb", ] @@ -3367,7 +3362,6 @@ dependencies = [ "rivet-runner-protocol", "rivet-types", "tracing", - "udb-util", "universaldb", ] @@ -4373,7 +4367,6 @@ dependencies = [ "tokio-tungstenite", "tracing", "tracing-subscriber", - "udb-util", "universaldb", "url", "uuid", @@ -4450,7 +4443,6 @@ dependencies = [ "tokio", "tower 0.5.2", "tracing", - "udb-util", "universaldb", "universalpubsub", "url", @@ -4560,7 +4552,6 @@ dependencies = [ "tracing", "tracing-logfmt", "tracing-subscriber", - "udb-util", "universaldb", "universalpubsub", "url", @@ -4710,7 +4701,7 @@ dependencies = [ "rivet-runner-protocol", "rivet-util", "serde", - "udb-util", + "universaldb", "utoipa", "versioned-data-util", ] @@ -4767,7 +4758,7 @@ version = "0.0.1" dependencies = [ "serde", "thiserror 1.0.69", - "udb-util", + "universaldb", "utoipa", "uuid", ] @@ -6335,20 +6326,6 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" -[[package]] -name = "udb-util" -version = "0.0.1" -dependencies = [ - "anyhow", - "async-trait", - "futures-util", - "lazy_static", - "rivet-metrics", - "tokio", - "tracing", - "universaldb", -] - [[package]] name = "uname" version = "0.1.1" @@ -6416,16 +6393,17 @@ dependencies = [ "rand 0.8.5", "rivet-config", "rivet-env", + "rivet-metrics", "rivet-pools", "rivet-test-deps-docker", "rocksdb", "serde", "tempfile", + "thiserror 1.0.69", "tokio", "tokio-postgres", "tracing", "tracing-subscriber", - "udb-util", "uuid", ] diff --git a/Cargo.toml b/Cargo.toml index 4ed65e5644..fd561bc4f5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [workspace] resolver = "2" -members = ["packages/common/api-builder","packages/common/api-client","packages/common/api-types","packages/common/api-util","packages/common/cache/build","packages/common/cache/result","packages/common/clickhouse-inserter","packages/common/clickhouse-user-query","packages/common/config","packages/common/env","packages/common/error/core","packages/common/error/macros","packages/common/gasoline/core","packages/common/gasoline/macros","packages/common/logs","packages/common/metrics","packages/common/pools","packages/common/runtime","packages/common/service-manager","packages/common/telemetry","packages/common/test-deps","packages/common/test-deps-docker","packages/common/types","packages/common/udb-util","packages/common/universaldb","packages/common/universalpubsub","packages/common/util/core","packages/common/util/id","packages/common/versioned-data-util","packages/core/actor-kv","packages/core/api-peer","packages/core/api-public","packages/core/bootstrap","packages/core/dump-openapi","packages/core/guard/core","packages/core/guard/server","packages/core/pegboard-gateway","packages/core/pegboard-runner-ws","packages/core/pegboard-serverless","packages/core/pegboard-tunnel","packages/core/workflow-worker","packages/infra/engine","packages/services/epoxy","packages/services/internal","packages/services/namespace","packages/services/pegboard","sdks/rust/api-full","sdks/rust/bare_gen","sdks/rust/data","sdks/rust/epoxy-protocol","sdks/rust/runner-protocol","sdks/rust/tunnel-protocol","sdks/rust/ups-protocol"] +members = ["packages/common/api-builder","packages/common/api-client","packages/common/api-types","packages/common/api-util","packages/common/cache/build","packages/common/cache/result","packages/common/clickhouse-inserter","packages/common/clickhouse-user-query","packages/common/config","packages/common/env","packages/common/error/core","packages/common/error/macros","packages/common/gasoline/core","packages/common/gasoline/macros","packages/common/logs","packages/common/metrics","packages/common/pools","packages/common/runtime","packages/common/service-manager","packages/common/telemetry","packages/common/test-deps","packages/common/test-deps-docker","packages/common/types","packages/common/universaldb","packages/common/universalpubsub","packages/common/util/core","packages/common/util/id","packages/common/versioned-data-util","packages/core/actor-kv","packages/core/api-peer","packages/core/api-public","packages/core/bootstrap","packages/core/dump-openapi","packages/core/guard/core","packages/core/guard/server","packages/core/pegboard-gateway","packages/core/pegboard-runner-ws","packages/core/pegboard-serverless","packages/core/pegboard-tunnel","packages/core/workflow-worker","packages/infra/engine","packages/services/epoxy","packages/services/internal","packages/services/namespace","packages/services/pegboard","sdks/rust/api-full","sdks/rust/bare_gen","sdks/rust/data","sdks/rust/epoxy-protocol","sdks/rust/runner-protocol","sdks/rust/tunnel-protocol","sdks/rust/ups-protocol"] [workspace.package] version = "0.0.1" @@ -318,9 +318,6 @@ path = "packages/common/test-deps-docker" [workspace.dependencies.rivet-types] path = "packages/common/types" -[workspace.dependencies.udb-util] -path = "packages/common/udb-util" - [workspace.dependencies.universaldb] path = "packages/common/universaldb" diff --git a/packages/common/gasoline/core/Cargo.toml b/packages/common/gasoline/core/Cargo.toml index 83223ecd37..ca0e51ae96 100644 --- a/packages/common/gasoline/core/Cargo.toml +++ b/packages/common/gasoline/core/Cargo.toml @@ -39,7 +39,6 @@ tokio-util.workspace = true tokio.workspace = true tracing-logfmt.workspace = true tracing-opentelemetry.workspace = true -udb-util.workspace = true universaldb.workspace = true tracing-subscriber = { workspace = true, features = ["env-filter"] } tracing.workspace = true diff --git a/packages/common/gasoline/core/src/ctx/workflow.rs b/packages/common/gasoline/core/src/ctx/workflow.rs index a8432f406a..7e9d0fc34e 100644 --- a/packages/common/gasoline/core/src/ctx/workflow.rs +++ b/packages/common/gasoline/core/src/ctx/workflow.rs @@ -671,34 +671,35 @@ impl WorkflowCtx { exec.execute(self).await } - /// Tests if the given error is unrecoverable. If it is, allows the user to run recovery code safely. - /// Should always be used when trying to handle activity errors manually. - #[tracing::instrument(skip_all)] - pub fn catch_unrecoverable(&mut self, res: Result) -> Result> { - match res { - Err(err) => { - // TODO: This should check .chain() for the error - match err.downcast::() { - Ok(inner_err) => { - // Despite "history diverged" errors being unrecoverable, they should not have be returned - // by this function because the state of the history is already messed up and no new - // workflow items should be run. - if !inner_err.is_recoverable() - && !matches!(inner_err, WorkflowError::HistoryDiverged(_)) - { - self.cursor.inc(); - - Ok(Err(inner_err.into())) - } else { - Err(inner_err.into()) - } - } - Err(err) => Err(err), - } - } - Ok(x) => Ok(Ok(x)), - } - } + // TODO: Replace with some method on WorkflowError + // /// Tests if the given error is unrecoverable. If it is, allows the user to run recovery code safely. + // /// Should always be used when trying to handle activity errors manually. + // #[tracing::instrument(skip_all)] + // pub fn catch_unrecoverable(&mut self, res: Result) -> Result> { + // match res { + // Err(err) => { + // // TODO: This should check .chain() for the error + // match err.downcast::() { + // Ok(inner_err) => { + // // Despite "history diverged" errors being unrecoverable, they should not have be returned + // // by this function because the state of the history is already messed up and no new + // // workflow items should be run. + // if !inner_err.is_recoverable() + // && !matches!(inner_err, WorkflowError::HistoryDiverged(_)) + // { + // self.cursor.inc(); + + // Ok(Err(inner_err.into())) + // } else { + // Err(inner_err.into()) + // } + // } + // Err(err) => Err(err), + // } + // } + // Ok(x) => Ok(Ok(x)), + // } + // } /// Creates a signal builder. pub fn signal(&mut self, body: T) -> builder::signal::SignalBuilder { diff --git a/packages/common/gasoline/core/src/db/kv/debug.rs b/packages/common/gasoline/core/src/db/kv/debug.rs index 7a556c1b58..50ffc08b0b 100644 --- a/packages/common/gasoline/core/src/db/kv/debug.rs +++ b/packages/common/gasoline/core/src/db/kv/debug.rs @@ -4,16 +4,16 @@ use std::{ result::Result::{Err, Ok}, }; -use anyhow::*; +use anyhow::{Context, Result, ensure}; use futures_util::{StreamExt, TryStreamExt}; use rivet_util::Id; use tracing::Instrument; -use udb_util::{FormalChunkedKey, FormalKey, SERIALIZABLE, SNAPSHOT, TxnExt, end_of_key_range}; +use universaldb::utils::{FormalChunkedKey, FormalKey, IsolationLevel::*, end_of_key_range}; use universaldb::{ - self as udb, - future::FdbValue, + RangeOption, options::{ConflictRangeType, StreamingMode}, tuple::{PackResult, TupleDepth, TupleUnpack}, + value::Value, }; use super::{DatabaseKv, keys, update_metric}; @@ -35,8 +35,8 @@ impl DatabaseKv { async fn get_workflows_inner( &self, workflow_ids: Vec, - tx: &udb::RetryableTransaction, - ) -> std::result::Result, udb::FdbBindingError> { + tx: &universaldb::RetryableTransaction, + ) -> Result> { let mut res = Vec::new(); // TODO: Parallelize @@ -70,91 +70,49 @@ impl DatabaseKv { silence_ts_entry, ) = tokio::try_join!( tx.get_ranges_keyvalues( - udb::RangeOption { + RangeOption { mode: StreamingMode::WantAll, ..(&tags_subspace).into() }, - SNAPSHOT, + Snapshot, ) - .map(|res| match res { - Ok(entry) => { - let key = self - .subspace - .unpack::(entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; - let v = serde_json::Value::String(key.v.clone()); + .map(|res| { + let key = self.subspace.unpack::(res?.key())?; + let v = serde_json::Value::String(key.v.clone()); - Ok((key.k, v)) - } - Err(err) => Err(Into::::into(err)), + Ok((key.k, v)) }) .try_collect::>(), - async { - tx.get(&self.subspace.pack(&name_key), SNAPSHOT) - .await - .map_err(Into::into) - }, - async { - tx.get(&self.subspace.pack(&create_ts_key), SNAPSHOT) - .await - .map_err(Into::into) - }, - async { - tx.get_ranges_keyvalues( - udb::RangeOption { - mode: StreamingMode::WantAll, - ..(&input_subspace).into() - }, - SNAPSHOT, - ) - .try_collect::>() - .await - .map_err(Into::into) - }, - async { - tx.get_ranges_keyvalues( - udb::RangeOption { - mode: StreamingMode::WantAll, - ..(&state_subspace).into() - }, - SNAPSHOT, - ) - .try_collect::>() - .await - .map_err(Into::into) - }, - async { - tx.get_ranges_keyvalues( - udb::RangeOption { - mode: StreamingMode::WantAll, - ..(&output_subspace).into() - }, - SNAPSHOT, - ) - .try_collect::>() - .await - .map_err(Into::into) - }, - async { - tx.get(&self.subspace.pack(&error_key), SNAPSHOT) - .await - .map_err(Into::into) - }, - async { - tx.get(&self.subspace.pack(&has_wake_condition_key), SNAPSHOT) - .await - .map_err(Into::into) - }, - async { - tx.get(&self.subspace.pack(&worker_instance_id_key), SNAPSHOT) - .await - .map_err(Into::into) - }, - async { - tx.get(&self.subspace.pack(&silence_ts_key), SNAPSHOT) - .await - .map_err(Into::into) - }, + tx.get(&self.subspace.pack(&name_key), Snapshot), + tx.get(&self.subspace.pack(&create_ts_key), Snapshot), + tx.get_ranges_keyvalues( + RangeOption { + mode: StreamingMode::WantAll, + ..(&input_subspace).into() + }, + Snapshot, + ) + .try_collect::>(), + tx.get_ranges_keyvalues( + RangeOption { + mode: StreamingMode::WantAll, + ..(&state_subspace).into() + }, + Snapshot, + ) + .try_collect::>(), + tx.get_ranges_keyvalues( + RangeOption { + mode: StreamingMode::WantAll, + ..(&output_subspace).into() + }, + Snapshot, + ) + .try_collect::>(), + tx.get(&self.subspace.pack(&error_key), Snapshot), + tx.get(&self.subspace.pack(&has_wake_condition_key), Snapshot), + tx.get(&self.subspace.pack(&worker_instance_id_key), Snapshot), + tx.get(&self.subspace.pack(&silence_ts_key), Snapshot), )?; let Some(create_ts_entry) = &create_ts_entry else { @@ -162,44 +120,26 @@ impl DatabaseKv { continue; }; - let create_ts = create_ts_key - .deserialize(&create_ts_entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let create_ts = create_ts_key.deserialize(&create_ts_entry)?; - let workflow_name = name_key - .deserialize(&name_entry.ok_or(udb::FdbBindingError::CustomError( - format!("key should exist: {name_key:?}").into(), - ))?) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let workflow_name = name_key.deserialize(&name_entry.context("key should exist")?)?; - let input = input_key - .combine(input_chunks) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let input = input_key.combine(input_chunks)?; let data = if state_chunks.is_empty() { serde_json::value::RawValue::NULL.to_owned() } else { - state_key - .combine(state_chunks) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? + state_key.combine(state_chunks)? }; let output = if output_chunks.is_empty() { None } else { - Some( - output_key - .combine(output_chunks) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ) + Some(output_key.combine(output_chunks)?) }; let error = if let Some(error_entry) = error_entry { - Some( - error_key - .deserialize(&error_entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ) + Some(error_key.deserialize(&error_entry)?) } else { None }; @@ -221,14 +161,9 @@ impl DatabaseKv { workflow_name, tags: serde_json::Value::Object(tags), create_ts, - input: serde_json::from_str(input.get()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - data: serde_json::from_str(data.get()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - output: output - .map(|x| serde_json::from_str(x.get())) - .transpose() - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + input: serde_json::from_str(input.get())?, + data: serde_json::from_str(data.get())?, + output: output.map(|x| serde_json::from_str(x.get())).transpose()?, error, state, }); @@ -241,8 +176,8 @@ impl DatabaseKv { async fn get_signals_inner( &self, signal_ids: Vec, - tx: &udb::RetryableTransaction, - ) -> std::result::Result, udb::FdbBindingError> { + tx: &universaldb::RetryableTransaction, + ) -> Result> { let mut res = Vec::new(); // TODO: Parallelize @@ -263,22 +198,22 @@ impl DatabaseKv { ack_ts_entry, silence_ts_entry, ) = tokio::try_join!( - tx.get(&self.subspace.pack(&name_key), SNAPSHOT), - tx.get(&self.subspace.pack(&workflow_id_key), SNAPSHOT), - tx.get(&self.subspace.pack(&create_ts_key), SNAPSHOT), + tx.get(&self.subspace.pack(&name_key), Snapshot), + tx.get(&self.subspace.pack(&workflow_id_key), Snapshot), + tx.get(&self.subspace.pack(&create_ts_key), Snapshot), async { tx.get_ranges_keyvalues( - udb::RangeOption { + RangeOption { mode: StreamingMode::WantAll, ..(&body_subspace).into() }, - SNAPSHOT, + Snapshot, ) .try_collect::>() .await }, - tx.get(&self.subspace.pack(&ack_ts_key), SNAPSHOT), - tx.get(&self.subspace.pack(&silence_ts_key), SNAPSHOT), + tx.get(&self.subspace.pack(&ack_ts_key), Snapshot), + tx.get(&self.subspace.pack(&silence_ts_key), Snapshot), )?; let Some(create_ts_entry) = &create_ts_entry else { @@ -286,36 +221,20 @@ impl DatabaseKv { continue; }; - let create_ts = create_ts_key - .deserialize(&create_ts_entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let create_ts = create_ts_key.deserialize(&create_ts_entry)?; - let signal_name = name_key - .deserialize(&name_entry.ok_or(udb::FdbBindingError::CustomError( - format!("key should exist: {name_key:?}").into(), - ))?) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let signal_name = name_key.deserialize(&name_entry.context("key should exist")?)?; - let body = body_key - .combine(body_chunks) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let body = body_key.combine(body_chunks)?; let workflow_id = if let Some(workflow_id_entry) = workflow_id_entry { - Some( - workflow_id_key - .deserialize(&workflow_id_entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ) + Some(workflow_id_key.deserialize(&workflow_id_entry)?) } else { None }; let ack_ts = if let Some(ack_ts_entry) = ack_ts_entry { - Some( - ack_ts_key - .deserialize(&ack_ts_entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ) + Some(ack_ts_key.deserialize(&ack_ts_entry)?) } else { None }; @@ -335,8 +254,7 @@ impl DatabaseKv { workflow_id, create_ts, ack_ts, - body: serde_json::from_str(body.get()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + body: serde_json::from_str(body.get())?, state, }); } @@ -345,7 +263,7 @@ impl DatabaseKv { } } -// NOTE: Most of the reads here are SNAPSHOT because we don't want this to conflict with the actual wf engine. +// NOTE: Most of the reads here are Snapshot because we don't want this to conflict with the actual wf engine. // Its just for debugging #[async_trait::async_trait] impl DatabaseDebug for DatabaseKv { @@ -353,7 +271,7 @@ impl DatabaseDebug for DatabaseKv { async fn get_workflows(&self, workflow_ids: Vec) -> Result> { self.pools .udb()? - .run(|tx, _mc| { + .run(|tx| { let workflow_ids = workflow_ids.clone(); async move { self.get_workflows_inner(workflow_ids, &tx).await } }) @@ -371,7 +289,7 @@ impl DatabaseDebug for DatabaseKv { // NOTE: this does a full scan of all keys under workflow/data and filters in memory self.pools .udb()? - .run(|tx, _mc| { + .run(|tx| { let name = name.clone(); async move { let mut workflow_ids = Vec::new(); @@ -381,11 +299,11 @@ impl DatabaseDebug for DatabaseKv { .subspace(&keys::workflow::DataSubspaceKey::new()); let mut stream = tx.get_ranges_keyvalues( - udb::RangeOption { + RangeOption { mode: StreamingMode::Iterator, ..(&data_subspace).into() }, - SNAPSHOT, + Snapshot, ); let mut current_workflow_id = None; @@ -394,10 +312,7 @@ impl DatabaseDebug for DatabaseKv { let mut state_matches = state.is_none() || state == Some(WorkflowState::Dead); while let Some(entry) = stream.try_next().await? { - let workflow_id = *self - .subspace - .unpack::(entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let workflow_id = *self.subspace.unpack::(entry.key())?; if let Some(curr) = current_workflow_id { if workflow_id != curr { @@ -431,9 +346,7 @@ impl DatabaseDebug for DatabaseKv { self.subspace.unpack::(entry.key()) { if let Some(name) = &name { - let workflow_name = name_key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let workflow_name = name_key.deserialize(entry.value())?; name_matches = &workflow_name == name; } @@ -499,7 +412,7 @@ impl DatabaseDebug for DatabaseKv { async fn silence_workflows(&self, workflow_ids: Vec) -> Result<()> { self.pools .udb()? - .run(|tx, _mc| { + .run(|tx| { let workflow_ids = workflow_ids.clone(); async move { @@ -524,15 +437,13 @@ impl DatabaseDebug for DatabaseKv { let error_key = keys::workflow::ErrorKey::new(workflow_id); let Some(name_entry) = - tx.get(&self.subspace.pack(&name_key), SERIALIZABLE).await? + tx.get(&self.subspace.pack(&name_key), Serializable).await? else { tracing::warn!(?workflow_id, "workflow not found"); continue; }; - let workflow_name = name_key - .deserialize(&name_entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let workflow_name = name_key.deserialize(&name_entry)?; let wake_conditions_subspace = self.subspace.subspace( &keys::wake::WorkflowWakeConditionKey::subspace_without_ts( @@ -553,107 +464,87 @@ impl DatabaseDebug for DatabaseKv { ) = tokio::try_join!( // Read sub workflow wake conditions tx.get_ranges_keyvalues( - udb::RangeOption { + RangeOption { mode: StreamingMode::WantAll, ..(&sub_workflow_wake_subspace).into() }, - SERIALIZABLE, + Serializable, ) - .map(|res| match res { - Ok(entry) => self - .subspace - .unpack::(entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into())), - Err(err) => Err(Into::::into(err)), - }) + .map(|res| self + .subspace + .unpack::(res?.key()) + .map_err(Into::into)) .try_collect::>(), // Read tags tx.get_ranges_keyvalues( - udb::RangeOption { + RangeOption { mode: StreamingMode::WantAll, ..(&tags_subspace).into() }, - SERIALIZABLE, + Serializable, ) - .map(|res| match res { - Ok(entry) => self - .subspace - .unpack::(entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into())), - Err(err) => Err(Into::::into(err)), - }) + .map(|res| self + .subspace + .unpack::(res?.key()) + .map_err(Into::into)) .try_collect::>(), // Read wake conditions tx.get_ranges_keyvalues( - udb::RangeOption { + RangeOption { mode: StreamingMode::WantAll, ..(&wake_conditions_subspace).into() }, - SNAPSHOT, + Snapshot, ) - .map(|res| match res { - Ok(entry) => Ok(( + .map(|res| { + let entry = res?; + + Ok(( entry.key().to_vec(), self.subspace - .unpack::(entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - )), - Err(err) => Err(Into::::into(err)), + .unpack::( + entry.key(), + )?, + )) }) .try_collect::>(), async { - tx.get(&self.subspace.pack(&worker_instance_id_key), SERIALIZABLE) + tx.get(&self.subspace.pack(&worker_instance_id_key), Serializable) .await - .map_err(Into::into) .map(|x| x.is_some()) }, async { tx.get_ranges_keyvalues( - udb::RangeOption { + RangeOption { mode: StreamingMode::WantAll, limit: Some(1), ..(&output_subspace).into() }, - SNAPSHOT, + Snapshot, ) .try_next() .await - .map_err(Into::into) .map(|x| x.is_some()) }, async { - tx.get(&self.subspace.pack(&has_wake_condition_key), SERIALIZABLE) + tx.get(&self.subspace.pack(&has_wake_condition_key), Serializable) .await - .map_err(Into::into) .map(|x| x.is_some()) }, async { - tx.get(&self.subspace.pack(&silence_ts_key), SERIALIZABLE) + tx.get(&self.subspace.pack(&silence_ts_key), Serializable) .await - .map_err(Into::into) .map(|x| x.is_some()) }, - async { - tx.get(&self.subspace.pack(&wake_sub_workflow_key), SERIALIZABLE) - .await - .map_err(Into::into) - }, - async { - tx.get(&self.subspace.pack(&error_key), SERIALIZABLE) - .await - .map_err(Into::into) - }, + tx.get(&self.subspace.pack(&wake_sub_workflow_key), Serializable), + tx.get(&self.subspace.pack(&error_key), Serializable), )?; if is_silenced { continue; } - if is_running { - return Err(udb::FdbBindingError::CustomError( - "cannot silence a running workflow".into(), - )); - } + ensure!(!is_running, "cannot silence a running workflow"); for key in sub_workflow_wake_keys { tracing::warn!( @@ -699,9 +590,7 @@ impl DatabaseDebug for DatabaseKv { // Clear sub workflow secondary idx if let Some(entry) = wake_sub_workflow_entry { - let sub_workflow_id = wake_sub_workflow_key - .deserialize(&entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let sub_workflow_id = wake_sub_workflow_key.deserialize(&entry)?; let sub_workflow_wake_key = keys::wake::SubWorkflowWakeKey::new(sub_workflow_id, workflow_id); @@ -720,29 +609,22 @@ impl DatabaseDebug for DatabaseKv { tx.set( &self.subspace.pack(&silence_ts_key), - &silence_ts_key - .serialize(rivet_util::timestamp::now()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &silence_ts_key.serialize(rivet_util::timestamp::now())?, ); // Clear metric let metric = if has_output { keys::metric::GaugeMetric::WorkflowComplete(workflow_name.clone()) } else if has_wake_condition { - let error = error_key - .deserialize(&error_entry.ok_or( - udb::FdbBindingError::CustomError( - format!("key should exist: {error_key:?}").into(), - ), - )?) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let error = + error_key.deserialize(&error_entry.context("key should exist")?)?; keys::metric::GaugeMetric::WorkflowDead(workflow_name.clone(), error) } else { keys::metric::GaugeMetric::WorkflowSleeping(workflow_name.clone()) }; - update_metric(&tx.subspace(self.subspace.clone()), Some(metric), None); + update_metric(&tx.with_subspace(self.subspace.clone()), Some(metric), None); } Ok(()) @@ -756,10 +638,10 @@ impl DatabaseDebug for DatabaseKv { async fn wake_workflows(&self, workflow_ids: Vec) -> Result<()> { self.pools .udb()? - .run(|tx, _mc| { + .run(|tx| { let workflow_ids = workflow_ids.clone(); async move { - let txs = tx.subspace(self.subspace.clone()); + let tx = tx.with_subspace(self.subspace.clone()); for workflow_id in workflow_ids { let name_key = keys::workflow::NameKey::new(workflow_id); @@ -780,38 +662,34 @@ impl DatabaseDebug for DatabaseKv { has_output, error, ) = tokio::try_join!( - txs.read(&name_key, SERIALIZABLE), - txs.exists(&worker_instance_id_key, SERIALIZABLE), - txs.exists(&has_wake_condition_key, SERIALIZABLE), - txs.exists(&silence_ts_key, SERIALIZABLE), + tx.read(&name_key, Serializable), + tx.exists(&worker_instance_id_key, Serializable), + tx.exists(&has_wake_condition_key, Serializable), + tx.exists(&silence_ts_key, Serializable), async { tx.get_ranges_keyvalues( - udb::RangeOption { + RangeOption { mode: StreamingMode::WantAll, limit: Some(1), ..(&output_subspace).into() }, - SNAPSHOT, + Snapshot, ) .try_next() .await .map_err(Into::into) .map(|x| x.is_some()) }, - txs.read_opt(&error_key, SERIALIZABLE), + tx.read_opt(&error_key, Serializable), )?; if is_running || is_silenced { continue; } - if has_output { - return Err(udb::FdbBindingError::CustomError( - "cannot silence a running workflow".into(), - )); - } + ensure!(!has_output, "cannot wake a completed workflow"); - txs.write( + tx.write( &keys::wake::WorkflowWakeConditionKey::new( workflow_name.clone(), workflow_id, @@ -820,16 +698,14 @@ impl DatabaseDebug for DatabaseKv { (), )?; - txs.write(&has_wake_condition_key, ())?; + tx.write(&has_wake_condition_key, ())?; if !has_wake_condition { update_metric( - &txs, + &tx, Some(keys::metric::GaugeMetric::WorkflowDead( workflow_name.clone(), - error.ok_or(udb::FdbBindingError::CustomError( - format!("key should exist: {error_key:?}").into(), - ))?, + error.context("key should exist")?, )), Some(keys::metric::GaugeMetric::WorkflowSleeping(workflow_name)), ); @@ -855,7 +731,7 @@ impl DatabaseDebug for DatabaseKv { ) -> Result> { self.pools .udb()? - .run(|tx, _mc| { + .run(|tx| { async move { let history_subspace = self.subspace @@ -872,7 +748,6 @@ impl DatabaseDebug for DatabaseKv { async { self.get_workflows(vec![workflow_id]) .await - .map_err(|x| udb::FdbBindingError::CustomError(x.into())) .map(|wfs| wfs.into_iter().next()) }, async { @@ -882,11 +757,11 @@ impl DatabaseDebug for DatabaseKv { WorkflowHistoryEventBuilder::new(Location::empty(), false); let mut stream = tx.get_ranges_keyvalues( - udb::RangeOption { + RangeOption { mode: StreamingMode::WantAll, ..(&history_subspace).into() }, - SERIALIZABLE, + Serializable, ); loop { @@ -897,8 +772,7 @@ impl DatabaseDebug for DatabaseKv { // Parse only the wf id and location of the current key let partial_key = self .subspace - .unpack::(entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + .unpack::(entry.key())?; if current_event.location != partial_key.location { if current_event.location.is_empty() { @@ -919,9 +793,7 @@ impl DatabaseDebug for DatabaseKv { events_by_location .entry(previous_event.location.root()) .or_default() - .push(Event::try_from(previous_event).map_err( - |x| udb::FdbBindingError::CustomError(x.into()), - )?); + .push(Event::try_from(previous_event)?); } } @@ -930,53 +802,41 @@ impl DatabaseDebug for DatabaseKv { .subspace .unpack::(entry.key()) { - let event_type = key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let event_type = key.deserialize(entry.value())?; current_event.event_type = Some(event_type); } else if let Ok(key) = self .subspace .unpack::(entry.key()) { - let version = key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let version = key.deserialize(entry.value())?; current_event.version = Some(version); } else if let Ok(key) = self .subspace .unpack::(entry.key()) { - let create_ts = key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let create_ts = key.deserialize(entry.value())?; current_event.create_ts = Some(create_ts); } else if let Ok(key) = self.subspace.unpack::(entry.key()) { - let name = key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let name = key.deserialize(entry.value())?; current_event.name = Some(name); } else if let Ok(key) = self .subspace .unpack::(entry.key()) { - let signal_id = key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let signal_id = key.deserialize(entry.value())?; current_event.signal_id = Some(signal_id); } else if let Ok(key) = self .subspace .unpack::(entry.key()) { - let sub_workflow_id = key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let sub_workflow_id = key.deserialize(entry.value())?; current_event.sub_workflow_id = Some(sub_workflow_id); } else if let Ok(_key) = self @@ -993,9 +853,7 @@ impl DatabaseDebug for DatabaseKv { .subspace .unpack::(entry.key()) { - let input_hash = key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let input_hash = key.deserialize(entry.value())?; current_event.input_hash = Some(input_hash); } else if let Ok(key) = @@ -1019,36 +877,28 @@ impl DatabaseDebug for DatabaseKv { .subspace .unpack::(entry.key()) { - let iteration = key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let iteration = key.deserialize(entry.value())?; current_event.iteration = Some(iteration); } else if let Ok(key) = self .subspace .unpack::(entry.key()) { - let deadline_ts = key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let deadline_ts = key.deserialize(entry.value())?; current_event.deadline_ts = Some(deadline_ts); } else if let Ok(key) = self .subspace .unpack::(entry.key()) { - let sleep_state = key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let sleep_state = key.deserialize(entry.value())?; current_event.sleep_state = Some(sleep_state); } else if let Ok(key) = self.subspace .unpack::(entry.key()) { - let inner_event_type = key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let inner_event_type = key.deserialize(entry.value())?; current_event.inner_event_type = Some(inner_event_type); } @@ -1060,9 +910,7 @@ impl DatabaseDebug for DatabaseKv { events_by_location .entry(current_event.location.root()) .or_default() - .push(Event::try_from(current_event).map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?); + .push(Event::try_from(current_event)?); } Ok(events_by_location) @@ -1077,7 +925,7 @@ impl DatabaseDebug for DatabaseKv { events.into_iter().flat_map(|(_, v)| v).collect::>(); flat_events.sort_by(|a, b| a.location.cmp(&b.location)); - Result::<_, udb::FdbBindingError>::Ok(Some(HistoryData { + Ok(Some(HistoryData { wf, events: flat_events, })) @@ -1092,7 +940,7 @@ impl DatabaseDebug for DatabaseKv { async fn get_signals(&self, signal_ids: Vec) -> Result> { self.pools .udb()? - .run(|tx, _mc| { + .run(|tx| { let signal_ids = signal_ids.clone(); async move { self.get_signals_inner(signal_ids, &tx).await } }) @@ -1112,7 +960,7 @@ impl DatabaseDebug for DatabaseKv { // NOTE: this does a full scan of all keys under signal/data and filters in memory self.pools .udb()? - .run(|tx, _mc| { + .run(|tx| { let name = name.clone(); let workflow_id = workflow_id.clone(); async move { @@ -1123,11 +971,11 @@ impl DatabaseDebug for DatabaseKv { .subspace(&keys::signal::DataSubspaceKey::new()); let mut stream = tx.get_ranges_keyvalues( - udb::RangeOption { + RangeOption { mode: StreamingMode::Iterator, ..(&data_subspace).into() }, - SNAPSHOT, + Snapshot, ); let mut current_signal_id = None; @@ -1136,10 +984,7 @@ impl DatabaseDebug for DatabaseKv { let mut state_matches = state.is_none() || state == Some(SignalState::Pending); while let Some(entry) = stream.try_next().await? { - let signal_id = *self - .subspace - .unpack::(entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let signal_id = *self.subspace.unpack::(entry.key())?; if let Some(curr) = current_signal_id { if signal_id != curr { @@ -1167,9 +1012,7 @@ impl DatabaseDebug for DatabaseKv { self.subspace.unpack::(entry.key()) { if let Some(name) = &name { - let signal_name = name_key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let signal_name = name_key.deserialize(entry.value())?; name_matches = &signal_name == name; } @@ -1178,9 +1021,8 @@ impl DatabaseDebug for DatabaseKv { .unpack::(entry.key()) { if let Some(workflow_id) = &workflow_id { - let signal_workflow_id = workflow_id_key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let signal_workflow_id = + workflow_id_key.deserialize(entry.value())?; workflow_id_matches = &signal_workflow_id == workflow_id; } @@ -1225,7 +1067,7 @@ impl DatabaseDebug for DatabaseKv { async fn silence_signals(&self, signal_ids: Vec) -> Result<()> { self.pools .udb()? - .run(|tx, _mc| { + .run(|tx| { let signal_ids = signal_ids.clone(); async move { @@ -1244,11 +1086,11 @@ impl DatabaseDebug for DatabaseKv { silence_ts_entry, ack_ts_entry, ) = tokio::try_join!( - tx.get(&self.subspace.pack(&signal_name_key), SERIALIZABLE), - tx.get(&self.subspace.pack(&create_ts_key), SERIALIZABLE), - tx.get(&self.subspace.pack(&workflow_id_key), SERIALIZABLE), - tx.get(&self.subspace.pack(&silence_ts_key), SERIALIZABLE), - tx.get(&self.subspace.pack(&ack_ts_key), SERIALIZABLE), + tx.get(&self.subspace.pack(&signal_name_key), Serializable), + tx.get(&self.subspace.pack(&create_ts_key), Serializable), + tx.get(&self.subspace.pack(&workflow_id_key), Serializable), + tx.get(&self.subspace.pack(&silence_ts_key), Serializable), + tx.get(&self.subspace.pack(&ack_ts_key), Serializable), )?; if silence_ts_entry.is_some() { @@ -1260,39 +1102,22 @@ impl DatabaseDebug for DatabaseKv { continue; }; - let signal_name = signal_name_key - .deserialize(&signal_name_entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let signal_name = signal_name_key.deserialize(&signal_name_entry)?; let create_ts = create_ts_key - .deserialize(&create_ts_entry.ok_or( - udb::FdbBindingError::CustomError( - format!("key should exist: {create_ts_key:?}").into(), - ), - )?) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + .deserialize(&create_ts_entry.context("key should exist")?)?; let workflow_id = workflow_id_key - .deserialize(&workflow_id_entry.ok_or( - udb::FdbBindingError::CustomError( - format!("key should exist: {workflow_id_key:?}").into(), - ), - )?) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + .deserialize(&workflow_id_entry.context("key should exist")?)?; let workflow_name_key = keys::workflow::NameKey::new(workflow_id); let workflow_name_entry = tx - .get(&self.subspace.pack(&workflow_name_key), SERIALIZABLE) + .get(&self.subspace.pack(&workflow_name_key), Serializable) .await?; let workflow_name = workflow_name_key - .deserialize(&workflow_name_entry.ok_or( - udb::FdbBindingError::CustomError( - format!("key should exist: {workflow_name_key:?}").into(), - ), - )?) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + .deserialize(&workflow_name_entry.context("key should exist")?)?; // Clear pending key let mut pending_signal_key = keys::workflow::PendingSignalKey::new( @@ -1314,14 +1139,12 @@ impl DatabaseDebug for DatabaseKv { tx.set( &self.subspace.pack(&silence_ts_key), - &silence_ts_key - .serialize(rivet_util::timestamp::now()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &silence_ts_key.serialize(rivet_util::timestamp::now())?, ); if ack_ts_entry.is_none() { update_metric( - &tx.subspace(self.subspace.clone()), + &tx.with_subspace(self.subspace.clone()), Some(keys::metric::GaugeMetric::SignalPending(signal_name)), None, ); @@ -1365,8 +1188,8 @@ struct WorkflowHistoryEventBuilder { name: Option, signal_id: Option, sub_workflow_id: Option, - input_chunks: Vec, - output_chunks: Vec, + input_chunks: Vec, + output_chunks: Vec, tags: Vec<(String, String)>, input_hash: Option>, errors: Vec, diff --git a/packages/common/gasoline/core/src/db/kv/keys/history.rs b/packages/common/gasoline/core/src/db/kv/keys/history.rs index 9b8a838fd1..e47b201c2b 100644 --- a/packages/common/gasoline/core/src/db/kv/keys/history.rs +++ b/packages/common/gasoline/core/src/db/kv/keys/history.rs @@ -2,7 +2,7 @@ use std::result::Result::Ok; use anyhow::*; use rivet_util::Id; -use udb_util::prelude::*; +use universaldb::prelude::*; use crate::history::{ event::{EventType, SleepState}, @@ -147,7 +147,7 @@ impl TuplePack for EventHistorySubspaceKey { // This ensures we are only reading events under the given location and not event data at the current // location - w.write_all(&[udb_util::codes::NESTED])?; + w.write_all(&[universaldb::utils::codes::NESTED])?; offset += 1; if let Some(idx) = self.idx { @@ -550,7 +550,7 @@ impl InputKey { Ok(value .get() .as_bytes() - .chunks(udb_util::CHUNK_SIZE) + .chunks(universaldb::utils::CHUNK_SIZE) .map(|x| x.to_vec()) .collect()) } @@ -569,7 +569,7 @@ impl FormalChunkedKey for InputKey { } } - fn combine(&self, chunks: Vec) -> Result { + fn combine(&self, chunks: Vec) -> Result { serde_json::value::RawValue::from_string(String::from_utf8( chunks .iter() @@ -666,7 +666,7 @@ impl OutputKey { Ok(value .get() .as_bytes() - .chunks(udb_util::CHUNK_SIZE) + .chunks(universaldb::utils::CHUNK_SIZE) .map(|x| x.to_vec()) .collect()) } @@ -685,7 +685,7 @@ impl FormalChunkedKey for OutputKey { } } - fn combine(&self, chunks: Vec) -> Result { + fn combine(&self, chunks: Vec) -> Result { serde_json::value::RawValue::from_string(String::from_utf8( chunks .iter() @@ -1301,8 +1301,7 @@ fn unpack_history_key<'de>( pub mod insert { use anyhow::Result; use rivet_util::Id; - use udb_util::{FormalChunkedKey, FormalKey}; - use universaldb as udb; + use universaldb::utils::{FormalChunkedKey, FormalKey}; use super::super::super::value_to_str; use crate::{ @@ -1314,8 +1313,8 @@ pub mod insert { }; pub fn common( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, event_type: EventType, @@ -1344,8 +1343,8 @@ pub mod insert { } pub fn signal_event( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, version: usize, @@ -1379,12 +1378,7 @@ pub mod insert { let signal_body_key = super::InputKey::new(workflow_id, location.clone()); // Write signal body - for (i, chunk) in signal_body_key - .split_ref(&body) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in signal_body_key.split_ref(&body)?.into_iter().enumerate() { let chunk_key = signal_body_key.chunk(i); tx.set(&subspace.pack(&chunk_key), &chunk); @@ -1394,8 +1388,8 @@ pub mod insert { } pub fn signal_send_event( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, version: usize, @@ -1430,12 +1424,7 @@ pub mod insert { let signal_body_key = super::InputKey::new(workflow_id, location.clone()); // Write signal body - for (i, chunk) in signal_body_key - .split_ref(&body) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in signal_body_key.split_ref(&body)?.into_iter().enumerate() { let chunk_key = signal_body_key.chunk(i); tx.set(&subspace.pack(&chunk_key), &chunk); @@ -1451,8 +1440,8 @@ pub mod insert { } pub fn sub_workflow_event( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, version: usize, @@ -1490,34 +1479,22 @@ pub mod insert { x.as_object() .ok_or_else(|| WorkflowError::InvalidTags("must be an object".to_string())) }) - .transpose() - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? + .transpose()? .into_iter() .flatten() .map(|(k, v)| Ok((k.clone(), value_to_str(v)?))) - .collect::>>() - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + .collect::>>()?; for (k, v) in &tags { // Write tag key let tag_key = super::TagKey::new(workflow_id, location.clone(), k.clone(), v.clone()); - tx.set( - &subspace.pack(&tag_key), - &tag_key - .serialize(()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ); + tx.set(&subspace.pack(&tag_key), &tag_key.serialize(())?); } let input_key = super::InputKey::new(workflow_id, location.clone()); // Write input - for (i, chunk) in input_key - .split_ref(&input) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in input_key.split_ref(&input)?.into_iter().enumerate() { let chunk_key = input_key.chunk(i); tx.set(&subspace.pack(&chunk_key), &chunk); @@ -1527,8 +1504,8 @@ pub mod insert { } pub fn activity_event( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, version: usize, @@ -1563,12 +1540,7 @@ pub mod insert { let input_key = super::InputKey::new(workflow_id, location.clone()); // Write input - for (i, chunk) in input_key - .split_ref(&input) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in input_key.split_ref(&input)?.into_iter().enumerate() { let chunk_key = input_key.chunk(i); tx.set(&subspace.pack(&chunk_key), &chunk); @@ -1579,12 +1551,7 @@ pub mod insert { let output_key = super::OutputKey::new(workflow_id, location.clone()); // Write output - for (i, chunk) in output_key - .split_ref(&output) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in output_key.split_ref(&output)?.into_iter().enumerate() { let chunk_key = output_key.chunk(i); tx.set(&subspace.pack(&chunk_key), &chunk); @@ -1605,8 +1572,8 @@ pub mod insert { } pub fn message_send_event( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, version: usize, @@ -1628,22 +1595,15 @@ pub mod insert { // Write tags let tags = tags .as_object() - .ok_or_else(|| WorkflowError::InvalidTags("must be an object".to_string())) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? + .ok_or_else(|| WorkflowError::InvalidTags("must be an object".to_string()))? .into_iter() .map(|(k, v)| Ok((k.clone(), value_to_str(v)?))) - .collect::>>() - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + .collect::>>()?; for (k, v) in &tags { // Write tag key let tag_key = super::TagKey::new(workflow_id, location.clone(), k.clone(), v.clone()); - tx.set( - &subspace.pack(&tag_key), - &tag_key - .serialize(()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ); + tx.set(&subspace.pack(&tag_key), &tag_key.serialize(())?); } let message_name_key = super::NameKey::new(workflow_id, location.clone()); @@ -1655,12 +1615,7 @@ pub mod insert { let body_key = super::InputKey::new(workflow_id, location.clone()); // Write body - for (i, chunk) in body_key - .split_ref(&body) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in body_key.split_ref(&body)?.into_iter().enumerate() { let chunk_key = body_key.chunk(i); tx.set(&subspace.pack(&chunk_key), &chunk); @@ -1670,8 +1625,8 @@ pub mod insert { } pub fn loop_event( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, version: usize, @@ -1699,12 +1654,7 @@ pub mod insert { let state_key = super::InputKey::new(workflow_id, location.clone()); // Write state - for (i, chunk) in state_key - .split_ref(&state) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in state_key.split_ref(&state)?.into_iter().enumerate() { let chunk_key = state_key.chunk(i); tx.set(&subspace.pack(&chunk_key), &chunk); @@ -1714,12 +1664,7 @@ pub mod insert { let output_key = super::OutputKey::new(workflow_id, location.clone()); // Write output - for (i, chunk) in output_key - .split_ref(&output) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in output_key.split_ref(&output)?.into_iter().enumerate() { let chunk_key = output_key.chunk(i); tx.set(&subspace.pack(&chunk_key), &chunk); @@ -1730,8 +1675,8 @@ pub mod insert { } pub fn update_loop_event( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, iteration: usize, @@ -1747,12 +1692,7 @@ pub mod insert { let state_key = super::InputKey::new(workflow_id, location.clone()); // Write state - for (i, chunk) in state_key - .split_ref(&state) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in state_key.split_ref(&state)?.into_iter().enumerate() { let chunk_key = state_key.chunk(i); tx.set(&subspace.pack(&chunk_key), &chunk); @@ -1762,12 +1702,7 @@ pub mod insert { let output_key = super::OutputKey::new(workflow_id, location.clone()); // Write output - for (i, chunk) in output_key - .split_ref(&output) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in output_key.split_ref(&output)?.into_iter().enumerate() { let chunk_key = output_key.chunk(i); tx.set(&subspace.pack(&chunk_key), &chunk); @@ -1778,8 +1713,8 @@ pub mod insert { } pub fn sleep_event( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, version: usize, @@ -1813,8 +1748,8 @@ pub mod insert { } pub fn update_sleep_event( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, sleep_state: SleepState, @@ -1829,8 +1764,8 @@ pub mod insert { } pub fn branch_event( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, version: usize, @@ -1850,8 +1785,8 @@ pub mod insert { } pub fn removed_event( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, version: usize, @@ -1887,8 +1822,8 @@ pub mod insert { } pub fn version_check_event( - subspace: &udb::tuple::Subspace, - tx: &udb::RetryableTransaction, + subspace: &universaldb::tuple::Subspace, + tx: &universaldb::RetryableTransaction, workflow_id: Id, location: &Location, version: usize, diff --git a/packages/common/gasoline/core/src/db/kv/keys/metric.rs b/packages/common/gasoline/core/src/db/kv/keys/metric.rs index 60697d670c..c20c60d2c7 100644 --- a/packages/common/gasoline/core/src/db/kv/keys/metric.rs +++ b/packages/common/gasoline/core/src/db/kv/keys/metric.rs @@ -1,5 +1,5 @@ use anyhow::*; -use udb_util::prelude::*; +use universaldb::prelude::*; #[derive(Debug, PartialEq, Eq)] pub enum GaugeMetric { diff --git a/packages/common/gasoline/core/src/db/kv/keys/signal.rs b/packages/common/gasoline/core/src/db/kv/keys/signal.rs index 150c34b681..04cb4437c5 100644 --- a/packages/common/gasoline/core/src/db/kv/keys/signal.rs +++ b/packages/common/gasoline/core/src/db/kv/keys/signal.rs @@ -2,7 +2,7 @@ use std::result::Result::Ok; use anyhow::*; use rivet_util::Id; -use udb_util::prelude::*; +use universaldb::prelude::*; pub struct BodyKey { signal_id: Id, @@ -17,7 +17,7 @@ impl BodyKey { Ok(value .get() .as_bytes() - .chunks(udb_util::CHUNK_SIZE) + .chunks(universaldb::utils::CHUNK_SIZE) .map(|x| x.to_vec()) .collect()) } @@ -34,7 +34,7 @@ impl FormalChunkedKey for BodyKey { } } - fn combine(&self, chunks: Vec) -> Result { + fn combine(&self, chunks: Vec) -> Result { serde_json::value::RawValue::from_string(String::from_utf8( chunks .iter() diff --git a/packages/common/gasoline/core/src/db/kv/keys/wake.rs b/packages/common/gasoline/core/src/db/kv/keys/wake.rs index 721f3714a4..9e71b3f43a 100644 --- a/packages/common/gasoline/core/src/db/kv/keys/wake.rs +++ b/packages/common/gasoline/core/src/db/kv/keys/wake.rs @@ -2,7 +2,7 @@ use std::result::Result::Ok; use anyhow::*; use rivet_util::Id; -use udb_util::prelude::*; +use universaldb::prelude::*; #[derive(Debug)] pub enum WakeCondition { diff --git a/packages/common/gasoline/core/src/db/kv/keys/worker_instance.rs b/packages/common/gasoline/core/src/db/kv/keys/worker_instance.rs index 851b1825bf..47b4a1103a 100644 --- a/packages/common/gasoline/core/src/db/kv/keys/worker_instance.rs +++ b/packages/common/gasoline/core/src/db/kv/keys/worker_instance.rs @@ -2,7 +2,7 @@ use std::result::Result::Ok; use anyhow::*; use rivet_util::Id; -use udb_util::prelude::*; +use universaldb::prelude::*; #[derive(Debug)] pub struct LastPingTsKey { diff --git a/packages/common/gasoline/core/src/db/kv/keys/workflow.rs b/packages/common/gasoline/core/src/db/kv/keys/workflow.rs index 4a5e01d8ef..63b166c42e 100644 --- a/packages/common/gasoline/core/src/db/kv/keys/workflow.rs +++ b/packages/common/gasoline/core/src/db/kv/keys/workflow.rs @@ -2,7 +2,7 @@ use std::result::Result::Ok; use anyhow::*; use rivet_util::Id; -use udb_util::prelude::*; +use universaldb::prelude::*; #[derive(Debug)] pub struct LeaseKey { @@ -159,7 +159,7 @@ impl InputKey { Ok(value .get() .as_bytes() - .chunks(udb_util::CHUNK_SIZE) + .chunks(universaldb::utils::CHUNK_SIZE) .map(|x| x.to_vec()) .collect()) } @@ -176,7 +176,7 @@ impl FormalChunkedKey for InputKey { } } - fn combine(&self, chunks: Vec) -> Result { + fn combine(&self, chunks: Vec) -> Result { serde_json::value::RawValue::from_string(String::from_utf8( chunks .iter() @@ -246,7 +246,7 @@ impl OutputKey { Ok(value .get() .as_bytes() - .chunks(udb_util::CHUNK_SIZE) + .chunks(universaldb::utils::CHUNK_SIZE) .map(|x| x.to_vec()) .collect()) } @@ -263,7 +263,7 @@ impl FormalChunkedKey for OutputKey { } } - fn combine(&self, chunks: Vec) -> Result { + fn combine(&self, chunks: Vec) -> Result { serde_json::value::RawValue::from_string(String::from_utf8( chunks .iter() @@ -333,7 +333,7 @@ impl StateKey { Ok(value .get() .as_bytes() - .chunks(udb_util::CHUNK_SIZE) + .chunks(universaldb::utils::CHUNK_SIZE) .map(|x| x.to_vec()) .collect()) } @@ -350,7 +350,7 @@ impl FormalChunkedKey for StateKey { } } - fn combine(&self, chunks: Vec) -> Result { + fn combine(&self, chunks: Vec) -> Result { serde_json::value::RawValue::from_string(String::from_utf8( chunks .iter() diff --git a/packages/common/gasoline/core/src/db/kv/mod.rs b/packages/common/gasoline/core/src/db/kv/mod.rs index 0321d3453c..9c2db70896 100644 --- a/packages/common/gasoline/core/src/db/kv/mod.rs +++ b/packages/common/gasoline/core/src/db/kv/mod.rs @@ -7,18 +7,18 @@ use std::{ time::Instant, }; +use anyhow::{Context, Result}; use futures_util::{StreamExt, TryStreamExt, stream::BoxStream}; use rivet_util::Id; use rivet_util::future::CustomInstrumentExt; use serde_json::json; use tracing::Instrument; -use udb_util::{ - FormalChunkedKey, FormalKey, SERIALIZABLE, SNAPSHOT, TxnExt, end_of_key_range, keys::*, +use universaldb::utils::{ + FormalChunkedKey, FormalKey, IsolationLevel::*, end_of_key_range, keys::*, }; use universaldb::{ - self as udb, - future::FdbValue, options::{ConflictRangeType, MutationType, StreamingMode}, + value::Value, }; use rivet_metrics::KeyValue; @@ -48,7 +48,7 @@ const WORKER_WAKE_SUBJECT: &str = "gasoline.worker.wake"; pub struct DatabaseKv { pools: rivet_pools::Pools, - subspace: udb_util::Subspace, + subspace: universaldb::utils::Subspace, } impl DatabaseKv { @@ -81,23 +81,21 @@ impl DatabaseKv { } } -// MARK: FDB Helpers +// MARK: UDB Helpers impl DatabaseKv { fn write_signal_wake_idxs( &self, workflow_id: Id, wake_signals: &[&str], - tx: &udb::RetryableTransaction, - ) -> Result<(), udb::FdbBindingError> { + tx: &universaldb::Transaction, + ) -> Result<()> { for signal_name in wake_signals { // Write to wake signals list let wake_signal_key = keys::workflow::WakeSignalKey::new(workflow_id, signal_name.to_string()); tx.set( &self.subspace.pack(&wake_signal_key), - &wake_signal_key - .serialize(()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &wake_signal_key.serialize(())?, ); } @@ -109,16 +107,14 @@ impl DatabaseKv { workflow_id: Id, workflow_name: &str, sub_workflow_id: Id, - tx: &udb::RetryableTransaction, - ) -> Result<(), udb::FdbBindingError> { + tx: &universaldb::Transaction, + ) -> Result<()> { let sub_workflow_wake_key = keys::wake::SubWorkflowWakeKey::new(sub_workflow_id, workflow_id); tx.set( &self.subspace.pack(&sub_workflow_wake_key), - &sub_workflow_wake_key - .serialize(workflow_name.to_string()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &sub_workflow_wake_key.serialize(workflow_name.to_string())?, ); Ok(()) @@ -131,8 +127,8 @@ impl DatabaseKv { signal_id: Id, signal_name: &str, body: &serde_json::value::RawValue, - tx: &udb::RetryableTransaction, - ) -> Result<(), udb::FdbBindingError> { + tx: &universaldb::Transaction, + ) -> Result<()> { tracing::debug!( ?ray_id, ?workflow_id, @@ -145,36 +141,25 @@ impl DatabaseKv { // Check if the workflow exists let Some(workflow_name_entry) = tx - .get(&self.subspace.pack(&workflow_name_key), SERIALIZABLE) + .get(&self.subspace.pack(&workflow_name_key), Serializable) .await? else { - return Err(udb::FdbBindingError::CustomError( - WorkflowError::WorkflowNotFound.into(), - )); + return Err(WorkflowError::WorkflowNotFound.into()); }; - let workflow_name = workflow_name_key - .deserialize(&workflow_name_entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let workflow_name = workflow_name_key.deserialize(&workflow_name_entry)?; // Write name let name_key = keys::signal::NameKey::new(signal_id); tx.set( &self.subspace.pack(&name_key), - &name_key - .serialize(signal_name.to_string()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &name_key.serialize(signal_name.to_string())?, ); let signal_body_key = keys::signal::BodyKey::new(signal_id); // Write signal body - for (i, chunk) in signal_body_key - .split_ref(body) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in signal_body_key.split_ref(body)?.into_iter().enumerate() { let chunk_key = signal_body_key.chunk(i); tx.set(&self.subspace.pack(&chunk_key), &chunk); @@ -186,36 +171,28 @@ impl DatabaseKv { tx.set( &self.subspace.pack(&pending_signal_key), - &pending_signal_key - .serialize(()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &pending_signal_key.serialize(())?, ); // Write create ts let create_ts_key = keys::signal::CreateTsKey::new(signal_id); tx.set( &self.subspace.pack(&create_ts_key), - &create_ts_key - .serialize(pending_signal_key.ts) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &create_ts_key.serialize(pending_signal_key.ts)?, ); // Write ray id let ray_id_key = keys::signal::RayIdKey::new(signal_id); tx.set( &self.subspace.pack(&ray_id_key), - &ray_id_key - .serialize(ray_id) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &ray_id_key.serialize(ray_id)?, ); // Write workflow id let workflow_id_key = keys::signal::WorkflowIdKey::new(signal_id); tx.set( &self.subspace.pack(&workflow_id_key), - &workflow_id_key - .serialize(workflow_id) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &workflow_id_key.serialize(workflow_id)?, ); let wake_signal_key = @@ -223,7 +200,7 @@ impl DatabaseKv { // If the workflow currently has a wake signal key for this signal, wake it if tx - .get(&self.subspace.pack(&wake_signal_key), SERIALIZABLE) + .get(&self.subspace.pack(&wake_signal_key), Serializable) .await? .is_some() { @@ -237,14 +214,12 @@ impl DatabaseKv { // Add wake condition for workflow tx.set( &self.subspace.pack(&wake_condition_key), - &wake_condition_key - .serialize(()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &wake_condition_key.serialize(())?, ); } update_metric( - &tx.subspace(self.subspace.clone()), + &tx.with_subspace(self.subspace.clone()), None, Some(keys::metric::GaugeMetric::SignalPending( signal_name.to_string(), @@ -262,15 +237,15 @@ impl DatabaseKv { tags: Option<&serde_json::Value>, input: &serde_json::value::RawValue, unique: bool, - tx: &udb::RetryableTransaction, - ) -> Result { - let txs = tx.subspace(self.subspace.clone()); + tx: &universaldb::Transaction, + ) -> Result { + let tx = tx.with_subspace(self.subspace.clone()); if unique { let empty_tags = json!({}); if let Some(existing_workflow_id) = self - .find_workflow_inner(workflow_name, tags.unwrap_or(&empty_tags), tx) + .find_workflow_inner(workflow_name, tags.unwrap_or(&empty_tags), &tx) .await? { tracing::debug!(?existing_workflow_id, "found existing workflow"); @@ -278,17 +253,17 @@ impl DatabaseKv { } } - txs.write( + tx.write( &keys::workflow::CreateTsKey::new(workflow_id), rivet_util::timestamp::now(), )?; - txs.write( + tx.write( &keys::workflow::NameKey::new(workflow_id), workflow_name.to_string(), )?; - txs.write(&keys::workflow::RayIdKey::new(workflow_id), ray_id)?; + tx.write(&keys::workflow::RayIdKey::new(workflow_id), ray_id)?; // Write tags let tags = tags @@ -296,17 +271,15 @@ impl DatabaseKv { x.as_object() .ok_or_else(|| WorkflowError::InvalidTags("must be an object".to_string())) }) - .transpose() - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? + .transpose()? .into_iter() .flatten() .map(|(k, v)| Ok((k.clone(), value_to_str(v)?))) - .collect::>>() - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + .collect::>>()?; for (k, v) in &tags { // Write tag key - txs.write( + tx.write( &keys::workflow::TagKey::new(workflow_id, k.clone(), v.clone()), (), )?; @@ -318,7 +291,7 @@ impl DatabaseKv { .map(|(k, v)| (k.clone(), v.clone())) .collect(); - txs.write( + tx.write( &keys::workflow::ByNameAndTagKey::new( workflow_name.to_string(), k.clone(), @@ -330,7 +303,7 @@ impl DatabaseKv { } // Write null key for the "by name and first tag" secondary index (all workflows have this) - txs.write( + tx.write( &keys::workflow::ByNameAndTagKey::null(workflow_name.to_string(), workflow_id), tags, )?; @@ -338,19 +311,14 @@ impl DatabaseKv { // Write input let input_key = keys::workflow::InputKey::new(workflow_id); - for (i, chunk) in input_key - .split_ref(input) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in input_key.split_ref(input)?.into_iter().enumerate() { let chunk_key = input_key.chunk(i); - txs.set(&self.subspace.pack(&chunk_key), &chunk); + tx.set(&self.subspace.pack(&chunk_key), &chunk); } // Write immediate wake condition - txs.write( + tx.write( &keys::wake::WorkflowWakeConditionKey::new( workflow_name.to_string(), workflow_id, @@ -359,11 +327,11 @@ impl DatabaseKv { (), )?; - txs.write(&keys::workflow::HasWakeConditionKey::new(workflow_id), ())?; + tx.write(&keys::workflow::HasWakeConditionKey::new(workflow_id), ())?; // Write metric update_metric( - &txs, + &tx, None, Some(keys::metric::GaugeMetric::WorkflowSleeping( workflow_name.to_string(), @@ -377,20 +345,14 @@ impl DatabaseKv { &self, workflow_name: &str, tags: &serde_json::Value, - tx: &udb::RetryableTransaction, - ) -> Result, udb::FdbBindingError> { + tx: &universaldb::Transaction, + ) -> Result> { // Convert to flat vec of strings let mut tag_iter = tags .as_object() - .ok_or_else(|| WorkflowError::InvalidTags("must be an object".to_string())) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? + .ok_or_else(|| WorkflowError::InvalidTags("must be an object".to_string()))? .iter() - .map(|(k, v)| { - Result::<_, udb::FdbBindingError>::Ok(( - k.clone(), - value_to_str(v).map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - )) - }); + .map(|(k, v)| Result::<_>::Ok((k.clone(), value_to_str(v)?))); let first_tag = tag_iter.next().transpose()?; let rest_of_tags = tag_iter.collect::, _>>()?; @@ -412,11 +374,11 @@ impl DatabaseKv { }; let mut stream = tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::Iterator, ..(&workflow_by_name_and_tag_subspace).into() }, - SERIALIZABLE, + Serializable, ); loop { @@ -427,13 +389,10 @@ impl DatabaseKv { // Unpack key let workflow_by_name_and_tag_key = self .subspace - .unpack::(&entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + .unpack::(&entry.key())?; // Deserialize value - let wf_rest_of_tags = workflow_by_name_and_tag_key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let wf_rest_of_tags = workflow_by_name_and_tag_key.deserialize(entry.value())?; // Compute intersection between wf tags and input let tags_match = rest_of_tags.iter().all(|(k, v)| { @@ -459,7 +418,7 @@ impl Database for DatabaseKv { async fn from_pools(pools: rivet_pools::Pools) -> anyhow::Result> { Ok(Arc::new(DatabaseKv { pools, - subspace: udb_util::Subspace::new(&(RIVET, GASOLINE, KV)), + subspace: universaldb::utils::Subspace::new(&(RIVET, GASOLINE, KV)), })) } @@ -496,7 +455,7 @@ impl Database for DatabaseKv { .pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| { + .run(|tx| { async move { let now = rivet_util::timestamp::now(); @@ -510,23 +469,21 @@ impl Database for DatabaseKv { // List all active leases let mut stream = tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&lease_subspace).into() }, - // Not SERIALIZABLE because we don't want this to conflict with other queries which write + // Not Serializable because we don't want this to conflict with other queries which write // leases - SNAPSHOT, + Snapshot, ); while let Some(lease_key_entry) = stream.try_next().await? { let lease_key = self .subspace - .unpack::(lease_key_entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; - let (workflow_name, worker_instance_id) = lease_key - .deserialize(lease_key_entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + .unpack::(lease_key_entry.key())?; + let (workflow_name, worker_instance_id) = + lease_key.deserialize(lease_key_entry.value())?; let last_ping_ts_key = keys::worker_instance::LastPingTsKey::new(worker_instance_id); @@ -539,15 +496,13 @@ impl Database for DatabaseKv { } else if let Some(last_ping_entry) = tx .get( &self.subspace.pack(&last_ping_ts_key), - // Not SERIALIZABLE because we don't want this to conflict - SNAPSHOT, + // Not Serializable because we don't want this to conflict + Snapshot, ) .await? { // Deserialize last ping value - let last_ping_ts = last_ping_ts_key - .deserialize(&last_ping_entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let last_ping_ts = last_ping_ts_key.deserialize(&last_ping_entry)?; // Update cache last_ping_cache.push((worker_instance_id, last_ping_ts)); @@ -566,7 +521,7 @@ impl Database for DatabaseKv { let silence_ts_key = keys::workflow::SilenceTsKey::new(lease_key.workflow_id); if tx - .get(&self.subspace.pack(&silence_ts_key), SERIALIZABLE) + .get(&self.subspace.pack(&silence_ts_key), Serializable) .await? .is_some() { @@ -596,13 +551,11 @@ impl Database for DatabaseKv { ); tx.set( &self.subspace.pack(&wake_condition_key), - &wake_condition_key - .serialize(()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &wake_condition_key.serialize(())?, ); update_metric( - &tx.subspace(self.subspace.clone()), + &tx.with_subspace(self.subspace.clone()), Some(keys::metric::GaugeMetric::WorkflowActive( workflow_name.to_string(), )), @@ -622,7 +575,8 @@ impl Database for DatabaseKv { } }) .custom_instrument(tracing::info_span!("clear_expired_leases_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; if expired_workflow_count != 0 { tracing::info!( @@ -644,13 +598,13 @@ impl Database for DatabaseKv { .pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| { + .run(|tx| { async move { - let txs = tx.subspace(self.subspace.clone()); + let tx = tx.with_subspace(self.subspace.clone()); // Read existing lock - let lock_expired = if let Some(lock_ts) = txs - .read_opt(&keys::worker_instance::MetricsLockKey::new(), SERIALIZABLE) + let lock_expired = if let Some(lock_ts) = tx + .read_opt(&keys::worker_instance::MetricsLockKey::new(), Serializable) .await? { lock_ts < rivet_util::timestamp::now() - METRICS_LOCK_TIMEOUT_MS @@ -659,9 +613,9 @@ impl Database for DatabaseKv { }; if lock_expired { - // Write to lock key. FDB transactions guarantee that if multiple workers are running this + // Write to lock key. UDB transactions guarantee that if multiple workers are running this // query at the same time only one will succeed which means only one will have the lock. - txs.write( + tx.write( &keys::worker_instance::MetricsLockKey::new(), rivet_util::timestamp::now(), )?; @@ -671,33 +625,37 @@ impl Database for DatabaseKv { } }) .custom_instrument(tracing::info_span!("acquire_lock_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; if acquired_lock { let entries = self .pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { - let txs = tx.subspace(self.subspace.clone()); + .run(|tx| async move { + let tx = tx.with_subspace(self.subspace.clone()); - let metrics_subspace = txs.subspace(&keys::metric::GaugeMetricKey::subspace()); - txs.get_ranges_keyvalues( - udb::RangeOption { + let metrics_subspace = self + .subspace + .subspace(&keys::metric::GaugeMetricKey::subspace()); + tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&metrics_subspace).into() }, - SERIALIZABLE, + Serializable, ) .map(|res| match res { - Ok(entry) => txs.read_entry::(&entry), + Ok(entry) => tx.read_entry::(&entry), Err(err) => Err(err.into()), }) .try_collect::>() .await }) .custom_instrument(tracing::info_span!("read_metrics_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; let mut total_workflow_counts: Vec<(String, usize)> = Vec::new(); @@ -779,14 +737,15 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { let metrics_lock_key = keys::worker_instance::MetricsLockKey::new(); tx.clear(&self.subspace.pack(&metrics_lock_key)); Ok(()) }) .custom_instrument(tracing::info_span!("clear_lock_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; } Ok(()) @@ -805,23 +764,22 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| { + .run(|tx| { async move { // Update worker instance ping let last_ping_ts_key = keys::worker_instance::LastPingTsKey::new(worker_instance_id); tx.set( &self.subspace.pack(&last_ping_ts_key), - &last_ping_ts_key - .serialize(rivet_util::timestamp::now()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &last_ping_ts_key.serialize(rivet_util::timestamp::now())?, ); Ok(()) } }) .custom_instrument(tracing::info_span!("update_worker_ping_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; Ok(()) } @@ -840,7 +798,7 @@ impl Database for DatabaseKv { .pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { self.dispatch_workflow_inner( ray_id, workflow_id, @@ -853,7 +811,8 @@ impl Database for DatabaseKv { .await }) .custom_instrument(tracing::info_span!("dispatch_workflow_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; self.wake_worker(); @@ -865,7 +824,7 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| { + .run(|tx| { let workflow_ids = workflow_ids.clone(); async move { futures_util::stream::iter(workflow_ids) @@ -889,56 +848,50 @@ impl Database for DatabaseKv { has_wake_condition_entry, ) = tokio::try_join!( tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&input_subspace).into() }, - SERIALIZABLE, + Serializable, ) .try_collect::>(), tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&state_subspace).into() }, - SERIALIZABLE, + Serializable, ) .try_collect::>(), tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&output_subspace).into() }, - SERIALIZABLE, + Serializable, ) .try_collect::>(), tx.get( &self.subspace.pack(&has_wake_condition_key), - SERIALIZABLE + Serializable ), )?; if input_chunks.is_empty() { Ok(None) } else { - let input = input_key - .combine(input_chunks) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let input = input_key.combine(input_chunks)?; let state = if state_chunks.is_empty() { serde_json::value::RawValue::NULL.to_owned() } else { - state_key.combine(state_chunks).map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })? + state_key.combine(state_chunks)? }; let output = if output_chunks.is_empty() { None } else { - Some(output_key.combine(output_chunks).map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?) + Some(output_key.combine(output_chunks)?) }; Ok(Some(WorkflowData { @@ -960,7 +913,7 @@ impl Database for DatabaseKv { }) .custom_instrument(tracing::info_span!("get_workflow_tx")) .await - .map_err(Into::into) + .map_err(WorkflowError::Udb) } /// Returns the first incomplete workflow with the given name and tags, first meaning the one with the @@ -978,9 +931,10 @@ impl Database for DatabaseKv { .pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { self.find_workflow_inner(workflow_name, tags, &tx).await }) + .run(|tx| async move { self.find_workflow_inner(workflow_name, tags, &tx).await }) .custom_instrument(tracing::info_span!("find_workflow_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; let dt = start_instant.elapsed().as_secs_f64(); metrics::FIND_WORKFLOWS_DURATION.record( @@ -1007,16 +961,14 @@ impl Database for DatabaseKv { .pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| { + .run(|tx| { let owned_filter = owned_filter.clone(); async move { let now = rivet_util::timestamp::now(); // All wake conditions with a timestamp after this timestamp will be pulled - let pull_before = now - + i64::try_from(self.worker_poll_interval().as_millis()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let pull_before = now + i64::try_from(self.worker_poll_interval().as_millis())?; // Pull all available wake conditions from all registered wf names let entries = futures_util::stream::iter(owned_filter) @@ -1044,24 +996,24 @@ impl Database for DatabaseKv { .to_vec(); tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(wake_subspace_start, wake_subspace_end).into() }, - // Must be a snapshot to not conflict with any new wake conditions being + // Must be a Snapshot to not conflict with any new wake conditions being // inserted - SNAPSHOT, + Snapshot, ) }) .flatten() - .map(|res| match res { - Ok(entry) => Ok(( + .map(|res| { + let entry = res?; + + anyhow::Ok(( entry.key().to_vec(), self.subspace - .unpack::(entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - )), - Err(err) => Err(Into::::into(err)), + .unpack::(entry.key())?, + )) }) .try_collect::>() .await?; @@ -1101,17 +1053,16 @@ impl Database for DatabaseKv { let lease_key_buf = self.subspace.pack(&lease_key); // Check lease - if tx.get(&lease_key_buf, SERIALIZABLE).await?.is_some() { - Result::<_, udb::FdbBindingError>::Ok(None) + if tx.get(&lease_key_buf, Serializable).await?.is_some() { + Result::<_>::Ok(None) } else { // Write lease tx.set( &lease_key_buf, - &lease_key - .serialize((workflow_name.clone(), worker_instance_id)) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?, + &lease_key.serialize(( + workflow_name.clone(), + worker_instance_id, + ))?, ); // Write worker instance id @@ -1119,15 +1070,11 @@ impl Database for DatabaseKv { keys::workflow::WorkerInstanceIdKey::new(workflow_id); tx.set( &self.subspace.pack(&worker_instance_id_key), - &worker_instance_id_key - .serialize(worker_instance_id) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?, + &worker_instance_id_key.serialize(worker_instance_id)?, ); update_metric( - &tx.subspace(self.subspace.clone()), + &tx.with_subspace(self.subspace.clone()), Some(keys::metric::GaugeMetric::WorkflowSleeping( workflow_name.clone(), )), @@ -1168,12 +1115,10 @@ impl Database for DatabaseKv { let wake_sub_workflow_key = keys::workflow::WakeSubWorkflowKey::new(*workflow_id); if let Some(entry) = tx - .get(&self.subspace.pack(&wake_sub_workflow_key), SERIALIZABLE) + .get(&self.subspace.pack(&wake_sub_workflow_key), Serializable) .await? { - let sub_workflow_id = wake_sub_workflow_key - .deserialize(&entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let sub_workflow_id = wake_sub_workflow_key.deserialize(&entry)?; let sub_workflow_wake_key = keys::wake::SubWorkflowWakeKey::new(sub_workflow_id, *workflow_id); @@ -1194,7 +1139,8 @@ impl Database for DatabaseKv { } }) .custom_instrument(tracing::info_span!("pull_workflows_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; let worker_instance_id_str = worker_instance_id.to_string(); let dt = start_instant.elapsed().as_secs_f64(); @@ -1223,7 +1169,7 @@ impl Database for DatabaseKv { .pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| { + .run(|tx| { let leased_workflows = leased_workflows.clone(); async move { @@ -1253,42 +1199,33 @@ impl Database for DatabaseKv { events, ) = tokio::try_join!( async { - tx.get(&self.subspace.pack(&create_ts_key), SERIALIZABLE) + tx.get(&self.subspace.pack(&create_ts_key), Serializable) .await - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - }) }, async { - tx.get(&self.subspace.pack(&ray_id_key), SERIALIZABLE) - .await - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - }) + tx.get(&self.subspace.pack(&ray_id_key), Serializable).await }, async { tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&input_subspace).into() }, - SERIALIZABLE, + Serializable, ) .try_collect::>() .await - .map_err(|x| udb::FdbBindingError::CustomError(x.into())) }, async { tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&state_subspace).into() }, - SERIALIZABLE, + Serializable, ) .try_collect::>() .await - .map_err(|x| udb::FdbBindingError::CustomError(x.into())) }, async { let mut events_by_location: HashMap> = @@ -1297,11 +1234,11 @@ impl Database for DatabaseKv { WorkflowHistoryEventBuilder::new(Location::empty()); let mut stream = tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&active_history_subspace).into() }, - SERIALIZABLE, + Serializable, ); loop { @@ -1313,11 +1250,8 @@ impl Database for DatabaseKv { let partial_key = self .subspace .unpack::( - entry.key(), - ) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?; + entry.key(), + )?; if current_event.location != partial_key.location { if current_event.location.is_empty() { @@ -1337,12 +1271,7 @@ impl Database for DatabaseKv { events_by_location .entry(previous_event.location.root()) .or_default() - .push( - Event::try_from(previous_event) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?, - ); + .push(Event::try_from(previous_event)?); } } @@ -1351,53 +1280,35 @@ impl Database for DatabaseKv { self.subspace.unpack::( entry.key(), ) { - let event_type = key - .deserialize(entry.value()) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?; + let event_type = key.deserialize(entry.value())?; current_event.event_type = Some(event_type); } else if let Ok(key) = self.subspace.unpack::( entry.key(), ) { - let version = key - .deserialize(entry.value()) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?; + let version = key.deserialize(entry.value())?; current_event.version = Some(version); } else if let Ok(key) = self.subspace.unpack::( entry.key(), ) { - let create_ts = key - .deserialize(entry.value()) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?; + let create_ts = key.deserialize(entry.value())?; current_event.create_ts = Some(create_ts); } else if let Ok(key) = self.subspace .unpack::(entry.key()) { - let name = key.deserialize(entry.value()).map_err( - |x| udb::FdbBindingError::CustomError(x.into()), - )?; + let name = key.deserialize(entry.value())?; current_event.name = Some(name); } else if let Ok(key) = self.subspace.unpack::( entry.key(), ) { - let signal_id = key - .deserialize(entry.value()) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?; + let signal_id = key.deserialize(entry.value())?; current_event.signal_id = Some(signal_id); } else if let Ok(key) = self @@ -1405,11 +1316,8 @@ impl Database for DatabaseKv { .unpack::( entry.key(), ) { - let sub_workflow_id = key - .deserialize(entry.value()) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?; + let sub_workflow_id = + key.deserialize(entry.value())?; current_event.sub_workflow_id = Some(sub_workflow_id); @@ -1429,11 +1337,7 @@ impl Database for DatabaseKv { self.subspace.unpack::( entry.key(), ) { - let input_hash = key - .deserialize(entry.value()) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?; + let input_hash = key.deserialize(entry.value())?; current_event.input_hash = Some(input_hash); } else if let Ok(_key) = @@ -1445,11 +1349,7 @@ impl Database for DatabaseKv { self.subspace.unpack::( entry.key(), ) { - let iteration = key - .deserialize(entry.value()) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?; + let iteration = key.deserialize(entry.value())?; current_event.iteration = Some(iteration); } else if let Ok(key) = self @@ -1457,11 +1357,7 @@ impl Database for DatabaseKv { .unpack::( entry.key(), ) { - let deadline_ts = key - .deserialize(entry.value()) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?; + let deadline_ts = key.deserialize(entry.value())?; current_event.deadline_ts = Some(deadline_ts); } else if let Ok(key) = self @@ -1469,11 +1365,7 @@ impl Database for DatabaseKv { .unpack::( entry.key(), ) { - let sleep_state = key - .deserialize(entry.value()) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?; + let sleep_state = key.deserialize(entry.value())?; current_event.sleep_state = Some(sleep_state); } else if let Ok(key) = self @@ -1481,11 +1373,8 @@ impl Database for DatabaseKv { .unpack::( entry.key(), ) { - let inner_event_type = key - .deserialize(entry.value()) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?; + let inner_event_type = + key.deserialize(entry.value())?; current_event.inner_event_type = Some(inner_event_type); @@ -1498,9 +1387,7 @@ impl Database for DatabaseKv { events_by_location .entry(current_event.location.root()) .or_default() - .push(Event::try_from(current_event).map_err( - |x| udb::FdbBindingError::CustomError(x.into()), - )?); + .push(Event::try_from(current_event)?); } Ok(events_by_location) @@ -1508,31 +1395,17 @@ impl Database for DatabaseKv { )?; let create_ts = create_ts_key - .deserialize(&create_ts_entry.ok_or( - udb::FdbBindingError::CustomError( - format!("key should exist: {create_ts_key:?}").into(), - ), - )?) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + .deserialize(&create_ts_entry.context("key should exist")?)?; let ray_id = ray_id_key - .deserialize(&ray_id_entry.ok_or( - udb::FdbBindingError::CustomError( - format!("key should exist: {ray_id_key:?}").into(), - ), - )?) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; - let input = input_key - .combine(input_chunks) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + .deserialize(&ray_id_entry.context("key should exist")?)?; + let input = input_key.combine(input_chunks)?; let state = if state_chunks.is_empty() { serde_json::value::RawValue::NULL.to_owned() } else { - state_key - .combine(state_chunks) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? + state_key.combine(state_chunks)? }; - Result::<_, udb::FdbBindingError>::Ok(PulledWorkflowData { + Result::<_>::Ok(PulledWorkflowData { workflow_id, workflow_name, create_ts, @@ -1552,7 +1425,8 @@ impl Database for DatabaseKv { } }) .custom_instrument(tracing::info_span!("pull_workflow_history_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; let dt2 = start_instant2.elapsed().as_secs_f64(); let dt = start_instant.elapsed().as_secs_f64(); @@ -1601,7 +1475,7 @@ impl Database for DatabaseKv { .pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| { + .run(|tx| { async move { let sub_workflow_wake_subspace = self .subspace @@ -1612,12 +1486,12 @@ impl Database for DatabaseKv { let wake_deadline_key = keys::workflow::WakeDeadlineKey::new(workflow_id); let mut stream = tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&sub_workflow_wake_subspace).into() }, - // NOTE: Must be serializable to conflict with `get_sub_workflow` - SERIALIZABLE, + // NOTE: Must be Serializable to conflict with `get_sub_workflow` + Serializable, ); let (wrote_to_wake_idx, tag_keys, wake_deadline_entry) = tokio::try_join!( @@ -1626,13 +1500,11 @@ impl Database for DatabaseKv { let mut wrote_to_wake_idx = false; while let Some(entry) = stream.try_next().await? { - let sub_workflow_wake_key = self - .subspace - .unpack::(&entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; - let workflow_name = sub_workflow_wake_key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let sub_workflow_wake_key = + self.subspace + .unpack::(&entry.key())?; + let workflow_name = + sub_workflow_wake_key.deserialize(entry.value())?; let wake_condition_key = keys::wake::WorkflowWakeConditionKey::new( workflow_name, @@ -1645,9 +1517,7 @@ impl Database for DatabaseKv { // Add wake condition for workflow tx.set( &self.subspace.pack(&wake_condition_key), - &wake_condition_key - .serialize(()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &wake_condition_key.serialize(())?, ); // Clear secondary index @@ -1656,33 +1526,23 @@ impl Database for DatabaseKv { wrote_to_wake_idx = true; } - Result::<_, udb::FdbBindingError>::Ok(wrote_to_wake_idx) + Ok(wrote_to_wake_idx) }, // Read tags - async { - tx.get_ranges_keyvalues( - udb::RangeOption { - mode: StreamingMode::WantAll, - ..(&tags_subspace).into() - }, - SERIALIZABLE, - ) - .map(|res| match res { - Ok(entry) => self - .subspace - .unpack::(entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into())), - Err(err) => Err(Into::::into(err)), - }) - .try_collect::>() - .await - .map_err(Into::into) - }, - async { - tx.get(&self.subspace.pack(&wake_deadline_key), SERIALIZABLE) - .await - .map_err(Into::into) - }, + tx.get_ranges_keyvalues( + universaldb::RangeOption { + mode: StreamingMode::WantAll, + ..(&tags_subspace).into() + }, + Serializable, + ) + .map(|res| { + self.subspace + .unpack::(res?.key()) + .map_err(anyhow::Error::from) + }) + .try_collect::>(), + tx.get(&self.subspace.pack(&wake_deadline_key), Serializable), )?; for key in tag_keys { @@ -1711,9 +1571,7 @@ impl Database for DatabaseKv { // reason this isn't immediately cleared in `pull_workflows` along with the rest of the // wake conditions is because it might be in the future. if let Some(raw) = wake_deadline_entry { - let deadline_ts = wake_deadline_key - .deserialize(&raw) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let deadline_ts = wake_deadline_key.deserialize(&raw)?; let wake_condition_key = keys::wake::WorkflowWakeConditionKey::new( workflow_name.to_string(), @@ -1732,12 +1590,7 @@ impl Database for DatabaseKv { // Write output let output_key = keys::workflow::OutputKey::new(workflow_id); - for (i, chunk) in output_key - .split_ref(output) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in output_key.split_ref(output)?.into_iter().enumerate() { let chunk_key = output_key.chunk(i); tx.set(&self.subspace.pack(&chunk_key), &chunk); @@ -1751,7 +1604,7 @@ impl Database for DatabaseKv { tx.clear(&self.subspace.pack(&worker_instance_id_key)); update_metric( - &tx.subspace(self.subspace.clone()), + &tx.with_subspace(self.subspace.clone()), Some(keys::metric::GaugeMetric::WorkflowActive( workflow_name.to_string(), )), @@ -1764,7 +1617,8 @@ impl Database for DatabaseKv { } }) .custom_instrument(tracing::info_span!("complete_workflows_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; // Wake worker again in case some other workflow was waiting for this one to complete if wrote_to_wake_idx { @@ -1796,12 +1650,12 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| { + .run(|tx| { async move { let wake_deadline_key = keys::workflow::WakeDeadlineKey::new(workflow_id); let wake_deadline_entry = tx - .get(&self.subspace.pack(&wake_deadline_key), SERIALIZABLE) + .get(&self.subspace.pack(&wake_deadline_key), Serializable) .await?; // Add immediate wake for workflow @@ -1813,9 +1667,7 @@ impl Database for DatabaseKv { ); tx.set( &self.subspace.pack(&wake_condition_key), - &wake_condition_key - .serialize(()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &wake_condition_key.serialize(())?, ); } @@ -1826,9 +1678,7 @@ impl Database for DatabaseKv { // reason this isn't immediately cleared in `pull_workflows` along with the rest of the // wake conditions is because it might be in the future. if let Some(raw) = wake_deadline_entry { - let deadline_ts = wake_deadline_key - .deserialize(&raw) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let deadline_ts = wake_deadline_key.deserialize(&raw)?; let wake_condition_key = keys::wake::WorkflowWakeConditionKey::new( workflow_name.to_string(), @@ -1850,17 +1700,13 @@ impl Database for DatabaseKv { // Add wake condition for workflow tx.set( &self.subspace.pack(&wake_condition_key), - &wake_condition_key - .serialize(()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &wake_condition_key.serialize(())?, ); // Write to wake deadline tx.set( &self.subspace.pack(&wake_deadline_key), - &wake_deadline_key - .serialize(deadline_ts) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &wake_deadline_key.serialize(deadline_ts)?, ); } @@ -1886,9 +1732,7 @@ impl Database for DatabaseKv { if has_wake_condition { tx.set( &self.subspace.pack(&has_wake_condition_key), - &has_wake_condition_key - .serialize(()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &has_wake_condition_key.serialize(())?, ); } else { tx.clear(&self.subspace.pack(&has_wake_condition_key)); @@ -1898,9 +1742,7 @@ impl Database for DatabaseKv { let error_key = keys::workflow::ErrorKey::new(workflow_id); tx.set( &self.subspace.pack(&error_key), - &error_key - .serialize(error.to_string()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &error_key.serialize(error.to_string())?, ); // Clear lease @@ -1911,7 +1753,7 @@ impl Database for DatabaseKv { tx.clear(&self.subspace.pack(&worker_instance_id_key)); update_metric( - &tx.subspace(self.subspace.clone()), + &tx.with_subspace(self.subspace.clone()), Some(keys::metric::GaugeMetric::WorkflowActive( workflow_name.to_string(), )), @@ -1929,7 +1771,8 @@ impl Database for DatabaseKv { } }) .custom_instrument(tracing::info_span!("commit_workflow_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; // Always wake the worker immediately again. This is an IMPORTANT implementation detail to prevent // race conditions with workflow sleep. Imagine the scenario: @@ -1975,12 +1818,12 @@ impl Database for DatabaseKv { .map(|x| x.to_string()) .collect::>(); - // Fetch signal from FDB + // Fetch signal from UDB let signal = self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| { + .run(|tx| { let owned_filter = owned_filter.clone(); async move { @@ -1997,14 +1840,14 @@ impl Database for DatabaseKv { ); tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, limit: Some(1), ..(&pending_signal_subspace).into() }, - // NOTE: This is serializable because any insert into this subspace + // NOTE: This is Serializable because any insert into this subspace // should cause a conflict and retry of this txn - SERIALIZABLE, + Serializable, ) }) .collect::>(); @@ -2013,15 +1856,12 @@ impl Database for DatabaseKv { let mut results = futures_util::future::try_join_all( streams.into_iter().map(|mut stream| async move { if let Some(entry) = stream.try_next().await? { - Result::<_, udb::FdbBindingError>::Ok(Some(( + Result::<_>::Ok(Some(( entry.key().to_vec(), self.subspace .unpack::( &entry.key(), - ) - .map_err(|x| { - udb::FdbBindingError::CustomError(x.into()) - })?, + )?, ))) } else { Ok(None) @@ -2058,13 +1898,11 @@ impl Database for DatabaseKv { )?; tx.set( &self.subspace.pack(&ack_ts_key), - &ack_ts_key - .serialize(rivet_util::timestamp::now()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &ack_ts_key.serialize(rivet_util::timestamp::now())?, ); update_metric( - &tx.subspace(self.subspace.clone()), + &tx.with_subspace(self.subspace.clone()), Some(keys::metric::GaugeMetric::SignalPending( signal_name.to_string(), )), @@ -2082,18 +1920,16 @@ impl Database for DatabaseKv { let chunks = tx .get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&body_subspace).into() }, - SERIALIZABLE, + Serializable, ) .try_collect::>() .await?; - let body = body_key - .combine(chunks) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let body = body_key.combine(chunks)?; // Insert history event keys::history::insert::signal_event( @@ -2106,8 +1942,7 @@ impl Database for DatabaseKv { signal_id, &signal_name, &body, - ) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + )?; Ok(Some(SignalData { signal_id, @@ -2142,7 +1977,8 @@ impl Database for DatabaseKv { } }) .custom_instrument(tracing::info_span!("pull_next_signal_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; Ok(signal) } @@ -2157,7 +1993,7 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| { + .run(|tx| { async move { let input_key = keys::workflow::InputKey::new(sub_workflow_id); let input_subspace = self.subspace.subspace(&input_key); @@ -2171,45 +2007,41 @@ impl Database for DatabaseKv { // Read input and output let (input_chunks, state_chunks, output_chunks, has_wake_condition_entry) = tokio::try_join!( tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&input_subspace).into() }, - SERIALIZABLE, + Serializable, ) .try_collect::>(), tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&state_subspace).into() }, - SERIALIZABLE, + Serializable, ) .try_collect::>(), tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&output_subspace).into() }, - SERIALIZABLE, + Serializable, ) .try_collect::>(), - tx.get(&self.subspace.pack(&has_wake_condition_key), SERIALIZABLE), + tx.get(&self.subspace.pack(&has_wake_condition_key), Serializable), )?; if input_chunks.is_empty() { Ok(None) } else { - let input = input_key - .combine(input_chunks) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let input = input_key.combine(input_chunks)?; let state = if state_chunks.is_empty() { serde_json::value::RawValue::NULL.to_owned() } else { - state_key - .combine(state_chunks) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? + state_key.combine(state_chunks)? }; let output = if output_chunks.is_empty() { @@ -2230,11 +2062,7 @@ impl Database for DatabaseKv { None } else { - Some( - output_key - .combine(output_chunks) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ) + Some(output_key.combine(output_chunks)?) }; Ok(Some(WorkflowData { @@ -2249,7 +2077,7 @@ impl Database for DatabaseKv { }) .custom_instrument(tracing::info_span!("get_sub_workflow_tx")) .await - .map_err(Into::into) + .map_err(WorkflowError::Udb) } #[tracing::instrument(skip_all)] @@ -2264,12 +2092,13 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { self.publish_signal_inner(ray_id, workflow_id, signal_id, signal_name, body, &tx) .await }) .custom_instrument(tracing::info_span!("publish_signal_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; self.wake_worker(); @@ -2292,7 +2121,7 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { self.publish_signal_inner( ray_id, to_workflow_id, @@ -2315,13 +2144,13 @@ impl Database for DatabaseKv { &signal_name, &body, to_workflow_id, - ) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + )?; Ok(()) }) .custom_instrument(tracing::info_span!("publish_signal_from_workflow_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; self.wake_worker(); @@ -2346,7 +2175,7 @@ impl Database for DatabaseKv { .pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { let sub_workflow_id = self .dispatch_workflow_inner( ray_id, @@ -2371,13 +2200,13 @@ impl Database for DatabaseKv { sub_workflow_name, tags, input, - ) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + )?; Ok(sub_workflow_id) }) .custom_instrument(tracing::info_span!("dispatch_sub_workflow_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; self.wake_worker(); @@ -2394,7 +2223,7 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| { + .run(|tx| { async move { let tags_subspace = self .subspace @@ -2403,18 +2232,16 @@ impl Database for DatabaseKv { // Read old tags let tag_keys = tx .get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&tags_subspace).into() }, - SERIALIZABLE, + Serializable, ) - .map(|res| match res { - Ok(entry) => self - .subspace - .unpack::(entry.key()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into())), - Err(err) => Err(Into::::into(err)), + .map(|res| { + self.subspace + .unpack::(res?.key()) + .map_err(anyhow::Error::from) }) .try_collect::>() .await?; @@ -2435,22 +2262,15 @@ impl Database for DatabaseKv { // Write new tags let tags = tags .as_object() - .ok_or_else(|| WorkflowError::InvalidTags("must be an object".to_string())) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? + .ok_or_else(|| WorkflowError::InvalidTags("must be an object".to_string()))? .into_iter() .map(|(k, v)| Ok((k.clone(), value_to_str(v)?))) - .collect::>>() - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + .collect::>>()?; for (k, v) in &tags { let tag_key = keys::workflow::TagKey::new(workflow_id, k.clone(), v.clone()); - tx.set( - &self.subspace.pack(&tag_key), - &tag_key - .serialize(()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ); + tx.set(&self.subspace.pack(&tag_key), &tag_key.serialize(())?); // Write new "by name and first tag" secondary index let by_name_and_tag_key = keys::workflow::ByNameAndTagKey::new( @@ -2466,9 +2286,7 @@ impl Database for DatabaseKv { .collect(); tx.set( &self.subspace.pack(&by_name_and_tag_key), - &by_name_and_tag_key - .serialize(rest_of_tags) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &by_name_and_tag_key.serialize(rest_of_tags)?, ); } @@ -2476,7 +2294,8 @@ impl Database for DatabaseKv { } }) .custom_instrument(tracing::info_span!("update_workflow_tags_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; Ok(()) } @@ -2490,17 +2309,12 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| { + .run(|tx| { async move { let state_key = keys::workflow::StateKey::new(workflow_id); // Write state - for (i, chunk) in state_key - .split_ref(&state) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .into_iter() - .enumerate() - { + for (i, chunk) in state_key.split_ref(&state)?.into_iter().enumerate() { let chunk_key = state_key.chunk(i); tx.set(&self.subspace.pack(&chunk_key), &chunk); @@ -2510,7 +2324,8 @@ impl Database for DatabaseKv { } }) .custom_instrument(tracing::info_span!("update_workflow_state_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; Ok(()) } @@ -2530,7 +2345,7 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { keys::history::insert::activity_event( &self.subspace, &tx, @@ -2542,13 +2357,13 @@ impl Database for DatabaseKv { &event_id.input_hash.to_be_bytes(), input, res, - ) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + )?; Ok(()) }) .custom_instrument(tracing::info_span!("commit_workflow_activity_event_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; Ok(()) } @@ -2567,7 +2382,7 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { keys::history::insert::message_send_event( &self.subspace, &tx, @@ -2578,13 +2393,13 @@ impl Database for DatabaseKv { tags, message_name, body, - ) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + )?; Ok(()) }) .custom_instrument(tracing::info_span!("commit_workflow_message_send_event_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; Ok(()) } @@ -2604,7 +2419,7 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { if iteration == 0 { keys::history::insert::loop_event( &self.subspace, @@ -2616,8 +2431,7 @@ impl Database for DatabaseKv { iteration, state, output, - ) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + )?; } else { keys::history::insert::update_loop_event( &self.subspace, @@ -2627,8 +2441,7 @@ impl Database for DatabaseKv { iteration, state, output, - ) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + )?; let active_history_subspace = self.subspace @@ -2653,11 +2466,11 @@ impl Database for DatabaseKv { )); let mut stream = tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&loop_events_subspace).into() }, - SERIALIZABLE, + Serializable, ); // Move all current events under this loop to the forgotten history @@ -2667,9 +2480,7 @@ impl Database for DatabaseKv { }; if !active_history_subspace.is_start_of(entry.key()) { - return Err(udb::FdbBindingError::CustomError( - udb::tuple::PackError::BadPrefix.into(), - )); + return Err(universaldb::tuple::PackError::BadPrefix.into()); } // Truncate tuple up to ACTIVE and replace it with FORGOTTEN @@ -2708,7 +2519,8 @@ impl Database for DatabaseKv { Ok(()) }) .custom_instrument(tracing::info_span!("commit_workflow_sleep_event_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; Ok(()) } @@ -2725,7 +2537,7 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { keys::history::insert::sleep_event( &self.subspace, &tx, @@ -2735,13 +2547,13 @@ impl Database for DatabaseKv { rivet_util::timestamp::now(), deadline_ts, SleepState::Normal, - ) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + )?; Ok(()) }) .custom_instrument(tracing::info_span!("commit_workflow_sleep_event_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; Ok(()) } @@ -2756,20 +2568,20 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { keys::history::insert::update_sleep_event( &self.subspace, &tx, from_workflow_id, location, state, - ) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + )?; Ok(()) }) .custom_instrument(tracing::info_span!("update_workflow_sleep_event_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; Ok(()) } @@ -2785,7 +2597,7 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { keys::history::insert::branch_event( &self.subspace, &tx, @@ -2793,13 +2605,13 @@ impl Database for DatabaseKv { location, version, rivet_util::timestamp::now(), - ) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + )?; Ok(()) }) .custom_instrument(tracing::info_span!("commit_workflow_branch_event_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; Ok(()) } @@ -2816,7 +2628,7 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { keys::history::insert::removed_event( &self.subspace, &tx, @@ -2826,13 +2638,13 @@ impl Database for DatabaseKv { rivet_util::timestamp::now(), event_type, event_name, - ) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + )?; Ok(()) }) .custom_instrument(tracing::info_span!("commit_workflow_removed_event_tx")) - .await?; + .await + .map_err(WorkflowError::Udb)?; Ok(()) } @@ -2848,7 +2660,7 @@ impl Database for DatabaseKv { self.pools .udb() .map_err(WorkflowError::PoolsGeneric)? - .run(|tx, _mc| async move { + .run(|tx| async move { keys::history::insert::version_check_event( &self.subspace, &tx, @@ -2856,22 +2668,22 @@ impl Database for DatabaseKv { location, version, rivet_util::timestamp::now(), - ) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + )?; Ok(()) }) .custom_instrument(tracing::info_span!( "commit_workflow_version_check_event_tx" )) - .await?; + .await + .map_err(WorkflowError::Udb)?; Ok(()) } } fn update_metric( - txs: &udb_util::TxnSubspace, + tx: &universaldb::Transaction, previous: Option, current: Option, ) { @@ -2880,7 +2692,7 @@ fn update_metric( } if let Some(previous) = previous { - txs.atomic_op( + tx.atomic_op( &keys::metric::GaugeMetricKey::new(previous), &(-1isize).to_le_bytes(), MutationType::Add, @@ -2888,7 +2700,7 @@ fn update_metric( } if let Some(current) = current { - txs.atomic_op( + tx.atomic_op( &keys::metric::GaugeMetricKey::new(current), &1usize.to_le_bytes(), MutationType::Add, @@ -2904,8 +2716,8 @@ struct WorkflowHistoryEventBuilder { name: Option, signal_id: Option, sub_workflow_id: Option, - input_chunks: Vec, - output_chunks: Vec, + input_chunks: Vec, + output_chunks: Vec, input_hash: Option>, error_count: usize, iteration: Option, diff --git a/packages/common/gasoline/core/src/error.rs b/packages/common/gasoline/core/src/error.rs index 7c79ae4617..7553046375 100644 --- a/packages/common/gasoline/core/src/error.rs +++ b/packages/common/gasoline/core/src/error.rs @@ -2,7 +2,6 @@ use std::time::{SystemTime, UNIX_EPOCH}; use rivet_util::Id; use tokio::time::Instant; -use universaldb as udb; use crate::ctx::common::RETRY_TIMEOUT_MS; @@ -134,8 +133,8 @@ pub enum WorkflowError { #[error("failed to deserialize event data: {0}")] DeserializeEventData(#[source] anyhow::Error), - #[error("fdb error: {0}")] - Fdb(#[from] udb::FdbBindingError), + #[error("udb error: {0}")] + Udb(#[source] anyhow::Error), #[error("pools error: {0}")] Pools(#[from] rivet_pools::Error), @@ -217,7 +216,7 @@ impl WorkflowError { } /// Any error that the workflow can continue on with its execution from. - pub fn is_recoverable(&self) -> bool { + pub(crate) fn is_recoverable(&self) -> bool { match self { WorkflowError::ActivityFailure(_, _) | WorkflowError::ActivityTimeout(_) diff --git a/packages/common/gasoline/core/src/history/location.rs b/packages/common/gasoline/core/src/history/location.rs index 3a7273b1be..d93f9d9141 100644 --- a/packages/common/gasoline/core/src/history/location.rs +++ b/packages/common/gasoline/core/src/history/location.rs @@ -158,9 +158,9 @@ impl Deref for Coordinate { } } -mod fdb { +mod udb { use super::Coordinate; - use udb_util::prelude::*; + use universaldb::prelude::*; impl TuplePack for Coordinate { fn pack( diff --git a/packages/common/gasoline/core/src/worker.rs b/packages/common/gasoline/core/src/worker.rs index b940811a4f..1223070763 100644 --- a/packages/common/gasoline/core/src/worker.rs +++ b/packages/common/gasoline/core/src/worker.rs @@ -257,11 +257,12 @@ impl Worker { // NOTE: No .in_current_span() because we want this to be a separate trace async move { if let Err(err) = ctx.run(current_span_ctx).await { - tracing::error!(?err, "unhandled workflow error"); + tracing::error!(?err, ?workflow_id, "unhandled workflow error"); sentry::with_scope( |scope| { scope.set_tag("error", err.to_string()); + scope.set_tag("workflow_id", workflow_id.to_string()); }, || { sentry::capture_message( diff --git a/packages/common/pools/Cargo.toml b/packages/common/pools/Cargo.toml index 113a5ddc00..f445b8cad3 100644 --- a/packages/common/pools/Cargo.toml +++ b/packages/common/pools/Cargo.toml @@ -10,8 +10,6 @@ anyhow.workspace = true async-nats.workspace = true clickhouse-inserter.workspace = true clickhouse.workspace = true -udb-util.workspace = true -universaldb.workspace = true futures-util.workspace = true governor.workspace = true hyper-tls.workspace = true @@ -20,8 +18,8 @@ lazy_static.workspace = true reqwest.workspace = true rivet-config.workspace = true rivet-metrics.workspace = true -universalpubsub.workspace = true rivet-util.workspace = true +serde.workspace = true tempfile.workspace = true thiserror.workspace = true tokio-native-tls.workspace = true @@ -30,9 +28,10 @@ tokio.workspace = true tracing-logfmt.workspace = true tracing-subscriber.workspace = true tracing.workspace = true +universaldb.workspace = true +universalpubsub.workspace = true url.workspace = true uuid.workspace = true -serde.workspace = true [dev-dependencies] divan.workspace = true diff --git a/packages/common/pools/src/db/udb.rs b/packages/common/pools/src/db/udb.rs index 7fff92efc8..9f6b0b855f 100644 --- a/packages/common/pools/src/db/udb.rs +++ b/packages/common/pools/src/db/udb.rs @@ -2,15 +2,14 @@ use std::{ops::Deref, sync::Arc}; use anyhow::*; use rivet_config::{Config, config}; -use universaldb as udb; #[derive(Clone)] pub struct UdbPool { - db: udb::Database, + db: universaldb::Database, } impl Deref for UdbPool { - type Target = udb::Database; + type Target = universaldb::Database; fn deref(&self) -> &Self::Target { &self.db @@ -21,18 +20,18 @@ impl Deref for UdbPool { pub async fn setup(config: Config) -> Result> { let db_driver = match config.database() { config::Database::Postgres(pg) => { - Arc::new(udb::driver::PostgresDatabaseDriver::new(pg.url.read().clone()).await?) - as udb::DatabaseDriverHandle + Arc::new(universaldb::driver::PostgresDatabaseDriver::new(pg.url.read().clone()).await?) + as universaldb::DatabaseDriverHandle } config::Database::FileSystem(fs) => { - Arc::new(udb::driver::RocksDbDatabaseDriver::new(fs.path.clone()).await?) - as udb::DatabaseDriverHandle + Arc::new(universaldb::driver::RocksDbDatabaseDriver::new(fs.path.clone()).await?) + as universaldb::DatabaseDriverHandle } }; tracing::debug!("udb started"); Ok(Some(UdbPool { - db: udb::Database::new(db_driver), + db: universaldb::Database::new(db_driver), })) } diff --git a/packages/common/types/Cargo.toml b/packages/common/types/Cargo.toml index cad21dfa05..0dde0d6af1 100644 --- a/packages/common/types/Cargo.toml +++ b/packages/common/types/Cargo.toml @@ -13,6 +13,6 @@ rivet-data.workspace = true rivet-runner-protocol.workspace = true rivet-util.workspace = true serde.workspace = true -udb-util.workspace = true +universaldb.workspace = true utoipa.workspace = true versioned-data-util.workspace = true diff --git a/packages/common/types/src/keys/pegboard/mod.rs b/packages/common/types/src/keys/pegboard/mod.rs index 1e3ff30358..406b2ae4b5 100644 --- a/packages/common/types/src/keys/pegboard/mod.rs +++ b/packages/common/types/src/keys/pegboard/mod.rs @@ -1,7 +1,7 @@ -use udb_util::prelude::*; +use universaldb::prelude::*; pub mod ns; -pub fn subspace() -> udb_util::Subspace { - udb_util::Subspace::new(&(RIVET, PEGBOARD)) +pub fn subspace() -> universaldb::utils::Subspace { + universaldb::utils::Subspace::new(&(RIVET, PEGBOARD)) } diff --git a/packages/common/types/src/keys/pegboard/ns.rs b/packages/common/types/src/keys/pegboard/ns.rs index fa04e89a6e..f513eae2fc 100644 --- a/packages/common/types/src/keys/pegboard/ns.rs +++ b/packages/common/types/src/keys/pegboard/ns.rs @@ -2,7 +2,7 @@ use std::result::Result::Ok; use anyhow::*; use gas::prelude::*; -use udb_util::prelude::*; +use universaldb::prelude::*; #[derive(Debug)] pub struct ServerlessDesiredSlotsKey { diff --git a/packages/common/udb-util/Cargo.toml b/packages/common/udb-util/Cargo.toml deleted file mode 100644 index e89a97a4ff..0000000000 --- a/packages/common/udb-util/Cargo.toml +++ /dev/null @@ -1,16 +0,0 @@ -[package] -name = "udb-util" -version.workspace = true -authors.workspace = true -license.workspace = true -edition.workspace = true - -[dependencies] -anyhow.workspace = true -async-trait.workspace = true -universaldb.workspace = true -futures-util.workspace = true -lazy_static.workspace = true -rivet-metrics.workspace = true -tokio.workspace = true -tracing.workspace = true diff --git a/packages/common/udb-util/src/ext.rs b/packages/common/udb-util/src/ext.rs deleted file mode 100644 index 0d5bcc9c0f..0000000000 --- a/packages/common/udb-util/src/ext.rs +++ /dev/null @@ -1,322 +0,0 @@ -use std::{ops::Deref, result::Result::Ok}; - -use anyhow::*; -use futures_util::TryStreamExt; -use universaldb::{ - self as udb, - options::{ConflictRangeType, MutationType, StreamingMode}, - tuple::{TuplePack, TupleUnpack}, -}; - -use crate::{FormalKey, Subspace, end_of_key_range}; - -pub trait TxnExt { - fn subspace<'a>(&'a self, subspace: Subspace) -> TxnSubspace<'a>; -} - -impl TxnExt for udb::Transaction { - fn subspace<'a>(&'a self, subspace: Subspace) -> TxnSubspace<'a> { - TxnSubspace { - tx: &self, - subspace, - } - } -} - -#[derive(Clone)] -pub struct TxnSubspace<'a> { - tx: &'a udb::Transaction, - subspace: Subspace, -} - -impl<'a> TxnSubspace<'a> { - pub fn subspace(&self, t: &T) -> Subspace { - self.subspace.subspace(t) - } - - pub fn pack(&self, t: &T) -> Vec { - self.subspace.pack(t) - } - - pub fn unpack<'de, T: TupleUnpack<'de>>( - &self, - key: &'de [u8], - ) -> Result { - self.subspace - .unpack(key) - .with_context(|| format!("failed unpacking key of {}", std::any::type_name::())) - .map_err(|x| udb::FdbBindingError::CustomError(x.into())) - } - - pub fn write( - &self, - key: &T, - value: T::Value, - ) -> Result<(), udb::FdbBindingError> { - self.tx.set( - &self.subspace.pack(key), - &key.serialize(value) - .with_context(|| { - format!( - "failed serializing key value of {}", - std::any::type_name::(), - ) - }) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, - ); - - Ok(()) - } - - pub async fn read<'de, T: FormalKey + TuplePack + TupleUnpack<'de>>( - &self, - key: &'de T, - snapshot: bool, - ) -> Result { - self.tx - .get(&self.subspace.pack(key), snapshot) - .await? - .read(key) - } - - pub async fn read_opt<'de, T: FormalKey + TuplePack + TupleUnpack<'de>>( - &self, - key: &'de T, - snapshot: bool, - ) -> Result, udb::FdbBindingError> { - self.tx - .get(&self.subspace.pack(key), snapshot) - .await? - .read_opt(key) - } - - pub async fn exists( - &self, - key: &T, - snapshot: bool, - ) -> Result { - Ok(self - .tx - .get(&self.subspace.pack(key), snapshot) - .await? - .is_some()) - } - - pub fn delete(&self, key: &T) { - self.tx.clear(&self.subspace.pack(key)); - } - - pub fn delete_key_subspace(&self, key: &T) { - self.tx - .clear_subspace_range(&self.subspace(&self.subspace.pack(key))); - } - - pub fn read_entry TupleUnpack<'de>>( - &self, - entry: &udb::future::FdbValue, - ) -> Result<(T, T::Value), udb::FdbBindingError> { - let key = self.unpack::(entry.key())?; - let value = key - .deserialize(entry.value()) - .with_context(|| { - format!( - "failed deserializing key value of {}", - std::any::type_name::() - ) - }) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; - - Ok((key, value)) - } - - pub async fn cherry_pick( - &self, - subspace: impl TuplePack + Send, - snapshot: bool, - ) -> Result { - T::cherry_pick(self, subspace, snapshot).await - } - - pub fn add_conflict_key( - &self, - key: &T, - conflict_type: ConflictRangeType, - ) -> Result<(), udb::FdbBindingError> { - let key_buf = self.subspace.pack(key); - - self.tx - .add_conflict_range(&key_buf, &end_of_key_range(&key_buf), conflict_type) - .map_err(Into::into) - } - - pub fn atomic_op<'de, T: FormalKey + TuplePack + TupleUnpack<'de>>( - &self, - key: &'de T, - param: &[u8], - op_type: MutationType, - ) { - self.tx.atomic_op(&self.subspace.pack(key), param, op_type) - } -} - -impl<'a> Deref for TxnSubspace<'a> { - type Target = udb::Transaction; - - fn deref(&self) -> &Self::Target { - self.tx - } -} - -pub trait SliceExt { - fn read<'de, T: FormalKey + TupleUnpack<'de>>( - &self, - key: &'de T, - ) -> Result; -} - -pub trait OptSliceExt { - fn read<'de, T: FormalKey + TupleUnpack<'de>>( - &self, - key: &'de T, - ) -> Result; - fn read_opt<'de, T: FormalKey + TupleUnpack<'de>>( - &self, - key: &'de T, - ) -> Result, udb::FdbBindingError>; -} - -impl SliceExt for udb::future::FdbSlice { - fn read<'de, T: FormalKey + TupleUnpack<'de>>( - &self, - key: &'de T, - ) -> Result { - key.deserialize(self) - .with_context(|| { - format!( - "failed deserializing key value of {}", - std::any::type_name::(), - ) - }) - .map_err(|x| udb::FdbBindingError::CustomError(x.into())) - } -} - -impl OptSliceExt for Option { - fn read<'de, T: FormalKey + TupleUnpack<'de>>( - &self, - key: &'de T, - ) -> Result { - key.deserialize(&self.as_ref().ok_or(udb::FdbBindingError::CustomError( - format!("key should exist: {}", std::any::type_name::()).into(), - ))?) - .with_context(|| { - format!( - "failed deserializing key value of {}", - std::any::type_name::(), - ) - }) - .map_err(|x| udb::FdbBindingError::CustomError(x.into())) - } - - fn read_opt<'de, T: FormalKey + TupleUnpack<'de>>( - &self, - key: &'de T, - ) -> Result, udb::FdbBindingError> { - if let Some(data) = self { - key.deserialize(data) - .map(Some) - .with_context(|| { - format!( - "failed deserializing key value of {}", - std::any::type_name::(), - ) - }) - .map_err(|x| udb::FdbBindingError::CustomError(x.into())) - } else { - Ok(None) - } - } -} - -#[async_trait::async_trait] -pub trait CherryPick { - type Output; - - async fn cherry_pick( - txs: &TxnSubspace<'_>, - subspace: S, - snapshot: bool, - ) -> Result; -} - -// Implements `CherryPick` for any tuple size -macro_rules! impl_tuple { - ($($args:ident),*) => { - #[async_trait::async_trait] - impl<$($args: FormalKey + for<'de> TupleUnpack<'de>),*> CherryPick for ($($args),*) - where - $($args::Value: Send),* - { - type Output = ($($args::Value),*); - - async fn cherry_pick( - txs: &TxnSubspace<'_>, - subspace: S, - snapshot: bool, - ) -> Result { - let subspace = txs.subspace(&subspace); - - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { - mode: StreamingMode::WantAll, - ..(&subspace).into() - }, - snapshot, - ); - - $( - #[allow(non_snake_case)] - let mut $args = None; - )* - - loop { - let Some(entry) = stream.try_next().await? else { - break; - }; - - $( - if let Ok(key) = txs.unpack::<$args>(entry.key()) { - if $args.is_some() { - return Err(udb::FdbBindingError::CustomError( - format!("{} already picked", std::any::type_name::<$args>()).into() - )); - } - - let value = key.read(entry.value())?; - $args = Some(value); - continue; - } - )* - } - - Ok(( - $( - $args.ok_or(udb::FdbBindingError::CustomError( - format!("key not found in cherry pick: {}", std::any::type_name::<$args>()).into(), - ))?, - )* - )) - } - } - } -} - -impl_tuple!(A, B); -impl_tuple!(A, B, C); -impl_tuple!(A, B, C, D); -impl_tuple!(A, B, C, D, E); -impl_tuple!(A, B, C, D, E, F); -impl_tuple!(A, B, C, D, E, F, G); -impl_tuple!(A, B, C, D, E, F, G, H); -impl_tuple!(A, B, C, D, E, F, G, H, I); -impl_tuple!(A, B, C, D, E, F, G, H, I, J); diff --git a/packages/common/universaldb/Cargo.toml b/packages/common/universaldb/Cargo.toml index 8b9277637e..7a86606a08 100644 --- a/packages/common/universaldb/Cargo.toml +++ b/packages/common/universaldb/Cargo.toml @@ -8,20 +8,21 @@ edition.workspace = true [dependencies] anyhow.workspace = true async-trait.workspace = true +deadpool-postgres.workspace = true +foundationdb-tuple.workspace = true futures-util.workspace = true lazy_static.workspace = true rand.workspace = true +rivet-metrics.workspace = true rocksdb.workspace = true -foundationdb-tuple.workspace = true serde.workspace = true -tokio.workspace = true +thiserror.workspace = true tokio-postgres.workspace = true +tokio.workspace = true tracing.workspace = true -deadpool-postgres.workspace = true uuid.workspace = true [dev-dependencies] -udb-util.workspace = true rivet-config.workspace = true rivet-env.workspace = true rivet-pools.workspace = true diff --git a/packages/common/universaldb/src/database.rs b/packages/common/universaldb/src/database.rs index 6288e2de8d..508a638e6b 100644 --- a/packages/common/universaldb/src/database.rs +++ b/packages/common/universaldb/src/database.rs @@ -1,12 +1,12 @@ use std::future::Future; +use anyhow::{Result, anyhow}; use futures_util::FutureExt; -use crate::{FdbBindingError, FdbResult, driver::Erased}; - use crate::{ - MaybeCommitted, RetryableTransaction, Transaction, driver::DatabaseDriverHandle, + driver::{DatabaseDriverHandle, Erased}, options::DatabaseOption, + transaction::{RetryableTransaction, Transaction}, }; #[derive(Clone)] @@ -20,32 +20,32 @@ impl Database { } /// Run a closure with automatic retry logic - pub async fn run<'a, F, Fut, T>(&'a self, closure: F) -> Result + pub async fn run<'a, F, Fut, T>(&'a self, closure: F) -> Result where - F: Fn(RetryableTransaction, MaybeCommitted) -> Fut + Send + Sync, - Fut: Future> + Send, + F: Fn(RetryableTransaction) -> Fut + Send + Sync, + Fut: Future> + Send, T: Send + 'a + 'static, { let closure = &closure; self.driver - .run(Box::new(|tx, mc| { - async move { closure(tx, mc).await.map(|value| Box::new(value) as Erased) }.boxed() + .run(Box::new(|tx| { + async move { closure(tx).await.map(|value| Box::new(value) as Erased) }.boxed() })) .await .and_then(|res| { - res.downcast::().map(|x| *x).map_err(|_| { - FdbBindingError::CustomError("failed to downcast `run` return type".into()) - }) + res.downcast::() + .map(|x| *x) + .map_err(|_| anyhow!("failed to downcast `run` return type")) }) } /// Creates a new txn instance. - pub fn create_trx(&self) -> FdbResult { + pub fn create_trx(&self) -> Result { self.driver.create_trx() } /// Set a database option - pub fn set_option(&self, opt: DatabaseOption) -> FdbResult<()> { + pub fn set_option(&self, opt: DatabaseOption) -> Result<()> { self.driver.set_option(opt) } } diff --git a/packages/common/universaldb/src/driver/mod.rs b/packages/common/universaldb/src/driver/mod.rs index bbf059ea0b..dd8fb11f3f 100644 --- a/packages/common/universaldb/src/driver/mod.rs +++ b/packages/common/universaldb/src/driver/mod.rs @@ -1,11 +1,14 @@ use std::{any::Any, future::Future, pin::Pin, sync::Arc}; +use anyhow::{Result, bail}; + use crate::{ - FdbBindingError, FdbError, FdbResult, KeySelector, RangeOption, RetryableTransaction, - Transaction, - future::{FdbSlice, FdbValues}, + key_selector::KeySelector, options::{ConflictRangeType, DatabaseOption, MutationType}, - types::{MaybeCommitted, TransactionCommitError, TransactionCommitted}, + range_option::RangeOption, + transaction::{RetryableTransaction, Transaction}, + utils::IsolationLevel, + value::{Slice, Value, Values}, }; mod postgres; @@ -20,20 +23,12 @@ pub type Erased = Box; pub type DatabaseDriverHandle = Arc; pub trait DatabaseDriver: Send + Sync { - fn create_trx(&self) -> FdbResult; + fn create_trx(&self) -> Result; fn run<'a>( &'a self, - closure: Box< - dyn Fn( - RetryableTransaction, - MaybeCommitted, - ) -> BoxFut<'a, Result> - + Send - + Sync - + 'a, - >, - ) -> BoxFut<'a, Result>; - fn set_option(&self, opt: DatabaseOption) -> FdbResult<()>; + closure: Box BoxFut<'a, Result> + Send + Sync + 'a>, + ) -> BoxFut<'a, Result>; + fn set_option(&self, opt: DatabaseOption) -> Result<()>; } pub trait TransactionDriver: Send + Sync { @@ -43,24 +38,24 @@ pub trait TransactionDriver: Send + Sync { fn get<'a>( &'a self, key: &[u8], - snapshot: bool, - ) -> Pin>> + Send + 'a>>; + isolation_level: IsolationLevel, + ) -> Pin>> + Send + 'a>>; fn get_key<'a>( &'a self, selector: &KeySelector<'a>, - snapshot: bool, - ) -> Pin> + Send + 'a>>; + isolation_level: IsolationLevel, + ) -> Pin> + Send + 'a>>; fn get_range<'a>( &'a self, opt: &RangeOption<'a>, iteration: usize, - snapshot: bool, - ) -> Pin> + Send + 'a>>; + isolation_level: IsolationLevel, + ) -> Pin> + Send + 'a>>; fn get_ranges_keyvalues<'a>( &'a self, opt: RangeOption<'a>, - snapshot: bool, - ) -> crate::future::FdbStream<'a, crate::future::FdbValue>; + isolation_level: IsolationLevel, + ) -> crate::value::Stream<'a, Value>; // Write operations fn set(&self, key: &[u8], value: &[u8]); @@ -68,9 +63,7 @@ pub trait TransactionDriver: Send + Sync { fn clear_range(&self, begin: &[u8], end: &[u8]); // Transaction management - fn commit( - self: Box, - ) -> Pin> + Send>>; + fn commit(self: Box) -> Pin> + Send>>; fn reset(&mut self); fn cancel(&self); fn add_conflict_range( @@ -78,18 +71,17 @@ pub trait TransactionDriver: Send + Sync { begin: &[u8], end: &[u8], conflict_type: ConflictRangeType, - ) -> FdbResult<()>; + ) -> Result<()>; fn get_estimated_range_size_bytes<'a>( &'a self, begin: &'a [u8], end: &'a [u8], - ) -> Pin> + Send + 'a>>; + ) -> Pin> + Send + 'a>>; // Helper for committing without consuming self (for database drivers that need it) - fn commit_owned(&self) -> Pin> + Send + '_>> { + fn commit_ref(&self) -> Pin> + Send + '_>> { Box::pin(async move { - // Default implementation returns error - drivers that need this should override - Err(FdbError::from_code(1510)) + bail!("`commit_ref` unimplemented"); }) } } diff --git a/packages/common/universaldb/src/driver/postgres/database.rs b/packages/common/universaldb/src/driver/postgres/database.rs index 458d1e128b..4b0d2cc40f 100644 --- a/packages/common/universaldb/src/driver/postgres/database.rs +++ b/packages/common/universaldb/src/driver/postgres/database.rs @@ -1,13 +1,15 @@ use std::sync::{Arc, Mutex}; +use anyhow::{Context, Result}; use deadpool_postgres::{Config, ManagerConfig, Pool, PoolConfig, RecyclingMethod, Runtime}; use tokio_postgres::NoTls; use crate::{ - FdbBindingError, FdbError, FdbResult, MaybeCommitted, RetryableTransaction, Transaction, + RetryableTransaction, Transaction, driver::{BoxFut, DatabaseDriver, Erased}, + error::DatabaseError, options::DatabaseOption, - utils::calculate_tx_retry_backoff, + utils::{MaybeCommitted, calculate_tx_retry_backoff}, }; use super::transaction::PostgresTransactionDriver; @@ -18,7 +20,7 @@ pub struct PostgresDatabaseDriver { } impl PostgresDatabaseDriver { - pub async fn new(connection_string: String) -> FdbResult { + pub async fn new(connection_string: String) -> Result { tracing::debug!(connection_string = ?connection_string, "Creating PostgresDatabaseDriver"); // Create deadpool config from connection string @@ -36,22 +38,19 @@ impl PostgresDatabaseDriver { // Create the pool let pool = config .create_pool(Some(Runtime::Tokio1), NoTls) - .map_err(|e| { - tracing::error!(error = ?e, "Failed to create Postgres pool"); - FdbError::from_code(1510) - })?; + .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.map_err(|e| { - tracing::error!(error = ?e, "Failed to get Postgres connection"); - FdbError::from_code(1510) - })?; + 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 - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to create btree_gist extension")?; // Create the KV table if it doesn't exist conn.execute( @@ -62,7 +61,7 @@ impl PostgresDatabaseDriver { &[], ) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to create kv table")?; // Create range_type type if it doesn't exist conn.execute( @@ -74,7 +73,7 @@ impl PostgresDatabaseDriver { &[], ) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to create range_type enum")?; // Create bytearange type if it doesn't exist conn.execute( @@ -89,7 +88,7 @@ impl PostgresDatabaseDriver { &[], ) .await - .map_err(|_| FdbError::from_code(1510))?; + .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 @@ -110,7 +109,7 @@ impl PostgresDatabaseDriver { &[], ) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to create conflict_ranges table")?; // Connection is automatically returned to the pool when dropped drop(conn); @@ -123,69 +122,60 @@ impl PostgresDatabaseDriver { } impl DatabaseDriver for PostgresDatabaseDriver { - fn create_trx(&self) -> FdbResult { + fn create_trx(&self) -> Result { // Pass the connection pool to the transaction driver - Ok(Transaction::new(Box::new(PostgresTransactionDriver::new( + Ok(Transaction::new(Arc::new(PostgresTransactionDriver::new( self.pool.clone(), )))) } fn run<'a>( &'a self, - closure: Box< - dyn Fn( - RetryableTransaction, - MaybeCommitted, - ) -> BoxFut<'a, Result> - + Send - + Sync - + 'a, - >, - ) -> BoxFut<'a, Result> { + closure: Box BoxFut<'a, Result> + Send + Sync + 'a>, + ) -> BoxFut<'a, Result> { Box::pin(async move { let mut maybe_committed = MaybeCommitted(false); let max_retries = *self.max_retries.lock().unwrap(); for attempt in 0..max_retries { let tx = self.create_trx()?; - let retryable = RetryableTransaction::new(tx); + let mut retryable = RetryableTransaction::new(tx); + retryable.maybe_committed = maybe_committed; // Execute transaction - let result = closure(retryable.clone(), maybe_committed).await; - let fdb_error = match result { - std::result::Result::Ok(res) => { - match retryable.inner.driver.commit_owned().await { - Ok(_) => return Ok(res), - Err(e) => e, - } - } - std::result::Result::Err(e) => { - if let Some(fdb_error) = e.get_fdb_error() { - fdb_error - } else { - return Err(e); - } - } + let error = match closure(retryable.clone()).await { + Ok(res) => match retryable.inner.driver.commit_ref().await { + Ok(_) => return Ok(res), + Err(e) => e, + }, + Err(e) => e, }; - // Handle retry or return error - if fdb_error.is_retryable() { - if fdb_error.is_maybe_committed() { - maybe_committed = MaybeCommitted(true); - } + let chain = error + .chain() + .find_map(|x| x.downcast_ref::()); + + if let Some(db_error) = chain { + // Handle retry or return error + if db_error.is_retryable() { + if db_error.is_maybe_committed() { + maybe_committed = MaybeCommitted(true); + } - let backoff_ms = calculate_tx_retry_backoff(attempt as usize); - tokio::time::sleep(tokio::time::Duration::from_millis(backoff_ms)).await; - } else { - return Err(FdbBindingError::from(fdb_error)); + let backoff_ms = calculate_tx_retry_backoff(attempt as usize); + tokio::time::sleep(tokio::time::Duration::from_millis(backoff_ms)).await; + continue; + } } + + return Err(error); } - Err(FdbBindingError::from(FdbError::from_code(1007))) // Retry limit exceeded + Err(DatabaseError::MaxRetriesReached.into()) }) } - fn set_option(&self, opt: DatabaseOption) -> FdbResult<()> { + fn set_option(&self, opt: DatabaseOption) -> Result<()> { match opt { DatabaseOption::TransactionRetryLimit(limit) => { *self.max_retries.lock().unwrap() = limit; diff --git a/packages/common/universaldb/src/driver/postgres/transaction.rs b/packages/common/universaldb/src/driver/postgres/transaction.rs index 93f3a62177..d275cfbbb5 100644 --- a/packages/common/universaldb/src/driver/postgres/transaction.rs +++ b/packages/common/universaldb/src/driver/postgres/transaction.rs @@ -4,15 +4,19 @@ use std::{ sync::{Arc, Mutex}, }; +use anyhow::{Context, Result}; use deadpool_postgres::Pool; use tokio::sync::{OnceCell, mpsc, oneshot}; use crate::{ - FdbError, FdbResult, KeySelector, RangeOption, TransactionCommitError, TransactionCommitted, driver::TransactionDriver, - future::{FdbSlice, FdbValues}, + error::DatabaseError, + key_selector::KeySelector, options::{ConflictRangeType, MutationType}, + range_option::RangeOption, tx_ops::{Operation, TransactionOperations}, + utils::IsolationLevel, + value::{KeyValue, Slice, Value, Values}, }; use super::transaction_task::{TransactionCommand, TransactionIsolationLevel, TransactionTask}; @@ -49,7 +53,7 @@ impl PostgresTransactionDriver { } /// Get or create the transaction task - async fn ensure_transaction(&self) -> FdbResult<&mpsc::Sender> { + async fn ensure_transaction(&self) -> Result<&mpsc::Sender> { self.tx_sender .get_or_try_init(|| async { let (sender, receiver) = mpsc::channel(100); @@ -62,16 +66,16 @@ impl PostgresTransactionDriver { ); tokio::spawn(task.run()); - Ok(sender) + anyhow::Ok(sender) }) .await - .map_err(|_: anyhow::Error| FdbError::from_code(1510)) + .context("failed to initialize postgres transaction task") } /// Get or create the snapshot transaction task /// This creates a separate REPEATABLE READ READ ONLY transaction /// to enforce reading from a consistent snapshot - async fn ensure_snapshot_transaction(&self) -> FdbResult<&mpsc::Sender> { + async fn ensure_snapshot_transaction(&self) -> Result<&mpsc::Sender> { self.snapshot_tx_sender .get_or_try_init(|| async { let (sender, receiver) = mpsc::channel(100); @@ -84,10 +88,10 @@ impl PostgresTransactionDriver { ); tokio::spawn(task.run()); - Ok(sender) + anyhow::Ok(sender) }) .await - .map_err(|_: anyhow::Error| FdbError::from_code(1510)) + .context("failed to initialize postgres transaction task") } } @@ -101,8 +105,8 @@ impl TransactionDriver for PostgresTransactionDriver { fn get<'a>( &'a self, key: &[u8], - snapshot: bool, - ) -> Pin>> + Send + 'a>> { + isolation_level: IsolationLevel, + ) -> Pin>> + Send + 'a>> { let key = key.to_vec(); Box::pin(async move { // Both snapshot and non-snapshot reads check local operations first @@ -113,7 +117,7 @@ impl TransactionDriver for PostgresTransactionDriver { }; ops.get_with_callback(&key, || async { - let tx_sender = if snapshot { + let tx_sender = if let IsolationLevel::Snapshot = isolation_level { self.ensure_snapshot_transaction().await? } else { self.ensure_transaction().await? @@ -127,12 +131,14 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; // Wait for response - let value = response_rx.await.map_err(|_| FdbError::from_code(1510))??; + let value = response_rx + .await + .context("failed to receive postgres response")??; - Ok(value) + Ok(value.map(Into::into)) }) .await }) @@ -141,8 +147,8 @@ impl TransactionDriver for PostgresTransactionDriver { fn get_key<'a>( &'a self, selector: &KeySelector<'a>, - snapshot: bool, - ) -> Pin> + Send + 'a>> { + isolation_level: IsolationLevel, + ) -> Pin> + Send + 'a>> { let selector = selector.clone(); Box::pin(async move { @@ -158,7 +164,7 @@ impl TransactionDriver for PostgresTransactionDriver { }; ops.get_key(&selector, || async { - let tx_sender = if snapshot { + let tx_sender = if let IsolationLevel::Snapshot = isolation_level { self.ensure_snapshot_transaction().await? } else { self.ensure_transaction().await? @@ -174,13 +180,15 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; // Wait for response - let result_key = response_rx.await.map_err(|_| FdbError::from_code(1510))??; + let result_key = response_rx + .await + .context("failed to receive postgres key selector response")??; // Return the key if found, or empty vector if not - Ok(result_key.unwrap_or_else(Vec::new)) + Ok(result_key.map(Into::into).unwrap_or_else(Slice::new)) }) .await }) @@ -190,8 +198,8 @@ impl TransactionDriver for PostgresTransactionDriver { &'a self, opt: &RangeOption<'a>, _iteration: usize, - snapshot: bool, - ) -> Pin> + Send + 'a>> { + isolation_level: IsolationLevel, + ) -> Pin> + Send + 'a>> { let opt = opt.clone(); Box::pin(async move { @@ -212,7 +220,7 @@ impl TransactionDriver for PostgresTransactionDriver { }; ops.get_range(&opt, || async { - let tx_sender = if snapshot { + let tx_sender = if let IsolationLevel::Snapshot = isolation_level { self.ensure_snapshot_transaction().await? } else { self.ensure_transaction().await? @@ -233,17 +241,19 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; // Wait for response - let keyvalues_data = response_rx.await.map_err(|_| FdbError::from_code(1510))??; + let keyvalues_data = response_rx + .await + .context("failed to receive postgres range response")??; let keyvalues: Vec<_> = keyvalues_data .into_iter() - .map(|(key, value)| crate::future::FdbKeyValue::new(key, value)) + .map(|(key, value)| KeyValue::new(key, value)) .collect(); - Ok(crate::future::FdbValues::new(keyvalues)) + Ok(Values::new(keyvalues)) }) .await }) @@ -252,16 +262,16 @@ impl TransactionDriver for PostgresTransactionDriver { fn get_ranges_keyvalues<'a>( &'a self, opt: RangeOption<'a>, - snapshot: bool, - ) -> crate::future::FdbStream<'a, crate::future::FdbValue> { + isolation_level: IsolationLevel, + ) -> crate::value::Stream<'a, Value> { use futures_util::{StreamExt, stream}; // Convert the range result into a stream let fut = async move { - match self.get_range(&opt, 1, snapshot).await { + match self.get_range(&opt, 1, isolation_level).await { Ok(values) => values .into_iter() - .map(|kv| Ok(crate::future::FdbValue::from_keyvalue(kv))) + .map(|kv| Ok(Value::from_keyvalue(kv))) .collect::>(), Err(e) => vec![Err(e)], } @@ -288,10 +298,7 @@ impl TransactionDriver for PostgresTransactionDriver { } } - fn commit( - self: Box, - ) -> Pin> + Send>> - { + fn commit(self: Box) -> Pin> + Send>> { Box::pin(async move { // Get operations and mark as committed let operations = { @@ -321,16 +328,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })?; + .context("failed to send postgres transaction command")?; response_rx .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })? - .map_err(TransactionCommitError::new)?; + .context("failed to receive postgres response")??; } Operation::Clear { key } => { let (response_tx, response_rx) = oneshot::channel(); @@ -340,16 +342,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })?; + .context("failed to send postgres transaction command")?; response_rx .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })? - .map_err(TransactionCommitError::new)?; + .context("failed to receive postgres response")??; } Operation::ClearRange { begin, end } => { let (response_tx, response_rx) = oneshot::channel(); @@ -360,16 +357,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })?; + .context("failed to send postgres transaction command")?; response_rx .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })? - .map_err(TransactionCommitError::new)?; + .context("failed to receive postgres response")??; } Operation::AtomicOp { key, @@ -385,16 +377,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })?; + .context("failed to send postgres transaction command")?; response_rx .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })? - .map_err(TransactionCommitError::new)?; + .context("failed to receive postgres response")??; } } } @@ -407,20 +394,15 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| TransactionCommitError::new(FdbError::from_code(1510)))?; + .context("failed to send postgres transaction command")?; // Wait for commit response response_rx .await - .map_err(|_| TransactionCommitError::new(FdbError::from_code(1510)))? - .map_err(TransactionCommitError::new)?; + .context("failed to receive postgres commit response")??; } else if !operations.operations().is_empty() { // We have operations but no transaction - create one just for commit - let tx_sender = self - .ensure_transaction() - .await - .map_err(TransactionCommitError::new)?; - + let tx_sender = self.ensure_transaction().await?; // Execute all operations for op in operations.operations() { match op { @@ -433,16 +415,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })?; + .context("failed to send postgres transaction command")?; response_rx .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })? - .map_err(TransactionCommitError::new)?; + .context("failed to receive postgres response")??; } Operation::Clear { key } => { let (response_tx, response_rx) = oneshot::channel(); @@ -452,16 +429,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })?; + .context("failed to send postgres transaction command")?; response_rx .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })? - .map_err(TransactionCommitError::new)?; + .context("failed to receive postgres response")??; } Operation::ClearRange { begin, end } => { let (response_tx, response_rx) = oneshot::channel(); @@ -472,16 +444,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })?; + .context("failed to send postgres transaction command")?; response_rx .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })? - .map_err(TransactionCommitError::new)?; + .context("failed to receive postgres response")??; } Operation::AtomicOp { key, @@ -497,16 +464,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })?; + .context("failed to send postgres transaction command")?; response_rx .await - .map_err(|_| { - TransactionCommitError::new(FdbError::from_code(1510)) - })? - .map_err(TransactionCommitError::new)?; + .context("failed to receive postgres response")??; } } } @@ -519,13 +481,12 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| TransactionCommitError::new(FdbError::from_code(1510)))?; + .context("failed to send postgres transaction command")?; // Wait for commit response response_rx .await - .map_err(|_| TransactionCommitError::new(FdbError::from_code(1510)))? - .map_err(TransactionCommitError::new)?; + .context("failed to receive postgres commit response")??; } Ok(()) @@ -554,7 +515,7 @@ impl TransactionDriver for PostgresTransactionDriver { begin: &[u8], end: &[u8], conflict_type: ConflictRangeType, - ) -> FdbResult<()> { + ) -> Result<()> { // For PostgreSQL, we implement conflict ranges using the conflict_ranges table // This ensures serializable isolation for the specified range @@ -585,31 +546,21 @@ impl TransactionDriver for PostgresTransactionDriver { // Try to send the add conflict range command // Since this is a synchronous method, we use try_send let (response_tx, _response_rx) = oneshot::channel(); - match tx_sender.try_send(TransactionCommand::AddConflictRange { - begin: begin_vec, - end: end_vec, - conflict_type, - response: response_tx, - }) { - Ok(_) => { - // Command sent successfully - // Note: We can't wait for the response in a sync method - // The actual conflict range acquisition will happen asynchronously - Ok(()) - } - Err(_) => { - // Channel is full or closed - // Return an error indicating we couldn't add the conflict range - Err(FdbError::from_code(1020)) // Transaction conflict error - } - } + tx_sender + .try_send(TransactionCommand::AddConflictRange { + begin: begin_vec, + end: end_vec, + conflict_type, + response: response_tx, + }) + .map_err(|_| DatabaseError::NotCommitted.into()) } fn get_estimated_range_size_bytes<'a>( &'a self, begin: &'a [u8], end: &'a [u8], - ) -> Pin> + Send + 'a>> { + ) -> Pin> + Send + 'a>> { let begin = begin.to_vec(); let end = end.to_vec(); @@ -625,16 +576,18 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres command")?; // Wait for response - let size = response_rx.await.map_err(|_| FdbError::from_code(1510))??; + let size = response_rx + .await + .context("failed to receive postgres size response")??; Ok(size) }) } - fn commit_owned(&self) -> Pin> + Send + '_>> { + fn commit_ref(&self) -> Pin> + Send + '_>> { Box::pin(async move { // Get operations and mark as committed let operations = { @@ -664,9 +617,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; - response_rx.await.map_err(|_| FdbError::from_code(1510))??; + response_rx + .await + .context("failed to receive postgres response")??; } Operation::Clear { key } => { let (response_tx, response_rx) = oneshot::channel(); @@ -676,9 +631,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; - response_rx.await.map_err(|_| FdbError::from_code(1510))??; + response_rx + .await + .context("failed to receive postgres response")??; } Operation::ClearRange { begin, end } => { let (response_tx, response_rx) = oneshot::channel(); @@ -689,9 +646,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; - response_rx.await.map_err(|_| FdbError::from_code(1510))??; + response_rx + .await + .context("failed to receive postgres response")??; } Operation::AtomicOp { key, @@ -707,9 +666,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; - response_rx.await.map_err(|_| FdbError::from_code(1510))??; + response_rx + .await + .context("failed to receive postgres response")??; } } } @@ -722,10 +683,12 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; // Wait for commit response - response_rx.await.map_err(|_| FdbError::from_code(1510))??; + response_rx + .await + .context("failed to receive postgres commit response")??; } else if !operations.operations().is_empty() { // We have operations but no transaction - create one just for commit let tx_sender = self.ensure_transaction().await?; @@ -742,9 +705,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; - response_rx.await.map_err(|_| FdbError::from_code(1510))??; + response_rx + .await + .context("failed to receive postgres response")??; } Operation::Clear { key } => { let (response_tx, response_rx) = oneshot::channel(); @@ -754,9 +719,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; - response_rx.await.map_err(|_| FdbError::from_code(1510))??; + response_rx + .await + .context("failed to receive postgres response")??; } Operation::ClearRange { begin, end } => { let (response_tx, response_rx) = oneshot::channel(); @@ -767,9 +734,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; - response_rx.await.map_err(|_| FdbError::from_code(1510))??; + response_rx + .await + .context("failed to receive postgres response")??; } Operation::AtomicOp { key, @@ -785,9 +754,11 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; - response_rx.await.map_err(|_| FdbError::from_code(1510))??; + response_rx + .await + .context("failed to receive postgres response")??; } } } @@ -800,10 +771,12 @@ impl TransactionDriver for PostgresTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send postgres transaction command")?; // Wait for commit response - response_rx.await.map_err(|_| FdbError::from_code(1510))??; + response_rx + .await + .context("failed to receive postgres commit response")??; } Ok(()) diff --git a/packages/common/universaldb/src/driver/postgres/transaction_task.rs b/packages/common/universaldb/src/driver/postgres/transaction_task.rs index 230acbc1d7..3587b791db 100644 --- a/packages/common/universaldb/src/driver/postgres/transaction_task.rs +++ b/packages/common/universaldb/src/driver/postgres/transaction_task.rs @@ -1,10 +1,11 @@ +use anyhow::{Result, anyhow}; use deadpool_postgres::Pool; use tokio::sync::{mpsc, oneshot}; use tokio_postgres::IsolationLevel; use crate::{ - FdbError, FdbResult, atomic::apply_atomic_op, + error::DatabaseError, options::{ConflictRangeType, MutationType}, versionstamp::substitute_versionstamp_if_incomplete, }; @@ -20,13 +21,13 @@ pub enum TransactionCommand { // Read operations Get { key: Vec, - response: oneshot::Sender>>>, + response: oneshot::Sender>>>, }, GetKey { key: Vec, or_equal: bool, offset: i32, - response: oneshot::Sender>>>, + response: oneshot::Sender>>>, }, GetRange { begin: Vec, @@ -37,45 +38,45 @@ pub enum TransactionCommand { end_offset: i32, limit: Option, reverse: bool, - response: oneshot::Sender, Vec)>>>, + response: oneshot::Sender, Vec)>>>, }, // Write operations Set { key: Vec, value: Vec, - response: oneshot::Sender>, + response: oneshot::Sender>, }, Clear { key: Vec, - response: oneshot::Sender>, + response: oneshot::Sender>, }, ClearRange { begin: Vec, end: Vec, - response: oneshot::Sender>, + response: oneshot::Sender>, }, AtomicOp { key: Vec, param: Vec, op_type: MutationType, - response: oneshot::Sender>, + response: oneshot::Sender>, }, // Transaction control Commit { has_conflict_ranges: bool, - response: oneshot::Sender>, + response: oneshot::Sender>, }, // Conflict ranges AddConflictRange { begin: Vec, end: Vec, conflict_type: ConflictRangeType, - response: oneshot::Sender>, + response: oneshot::Sender>, }, GetEstimatedRangeSize { begin: Vec, end: Vec, - response: oneshot::Sender>, + response: oneshot::Sender>, }, } @@ -117,34 +118,44 @@ impl TransactionTask { while let Some(cmd) = self.receiver.recv().await { match cmd { TransactionCommand::Get { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::GetKey { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::GetRange { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::Set { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::Clear { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::ClearRange { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::AtomicOp { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::Commit { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::AddConflictRange { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::GetEstimatedRangeSize { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } } } @@ -176,34 +187,44 @@ impl TransactionTask { while let Some(cmd) = self.receiver.recv().await { match cmd { TransactionCommand::Get { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::GetKey { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::GetRange { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::Set { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::Clear { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::ClearRange { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::AtomicOp { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::Commit { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::AddConflictRange { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } TransactionCommand::GetEstimatedRangeSize { response, .. } => { - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); } } } @@ -262,7 +283,8 @@ impl TransactionTask { _ => { // For other offset values, we need more complex logic // This is a simplified fallback that may not handle all cases perfectly - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); continue; } }; @@ -373,10 +395,12 @@ impl TransactionTask { if let TransactionIsolationLevel::RepeatableReadReadOnly = self.isolation_level { tracing::error!("cannot set in read only txn"); - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = + response.send(Err(anyhow!("postgres transaction connection failed"))); continue; }; + // TODO: versionstamps need to be calculated on the sql side, not in rust let value = substitute_versionstamp_if_incomplete(value, 0); let query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2"; @@ -395,7 +419,8 @@ impl TransactionTask { if let TransactionIsolationLevel::RepeatableReadReadOnly = self.isolation_level { tracing::error!("cannot set in read only txn"); - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = + response.send(Err(anyhow!("postgres transaction connection failed"))); continue; }; @@ -419,7 +444,8 @@ impl TransactionTask { if let TransactionIsolationLevel::RepeatableReadReadOnly = self.isolation_level { tracing::error!("cannot clear range in read only txn"); - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = + response.send(Err(anyhow!("postgres transaction connection failed"))); continue; }; @@ -453,7 +479,8 @@ impl TransactionTask { if let TransactionIsolationLevel::RepeatableReadReadOnly = self.isolation_level { tracing::error!("cannot apply atomic op in read only txn"); - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = + response.send(Err(anyhow!("postgres transaction connection failed"))); continue; }; @@ -513,7 +540,8 @@ impl TransactionTask { self.isolation_level { tracing::error!("cannot release conflict ranges in read only txn"); - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = response + .send(Err(anyhow!("postgres transaction connection failed"))); continue; }; @@ -540,7 +568,8 @@ impl TransactionTask { if let TransactionIsolationLevel::RepeatableReadReadOnly = self.isolation_level { tracing::error!("cannot add conflict range in read only txn"); - let _ = response.send(Err(FdbError::from_code(1510))); + let _ = + response.send(Err(anyhow!("postgres transaction connection failed"))); continue; }; @@ -605,26 +634,26 @@ impl TransactionTask { } } -/// Maps PostgreSQL errors to FdbError codes -fn map_postgres_error(err: tokio_postgres::Error) -> FdbError { +/// Maps PostgreSQL error to DatabaseError +fn map_postgres_error(err: tokio_postgres::Error) -> anyhow::Error { 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 - FdbError::from_code(1020) + DatabaseError::NotCommitted.into() } else if error_str.contains("serialization failure") || error_str.contains("could not serialize") || error_str.contains("deadlock detected") { // Retryable - transaction conflict - FdbError::from_code(1020) + 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 - FdbError::from_code(1020) + DatabaseError::NotCommitted.into() } else { tracing::error!(%err, "postgres error"); // Non-retryable error - FdbError::from_code(1510) + anyhow::Error::new(err) } } diff --git a/packages/common/universaldb/src/driver/rocksdb/conflict_range_tracker.rs b/packages/common/universaldb/src/driver/rocksdb/conflict_range_tracker.rs index 22e2d20183..4f33b395b2 100644 --- a/packages/common/universaldb/src/driver/rocksdb/conflict_range_tracker.rs +++ b/packages/common/universaldb/src/driver/rocksdb/conflict_range_tracker.rs @@ -1,7 +1,9 @@ use std::collections::HashMap; use std::sync::{Arc, RwLock}; -use crate::{FdbError, FdbResult}; +use anyhow::Result; + +use crate::error::DatabaseError; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct TransactionId(u64); @@ -58,7 +60,7 @@ impl ConflictRangeTracker { begin: &[u8], end: &[u8], is_write: bool, - ) -> FdbResult<()> { + ) -> Result<()> { let new_range = ConflictRange { begin: begin.to_vec(), end: end.to_vec(), @@ -77,7 +79,7 @@ impl ConflictRangeTracker { for existing_range in ranges { if new_range.conflicts_with(existing_range) { // Found a conflict - return retryable error - return Err(FdbError::from_code(1020)); + return Err(DatabaseError::NotCommitted.into()); } } } @@ -92,7 +94,7 @@ impl ConflictRangeTracker { begin: &[u8], end: &[u8], is_write: bool, - ) -> FdbResult<()> { + ) -> Result<()> { // First check for conflicts self.check_conflict(tx_id, begin, end, is_write)?; @@ -200,9 +202,12 @@ mod tests { // Try to add overlapping write range for tx2 - should conflict let result = tracker.add_range(tx2, b"b", b"d", true); assert!(result.is_err()); - // Check for conflict error code 1020 + // Check for conflict error if let Err(e) = result { - assert_eq!(e.code(), 1020); + assert!(matches!( + e.downcast::().unwrap(), + DatabaseError::NotCommitted + )); } } @@ -220,9 +225,12 @@ mod tests { // Try to add overlapping write range for tx2 - should conflict let result = tracker.add_range(tx2, b"b", b"d", true); assert!(result.is_err()); - // Check for conflict error code 1020 + // Check for conflict error if let Err(e) = result { - assert_eq!(e.code(), 1020); + assert!(matches!( + e.downcast::().unwrap(), + DatabaseError::NotCommitted + )); } } @@ -264,9 +272,12 @@ mod tests { // Try to add overlapping read range for tx2 - should conflict let result = tracker.add_range(tx2, b"b", b"d", false); assert!(result.is_err()); - // Check for conflict error code 1020 + // Check for conflict error if let Err(e) = result { - assert_eq!(e.code(), 1020); + assert!(matches!( + e.downcast::().unwrap(), + DatabaseError::NotCommitted + )); } } diff --git a/packages/common/universaldb/src/driver/rocksdb/database.rs b/packages/common/universaldb/src/driver/rocksdb/database.rs index 745fb752cb..28575beefd 100644 --- a/packages/common/universaldb/src/driver/rocksdb/database.rs +++ b/packages/common/universaldb/src/driver/rocksdb/database.rs @@ -3,13 +3,15 @@ use std::{ sync::{Arc, Mutex}, }; +use anyhow::{Context, Result}; use rocksdb::{OptimisticTransactionDB, Options}; use crate::{ - FdbBindingError, FdbError, FdbResult, MaybeCommitted, RetryableTransaction, Transaction, + RetryableTransaction, Transaction, driver::{BoxFut, DatabaseDriver, Erased}, + error::DatabaseError, options::DatabaseOption, - utils::calculate_tx_retry_backoff, + utils::{MaybeCommitted, calculate_tx_retry_backoff}, }; use super::{conflict_range_tracker::ConflictRangeTracker, transaction::RocksDbTransactionDriver}; @@ -21,9 +23,9 @@ pub struct RocksDbDatabaseDriver { } impl RocksDbDatabaseDriver { - pub async fn new(db_path: PathBuf) -> FdbResult { + pub async fn new(db_path: PathBuf) -> Result { // Create directory if it doesn't exist - std::fs::create_dir_all(&db_path).map_err(|_| FdbError::from_code(1510))?; + std::fs::create_dir_all(&db_path).context("failed to create database directory")?; // Configure RocksDB options let mut opts = Options::default(); @@ -33,8 +35,7 @@ impl RocksDbDatabaseDriver { opts.set_max_total_wal_size(64 * 1024 * 1024); // 64MB // Open the OptimisticTransactionDB - let db = - OptimisticTransactionDB::open(&opts, db_path).map_err(|_| FdbError::from_code(1510))?; + let db = OptimisticTransactionDB::open(&opts, db_path).context("failed to open rocksdb")?; Ok(RocksDbDatabaseDriver { db: Arc::new(db), @@ -45,8 +46,8 @@ impl RocksDbDatabaseDriver { } impl DatabaseDriver for RocksDbDatabaseDriver { - fn create_trx(&self) -> FdbResult { - Ok(Transaction::new(Box::new(RocksDbTransactionDriver::new( + fn create_trx(&self) -> Result { + Ok(Transaction::new(Arc::new(RocksDbTransactionDriver::new( self.db.clone(), self.conflict_tracker.clone(), )))) @@ -54,61 +55,51 @@ impl DatabaseDriver for RocksDbDatabaseDriver { fn run<'a>( &'a self, - closure: Box< - dyn Fn( - RetryableTransaction, - MaybeCommitted, - ) -> BoxFut<'a, Result> - + Send - + Sync - + 'a, - >, - ) -> BoxFut<'a, Result> { + closure: Box BoxFut<'a, Result> + Send + Sync + 'a>, + ) -> BoxFut<'a, Result> { Box::pin(async move { let mut maybe_committed = MaybeCommitted(false); let max_retries = *self.max_retries.lock().unwrap(); for attempt in 0..max_retries { let tx = self.create_trx()?; - let retryable = RetryableTransaction::new(tx); + let mut retryable = RetryableTransaction::new(tx); + retryable.maybe_committed = maybe_committed; // Execute transaction - let result = closure(retryable.clone(), maybe_committed).await; - let fdb_error = match result { - std::result::Result::Ok(res) => { - match retryable.inner.driver.commit_owned().await { - Ok(_) => return Ok(res), - Err(e) => e, - } - } - std::result::Result::Err(e) => { - if let Some(fdb_error) = e.get_fdb_error() { - fdb_error - } else { - return Err(e); - } - } + let error = match closure(retryable.clone()).await { + Ok(res) => match retryable.inner.driver.commit_ref().await { + Ok(_) => return Ok(res), + Err(e) => e, + }, + Err(e) => e, }; - // Handle retry or return error - if fdb_error.is_retryable() { - if fdb_error.is_maybe_committed() { - maybe_committed = MaybeCommitted(true); - } + let chain = error + .chain() + .find_map(|x| x.downcast_ref::()); - let backoff_ms = calculate_tx_retry_backoff(attempt as usize); - tokio::time::sleep(tokio::time::Duration::from_millis(backoff_ms)).await; - } else { - return Err(FdbBindingError::from(fdb_error)); + if let Some(db_error) = chain { + // Handle retry or return error + if db_error.is_retryable() { + if db_error.is_maybe_committed() { + maybe_committed = MaybeCommitted(true); + } + + let backoff_ms = calculate_tx_retry_backoff(attempt as usize); + tokio::time::sleep(tokio::time::Duration::from_millis(backoff_ms)).await; + continue; + } } + + return Err(error); } - // Max retries exceeded - Err(FdbBindingError::from(FdbError::from_code(1007))) + Err(DatabaseError::MaxRetriesReached.into()) }) } - fn set_option(&self, opt: DatabaseOption) -> FdbResult<()> { + fn set_option(&self, opt: DatabaseOption) -> Result<()> { match opt { DatabaseOption::TransactionRetryLimit(limit) => { *self.max_retries.lock().unwrap() = limit; diff --git a/packages/common/universaldb/src/driver/rocksdb/transaction.rs b/packages/common/universaldb/src/driver/rocksdb/transaction.rs index 9798d1d369..82de40ca56 100644 --- a/packages/common/universaldb/src/driver/rocksdb/transaction.rs +++ b/packages/common/universaldb/src/driver/rocksdb/transaction.rs @@ -4,15 +4,19 @@ use std::{ sync::{Arc, Mutex}, }; +use anyhow::{Context, Result, anyhow}; use rocksdb::OptimisticTransactionDB; use tokio::sync::{OnceCell, mpsc, oneshot}; use crate::{ - FdbError, FdbResult, KeySelector, RangeOption, TransactionCommitError, TransactionCommitted, driver::TransactionDriver, - future::{FdbSlice, FdbValues}, + error::DatabaseError, + key_selector::KeySelector, options::{ConflictRangeType, MutationType}, + range_option::RangeOption, tx_ops::TransactionOperations, + utils::IsolationLevel, + value::{Slice, Value, Values}, }; use super::{ @@ -63,7 +67,7 @@ impl RocksDbTransactionDriver { } /// Get or create the transaction task for non-snapshot operations - async fn ensure_transaction(&self) -> FdbResult<&mpsc::Sender> { + async fn ensure_transaction(&self) -> Result<&mpsc::Sender> { self.tx_sender .get_or_try_init(|| async { let (sender, receiver) = mpsc::channel(100); @@ -76,14 +80,14 @@ impl RocksDbTransactionDriver { ); tokio::spawn(task.run()); - Ok(sender) + anyhow::Ok(sender) }) .await - .map_err(|_: anyhow::Error| FdbError::from_code(1510)) + .context("failed to initialize transaction task") } /// Get or create the transaction task for snapshot operations - async fn ensure_snapshot_transaction(&self) -> FdbResult<&mpsc::Sender> { + async fn ensure_snapshot_transaction(&self) -> Result<&mpsc::Sender> { self.snapshot_tx_sender .get_or_try_init(|| async { let (sender, receiver) = mpsc::channel(100); @@ -96,10 +100,10 @@ impl RocksDbTransactionDriver { ); tokio::spawn(task.run()); - Ok(sender) + anyhow::Ok(sender) }) .await - .map_err(|_: anyhow::Error| FdbError::from_code(1510)) + .context("failed to initialize transaction task") } } @@ -120,8 +124,8 @@ impl TransactionDriver for RocksDbTransactionDriver { fn get<'a>( &'a self, key: &[u8], - snapshot: bool, - ) -> Pin>> + Send + 'a>> { + isolation_level: IsolationLevel, + ) -> Pin>> + Send + 'a>> { let key = key.to_vec(); Box::pin(async move { // Both snapshot and non-snapshot reads check local operations first @@ -132,7 +136,7 @@ impl TransactionDriver for RocksDbTransactionDriver { }; ops.get_with_callback(&key, || async { - if snapshot { + if let IsolationLevel::Snapshot = isolation_level { // For snapshot reads, don't add conflict ranges let tx_sender = self.ensure_snapshot_transaction().await?; @@ -144,10 +148,12 @@ impl TransactionDriver for RocksDbTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send transaction command")?; // Wait for response - let value = response_rx.await.map_err(|_| FdbError::from_code(1510))??; + let value = response_rx + .await + .context("failed to receive transaction response")??; Ok(value) } else { @@ -169,10 +175,12 @@ impl TransactionDriver for RocksDbTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send transaction command")?; // Wait for response - let value = response_rx.await.map_err(|_| FdbError::from_code(1510))??; + let value = response_rx + .await + .context("failed to receive transaction response")??; Ok(value) } @@ -184,8 +192,8 @@ impl TransactionDriver for RocksDbTransactionDriver { fn get_key<'a>( &'a self, selector: &KeySelector<'a>, - snapshot: bool, - ) -> Pin> + Send + 'a>> { + isolation_level: IsolationLevel, + ) -> Pin> + Send + 'a>> { let selector = selector.clone(); Box::pin(async move { @@ -201,7 +209,7 @@ impl TransactionDriver for RocksDbTransactionDriver { }; ops.get_key(&selector, || async { - let tx_sender = if snapshot { + let tx_sender = if let IsolationLevel::Snapshot = isolation_level { self.ensure_snapshot_transaction().await? } else { self.ensure_transaction().await? @@ -217,13 +225,15 @@ impl TransactionDriver for RocksDbTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send commit command")?; // Wait for response - let result_key = response_rx.await.map_err(|_| FdbError::from_code(1510))??; + let result_key = response_rx + .await + .context("failed to receive key selector response")??; // Return the key if found, or empty vector if not - Ok(result_key.unwrap_or_else(Vec::new)) + Ok(result_key.unwrap_or_else(Slice::new)) }) .await }) @@ -233,8 +243,8 @@ impl TransactionDriver for RocksDbTransactionDriver { &'a self, opt: &RangeOption<'a>, iteration: usize, - snapshot: bool, - ) -> Pin> + Send + 'a>> { + isolation_level: IsolationLevel, + ) -> Pin> + Send + 'a>> { // Extract fields from RangeOption for the async closure let opt = opt.clone(); let begin_selector = opt.begin.clone(); @@ -251,7 +261,7 @@ impl TransactionDriver for RocksDbTransactionDriver { }; ops.get_range(&opt, || async { - if snapshot { + if let IsolationLevel::Snapshot = isolation_level { // For snapshot reads, don't add conflict ranges let tx_sender = self.ensure_snapshot_transaction().await?; @@ -271,10 +281,12 @@ impl TransactionDriver for RocksDbTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send transaction command")?; // Wait for response - let values = response_rx.await.map_err(|_| FdbError::from_code(1510))??; + let values = response_rx + .await + .context("failed to receive range response")??; Ok(values) } else { @@ -304,10 +316,12 @@ impl TransactionDriver for RocksDbTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send transaction command")?; // Wait for response - let values = response_rx.await.map_err(|_| FdbError::from_code(1510))??; + let values = response_rx + .await + .context("failed to receive range response")??; Ok(values) } @@ -319,8 +333,8 @@ impl TransactionDriver for RocksDbTransactionDriver { fn get_ranges_keyvalues<'a>( &'a self, opt: RangeOption<'a>, - snapshot: bool, - ) -> crate::future::FdbStream<'a, crate::future::FdbValue> { + isolation_level: IsolationLevel, + ) -> crate::value::Stream<'a, Value> { use futures_util::StreamExt; // Extract the selectors from RangeOption, same as get_range does @@ -332,7 +346,7 @@ impl TransactionDriver for RocksDbTransactionDriver { Box::pin( futures_util::stream::once(async move { // Get the transaction sender based on snapshot mode - let tx_sender = if snapshot { + let tx_sender = if let IsolationLevel::Snapshot = isolation_level { match self.ensure_snapshot_transaction().await { Ok(sender) => sender, Err(e) => return futures_util::stream::iter(vec![Err(e)]), @@ -360,26 +374,25 @@ impl TransactionDriver for RocksDbTransactionDriver { }) .await { - return futures_util::stream::iter(vec![Err(FdbError::from_code(1510))]); + return futures_util::stream::iter(vec![Err(anyhow!( + "failed to send stream command" + ))]); } match response_rx.await { Ok(Ok(result)) => { - // Convert to FdbValues for the stream + // Convert to Values for the stream let values: Vec<_> = result .iter() - .map(|kv| { - Ok(crate::future::FdbValue::new( - kv.key().to_vec(), - kv.value().to_vec(), - )) - }) + .map(|kv| Ok(Value::new(kv.key().to_vec(), kv.value().to_vec()))) .collect(); futures_util::stream::iter(values) } Ok(Err(e)) => futures_util::stream::iter(vec![Err(e)]), - Err(_) => futures_util::stream::iter(vec![Err(FdbError::from_code(1510))]), + Err(_) => futures_util::stream::iter(vec![Err(anyhow!( + "failed to receive stream response" + ))]), } }) .flatten(), @@ -422,16 +435,13 @@ impl TransactionDriver for RocksDbTransactionDriver { state.operations.clear_range(begin, end); } - fn commit( - self: Box, - ) -> Pin> + Send>> - { + fn commit(self: Box) -> Pin> + Send>> { Box::pin(async move { // Get the operations and conflict ranges to commit let operations = { let mut state = self.state.lock().unwrap(); if state.committed { - return Err(TransactionCommitError::new(FdbError::from_code(2017))); + return Err(DatabaseError::UsedDuringCommit.into()); } state.committed = true; @@ -439,10 +449,7 @@ impl TransactionDriver for RocksDbTransactionDriver { }; // Get the transaction sender - let tx_sender = self - .ensure_transaction() - .await - .map_err(|e| TransactionCommitError::new(e))?; + let tx_sender = self.ensure_transaction().await.map_err(|e| e)?; // Send commit command with operations and conflict ranges let (response_tx, response_rx) = oneshot::channel(); @@ -452,19 +459,19 @@ impl TransactionDriver for RocksDbTransactionDriver { response: response_tx, }) .await - .map_err(|_| TransactionCommitError::new(FdbError::from_code(1510)))?; + .context("failed to send commit command")?; // Wait for response let result = response_rx .await - .map_err(|_| TransactionCommitError::new(FdbError::from_code(1510)))?; + .context("failed to receive commit response")?; // Release conflict ranges after successful commit if result.is_ok() { self.conflict_tracker.release_transaction(self.tx_id); } - result.map_err(|e| TransactionCommitError::new(e)) + result.map_err(|e| e) }) } @@ -500,7 +507,7 @@ impl TransactionDriver for RocksDbTransactionDriver { begin: &[u8], end: &[u8], conflict_type: ConflictRangeType, - ) -> FdbResult<()> { + ) -> Result<()> { // Determine if this is a write conflict range let is_write = match conflict_type { ConflictRangeType::Write => true, @@ -523,7 +530,7 @@ impl TransactionDriver for RocksDbTransactionDriver { &'a self, begin: &'a [u8], end: &'a [u8], - ) -> Pin> + Send + 'a>> { + ) -> Pin> + Send + 'a>> { let begin = begin.to_vec(); let end = end.to_vec(); @@ -539,22 +546,24 @@ impl TransactionDriver for RocksDbTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send commit command")?; // Wait for response - let size = response_rx.await.map_err(|_| FdbError::from_code(1510))??; + let size = response_rx + .await + .context("failed to receive size response")??; Ok(size) }) } - fn commit_owned(&self) -> Pin> + Send + '_>> { + fn commit_ref(&self) -> Pin> + Send + '_>> { Box::pin(async move { // Get the operations to commit let operations = { let mut state = self.state.lock().unwrap(); if state.committed { - return Err(FdbError::from_code(2017)); + return Err(DatabaseError::UsedDuringCommit.into()); } state.committed = true; @@ -572,10 +581,12 @@ impl TransactionDriver for RocksDbTransactionDriver { response: response_tx, }) .await - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to send commit command")?; // Wait for response - let result = response_rx.await.map_err(|_| FdbError::from_code(1510))?; + let result = response_rx + .await + .context("failed to receive commit response")?; // Release conflict ranges after successful commit if result.is_ok() { diff --git a/packages/common/universaldb/src/driver/rocksdb/transaction_task.rs b/packages/common/universaldb/src/driver/rocksdb/transaction_task.rs index 6f00173fcc..9d14e7693c 100644 --- a/packages/common/universaldb/src/driver/rocksdb/transaction_task.rs +++ b/packages/common/universaldb/src/driver/rocksdb/transaction_task.rs @@ -1,27 +1,29 @@ use std::sync::Arc; +use anyhow::{Context, Result, bail}; use rocksdb::{ OptimisticTransactionDB, ReadOptions, Transaction as RocksDbTransaction, WriteOptions, }; use tokio::sync::{mpsc, oneshot}; use crate::{ - FdbError, FdbResult, KeySelector, TransactionCommitted, atomic::apply_atomic_op, - future::{FdbKeyValue, FdbSlice, FdbValues}, + error::DatabaseError, + key_selector::KeySelector, tx_ops::{Operation, TransactionOperations}, + value::{KeyValue, Slice, Values}, }; pub enum TransactionCommand { Get { key: Vec, - response: oneshot::Sender>>, + response: oneshot::Sender>>, }, GetKey { key: Vec, or_equal: bool, offset: i32, - response: oneshot::Sender>>, + response: oneshot::Sender>>, }, GetRange { begin_key: Vec, @@ -33,16 +35,16 @@ pub enum TransactionCommand { limit: Option, reverse: bool, iteration: usize, - response: oneshot::Sender>, + response: oneshot::Sender>, }, Commit { operations: TransactionOperations, - response: oneshot::Sender>, + response: oneshot::Sender>, }, GetEstimatedRangeSize { begin: Vec, end: Vec, - response: oneshot::Sender>, + response: oneshot::Sender>, }, Cancel, } @@ -138,16 +140,15 @@ impl TransactionTask { self.db.transaction_opt(&write_opts, &txn_opts) } - async fn handle_get(&mut self, key: &[u8]) -> FdbResult> { + async fn handle_get(&mut self, key: &[u8]) -> Result> { let txn = self.create_transaction(); let read_opts = ReadOptions::default(); - match txn.get_opt(key, &read_opts) { - Ok(Some(value)) => Ok(Some(value)), - Ok(None) => Ok(None), - Err(_) => Err(FdbError::from_code(1510)), - } + Ok(txn + .get_opt(key, &read_opts) + .context("failed to read key from rocksdb")? + .map(|v| v.into())) } async fn handle_get_key( @@ -155,7 +156,7 @@ impl TransactionTask { key: &[u8], or_equal: bool, offset: i32, - ) -> FdbResult> { + ) -> Result> { let txn = self.create_transaction(); let read_opts = ReadOptions::default(); @@ -174,14 +175,9 @@ impl TransactionTask { read_opts, ); for item in iter { - match item { - Ok((k, _v)) => { - return Ok(Some(k.to_vec())); - } - Err(_) => { - return Err(FdbError::from_code(1510)); - } - } + let (k, _v) = + item.context("failed to iterate rocksdb for first_greater_or_equal")?; + return Ok(Some(k.to_vec().into())); } Ok(None) } @@ -192,18 +188,13 @@ impl TransactionTask { read_opts, ); for item in iter { - match item { - Ok((k, _v)) => { - // Skip if it's the exact key - if k.as_ref() == key { - continue; - } - return Ok(Some(k.to_vec())); - } - Err(_) => { - return Err(FdbError::from_code(1510)); - } + let (k, _v) = + item.context("failed to iterate rocksdb for first_greater_than")?; + // Skip if it's the exact key + if k.as_ref() == key { + continue; } + return Ok(Some(k.to_vec().into())); } Ok(None) } @@ -216,16 +207,10 @@ impl TransactionTask { ); for item in iter { - match item { - Ok((k, _v)) => { - // We want strictly less than - if k.as_ref() < key { - return Ok(Some(k.to_vec())); - } - } - Err(_) => { - return Err(FdbError::from_code(1510)); - } + let (k, _v) = item.context("failed to iterate rocksdb for last_less_than")?; + // We want strictly less than + if k.as_ref() < key { + return Ok(Some(k.to_vec().into())); } } Ok(None) @@ -239,23 +224,18 @@ impl TransactionTask { ); for item in iter { - match item { - Ok((k, _v)) => { - // We want less than or equal - if k.as_ref() <= key { - return Ok(Some(k.to_vec())); - } - } - Err(_) => { - return Err(FdbError::from_code(1510)); - } + let (k, _v) = + item.context("failed to iterate rocksdb for last_less_or_equal")?; + // We want less than or equal + if k.as_ref() <= key { + return Ok(Some(k.to_vec().into())); } } Ok(None) } _ => { // For other offset values, return an error - Err(FdbError::from_code(1510)) + bail!("invalid key selector offset") } } } @@ -266,7 +246,7 @@ impl TransactionTask { txn: &RocksDbTransaction, selector: &KeySelector<'_>, _read_opts: &ReadOptions, - ) -> FdbResult> { + ) -> Result> { let key = selector.key(); let offset = selector.offset(); let or_equal = selector.or_equal(); @@ -285,14 +265,10 @@ impl TransactionTask { let mut keys: Vec> = Vec::new(); for item in iter { - match item { - Ok((k, _v)) => { - keys.push(k.to_vec()); - if keys.len() > (offset.abs() + 1) as usize { - break; - } - } - Err(_) => return Err(FdbError::from_code(1510)), + let (k, _v) = item.context("failed to iterate rocksdb for key selector")?; + keys.push(k.to_vec()); + if keys.len() > (offset.abs() + 1) as usize { + break; } } @@ -326,10 +302,7 @@ impl TransactionTask { } } - async fn handle_commit( - &mut self, - operations: TransactionOperations, - ) -> FdbResult { + async fn handle_commit(&mut self, operations: TransactionOperations) -> Result<()> { // Create a new transaction for this commit let txn = self.create_transaction(); @@ -346,10 +319,11 @@ impl TransactionTask { ); txn.put(key, &value) - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to set key in rocksdb")?; } Operation::Clear { key } => { - txn.delete(key).map_err(|_| FdbError::from_code(1510))?; + txn.delete(key) + .context("failed to delete key from rocksdb")?; } Operation::ClearRange { begin, end } => { // RocksDB doesn't have a native clear_range, so we need to iterate and delete @@ -360,15 +334,12 @@ impl TransactionTask { ); for item in iter { - match item { - Ok((k, _v)) => { - if k.as_ref() >= end.as_slice() { - break; - } - txn.delete(&k).map_err(|_| FdbError::from_code(1510))?; - } - Err(_) => return Err(FdbError::from_code(1510)), + let (k, _v) = item.context("failed to iterate rocksdb for clear range")?; + if k.as_ref() >= end.as_slice() { + break; } + txn.delete(&k) + .context("failed to delete key in range from rocksdb")?; } } Operation::AtomicOp { @@ -380,7 +351,7 @@ impl TransactionTask { let read_opts = ReadOptions::default(); let current_value = txn .get_opt(key, &read_opts) - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to get current value for atomic operation")?; // Apply the atomic operation let current_slice = current_value.as_deref(); @@ -389,9 +360,10 @@ impl TransactionTask { // Store the result if let Some(new_value) = &new_value { txn.put(key, new_value) - .map_err(|_| FdbError::from_code(1510))?; + .context("failed to set atomic operation result")?; } else { - txn.delete(key).map_err(|_| FdbError::from_code(1510))?; + txn.delete(key) + .context("failed to delete key after atomic operation")?; } } } @@ -407,10 +379,10 @@ impl TransactionTask { Err(e) => { // Check if this is a conflict error if e.to_string().contains("conflict") { - // Return retryable error code 1020 - Err(FdbError::from_code(1020)) + // Return retryable error + Err(DatabaseError::NotCommitted.into()) } else { - Err(FdbError::from_code(1510)) + Err(e.into()) } } } @@ -427,7 +399,7 @@ impl TransactionTask { limit: Option, reverse: bool, _iteration: usize, - ) -> FdbResult { + ) -> Result { let txn = self.create_transaction(); let read_opts = ReadOptions::default(); @@ -449,20 +421,16 @@ impl TransactionTask { let limit = limit.unwrap_or(usize::MAX); for item in iter { - match item { - Ok((k, v)) => { - // Check if we've reached the end key - if k.as_ref() >= resolved_end.as_slice() { - break; - } + let (k, v) = item.context("failed to iterate rocksdb for get range")?; + // Check if we've reached the end key + if k.as_ref() >= resolved_end.as_slice() { + break; + } - results.push(FdbKeyValue::new(k.to_vec(), v.to_vec())); + results.push(KeyValue::new(k.to_vec(), v.to_vec())); - if results.len() >= limit { - break; - } - } - Err(_) => return Err(FdbError::from_code(1510)), + if results.len() >= limit { + break; } } @@ -471,7 +439,7 @@ impl TransactionTask { results.reverse(); } - Ok(FdbValues::new(results)) + Ok(Values::new(results)) } fn resolve_key_selector_for_range( @@ -480,7 +448,7 @@ impl TransactionTask { key: &[u8], or_equal: bool, offset: i32, - ) -> FdbResult> { + ) -> Result> { // Based on PostgreSQL's interpretation: // (false, 1) => first_greater_or_equal // (true, 1) => first_greater_than @@ -497,14 +465,10 @@ impl TransactionTask { read_opts, ); for item in iter { - match item { - Ok((k, _v)) => { - return Ok(k.to_vec()); - } - Err(_) => { - return Err(FdbError::from_code(1510)); - } - } + let (k, _v) = item.context( + "failed to iterate rocksdb for range selector first_greater_or_equal", + )?; + return Ok(k.to_vec()); } // If no key found, return a key that will make the range empty Ok(vec![0xff; 255]) @@ -516,18 +480,14 @@ impl TransactionTask { read_opts, ); for item in iter { - match item { - Ok((k, _v)) => { - // Skip if it's the exact key - if k.as_ref() == key { - continue; - } - return Ok(k.to_vec()); - } - Err(_) => { - return Err(FdbError::from_code(1510)); - } + let (k, _v) = item.context( + "failed to iterate rocksdb for range selector first_greater_than", + )?; + // Skip if it's the exact key + if k.as_ref() == key { + continue; } + return Ok(k.to_vec()); } // If no key found, return a key that will make the range empty Ok(vec![0xff; 255]) @@ -540,11 +500,7 @@ impl TransactionTask { } } - async fn handle_get_estimated_range_size( - &mut self, - begin: &[u8], - end: &[u8], - ) -> FdbResult { + async fn handle_get_estimated_range_size(&mut self, begin: &[u8], end: &[u8]) -> Result { let range = rocksdb::Range::new(begin, end); Ok(self diff --git a/packages/common/universaldb/src/error.rs b/packages/common/universaldb/src/error.rs new file mode 100644 index 0000000000..04451239e5 --- /dev/null +++ b/packages/common/universaldb/src/error.rs @@ -0,0 +1,32 @@ +#[derive(thiserror::Error, Debug)] +pub enum DatabaseError { + #[error("transaction not committed due to conflict with another transaction")] + NotCommitted, + + // TODO: Implement in rocksdb and postgres drivers + #[error("transaction is too old to perform reads or be committed")] + TransactionTooOld, + + #[error("max number of transaction retries reached")] + MaxRetriesReached, + + #[error("operation issued while a commit was outstanding")] + UsedDuringCommit, + // #[error(transparent)] + // Custom(Box), +} + +impl DatabaseError { + pub fn is_retryable(&self) -> bool { + use DatabaseError::*; + + match self { + NotCommitted | TransactionTooOld | MaxRetriesReached => true, + _ => false, + } + } + + pub fn is_maybe_committed(&self) -> bool { + false + } +} diff --git a/packages/common/universaldb/src/future.rs b/packages/common/universaldb/src/future.rs deleted file mode 100644 index 56497cd845..0000000000 --- a/packages/common/universaldb/src/future.rs +++ /dev/null @@ -1,129 +0,0 @@ -use std::pin::Pin; - -use crate::FdbError; -use futures_util::Stream; - -pub type FdbSlice = Vec; - -#[derive(Debug, Clone)] -pub struct FdbValue(FdbKeyValue); - -impl FdbValue { - pub fn new(key: Vec, value: Vec) -> Self { - FdbValue(FdbKeyValue::new(key, value)) - } - - pub fn from_keyvalue(kv: FdbKeyValue) -> Self { - FdbValue(kv) - } - - pub fn key(&self) -> &[u8] { - self.0.key() - } - - pub fn value(&self) -> &[u8] { - self.0.value() - } - - pub fn into_parts(self) -> (Vec, Vec) { - self.0.into_parts() - } -} - -// FdbValues wraps a Vec to match FoundationDB API -#[derive(Debug, Clone)] -pub struct FdbValues { - values: Vec, - more: bool, -} - -impl FdbValues { - pub fn new(values: Vec) -> Self { - FdbValues { - values, - more: false, - } - } - - pub fn with_more(values: Vec, more: bool) -> Self { - FdbValues { values, more } - } - - pub fn more(&self) -> bool { - self.more - } - - pub fn into_vec(self) -> Vec { - self.values - } - - pub fn len(&self) -> usize { - self.values.len() - } - - pub fn is_empty(&self) -> bool { - self.values.is_empty() - } - - pub fn iter(&self) -> std::slice::Iter<'_, FdbKeyValue> { - self.values.iter() - } - - pub fn into_iter(self) -> std::vec::IntoIter { - self.values.into_iter() - } -} - -// impl Deref for FdbValues { -// type Target = [FdbKeyValue]; -// fn deref(&self) -> &Self::Target { -// &self.values -// } -// } -// impl AsRef<[FdbKeyValue]> for FdbValues { -// fn as_ref(&self) -> &[FdbKeyValue] { -// self.deref() -// } -// } - -// KeyValue type with key() and value() methods -#[derive(Debug, Clone)] -pub struct FdbKeyValue { - key: Vec, - value: Vec, -} - -impl FdbKeyValue { - pub fn new(key: Vec, value: Vec) -> Self { - FdbKeyValue { key, value } - } - - pub fn key(&self) -> &[u8] { - &self.key - } - - pub fn value(&self) -> &[u8] { - &self.value - } - - pub fn into_parts(self) -> (Vec, Vec) { - (self.key, self.value) - } - - pub fn to_value(self) -> FdbValue { - FdbValue::from_keyvalue(self) - } - - pub fn value_ref(&self) -> FdbValue { - FdbValue::from_keyvalue(self.clone()) - } -} - -// Stream type for range queries - generic over item type -pub type FdbStream<'a, T = FdbKeyValue> = - Pin> + Send + 'a>>; - -// UNIMPLEMENTED: -pub type FdbAddress = (); -pub type FdbAddresses = (); -pub type FdbValuesIter = (); diff --git a/packages/common/universaldb/src/inherited/README.md b/packages/common/universaldb/src/inherited/README.md deleted file mode 100644 index 8a554cd3c8..0000000000 --- a/packages/common/universaldb/src/inherited/README.md +++ /dev/null @@ -1,7 +0,0 @@ -# Inherited source - -Files from this folder were copied from the foundationdbrs crate. Since this repo doesn't directly use FDB, we only copied parts that we needed for the UniversalDB api to work. - -This removes the dependency on foundationdb-sys, allowing static compilation with musl. - -Origin: https://github.com/foundationdb-rs/foundationdb-rs at 34955a582e964c42c68717b03f97fd0ea3b3cc02 \ No newline at end of file diff --git a/packages/common/universaldb/src/inherited/error.rs b/packages/common/universaldb/src/inherited/error.rs deleted file mode 100644 index 8ac63b35e2..0000000000 --- a/packages/common/universaldb/src/inherited/error.rs +++ /dev/null @@ -1,1313 +0,0 @@ -// Copyright 2018 foundationdb-rs developers, https://github.com/Clikengo/foundationdb-rs/graphs/contributors -// Copyright 2013-2018 Apple, Inc and the FoundationDB project authors. -// -// Licensed under the Apache License, Version 2.0, or the MIT license , at your option. This file may not be -// copied, modified, or distributed except according to those terms. - -//! Error types for the Fdb crate - -#![allow(non_upper_case_globals)] - -use crate::options; -use crate::tuple::PackError; -use std::fmt; -use std::fmt::{Debug, Display, Formatter}; - -#[allow(non_camel_case_types)] -pub type fdb_error_t = ::std::os::raw::c_int; -#[allow(non_camel_case_types)] -pub type fdb_bool_t = ::std::os::raw::c_int; - -pub const success: i32 = 0; -pub const end_of_stream: i32 = 1; -pub const operation_failed: i32 = 1000; -pub const wrong_shard_server: i32 = 1001; -pub const operation_obsolete: i32 = 1002; -pub const cold_cache_server: i32 = 1003; -pub const timed_out: i32 = 1004; -pub const coordinated_state_conflict: i32 = 1005; -pub const all_alternatives_failed: i32 = 1006; -pub const transaction_too_old: i32 = 1007; -pub const no_more_servers: i32 = 1008; -pub const future_version: i32 = 1009; -pub const movekeys_conflict: i32 = 1010; -pub const tlog_stopped: i32 = 1011; -pub const server_request_queue_full: i32 = 1012; -pub const not_committed: i32 = 1020; -pub const commit_unknown_result: i32 = 1021; -pub const commit_unknown_result_fatal: i32 = 1022; -pub const transaction_cancelled: i32 = 1025; -pub const connection_failed: i32 = 1026; -pub const coordinators_changed: i32 = 1027; -pub const new_coordinators_timed_out: i32 = 1028; -pub const watch_cancelled: i32 = 1029; -pub const request_maybe_delivered: i32 = 1030; -pub const transaction_timed_out: i32 = 1031; -pub const too_many_watches: i32 = 1032; -pub const locality_information_unavailable: i32 = 1033; -pub const watches_disabled: i32 = 1034; -pub const default_error_or: i32 = 1035; -pub const accessed_unreadable: i32 = 1036; -pub const process_behind: i32 = 1037; -pub const database_locked: i32 = 1038; -pub const cluster_version_changed: i32 = 1039; -pub const external_client_already_loaded: i32 = 1040; -pub const lookup_failed: i32 = 1041; -pub const commit_proxy_memory_limit_exceeded: i32 = 1042; -pub const shutdown_in_progress: i32 = 1043; -pub const serialization_failed: i32 = 1044; -pub const connection_unreferenced: i32 = 1048; -pub const connection_idle: i32 = 1049; -pub const disk_adapter_reset: i32 = 1050; -pub const batch_transaction_throttled: i32 = 1051; -pub const dd_cancelled: i32 = 1052; -pub const dd_not_found: i32 = 1053; -pub const wrong_connection_file: i32 = 1054; -pub const version_already_compacted: i32 = 1055; -pub const local_config_changed: i32 = 1056; -pub const failed_to_reach_quorum: i32 = 1057; -pub const unsupported_format_version: i32 = 1058; -pub const unknown_change_feed: i32 = 1059; -pub const change_feed_not_registered: i32 = 1060; -pub const granule_assignment_conflict: i32 = 1061; -pub const change_feed_cancelled: i32 = 1062; -pub const blob_granule_file_load_error: i32 = 1063; -pub const blob_granule_transaction_too_old: i32 = 1064; -pub const blob_manager_replaced: i32 = 1065; -pub const change_feed_popped: i32 = 1066; -pub const remote_kvs_cancelled: i32 = 1067; -pub const page_header_wrong_page_id: i32 = 1068; -pub const page_header_checksum_failed: i32 = 1069; -pub const page_header_version_not_supported: i32 = 1070; -pub const page_encoding_not_supported: i32 = 1071; -pub const page_decoding_failed: i32 = 1072; -pub const unexpected_encoding_type: i32 = 1073; -pub const encryption_key_not_found: i32 = 1074; -pub const data_move_cancelled: i32 = 1075; -pub const data_move_dest_team_not_found: i32 = 1076; -pub const blob_worker_full: i32 = 1077; -pub const grv_proxy_memory_limit_exceeded: i32 = 1078; -pub const blob_granule_request_failed: i32 = 1079; -pub const storage_too_many_feed_streams: i32 = 1080; -pub const storage_engine_not_initialized: i32 = 1081; -pub const unknown_storage_engine: i32 = 1082; -pub const duplicate_snapshot_request: i32 = 1083; -pub const dd_config_changed: i32 = 1084; -pub const consistency_check_urgent_task_failed: i32 = 1085; -pub const data_move_conflict: i32 = 1086; -pub const consistency_check_urgent_duplicate_request: i32 = 1087; -pub const broken_promise: i32 = 1100; -pub const operation_cancelled: i32 = 1101; -pub const future_released: i32 = 1102; -pub const connection_leaked: i32 = 1103; -pub const never_reply: i32 = 1104; -pub const retry: i32 = 1105; -pub const recruitment_failed: i32 = 1200; -pub const move_to_removed_server: i32 = 1201; -pub const worker_removed: i32 = 1202; -pub const cluster_recovery_failed: i32 = 1203; -pub const master_max_versions_in_flight: i32 = 1204; -pub const tlog_failed: i32 = 1205; -pub const worker_recovery_failed: i32 = 1206; -pub const please_reboot: i32 = 1207; -pub const please_reboot_delete: i32 = 1208; -pub const commit_proxy_failed: i32 = 1209; -pub const resolver_failed: i32 = 1210; -pub const server_overloaded: i32 = 1211; -pub const backup_worker_failed: i32 = 1212; -pub const tag_throttled: i32 = 1213; -pub const grv_proxy_failed: i32 = 1214; -pub const dd_tracker_cancelled: i32 = 1215; -pub const failed_to_progress: i32 = 1216; -pub const invalid_cluster_id: i32 = 1217; -pub const restart_cluster_controller: i32 = 1218; -pub const please_reboot_kv_store: i32 = 1219; -pub const incompatible_software_version: i32 = 1220; -pub const audit_storage_failed: i32 = 1221; -pub const audit_storage_exceeded_request_limit: i32 = 1222; -pub const proxy_tag_throttled: i32 = 1223; -pub const key_value_store_deadline_exceeded: i32 = 1224; -pub const storage_quota_exceeded: i32 = 1225; -pub const audit_storage_error: i32 = 1226; -pub const master_failed: i32 = 1227; -pub const test_failed: i32 = 1228; -pub const retry_clean_up_datamove_tombstone_added: i32 = 1229; -pub const persist_new_audit_metadata_error: i32 = 1230; -pub const cancel_audit_storage_failed: i32 = 1231; -pub const audit_storage_cancelled: i32 = 1232; -pub const location_metadata_corruption: i32 = 1233; -pub const audit_storage_task_outdated: i32 = 1234; -pub const transaction_throttled_hot_shard: i32 = 1235; -pub const storage_replica_comparison_error: i32 = 1236; -pub const unreachable_storage_replica: i32 = 1237; -pub const bulkload_task_failed: i32 = 1238; -pub const bulkload_task_outdated: i32 = 1239; -pub const range_lock_failed: i32 = 1241; -pub const transaction_rejected_range_locked: i32 = 1242; -pub const bulkdump_task_failed: i32 = 1243; -pub const bulkdump_task_outdated: i32 = 1244; -pub const bulkload_fileset_invalid_filepath: i32 = 1245; -pub const bulkload_manifest_decode_error: i32 = 1246; -pub const range_lock_reject: i32 = 1247; -pub const range_unlock_reject: i32 = 1248; -pub const bulkload_dataset_not_cover_required_range: i32 = 1249; -pub const platform_error: i32 = 1500; -pub const large_alloc_failed: i32 = 1501; -pub const performance_counter_error: i32 = 1502; -pub const bad_allocator: i32 = 1503; -pub const io_error: i32 = 1510; -pub const file_not_found: i32 = 1511; -pub const bind_failed: i32 = 1512; -pub const file_not_readable: i32 = 1513; -pub const file_not_writable: i32 = 1514; -pub const no_cluster_file_found: i32 = 1515; -pub const file_too_large: i32 = 1516; -pub const non_sequential_op: i32 = 1517; -pub const http_bad_response: i32 = 1518; -pub const http_not_accepted: i32 = 1519; -pub const checksum_failed: i32 = 1520; -pub const io_timeout: i32 = 1521; -pub const file_corrupt: i32 = 1522; -pub const http_request_failed: i32 = 1523; -pub const http_auth_failed: i32 = 1524; -pub const http_bad_request_id: i32 = 1525; -pub const rest_invalid_uri: i32 = 1526; -pub const rest_invalid_rest_client_knob: i32 = 1527; -pub const rest_connectpool_key_not_found: i32 = 1528; -pub const lock_file_failure: i32 = 1529; -pub const rest_unsupported_protocol: i32 = 1530; -pub const rest_malformed_response: i32 = 1531; -pub const rest_max_base_cipher_len: i32 = 1532; -pub const resource_not_found: i32 = 1533; -pub const client_invalid_operation: i32 = 2000; -pub const commit_read_incomplete: i32 = 2002; -pub const test_specification_invalid: i32 = 2003; -pub const key_outside_legal_range: i32 = 2004; -pub const inverted_range: i32 = 2005; -pub const invalid_option_value: i32 = 2006; -pub const invalid_option: i32 = 2007; -pub const network_not_setup: i32 = 2008; -pub const network_already_setup: i32 = 2009; -pub const read_version_already_set: i32 = 2010; -pub const version_invalid: i32 = 2011; -pub const range_limits_invalid: i32 = 2012; -pub const invalid_database_name: i32 = 2013; -pub const attribute_not_found: i32 = 2014; -pub const future_not_set: i32 = 2015; -pub const future_not_error: i32 = 2016; -pub const used_during_commit: i32 = 2017; -pub const invalid_mutation_type: i32 = 2018; -pub const attribute_too_large: i32 = 2019; -pub const transaction_invalid_version: i32 = 2020; -pub const no_commit_version: i32 = 2021; -pub const environment_variable_network_option_failed: i32 = 2022; -pub const transaction_read_only: i32 = 2023; -pub const invalid_cache_eviction_policy: i32 = 2024; -pub const network_cannot_be_restarted: i32 = 2025; -pub const blocked_from_network_thread: i32 = 2026; -pub const invalid_config_db_range_read: i32 = 2027; -pub const invalid_config_db_key: i32 = 2028; -pub const invalid_config_path: i32 = 2029; -pub const mapper_bad_index: i32 = 2030; -pub const mapper_no_such_key: i32 = 2031; -pub const mapper_bad_range_decriptor: i32 = 2032; -pub const quick_get_key_values_has_more: i32 = 2033; -pub const quick_get_value_miss: i32 = 2034; -pub const quick_get_key_values_miss: i32 = 2035; -pub const blob_granule_no_ryw: i32 = 2036; -pub const blob_granule_not_materialized: i32 = 2037; -pub const get_mapped_key_values_has_more: i32 = 2038; -pub const get_mapped_range_reads_your_writes: i32 = 2039; -pub const checkpoint_not_found: i32 = 2040; -pub const key_not_tuple: i32 = 2041; -pub const value_not_tuple: i32 = 2042; -pub const mapper_not_tuple: i32 = 2043; -pub const invalid_checkpoint_format: i32 = 2044; -pub const invalid_throttle_quota_value: i32 = 2045; -pub const failed_to_create_checkpoint: i32 = 2046; -pub const failed_to_restore_checkpoint: i32 = 2047; -pub const failed_to_create_checkpoint_shard_metadata: i32 = 2048; -pub const address_parse_error: i32 = 2049; -pub const incompatible_protocol_version: i32 = 2100; -pub const transaction_too_large: i32 = 2101; -pub const key_too_large: i32 = 2102; -pub const value_too_large: i32 = 2103; -pub const connection_string_invalid: i32 = 2104; -pub const address_in_use: i32 = 2105; -pub const invalid_local_address: i32 = 2106; -pub const tls_error: i32 = 2107; -pub const unsupported_operation: i32 = 2108; -pub const too_many_tags: i32 = 2109; -pub const tag_too_long: i32 = 2110; -pub const too_many_tag_throttles: i32 = 2111; -pub const special_keys_cross_module_read: i32 = 2112; -pub const special_keys_no_module_found: i32 = 2113; -pub const special_keys_write_disabled: i32 = 2114; -pub const special_keys_no_write_module_found: i32 = 2115; -pub const special_keys_cross_module_clear: i32 = 2116; -pub const special_keys_api_failure: i32 = 2117; -pub const client_lib_invalid_metadata: i32 = 2118; -pub const client_lib_already_exists: i32 = 2119; -pub const client_lib_not_found: i32 = 2120; -pub const client_lib_not_available: i32 = 2121; -pub const client_lib_invalid_binary: i32 = 2122; -pub const no_external_client_provided: i32 = 2123; -pub const all_external_clients_failed: i32 = 2124; -pub const incompatible_client: i32 = 2125; -pub const tenant_name_required: i32 = 2130; -pub const tenant_not_found: i32 = 2131; -pub const tenant_already_exists: i32 = 2132; -pub const tenant_not_empty: i32 = 2133; -pub const invalid_tenant_name: i32 = 2134; -pub const tenant_prefix_allocator_conflict: i32 = 2135; -pub const tenants_disabled: i32 = 2136; -pub const illegal_tenant_access: i32 = 2138; -pub const invalid_tenant_group_name: i32 = 2139; -pub const invalid_tenant_configuration: i32 = 2140; -pub const cluster_no_capacity: i32 = 2141; -pub const tenant_removed: i32 = 2142; -pub const invalid_tenant_state: i32 = 2143; -pub const tenant_locked: i32 = 2144; -pub const invalid_cluster_name: i32 = 2160; -pub const invalid_metacluster_operation: i32 = 2161; -pub const cluster_already_exists: i32 = 2162; -pub const cluster_not_found: i32 = 2163; -pub const cluster_not_empty: i32 = 2164; -pub const cluster_already_registered: i32 = 2165; -pub const metacluster_no_capacity: i32 = 2166; -pub const management_cluster_invalid_access: i32 = 2167; -pub const tenant_creation_permanently_failed: i32 = 2168; -pub const cluster_removed: i32 = 2169; -pub const cluster_restoring: i32 = 2170; -pub const invalid_data_cluster: i32 = 2171; -pub const metacluster_mismatch: i32 = 2172; -pub const conflicting_restore: i32 = 2173; -pub const invalid_metacluster_configuration: i32 = 2174; -pub const unsupported_metacluster_version: i32 = 2175; -pub const api_version_unset: i32 = 2200; -pub const api_version_already_set: i32 = 2201; -pub const api_version_invalid: i32 = 2202; -pub const api_version_not_supported: i32 = 2203; -pub const api_function_missing: i32 = 2204; -pub const exact_mode_without_limits: i32 = 2210; -pub const invalid_tuple_data_type: i32 = 2250; -pub const invalid_tuple_index: i32 = 2251; -pub const key_not_in_subspace: i32 = 2252; -pub const manual_prefixes_not_enabled: i32 = 2253; -pub const prefix_in_partition: i32 = 2254; -pub const cannot_open_root_directory: i32 = 2255; -pub const directory_already_exists: i32 = 2256; -pub const directory_does_not_exist: i32 = 2257; -pub const parent_directory_does_not_exist: i32 = 2258; -pub const mismatched_layer: i32 = 2259; -pub const invalid_directory_layer_metadata: i32 = 2260; -pub const cannot_move_directory_between_partitions: i32 = 2261; -pub const cannot_use_partition_as_subspace: i32 = 2262; -pub const incompatible_directory_version: i32 = 2263; -pub const directory_prefix_not_empty: i32 = 2264; -pub const directory_prefix_in_use: i32 = 2265; -pub const invalid_destination_directory: i32 = 2266; -pub const cannot_modify_root_directory: i32 = 2267; -pub const invalid_uuid_size: i32 = 2268; -pub const invalid_versionstamp_size: i32 = 2269; -pub const backup_error: i32 = 2300; -pub const restore_error: i32 = 2301; -pub const backup_duplicate: i32 = 2311; -pub const backup_unneeded: i32 = 2312; -pub const backup_bad_block_size: i32 = 2313; -pub const backup_invalid_url: i32 = 2314; -pub const backup_invalid_info: i32 = 2315; -pub const backup_cannot_expire: i32 = 2316; -pub const backup_auth_missing: i32 = 2317; -pub const backup_auth_unreadable: i32 = 2318; -pub const backup_does_not_exist: i32 = 2319; -pub const backup_not_filterable_with_key_ranges: i32 = 2320; -pub const backup_not_overlapped_with_keys_filter: i32 = 2321; -pub const bucket_not_in_url: i32 = 2322; -pub const backup_parse_s3_response_failure: i32 = 2323; -pub const restore_invalid_version: i32 = 2361; -pub const restore_corrupted_data: i32 = 2362; -pub const restore_missing_data: i32 = 2363; -pub const restore_duplicate_tag: i32 = 2364; -pub const restore_unknown_tag: i32 = 2365; -pub const restore_unknown_file_type: i32 = 2366; -pub const restore_unsupported_file_version: i32 = 2367; -pub const restore_bad_read: i32 = 2368; -pub const restore_corrupted_data_padding: i32 = 2369; -pub const restore_destination_not_empty: i32 = 2370; -pub const restore_duplicate_uid: i32 = 2371; -pub const task_invalid_version: i32 = 2381; -pub const task_interrupted: i32 = 2382; -pub const invalid_encryption_key_file: i32 = 2383; -pub const blob_restore_missing_logs: i32 = 2384; -pub const blob_restore_corrupted_logs: i32 = 2385; -pub const blob_restore_invalid_manifest_url: i32 = 2386; -pub const blob_restore_corrupted_manifest: i32 = 2387; -pub const blob_restore_missing_manifest: i32 = 2388; -pub const blob_migrator_replaced: i32 = 2389; -pub const key_not_found: i32 = 2400; -pub const json_malformed: i32 = 2401; -pub const json_eof_expected: i32 = 2402; -pub const snap_disable_tlog_pop_failed: i32 = 2500; -pub const snap_storage_failed: i32 = 2501; -pub const snap_tlog_failed: i32 = 2502; -pub const snap_coord_failed: i32 = 2503; -pub const snap_enable_tlog_pop_failed: i32 = 2504; -pub const snap_path_not_whitelisted: i32 = 2505; -pub const snap_not_fully_recovered_unsupported: i32 = 2506; -pub const snap_log_anti_quorum_unsupported: i32 = 2507; -pub const snap_with_recovery_unsupported: i32 = 2508; -pub const snap_invalid_uid_string: i32 = 2509; -pub const encrypt_ops_error: i32 = 2700; -pub const encrypt_header_metadata_mismatch: i32 = 2701; -pub const encrypt_key_not_found: i32 = 2702; -pub const encrypt_key_ttl_expired: i32 = 2703; -pub const encrypt_header_authtoken_mismatch: i32 = 2704; -pub const encrypt_update_cipher: i32 = 2705; -pub const encrypt_invalid_id: i32 = 2706; -pub const encrypt_keys_fetch_failed: i32 = 2707; -pub const encrypt_invalid_kms_config: i32 = 2708; -pub const encrypt_unsupported: i32 = 2709; -pub const encrypt_mode_mismatch: i32 = 2710; -pub const encrypt_key_check_value_mismatch: i32 = 2711; -pub const encrypt_max_base_cipher_len: i32 = 2712; -pub const unknown_error: i32 = 4000; -pub const internal_error: i32 = 4100; -pub const not_implemented: i32 = 4200; -pub const permission_denied: i32 = 6000; -pub const unauthorized_attempt: i32 = 6001; -pub const digital_signature_ops_error: i32 = 6002; -pub const authorization_token_verify_failed: i32 = 6003; -pub const pkey_decode_error: i32 = 6004; -pub const pkey_encode_error: i32 = 6005; -pub const grpc_error: i32 = 7000; - -pub fn fdb_get_error(code: fdb_error_t) -> &'static str { - if code == success { - "Success" - } else if code == end_of_stream { - "End of stream" - } else if code == operation_failed { - "Operation failed" - } else if code == wrong_shard_server { - "Shard is not available from this server" - } else if code == operation_obsolete { - "Operation result no longer necessary" - } else if code == cold_cache_server { - "Cache server is not warm for this range" - } else if code == timed_out { - "Operation timed out" - } else if code == coordinated_state_conflict { - "Conflict occurred while changing coordination information" - } else if code == all_alternatives_failed { - "All alternatives failed" - } else if code == transaction_too_old { - "Transaction is too old to perform reads or be committed" - } else if code == no_more_servers { - "Not enough physical servers available" - } else if code == future_version { - "Request for future version" - } else if code == movekeys_conflict { - "Conflicting attempts to change data distribution" - } else if code == tlog_stopped { - "TLog stopped" - } else if code == server_request_queue_full { - "Server request queue is full" - } else if code == not_committed { - "Transaction not committed due to conflict with another transaction" - } else if code == commit_unknown_result { - "Transaction may or may not have committed" - } else if code == commit_unknown_result_fatal { - "Idempotency id for transaction may have expired, so the commit status of the transaction cannot be determined" - } else if code == transaction_cancelled { - "Operation aborted because the transaction was cancelled" - } else if code == connection_failed { - "Network connection failed" - } else if code == coordinators_changed { - "Coordination servers have changed" - } else if code == new_coordinators_timed_out { - "New coordination servers did not respond in a timely way" - } else if code == watch_cancelled { - "Watch cancelled because storage server watch limit exceeded" - } else if code == request_maybe_delivered { - "Request may or may not have been delivered" - } else if code == transaction_timed_out { - "Operation aborted because the transaction timed out" - } else if code == too_many_watches { - "Too many watches currently set" - } else if code == locality_information_unavailable { - "Locality information not available" - } else if code == watches_disabled { - "Watches cannot be set if read your writes is disabled" - } else if code == default_error_or { - "Default error for an ErrorOr object" - } else if code == accessed_unreadable { - "Read or wrote an unreadable key" - } else if code == process_behind { - "Storage process does not have recent mutations" - } else if code == database_locked { - "Database is locked" - } else if code == cluster_version_changed { - "The protocol version of the cluster has changed" - } else if code == external_client_already_loaded { - "External client has already been loaded" - } else if code == lookup_failed { - "DNS lookup failed" - } else if code == commit_proxy_memory_limit_exceeded { - "CommitProxy commit memory limit exceeded" - } else if code == shutdown_in_progress { - "Operation no longer supported due to shutdown" - } else if code == serialization_failed { - "Failed to deserialize an object" - } else if code == connection_unreferenced { - "No peer references for connection" - } else if code == connection_idle { - "Connection closed after idle timeout" - } else if code == disk_adapter_reset { - "The disk queue adapter reset" - } else if code == batch_transaction_throttled { - "Batch GRV request rate limit exceeded" - } else if code == dd_cancelled { - "Data distribution components cancelled" - } else if code == dd_not_found { - "Data distributor not found" - } else if code == wrong_connection_file { - "Connection file mismatch" - } else if code == version_already_compacted { - "The requested changes have been compacted away" - } else if code == local_config_changed { - "Local configuration file has changed. Restart and apply these changes" - } else if code == failed_to_reach_quorum { - "Failed to reach quorum from configuration database nodes. Retry sending these requests" - } else if code == unsupported_format_version { - "Format version not supported" - } else if code == unknown_change_feed { - "Change feed not found" - } else if code == change_feed_not_registered { - "Change feed not registered" - } else if code == granule_assignment_conflict { - "Conflicting attempts to assign blob granules" - } else if code == change_feed_cancelled { - "Change feed was cancelled" - } else if code == blob_granule_file_load_error { - "Error loading a blob file during granule materialization" - } else if code == blob_granule_transaction_too_old { - "Read version is older than blob granule history supports" - } else if code == blob_manager_replaced { - "This blob manager has been replaced." - } else if code == change_feed_popped { - "Tried to read a version older than what has been popped from the change feed" - } else if code == remote_kvs_cancelled { - "The remote key-value store is cancelled" - } else if code == page_header_wrong_page_id { - "Page header does not match location on disk" - } else if code == page_header_checksum_failed { - "Page header checksum failed" - } else if code == page_header_version_not_supported { - "Page header version is not supported" - } else if code == page_encoding_not_supported { - "Page encoding type is not supported or not valid" - } else if code == page_decoding_failed { - "Page content decoding failed" - } else if code == unexpected_encoding_type { - "Page content decoding failed" - } else if code == encryption_key_not_found { - "Encryption key not found" - } else if code == data_move_cancelled { - "Data move was cancelled" - } else if code == data_move_dest_team_not_found { - "Dest team was not found for data move" - } else if code == blob_worker_full { - "Blob worker cannot take on more granule assignments" - } else if code == grv_proxy_memory_limit_exceeded { - "GetReadVersion proxy memory limit exceeded" - } else if code == blob_granule_request_failed { - "BlobGranule request failed" - } else if code == storage_too_many_feed_streams { - "Too many feed streams to a single storage server" - } else if code == storage_engine_not_initialized { - "Storage engine was never successfully initialized." - } else if code == unknown_storage_engine { - "Storage engine type is not recognized." - } else if code == duplicate_snapshot_request { - "A duplicate snapshot request has been sent, the old request is discarded." - } else if code == dd_config_changed { - "DataDistribution configuration changed." - } else if code == consistency_check_urgent_task_failed { - "Consistency check urgent task is failed" - } else if code == data_move_conflict { - "Data move conflict in SS" - } else if code == consistency_check_urgent_duplicate_request { - "Consistency check urgent got a duplicate request" - } else if code == broken_promise { - "Broken promise" - } else if code == operation_cancelled { - "Asynchronous operation cancelled" - } else if code == future_released { - "Future has been released" - } else if code == connection_leaked { - "Connection object leaked" - } else if code == never_reply { - "Never reply to the request" - } else if code == retry { - "Retry operation" - } - // Be careful, catching this will delete the data of a storage server or tlog permanently - else if code == recruitment_failed { - "Recruitment of a server failed" - } else if code == move_to_removed_server { - "Attempt to move keys to a storage server that was removed" - } - // Be careful, catching this will delete the data of a storage server or tlog permanently - else if code == worker_removed { - "Normal worker shut down" - } else if code == cluster_recovery_failed { - "Cluster recovery failed" - } else if code == master_max_versions_in_flight { - "Master hit maximum number of versions in flight" - } - // similar to tlog_stopped, but the tlog has actually died - else if code == tlog_failed { - "Cluster recovery terminating because a TLog failed" - } else if code == worker_recovery_failed { - "Recovery of a worker process failed" - } else if code == please_reboot { - "Reboot of server process requested" - } else if code == please_reboot_delete { - "Reboot of server process requested, with deletion of state" - } else if code == commit_proxy_failed { - "Master terminating because a CommitProxy failed" - } else if code == resolver_failed { - "Cluster recovery terminating because a Resolver failed" - } else if code == server_overloaded { - "Server is under too much load and cannot respond" - } else if code == backup_worker_failed { - "Cluster recovery terminating because a backup worker failed" - } else if code == tag_throttled { - "Transaction tag is being throttled" - } else if code == grv_proxy_failed { - "Cluster recovery terminating because a GRVProxy failed" - } else if code == dd_tracker_cancelled { - "The data distribution tracker has been cancelled" - } else if code == failed_to_progress { - "Process has failed to make sufficient progress" - } else if code == invalid_cluster_id { - "Attempted to join cluster with a different cluster ID" - } else if code == restart_cluster_controller { - "Restart cluster controller process" - } else if code == please_reboot_kv_store { - "Need to reboot the storage engine" - } else if code == incompatible_software_version { - "Current software does not support database format" - } else if code == audit_storage_failed { - "Validate storage consistency operation failed" - } else if code == audit_storage_exceeded_request_limit { - "Exceeded the max number of allowed concurrent audit storage requests" - } else if code == proxy_tag_throttled { - "Exceeded maximum proxy tag throttling duration" - } else if code == key_value_store_deadline_exceeded { - "Exceeded maximum time allowed to read or write." - } else if code == storage_quota_exceeded { - "Exceeded the maximum storage quota allocated to the tenant." - } else if code == audit_storage_error { - "Found data corruption" - } else if code == master_failed { - "Cluster recovery terminating because master has failed" - } else if code == test_failed { - "Test failed" - } else if code == retry_clean_up_datamove_tombstone_added { - "Need background datamove cleanup" - } else if code == persist_new_audit_metadata_error { - "Persist new audit metadata error" - } else if code == cancel_audit_storage_failed { - "Failed to cancel an audit" - } else if code == audit_storage_cancelled { - "Audit has been cancelled" - } else if code == location_metadata_corruption { - "Found location metadata corruption" - } else if code == audit_storage_task_outdated { - "Audit task is scheduled by an outdated DD" - } else if code == transaction_throttled_hot_shard { - "Transaction throttled due to hot shard" - } else if code == storage_replica_comparison_error { - "Storage replicas not consistent" - } else if code == unreachable_storage_replica { - "Storage replica cannot be reached" - } else if code == bulkload_task_failed { - "Bulk loading task failed" - } else if code == bulkload_task_outdated { - "Bulk loading task outdated" - } else if code == range_lock_failed { - "Lock range failed" - } else if code == transaction_rejected_range_locked { - "Transaction rejected due to range lock" - } else if code == bulkdump_task_failed { - "Bulk dumping task failed" - } else if code == bulkdump_task_outdated { - "Bulk dumping task outdated" - } else if code == bulkload_fileset_invalid_filepath { - "Bulkload fileset provides invalid filepath" - } else if code == bulkload_manifest_decode_error { - "Bulkload manifest string is failed to decode" - } else if code == range_lock_reject { - "Range lock is rejected" - } else if code == range_unlock_reject { - "Range unlock is rejected" - } else if code == bulkload_dataset_not_cover_required_range { - "Bulkload dataset does not cover the required range" - } - // 15xx Platform errors - else if code == platform_error { - "Platform error" - } else if code == large_alloc_failed { - "Large block allocation failed" - } else if code == performance_counter_error { - "QueryPerformanceCounter error" - } else if code == bad_allocator { - "Null allocator was used to allocate memory" - } else if code == io_error { - "Disk i/o operation failed" - } else if code == file_not_found { - "File not found" - } else if code == bind_failed { - "Unable to bind to network" - } else if code == file_not_readable { - "File could not be read" - } else if code == file_not_writable { - "File could not be written" - } else if code == no_cluster_file_found { - "No cluster file found in current directory or default location" - } else if code == file_too_large { - "File too large to be read" - } else if code == non_sequential_op { - "Non sequential file operation not allowed" - } else if code == http_bad_response { - "HTTP response was badly formed" - } else if code == http_not_accepted { - "HTTP request not accepted" - } else if code == checksum_failed { - "A data checksum failed" - } else if code == io_timeout { - "A disk IO operation failed to complete in a timely manner" - } else if code == file_corrupt { - "A structurally corrupt data file was detected" - } else if code == http_request_failed { - "HTTP response code not received or indicated failure" - } else if code == http_auth_failed { - "HTTP request failed due to bad credentials" - } else if code == http_bad_request_id { - "HTTP response contained an unexpected X-Request-ID header" - } else if code == rest_invalid_uri { - "Invalid REST URI" - } else if code == rest_invalid_rest_client_knob { - "Invalid RESTClient knob" - } else if code == rest_connectpool_key_not_found { - "ConnectKey not found in connection pool" - } else if code == lock_file_failure { - "Unable to lock the file" - } else if code == rest_unsupported_protocol { - "Unsupported REST protocol" - } else if code == rest_malformed_response { - "Malformed REST response" - } else if code == rest_max_base_cipher_len { - "Max BaseCipher length violation" - } else if code == resource_not_found { - "Requested resource was not found" - } - // 2xxx Attempt (presumably by a _client_) to do something illegal. If an error is known to - // be internally caused, it should be 41xx - else if code == client_invalid_operation { - "Invalid API call" - } else if code == commit_read_incomplete { - "Commit with incomplete read" - } else if code == test_specification_invalid { - "Invalid test specification" - } else if code == key_outside_legal_range { - "Key outside legal range" - } else if code == inverted_range { - "Range begin key larger than end key" - } else if code == invalid_option_value { - "Option set with an invalid value" - } else if code == invalid_option { - "Option not valid in this context" - } else if code == network_not_setup { - "Action not possible before the network is configured" - } else if code == network_already_setup { - "Network can be configured only once" - } else if code == read_version_already_set { - "Transaction already has a read version set" - } else if code == version_invalid { - "Version not valid" - } else if code == range_limits_invalid { - "Range limits not valid" - } else if code == invalid_database_name { - "Database name must be 'DB'" - } else if code == attribute_not_found { - "Attribute not found" - } else if code == future_not_set { - "Future not ready" - } else if code == future_not_error { - "Future not an error" - } else if code == used_during_commit { - "Operation issued while a commit was outstanding" - } else if code == invalid_mutation_type { - "Unrecognized atomic mutation type" - } else if code == attribute_too_large { - "Attribute too large for type int" - } else if code == transaction_invalid_version { - "Transaction does not have a valid commit version" - } else if code == no_commit_version { - "Transaction is read-only and therefore does not have a commit version" - } else if code == environment_variable_network_option_failed { - "Environment variable network option could not be set" - } else if code == transaction_read_only { - "Attempted to commit a transaction specified as read-only" - } else if code == invalid_cache_eviction_policy { - "Invalid cache eviction policy, only random and lru are supported" - } else if code == network_cannot_be_restarted { - "Network can only be started once" - } else if code == blocked_from_network_thread { - "Detected a deadlock in a callback called from the network thread" - } else if code == invalid_config_db_range_read { - "Invalid configuration database range read" - } else if code == invalid_config_db_key { - "Invalid configuration database key provided" - } else if code == invalid_config_path { - "Invalid configuration path" - } else if code == mapper_bad_index { - "The index in K[] or V[] is not a valid number or out of range" - } else if code == mapper_no_such_key { - "A mapped key is not set in database" - } else if code == mapper_bad_range_decriptor { - "\"{...}\" must be the last element of the mapper tuple" - } else if code == quick_get_key_values_has_more { - "One of the mapped range queries is too large" - } else if code == quick_get_value_miss { - "Found a mapped key that is not served in the same SS" - } else if code == quick_get_key_values_miss { - "Found a mapped range that is not served in the same SS" - } else if code == blob_granule_no_ryw { - "Blob Granule Read Transactions must be specified as ryw-disabled" - } else if code == blob_granule_not_materialized { - "Blob Granule Read was not materialized" - } else if code == get_mapped_key_values_has_more { - "getMappedRange does not support continuation for now" - } else if code == get_mapped_range_reads_your_writes { - "getMappedRange tries to read data that were previously written in the transaction" - } else if code == checkpoint_not_found { - "Checkpoint not found" - } else if code == key_not_tuple { - "The key cannot be parsed as a tuple" - } else if code == value_not_tuple { - "The value cannot be parsed as a tuple" - } else if code == mapper_not_tuple { - "The mapper cannot be parsed as a tuple" - } else if code == invalid_checkpoint_format { - "Invalid checkpoint format" - } else if code == invalid_throttle_quota_value { - "Invalid quota value. Note that reserved_throughput cannot exceed total_throughput" - } else if code == failed_to_create_checkpoint { - "Failed to create a checkpoint" - } else if code == failed_to_restore_checkpoint { - "Failed to restore a checkpoint" - } else if code == failed_to_create_checkpoint_shard_metadata { - "Failed to dump shard metadata for a checkpoint to a sst file" - } else if code == address_parse_error { - "Failed to parse address" - } else if code == incompatible_protocol_version { - "Incompatible protocol version" - } else if code == transaction_too_large { - "Transaction exceeds byte limit" - } else if code == key_too_large { - "Key length exceeds limit" - } else if code == value_too_large { - "Value length exceeds limit" - } else if code == connection_string_invalid { - "Connection string invalid" - } else if code == address_in_use { - "Local address in use" - } else if code == invalid_local_address { - "Invalid local address" - } else if code == tls_error { - "TLS error" - } else if code == unsupported_operation { - "Operation is not supported" - } else if code == too_many_tags { - "Too many tags set on transaction" - } else if code == tag_too_long { - "Tag set on transaction is too long" - } else if code == too_many_tag_throttles { - "Too many tag throttles have been created" - } else if code == special_keys_cross_module_read { - "Special key space range read crosses modules. Refer to the `special_key_space_relaxed' transaction option for more details." - } else if code == special_keys_no_module_found { - "Special key space range read does not intersect a module. Refer to the `special_key_space_relaxed' transaction option for more details." - } else if code == special_keys_write_disabled { - "Special Key space is not allowed to write by default. Refer to the `special_key_space_enable_writes` transaction option for more details." - } else if code == special_keys_no_write_module_found { - "Special key space key or keyrange in set or clear does not intersect a module" - } else if code == special_keys_cross_module_clear { - "Special key space clear crosses modules" - } else if code == special_keys_api_failure { - "Api call through special keys failed. For more information, call get on special key 0xff0xff/error_message to get a json string of the error message." - } else if code == client_lib_invalid_metadata { - "Invalid client library metadata." - } else if code == client_lib_already_exists { - "Client library with same identifier already exists on the cluster." - } else if code == client_lib_not_found { - "Client library for the given identifier not found." - } else if code == client_lib_not_available { - "Client library exists, but is not available for download." - } else if code == client_lib_invalid_binary { - "Invalid client library binary." - } else if code == no_external_client_provided { - "No external client library provided." - } else if code == all_external_clients_failed { - "All external clients have failed." - } else if code == incompatible_client { - "None of the available clients match the protocol version of the cluster." - } else if code == tenant_name_required { - "Tenant name must be specified to access data in the cluster" - } else if code == tenant_not_found { - "Tenant does not exist" - } else if code == tenant_already_exists { - "A tenant with the given name already exists" - } else if code == tenant_not_empty { - "Cannot delete a non-empty tenant" - } else if code == invalid_tenant_name { - "Tenant name cannot begin with \\xff" - } else if code == tenant_prefix_allocator_conflict { - "The database already has keys stored at the prefix allocated for the tenant" - } else if code == tenants_disabled { - "Tenants have been disabled in the cluster" - } else if code == illegal_tenant_access { - "Illegal tenant access" - } else if code == invalid_tenant_group_name { - "Tenant group name cannot begin with \\xff" - } else if code == invalid_tenant_configuration { - "Tenant configuration is invalid" - } else if code == cluster_no_capacity { - "Cluster does not have capacity to perform the specified operation" - } else if code == tenant_removed { - "The tenant was removed" - } else if code == invalid_tenant_state { - "Operation cannot be applied to tenant in its current state" - } else if code == tenant_locked { - "Tenant is locked" - } else if code == invalid_cluster_name { - "Data cluster name cannot begin with \\xff" - } else if code == invalid_metacluster_operation { - "Metacluster operation performed on non-metacluster" - } else if code == cluster_already_exists { - "A data cluster with the given name already exists" - } else if code == cluster_not_found { - "Data cluster does not exist" - } else if code == cluster_not_empty { - "Cluster must be empty" - } else if code == cluster_already_registered { - "Data cluster is already registered with a metacluster" - } else if code == metacluster_no_capacity { - "Metacluster does not have capacity to create new tenants" - } else if code == management_cluster_invalid_access { - "Standard transactions cannot be run against the management cluster" - } else if code == tenant_creation_permanently_failed { - "The tenant creation did not complete in a timely manner and has permanently failed" - } else if code == cluster_removed { - "The cluster is being removed from the metacluster" - } else if code == cluster_restoring { - "The cluster is being restored to the metacluster" - } else if code == invalid_data_cluster { - "The data cluster being restored has no record of its metacluster" - } else if code == metacluster_mismatch { - "The cluster does not have the expected name or is associated with a different metacluster" - } else if code == conflicting_restore { - "Another restore is running for the same data cluster" - } else if code == invalid_metacluster_configuration { - "Metacluster configuration is invalid" - } else if code == unsupported_metacluster_version { - "Client is not compatible with the metacluster" - } - // 2200 - errors from bindings and official APIs - else if code == api_version_unset { - "API version is not set" - } else if code == api_version_already_set { - "API version may be set only once" - } else if code == api_version_invalid { - "API version not valid" - } else if code == api_version_not_supported { - "API version not supported" - } else if code == api_function_missing { - "Failed to load a required FDB API function." - } else if code == exact_mode_without_limits { - "EXACT streaming mode requires limits, but none were given" - } else if code == invalid_tuple_data_type { - "Unrecognized data type in packed tuple" - } else if code == invalid_tuple_index { - "Tuple does not have element at specified index" - } else if code == key_not_in_subspace { - "Cannot unpack key that is not in subspace" - } else if code == manual_prefixes_not_enabled { - "Cannot specify a prefix unless manual prefixes are enabled" - } else if code == prefix_in_partition { - "Cannot specify a prefix in a partition" - } else if code == cannot_open_root_directory { - "Root directory cannot be opened" - } else if code == directory_already_exists { - "Directory already exists" - } else if code == directory_does_not_exist { - "Directory does not exist" - } else if code == parent_directory_does_not_exist { - "Directory's parent does not exist" - } else if code == mismatched_layer { - "Directory has already been created with a different layer string" - } else if code == invalid_directory_layer_metadata { - "Invalid directory layer metadata" - } else if code == cannot_move_directory_between_partitions { - "Directory cannot be moved between partitions" - } else if code == cannot_use_partition_as_subspace { - "Directory partition cannot be used as subspace" - } else if code == incompatible_directory_version { - "Directory layer was created with an incompatible version" - } else if code == directory_prefix_not_empty { - "Database has keys stored at the prefix chosen by the automatic prefix allocator" - } else if code == directory_prefix_in_use { - "Directory layer already has a conflicting prefix" - } else if code == invalid_destination_directory { - "Target directory is invalid" - } else if code == cannot_modify_root_directory { - "Root directory cannot be modified" - } else if code == invalid_uuid_size { - "UUID is not sixteen bytes" - } else if code == invalid_versionstamp_size { - "Versionstamp is not exactly twelve bytes" - } - // 2300 - backup and restore errors - else if code == backup_error { - "Backup error" - } else if code == restore_error { - "Restore error" - } else if code == backup_duplicate { - "Backup duplicate request" - } else if code == backup_unneeded { - "Backup unneeded request" - } else if code == backup_bad_block_size { - "Backup file block size too small" - } else if code == backup_invalid_url { - "Backup Container URL invalid" - } else if code == backup_invalid_info { - "Backup Container info invalid" - } else if code == backup_cannot_expire { - "Cannot expire requested data from backup without violating minimum restorability" - } else if code == backup_auth_missing { - "Cannot find authentication details (such as a password or secret key) for the specified Backup Container URL" - } else if code == backup_auth_unreadable { - "Cannot read or parse one or more sources of authentication information for Backup Container URLs" - } else if code == backup_does_not_exist { - "Backup does not exist" - } else if code == backup_not_filterable_with_key_ranges { - "Backup before 6.3 cannot be filtered with key ranges" - } else if code == backup_not_overlapped_with_keys_filter { - "Backup key ranges doesn't overlap with key ranges filter" - } else if code == bucket_not_in_url { - "bucket is not in the URL for backup" - } else if code == backup_parse_s3_response_failure { - "cannot parse s3 response properly" - } else if code == restore_invalid_version { - "Invalid restore version" - } else if code == restore_corrupted_data { - "Corrupted backup data" - } else if code == restore_missing_data { - "Missing backup data" - } else if code == restore_duplicate_tag { - "Restore duplicate request" - } else if code == restore_unknown_tag { - "Restore tag does not exist" - } else if code == restore_unknown_file_type { - "Unknown backup/restore file type" - } else if code == restore_unsupported_file_version { - "Unsupported backup file version" - } else if code == restore_bad_read { - "Unexpected number of bytes read" - } else if code == restore_corrupted_data_padding { - "Backup file has unexpected padding bytes" - } else if code == restore_destination_not_empty { - "Attempted to restore into a non-empty destination database" - } else if code == restore_duplicate_uid { - "Attempted to restore using a UID that had been used for an aborted restore" - } else if code == task_invalid_version { - "Invalid task version" - } else if code == task_interrupted { - "Task execution stopped due to timeout, abort, or completion by another worker" - } else if code == invalid_encryption_key_file { - "The provided encryption key file has invalid contents" - } else if code == blob_restore_missing_logs { - "Missing mutation logs" - } else if code == blob_restore_corrupted_logs { - "Corrupted mutation logs" - } else if code == blob_restore_invalid_manifest_url { - "Invalid manifest URL" - } else if code == blob_restore_corrupted_manifest { - "Corrupted manifest" - } else if code == blob_restore_missing_manifest { - "Missing manifest" - } else if code == blob_migrator_replaced { - "Blob migrator is replaced" - } else if code == key_not_found { - "Expected key is missing" - } else if code == json_malformed { - "JSON string was malformed" - } else if code == json_eof_expected { - "JSON string did not terminate where expected" - } - // 2500 - disk snapshot based backup errors - else if code == snap_disable_tlog_pop_failed { - "Failed to disable tlog pops" - } else if code == snap_storage_failed { - "Failed to snapshot storage nodes" - } else if code == snap_tlog_failed { - "Failed to snapshot TLog nodes" - } else if code == snap_coord_failed { - "Failed to snapshot coordinator nodes" - } else if code == snap_enable_tlog_pop_failed { - "Failed to enable tlog pops" - } else if code == snap_path_not_whitelisted { - "Snapshot create binary path not whitelisted" - } else if code == snap_not_fully_recovered_unsupported { - "Unsupported when the cluster is not fully recovered" - } else if code == snap_log_anti_quorum_unsupported { - "Unsupported when log anti quorum is configured" - } else if code == snap_with_recovery_unsupported { - "Cluster recovery during snapshot operation not supported" - } else if code == snap_invalid_uid_string { - "The given uid string is not a 32-length hex string" - } - // 27XX - Encryption operations errors - else if code == encrypt_ops_error { - "Encryption operation error" - } else if code == encrypt_header_metadata_mismatch { - "Encryption header metadata mismatch" - } else if code == encrypt_key_not_found { - "Expected encryption key is missing" - } else if code == encrypt_key_ttl_expired { - "Expected encryption key TTL has expired" - } else if code == encrypt_header_authtoken_mismatch { - "Encryption header authentication token mismatch" - } else if code == encrypt_update_cipher { - "Attempt to update encryption cipher key" - } else if code == encrypt_invalid_id { - "Invalid encryption cipher details" - } else if code == encrypt_keys_fetch_failed { - "Encryption keys fetch from external KMS failed" - } else if code == encrypt_invalid_kms_config { - "Invalid encryption/kms configuration: discovery-url, validation-token, endpoint etc." - } else if code == encrypt_unsupported { - "Encryption not supported" - } else if code == encrypt_mode_mismatch { - "Encryption mode mismatch with configuration" - } else if code == encrypt_key_check_value_mismatch { - "Encryption key-check-value mismatch" - } else if code == encrypt_max_base_cipher_len { - "Max BaseCipher buffer length violation" - } - // 4xxx Internal errors (those that should be generated only by bugs) are decimal 4xxx - // C++ exception not of type Error - else if code == unknown_error { - "An unknown error occurred" - } else if code == internal_error { - "An internal error occurred" - } else if code == not_implemented { - "Not implemented yet" - } - // 6xxx Authorization and authentication error codes - else if code == permission_denied { - "Client tried to access unauthorized data" - } else if code == unauthorized_attempt { - "A untrusted client tried to send a message to a private endpoint" - } else if code == digital_signature_ops_error { - "Digital signature operation error" - } else if code == authorization_token_verify_failed { - "Failed to verify authorization token" - } else if code == pkey_decode_error { - "Failed to decode public/private key" - } else if code == pkey_encode_error { - "Failed to encode public/private key" - } - // gRPC error - else if code == grpc_error { - "gRPC Error" - } else { - "Unknown error" - } -} - -pub fn fdb_error_predicate(predicate_test: options::ErrorPredicate, code: fdb_error_t) -> bool { - if predicate_test == options::ErrorPredicate::Retryable { - return fdb_error_predicate(options::ErrorPredicate::MaybeCommitted, code) - || fdb_error_predicate(options::ErrorPredicate::RetryableNotCommitted, code); - } - if predicate_test == options::ErrorPredicate::MaybeCommitted { - return code == commit_unknown_result || code == cluster_version_changed; - } - if predicate_test == options::ErrorPredicate::RetryableNotCommitted { - return code == not_committed - || code == transaction_too_old - || code == future_version - || code == database_locked - || code == grv_proxy_memory_limit_exceeded - || code == commit_proxy_memory_limit_exceeded - || code == transaction_throttled_hot_shard - || code == batch_transaction_throttled - || code == process_behind - || code == tag_throttled - || code == proxy_tag_throttled - || code == transaction_rejected_range_locked; - } - - false -} - -/// Error returned when attempting to access metrics on a transaction that wasn't created with metrics instrumentation. -/// -/// This error occurs when calling methods like `set_custom_metric` or `increment_custom_metric` on a -/// transaction that was created without metrics instrumentation (i.e., using `create_trx` instead of -/// `create_instrumented_trx`). -#[derive(Debug)] -pub struct TransactionMetricsNotFound; - -impl std::fmt::Display for TransactionMetricsNotFound { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "Transaction metrics not found") - } -} - -impl std::error::Error for TransactionMetricsNotFound {} - -/// The Standard Error type of FoundationDB -#[derive(Debug, Copy, Clone)] -pub struct FdbError { - /// The FoundationDB error code - error_code: i32, -} - -impl FdbError { - /// Converts from a raw foundationDB error code - pub fn from_code(error_code: fdb_error_t) -> Self { - Self { error_code } - } - - pub fn message(self) -> &'static str { - fdb_get_error(self.error_code) - } - - fn is_error_predicate(self, predicate: options::ErrorPredicate) -> bool { - fdb_error_predicate(predicate, self.error_code) - } - - /// Indicates the transaction may have succeeded, though not in a way the system can verify. - pub fn is_maybe_committed(self) -> bool { - self.is_error_predicate(options::ErrorPredicate::MaybeCommitted) - } - - /// Indicates the operations in the transactions should be retried because of transient error. - pub fn is_retryable(self) -> bool { - self.is_error_predicate(options::ErrorPredicate::Retryable) - } - - /// Indicates the transaction has not committed, though in a way that can be retried. - pub fn is_retryable_not_committed(self) -> bool { - self.is_error_predicate(options::ErrorPredicate::RetryableNotCommitted) - } - - /// Raw foundationdb error code - pub fn code(self) -> i32 { - self.error_code - } -} - -impl fmt::Display for FdbError { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - std::fmt::Display::fmt(&self.message(), f) - } -} - -impl std::error::Error for FdbError {} - -/// Alias for `Result<..., FdbError>` -pub type FdbResult = Result; - -/// This error represent all errors that can be throwed by `db.run`. -/// Layer developers may use the `CustomError`. -pub enum FdbBindingError { - NonRetryableFdbError(FdbError), - PackError(PackError), - /// A reference to the `RetryableTransaction` has been kept - ReferenceToTransactionKept, - /// A custom error that layer developers can use - CustomError(Box), - /// Error returned when attempting to access metrics on a transaction that wasn't created with metrics instrumentation - TransactionMetricsNotFound, -} - -impl FdbBindingError { - /// Returns the underlying `FdbError`, if any. - pub fn get_fdb_error(&self) -> Option { - match *self { - Self::NonRetryableFdbError(error) => Some(error), - Self::CustomError(ref error) => { - if let Some(e) = error.downcast_ref::() { - Some(*e) - } else if let Some(e) = error.downcast_ref::() { - e.get_fdb_error() - } else { - None - } - } - _ => None, - } - } -} - -impl From for FdbBindingError { - fn from(e: FdbError) -> Self { - Self::NonRetryableFdbError(e) - } -} - -impl From for FdbBindingError { - fn from(_e: TransactionMetricsNotFound) -> Self { - Self::TransactionMetricsNotFound - } -} - -impl FdbBindingError { - /// create a new custom error - pub fn new_custom_error(e: Box) -> Self { - Self::CustomError(e) - } -} - -impl Debug for FdbBindingError { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - match self { - FdbBindingError::NonRetryableFdbError(err) => write!(f, "{err:?}"), - FdbBindingError::PackError(err) => write!(f, "{err:?}"), - FdbBindingError::ReferenceToTransactionKept => { - write!(f, "Reference to transaction kept") - } - FdbBindingError::CustomError(err) => write!(f, "{err:?}"), - FdbBindingError::TransactionMetricsNotFound => { - write!(f, "Transaction metrics not found") - } - } - } -} - -impl Display for FdbBindingError { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - std::fmt::Debug::fmt(&self, f) - } -} - -impl std::error::Error for FdbBindingError {} diff --git a/packages/common/universaldb/src/inherited/mod.rs b/packages/common/universaldb/src/inherited/mod.rs deleted file mode 100644 index 81ff533efb..0000000000 --- a/packages/common/universaldb/src/inherited/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -pub mod error; -pub mod keyselector; -pub mod options; -pub mod rangeoption; diff --git a/packages/common/universaldb/src/inherited/keyselector.rs b/packages/common/universaldb/src/key_selector.rs similarity index 82% rename from packages/common/universaldb/src/inherited/keyselector.rs rename to packages/common/universaldb/src/key_selector.rs index 4ccb0c468c..18d7fabb45 100644 --- a/packages/common/universaldb/src/inherited/keyselector.rs +++ b/packages/common/universaldb/src/key_selector.rs @@ -1,11 +1,3 @@ -// Copyright 2018 foundationdb-rs developers, https://github.com/Clikengo/foundationdb-rs/graphs/contributors -// Copyright 2013-2018 Apple, Inc and the FoundationDB project authors. -// -// Licensed under the Apache License, Version 2.0, or the MIT license , at your option. This file may not be -// copied, modified, or distributed except according to those terms. - //! A `KeySelector` identifies a particular key in the database. use crate::tuple::Bytes; @@ -29,8 +21,6 @@ use std::borrow::Cow; /// - `last_less_or_equal` /// - `first_greater_than` /// - `first_greater_or_equal` -/// -/// A dedicated [example](https://github.com/foundationdb-rs/foundationdb-rs/blob/main/foundationdb/examples/key_selectors.rs) is available on Github. #[derive(Clone, Debug)] pub struct KeySelector<'a> { key: Bytes<'a>, diff --git a/packages/common/universaldb/src/lib.rs b/packages/common/universaldb/src/lib.rs index b4dc229153..96260a177d 100644 --- a/packages/common/universaldb/src/lib.rs +++ b/packages/common/universaldb/src/lib.rs @@ -1,27 +1,24 @@ pub(crate) mod atomic; mod database; pub mod driver; -pub mod future; -pub mod inherited; +pub mod error; +pub mod key_selector; +mod metrics; +pub mod options; +pub mod prelude; +pub mod range_option; mod transaction; pub(crate) mod tx_ops; -mod types; pub mod utils; +pub mod value; pub mod versionstamp; -// Export UDB-specific types pub use database::Database; pub use driver::DatabaseDriverHandle; +pub use key_selector::KeySelector; +pub use range_option::RangeOption; +pub use transaction::{RetryableTransaction, Transaction}; +pub use utils::{Subspace, calculate_tx_retry_backoff}; // Re-export FDB types pub use foundationdb_tuple as tuple; -pub use future::{FdbKeyValue, FdbValue}; -pub use inherited::options; -pub use inherited::{ - error::FdbBindingError, error::FdbError, error::FdbResult, keyselector::KeySelector, - rangeoption::RangeOption, -}; -pub use options::DatabaseOption; -pub use transaction::{RetryableTransaction, Transaction}; -pub use types::*; -pub use utils::calculate_tx_retry_backoff; diff --git a/packages/common/udb-util/src/metrics.rs b/packages/common/universaldb/src/metrics.rs similarity index 94% rename from packages/common/udb-util/src/metrics.rs rename to packages/common/universaldb/src/metrics.rs index 95a8f36811..fd9aa61251 100644 --- a/packages/common/udb-util/src/metrics.rs +++ b/packages/common/universaldb/src/metrics.rs @@ -4,7 +4,7 @@ use rivet_metrics::{ }; lazy_static::lazy_static! { - static ref METER: Meter = meter("rivet-udb-util"); + static ref METER: Meter = meter("rivet-universaldb"); /// Has no expected attributes pub static ref PING_DURATION: Histogram = METER.f64_histogram("rivet_udb_ping_duration") diff --git a/packages/common/universaldb/src/inherited/options.rs b/packages/common/universaldb/src/options.rs similarity index 77% rename from packages/common/universaldb/src/inherited/options.rs rename to packages/common/universaldb/src/options.rs index 4d4492227a..64442cd9ea 100644 --- a/packages/common/universaldb/src/inherited/options.rs +++ b/packages/common/universaldb/src/options.rs @@ -1,160 +1,3 @@ -//! IMPORTANT: This file is generated by foundationdb-gen for fdb version 7.3. Anything referencing -// foundationdb_sys was stripped out. - -#[derive(Clone, Debug)] -#[non_exhaustive] -pub enum NetworkOption { - /// IP:PORT - /// - /// Deprecated - LocalAddress(String), - /// path to cluster file - /// - /// Deprecated - ClusterFile(String), - /// path to output directory (or NULL for current working directory) - /// - /// Enables trace output to a file in a directory of the clients choosing - TraceEnable(String), - /// max size of a single trace output file - /// - /// Sets the maximum size in bytes of a single trace output file. This value should be in the range ``[0, INT64_MAX]``. If the value is set to 0, there is no limit on individual file size. The default is a maximum size of 10,485,760 bytes. - TraceRollSize(i32), - /// max total size of trace files - /// - /// Sets the maximum size of all the trace output files put together. This value should be in the range ``[0, INT64_MAX]``. If the value is set to 0, there is no limit on the total size of the files. The default is a maximum size of 104,857,600 bytes. If the default roll size is used, this means that a maximum of 10 trace files will be written at a time. - TraceMaxLogsSize(i32), - /// value of the LogGroup attribute - /// - /// Sets the 'LogGroup' attribute with the specified value for all events in the trace output files. The default log group is 'default'. - TraceLogGroup(String), - /// Format of trace files - /// - /// Select the format of the log files. xml (the default) and json are supported. - TraceFormat(String), - /// Trace clock source - /// - /// Select clock source for trace files. now (the default) or realtime are supported. - TraceClockSource(String), - /// The identifier that will be part of all trace file names - /// - /// Once provided, this string will be used to replace the port/PID in the log file names. - TraceFileIdentifier(String), - /// Use the same base trace file name for all client threads as it did before version 7.2. The current default behavior is to use distinct trace file names for client threads by including their version and thread index. - TraceShareAmongClientThreads, - /// Initialize trace files on network setup, determine the local IP later. Otherwise tracing is initialized when opening the first database. - TraceInitializeOnSetup, - /// Append this suffix to partially written log files. When a log file is complete, it is renamed to remove the suffix. No separator is added between the file and the suffix. If you want to add a file extension, you should include the separator - e.g. '.tmp' instead of 'tmp' to add the 'tmp' extension. - /// - /// Set file suffix for partially written log files. - TracePartialFileSuffix(String), - /// knob_name=knob_value - /// - /// Set internal tuning or debugging knobs - Knob(String), - /// file path or linker-resolved name - /// - /// Deprecated - TLSPlugin(String), - /// certificates - /// - /// Set the certificate chain - TLSCertBytes(Vec), - /// file path - /// - /// Set the file from which to load the certificate chain - TLSCertPath(String), - /// key - /// - /// Set the private key corresponding to your own certificate - TLSKeyBytes(Vec), - /// file path - /// - /// Set the file from which to load the private key corresponding to your own certificate - TLSKeyPath(String), - /// verification pattern - /// - /// Set the peer certificate field verification criteria - TLSVerifyPeers(Vec), - BuggifyEnable, - BuggifyDisable, - /// probability expressed as a percentage between 0 and 100 - /// - /// Set the probability of a BUGGIFY section being active for the current execution. Only applies to code paths first traversed AFTER this option is changed. - BuggifySectionActivatedProbability(i32), - /// probability expressed as a percentage between 0 and 100 - /// - /// Set the probability of an active BUGGIFY section being fired - BuggifySectionFiredProbability(i32), - /// ca bundle - /// - /// Set the ca bundle - TLSCaBytes(Vec), - /// file path - /// - /// Set the file from which to load the certificate authority bundle - TLSCaPath(String), - /// key passphrase - /// - /// Set the passphrase for encrypted private key. Password should be set before setting the key for the password to be used. - TLSPassword(String), - /// Disables the multi-version client API and instead uses the local client directly. Must be set before setting up the network. - DisableMultiVersionClientApi, - /// If set, callbacks from external client libraries can be called from threads created by the FoundationDB client library. Otherwise, callbacks will be called from either the thread used to add the callback or the network thread. Setting this option can improve performance when connected using an external client, but may not be safe to use in all environments. Must be set before setting up the network. WARNING: This feature is considered experimental at this time. - CallbacksOnExternalThreads, - /// path to client library - /// - /// Adds an external client library for use by the multi-version client API. Must be set before setting up the network. - ExternalClientLibrary(String), - /// path to directory containing client libraries - /// - /// Searches the specified path for dynamic libraries and adds them to the list of client libraries for use by the multi-version client API. Must be set before setting up the network. - ExternalClientDirectory(String), - /// Prevents connections through the local client, allowing only connections through externally loaded client libraries. - DisableLocalClient, - /// Number of client threads to be spawned. Each cluster will be serviced by a single client thread. - /// - /// Spawns multiple worker threads for each version of the client that is loaded. Setting this to a number greater than one implies disable_local_client. - ClientThreadsPerVersion(i32), - /// path to client library - /// - /// Adds an external client library to be used with a future version protocol. This option can be used testing purposes only! - FutureVersionClientLibrary(String), - /// Retain temporary external client library copies that are created for enabling multi-threading. - RetainClientLibraryCopies, - /// Ignore the failure to initialize some of the external clients - IgnoreExternalClientFailures, - /// Fail with an error if there is no client matching the server version the client is connecting to - FailIncompatibleClient, - /// Disables logging of client statistics, such as sampled transaction activity. - DisableClientStatisticsLogging, - /// Deprecated - EnableSlowTaskProfiling, - /// Enables debugging feature to perform run loop profiling. Requires trace logging to be enabled. WARNING: this feature is not recommended for use in production. - EnableRunLoopProfiling, - /// Prevents the multi-version client API from being disabled, even if no external clients are configured. This option is required to use GRV caching. - DisableClientBypass, - /// Enable client buggify - will make requests randomly fail (intended for client testing) - ClientBuggifyEnable, - /// Disable client buggify - ClientBuggifyDisable, - /// probability expressed as a percentage between 0 and 100 - /// - /// Set the probability of a CLIENT_BUGGIFY section being active for the current execution. - ClientBuggifySectionActivatedProbability(i32), - /// probability expressed as a percentage between 0 and 100 - /// - /// Set the probability of an active CLIENT_BUGGIFY section being fired. A section will only fire if it was activated - ClientBuggifySectionFiredProbability(i32), - /// Distributed tracer type. Choose from none, log_file, or network_lossy - /// - /// Set a tracer to run on the client. Should be set to the same value as the tracer set on the server. - DistributedClientTracer(String), - /// Client directory for temporary files. - /// - /// Sets the directory for storing temporary files created by FDB client, such as temporary copies of client libraries. Defaults to /tmp - ClientTmpDir(String), -} #[derive(Clone, Debug)] #[non_exhaustive] pub enum DatabaseOption { diff --git a/packages/common/universaldb/src/prelude.rs b/packages/common/universaldb/src/prelude.rs new file mode 100644 index 0000000000..a5c1c897fa --- /dev/null +++ b/packages/common/universaldb/src/prelude.rs @@ -0,0 +1,8 @@ +pub use crate::{ + key_selector::KeySelector, + options::StreamingMode, + range_option::RangeOption, + tuple::{PackError, PackResult, TupleDepth, TuplePack, TupleUnpack, VersionstampOffset}, + utils::{FormalChunkedKey, FormalKey, IsolationLevel::*, OptSliceExt, SliceExt, keys::*}, + value::Value, +}; diff --git a/packages/common/universaldb/src/inherited/rangeoption.rs b/packages/common/universaldb/src/range_option.rs similarity index 96% rename from packages/common/universaldb/src/inherited/rangeoption.rs rename to packages/common/universaldb/src/range_option.rs index 83c4fa8023..f30952f321 100644 --- a/packages/common/universaldb/src/inherited/rangeoption.rs +++ b/packages/common/universaldb/src/range_option.rs @@ -15,8 +15,8 @@ use std::{ ops::{Range, RangeInclusive}, }; -use super::{keyselector::*, options}; -use crate::{future::*, tuple::Subspace}; +use super::{key_selector::KeySelector, options}; +use crate::{tuple::Subspace, value::Values}; /// `RangeOption` represents a query parameters for range scan query. #[derive(Debug, Clone)] @@ -47,7 +47,7 @@ impl RangeOption<'_> { self } - pub fn next_range(mut self, kvs: &FdbValues) -> Option { + pub fn next_range(mut self, kvs: &Values) -> Option { if !kvs.more() { return None; } diff --git a/packages/common/universaldb/src/transaction.rs b/packages/common/universaldb/src/transaction.rs index 2774395ff7..991c069db5 100644 --- a/packages/common/universaldb/src/transaction.rs +++ b/packages/common/universaldb/src/transaction.rs @@ -1,62 +1,225 @@ use std::{future::Future, ops::Deref, pin::Pin, sync::Arc}; +use anyhow::{Context, Result}; +use futures_util::StreamExt; + use crate::{ - FdbResult, KeySelector, RangeOption, driver::TransactionDriver, - future::{FdbSlice, FdbValues}, + key_selector::KeySelector, options::{ConflictRangeType, MutationType}, - tuple::Subspace, - types::{TransactionCommitError, TransactionCommitted}, + range_option::RangeOption, + tuple::{self, TuplePack, TupleUnpack}, + utils::{ + CherryPick, FormalKey, IsolationLevel, MaybeCommitted, OptSliceExt, Subspace, + end_of_key_range, + }, + value::{Slice, Value, Values}, }; +#[derive(Clone)] pub struct Transaction { - pub(crate) driver: Box, + pub(crate) driver: Arc, + subspace: Subspace, } impl Transaction { - pub(crate) fn new(driver: Box) -> Self { - Transaction { driver: driver } + pub(crate) fn new(driver: Arc) -> Self { + Transaction { + driver: driver, + subspace: tuple::Subspace::all().into(), + } } - pub fn atomic_op(&self, key: &[u8], param: &[u8], op_type: MutationType) { - self.driver.atomic_op(key, param, op_type) + /// Creates a new transaction instance with the provided subspace. + pub fn with_subspace(&self, subspace: Subspace) -> Self { + Transaction { + driver: self.driver.clone(), + subspace, + } } - // Read operations + pub fn informal(&self) -> InformalTransaction<'_> { + InformalTransaction { inner: self } + } + + pub fn pack(&self, t: &T) -> Vec { + self.subspace.pack(t) + } + + /// Unpacks a key based on the subspace of this transaction. + pub fn unpack<'de, T: TupleUnpack<'de>>(&self, key: &'de [u8]) -> Result { + self.subspace + .unpack(key) + .with_context(|| format!("failed unpacking key of {}", std::any::type_name::())) + } + + pub fn write(&self, key: &T, value: T::Value) -> Result<()> { + self.driver.set( + &self.subspace.pack(key), + &key.serialize(value).with_context(|| { + format!( + "failed serializing key value of {}", + std::any::type_name::(), + ) + })?, + ); + + Ok(()) + } + + pub async fn read<'de, T: FormalKey + TuplePack + TupleUnpack<'de>>( + &self, + key: &'de T, + isolation_level: IsolationLevel, + ) -> Result { + self.driver + .get(&self.subspace.pack(key), isolation_level) + .await? + .read(key) + } + + pub async fn read_opt<'de, T: FormalKey + TuplePack + TupleUnpack<'de>>( + &self, + key: &'de T, + isolation_level: IsolationLevel, + ) -> Result> { + self.driver + .get(&self.subspace.pack(key), isolation_level) + .await? + .read_opt(key) + } + + pub async fn exists( + &self, + key: &T, + isolation_level: IsolationLevel, + ) -> Result { + Ok(self + .driver + .get(&self.subspace.pack(key), isolation_level) + .await? + .is_some()) + } + + pub fn delete(&self, key: &T) { + self.driver.clear(&self.subspace.pack(key)); + } + + pub fn delete_key_subspace(&self, key: &T) { + self.informal() + .clear_subspace_range(&self.subspace.subspace(&self.subspace.pack(key))); + } + + pub fn read_entry TupleUnpack<'de>>( + &self, + entry: &Value, + ) -> Result<(T, T::Value)> { + let key = self.unpack::(entry.key())?; + let value = key.deserialize(entry.value()).with_context(|| { + format!( + "failed deserializing key value of {}", + std::any::type_name::() + ) + })?; + + Ok((key, value)) + } + + pub async fn cherry_pick( + &self, + subspace: impl TuplePack + Send, + isolation_level: IsolationLevel, + ) -> Result { + T::cherry_pick(self, subspace, isolation_level).await + } + + pub fn add_conflict_key( + &self, + key: &T, + conflict_type: ConflictRangeType, + ) -> Result<()> { + let key_buf = self.subspace.pack(key); + + self.driver + .add_conflict_range(&key_buf, &end_of_key_range(&key_buf), conflict_type) + .map_err(Into::into) + } + + pub fn atomic_op<'de, T: FormalKey + TuplePack + TupleUnpack<'de>>( + &self, + key: &'de T, + param: &[u8], + op_type: MutationType, + ) { + self.driver + .atomic_op(&self.subspace.pack(key), param, op_type) + } + + pub fn read_range<'a>( + &'a self, + opt: RangeOption<'a>, + isolation_level: IsolationLevel, + ) -> crate::value::Stream<'a, Value> { + let opt = RangeOption { + begin: KeySelector::new( + [self.subspace.bytes(), opt.begin.key()].concat().into(), + opt.begin.or_equal(), + opt.begin.offset(), + ), + end: KeySelector::new( + [self.subspace.bytes(), opt.end.key()].concat().into(), + opt.end.or_equal(), + opt.end.offset(), + ), + ..opt + }; + self.driver.get_ranges_keyvalues(opt, isolation_level) + } + + pub fn read_entries<'a, T: FormalKey + for<'de> TupleUnpack<'de>>( + &'a self, + opt: RangeOption<'a>, + isolation_level: IsolationLevel, + ) -> impl futures_util::Stream> { + self.driver + .get_ranges_keyvalues(opt, isolation_level) + .map(|res| self.read_entry(&res?)) + } + + // ==== TODO: Remove. all of these should only be used via `tx.informal()` ==== pub fn get<'a>( &'a self, key: &[u8], - snapshot: bool, - ) -> impl Future>> + 'a { - self.driver.get(key, snapshot) + isolation_level: IsolationLevel, + ) -> impl Future>> + 'a { + self.driver.get(key, isolation_level) } pub fn get_key<'a>( &'a self, selector: &KeySelector<'a>, - snapshot: bool, - ) -> impl Future> + 'a { - self.driver.get_key(selector, snapshot) + isolation_level: IsolationLevel, + ) -> impl Future> + 'a { + self.driver.get_key(selector, isolation_level) } pub fn get_range<'a>( &'a self, opt: &RangeOption<'a>, iteration: usize, - snapshot: bool, - ) -> impl Future> + 'a { - self.driver.get_range(opt, iteration, snapshot) + isolation_level: IsolationLevel, + ) -> impl Future> + 'a { + self.driver.get_range(opt, iteration, isolation_level) } pub fn get_ranges_keyvalues<'a>( &'a self, opt: RangeOption<'a>, - snapshot: bool, - ) -> crate::future::FdbStream<'a, crate::future::FdbValue> { - self.driver.get_ranges_keyvalues(opt, snapshot) + isolation_level: IsolationLevel, + ) -> crate::value::Stream<'a, Value> { + self.driver.get_ranges_keyvalues(opt, isolation_level) } - // Write operations pub fn set(&self, key: &[u8], value: &[u8]) { self.driver.set(key, value) } @@ -69,17 +232,9 @@ impl Transaction { self.driver.clear_range(begin, end) } - /// Clear all keys in a subspace range pub fn clear_subspace_range(&self, subspace: &Subspace) { let (begin, end) = subspace.range(); - self.clear_range(&begin, &end); - } - - pub fn commit( - self: Box, - ) -> Pin> + Send>> - { - self.driver.commit() + self.driver.clear_range(&begin, &end); } pub fn cancel(&self) { @@ -91,7 +246,7 @@ impl Transaction { begin: &[u8], end: &[u8], conflict_type: ConflictRangeType, - ) -> FdbResult<()> { + ) -> Result<()> { self.driver.add_conflict_range(begin, end, conflict_type) } @@ -99,23 +254,119 @@ impl Transaction { &'a self, begin: &'a [u8], end: &'a [u8], - ) -> Pin> + Send + 'a>> { + ) -> Pin> + Send + 'a>> { self.driver.get_estimated_range_size_bytes(begin, end) } } +pub struct InformalTransaction<'t> { + inner: &'t Transaction, +} + +impl<'t> InformalTransaction<'t> { + pub fn atomic_op(&self, key: &[u8], param: &[u8], op_type: MutationType) { + self.inner.driver.atomic_op(key, param, op_type) + } + + // Read operations + pub fn get<'a>( + &'a self, + key: &[u8], + isolation_level: IsolationLevel, + ) -> impl Future>> + 'a { + self.inner.driver.get(key, isolation_level) + } + + pub fn get_key<'a>( + &'a self, + selector: &KeySelector<'a>, + isolation_level: IsolationLevel, + ) -> impl Future> + 'a { + self.inner.driver.get_key(selector, isolation_level) + } + + pub fn get_range<'a>( + &'a self, + opt: &RangeOption<'a>, + iteration: usize, + isolation_level: IsolationLevel, + ) -> impl Future> + 'a { + self.inner.driver.get_range(opt, iteration, isolation_level) + } + + pub fn get_ranges_keyvalues<'a>( + &'a self, + opt: RangeOption<'a>, + isolation_level: IsolationLevel, + ) -> crate::value::Stream<'a, Value> { + self.inner.driver.get_ranges_keyvalues(opt, isolation_level) + } + + // Write operations + pub fn set(&self, key: &[u8], value: &[u8]) { + self.inner.driver.set(key, value) + } + + pub fn clear(&self, key: &[u8]) { + self.inner.driver.clear(key) + } + + pub fn clear_range(&self, begin: &[u8], end: &[u8]) { + self.inner.driver.clear_range(begin, end) + } + + /// Clear all keys in a subspace range + pub fn clear_subspace_range(&self, subspace: &Subspace) { + let (begin, end) = subspace.range(); + self.inner.driver.clear_range(&begin, &end); + } + + // pub fn commit(self: Box) -> Pin> + Send>> { + // self.inner.driver.commit() + // } + + pub fn cancel(&self) { + self.inner.driver.cancel() + } + + pub fn add_conflict_range( + &self, + begin: &[u8], + end: &[u8], + conflict_type: ConflictRangeType, + ) -> Result<()> { + self.inner + .driver + .add_conflict_range(begin, end, conflict_type) + } + + pub fn get_estimated_range_size_bytes<'a>( + &'a self, + begin: &'a [u8], + end: &'a [u8], + ) -> Pin> + Send + 'a>> { + self.inner.driver.get_estimated_range_size_bytes(begin, end) + } +} + /// Retryable transaction wrapper #[derive(Clone)] pub struct RetryableTransaction { - pub(crate) inner: Arc, + pub(crate) inner: Transaction, + pub(crate) maybe_committed: MaybeCommitted, } impl RetryableTransaction { pub fn new(transaction: Transaction) -> Self { RetryableTransaction { - inner: Arc::new(transaction), + inner: transaction, + maybe_committed: MaybeCommitted(false), } } + + pub fn maybe_committed(&self) -> MaybeCommitted { + self.maybe_committed + } } impl Deref for RetryableTransaction { @@ -125,11 +376,3 @@ impl Deref for RetryableTransaction { &self.inner } } - -impl RetryableTransaction { - /// Clear all keys in a subspace range - pub fn clear_subspace_range(&self, subspace: &Subspace) { - let (begin, end) = subspace.range(); - self.inner.clear_range(&begin, &end); - } -} diff --git a/packages/common/universaldb/src/tx_ops.rs b/packages/common/universaldb/src/tx_ops.rs index 52d4108a31..108619932d 100644 --- a/packages/common/universaldb/src/tx_ops.rs +++ b/packages/common/universaldb/src/tx_ops.rs @@ -1,10 +1,14 @@ +use std::collections::BTreeMap; + +use anyhow::Result; + use crate::{ - FdbResult, KeySelector, RangeOption, atomic::apply_atomic_op, - future::{FdbKeyValue, FdbSlice, FdbValues}, + key_selector::KeySelector, options::{ConflictRangeType, MutationType}, + range_option::RangeOption, + value::{KeyValue, Slice, Values}, }; -use std::collections::BTreeMap; #[derive(Debug, Clone)] pub enum Operation { @@ -134,14 +138,14 @@ impl TransactionOperations { &self, key: &[u8], get_from_db: F, - ) -> FdbResult> + ) -> Result> where F: FnOnce() -> Fut, - Fut: std::future::Future>>, + Fut: std::future::Future>>, { // Check local operations first match self.get(key) { - GetOutput::Value(value) => Ok(Some(value)), + GetOutput::Value(value) => Ok(Some(value.into())), GetOutput::Cleared => Ok(None), GetOutput::None => { // Fall back to database @@ -154,7 +158,12 @@ impl TransactionOperations { // Apply all atomic operations in order for (param, op_type) in atomic_ops { - result_value = apply_atomic_op(result_value.as_deref(), ¶m, op_type); + result_value = apply_atomic_op( + result_value.as_ref().map(|x| x.as_slice()), + ¶m, + op_type, + ) + .map(Into::into); } Ok(result_value) @@ -162,14 +171,10 @@ impl TransactionOperations { } } - pub async fn get_key( - &self, - selector: &KeySelector<'_>, - get_from_db: F, - ) -> FdbResult + pub async fn get_key(&self, selector: &KeySelector<'_>, get_from_db: F) -> Result where F: FnOnce() -> Fut, - Fut: std::future::Future>, + Fut: std::future::Future>, { // Get the database result first let db_key = get_from_db().await?; @@ -252,31 +257,27 @@ impl TransactionOperations { if db_key.as_slice() < local.as_slice() { Ok(db_key) } else { - Ok(local) + Ok(local.into()) } } else { // Return the larger key if db_key.as_slice() > local.as_slice() { Ok(db_key) } else { - Ok(local) + Ok(local.into()) } } } - (Some(local), _) => Ok(local), + (Some(local), _) => Ok(local.into()), (None, false) => Ok(db_key), - (None, true) => Ok(vec![]), + (None, true) => Ok(vec![].into()), } } - pub async fn get_range( - &self, - opt: &RangeOption<'_>, - get_from_db: F, - ) -> FdbResult + pub async fn get_range(&self, opt: &RangeOption<'_>, get_from_db: F) -> Result where F: FnOnce() -> Fut, - Fut: std::future::Future>, + Fut: std::future::Future>, { // Get database results let db_values = get_from_db().await?; @@ -349,10 +350,10 @@ impl TransactionOperations { let limit = opt.limit.unwrap_or(usize::MAX); for (key, value) in result_map.into_iter().take(limit) { - keyvalues.push(FdbKeyValue::new(key, value)); + keyvalues.push(KeyValue::new(key, value)); } - Ok(FdbValues::new(keyvalues)) + Ok(Values::new(keyvalues)) } pub fn clear_all(&mut self) { diff --git a/packages/common/universaldb/src/types.rs b/packages/common/universaldb/src/types.rs deleted file mode 100644 index b2b3dff74a..0000000000 --- a/packages/common/universaldb/src/types.rs +++ /dev/null @@ -1,49 +0,0 @@ -use std::{fmt, ops::Deref}; - -use crate::FdbError; - -pub struct TransactionCommitError { - pub err: FdbError, -} - -impl TransactionCommitError { - pub fn new(err: FdbError) -> Self { - Self { err } - } - - pub fn code(&self) -> i32 { - self.err.code() - } -} - -impl Deref for TransactionCommitError { - type Target = FdbError; - fn deref(&self) -> &FdbError { - &self.err - } -} - -impl From for FdbError { - fn from(tce: TransactionCommitError) -> FdbError { - tce.err - } -} - -impl fmt::Debug for TransactionCommitError { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "TransactionCommitError({})", self.err) - } -} - -impl fmt::Display for TransactionCommitError { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - self.err.fmt(f) - } -} - -pub type TransactionCommitted = (); -pub type TransactionCancelled = (); - -/// Indicates the transaction might have committed -#[derive(Debug, Clone, Copy)] -pub struct MaybeCommitted(pub bool); diff --git a/packages/common/universaldb/src/utils.rs b/packages/common/universaldb/src/utils.rs deleted file mode 100644 index 811a0d1668..0000000000 --- a/packages/common/universaldb/src/utils.rs +++ /dev/null @@ -1,36 +0,0 @@ -use crate::{KeySelector, RangeOption, future::FdbValues}; - -pub fn calculate_tx_retry_backoff(attempt: usize) -> u64 { - // TODO: Update this to mirror fdb 1:1: - // https://github.com/apple/foundationdb/blob/21407341d9b49e1d343514a7a5f395bd5f232079/fdbclient/NativeAPI.actor.cpp#L3162 - - let base_backoff_ms = 2_u64.pow((attempt as u32).min(10)) * 10; - - let jitter_ms = rand::random::() % 100; - - base_backoff_ms + jitter_ms -} - -pub fn next_range<'a>(mut range: RangeOption<'a>, kvs: &'a FdbValues) -> Option> { - if !kvs.more() { - return None; - } - - let last = kvs.iter().last()?; - let last_key = last.key(); - - if let Some(limit) = range.limit.as_mut() { - *limit = limit.saturating_sub(kvs.len()); - if *limit == 0 { - return None; - } - } - - if range.reverse { - range.end = KeySelector::first_greater_or_equal(last_key); - } else { - range.begin = KeySelector::first_greater_than(last_key); - } - - Some(range) -} diff --git a/packages/common/universaldb/src/utils/cherry_pick.rs b/packages/common/universaldb/src/utils/cherry_pick.rs new file mode 100644 index 0000000000..51a6262fbc --- /dev/null +++ b/packages/common/universaldb/src/utils/cherry_pick.rs @@ -0,0 +1,88 @@ +use anyhow::{Context, Result, ensure}; +use futures_util::TryStreamExt; + +use crate::{ + options::StreamingMode, + transaction::Transaction, + tuple::{TuplePack, TupleUnpack}, + utils::{FormalKey, IsolationLevel, Subspace}, +}; + +#[async_trait::async_trait] +pub trait CherryPick { + type Output; + + async fn cherry_pick( + tx: &Transaction, + subspace: S, + isolation_level: IsolationLevel, + ) -> Result; +} + +// Implements `CherryPick` for any tuple size +macro_rules! impl_tuple { + ($($args:ident),*) => { + #[async_trait::async_trait] + impl<$($args: FormalKey + for<'de> TupleUnpack<'de>),*> CherryPick for ($($args),*) + where + $($args::Value: Send),* + { + type Output = ($($args::Value),*); + + async fn cherry_pick( + tx: &Transaction, + subspace: S, + isolation_level: IsolationLevel, + ) -> Result { + let tx = tx.with_subspace(Subspace::new(&subspace)); + + let mut stream = tx.read_range( + $crate::range_option::RangeOption { + mode: StreamingMode::WantAll, + ..(&Subspace::all()).into() + }, + isolation_level, + ); + + $( + #[allow(non_snake_case)] + let mut $args = None; + )* + + loop { + let Some(entry) = stream.try_next().await? else { + break; + }; + + $( + if let Ok(key) = tx.unpack::<$args>(entry.key()) { + ensure!($args.is_none(), "{} already picked", std::any::type_name::<$args>()); + + let value = key.deserialize(entry.value())?; + $args = Some(value); + continue; + } + )* + } + + Ok(( + $( + $args.with_context(|| { + format!("key not found in cherry pick: {}", std::any::type_name::<$args>()) + })?, + )* + )) + } + } + } +} + +impl_tuple!(A, B); +impl_tuple!(A, B, C); +impl_tuple!(A, B, C, D); +impl_tuple!(A, B, C, D, E); +impl_tuple!(A, B, C, D, E, F); +impl_tuple!(A, B, C, D, E, F, G); +impl_tuple!(A, B, C, D, E, F, G, H); +impl_tuple!(A, B, C, D, E, F, G, H, I); +impl_tuple!(A, B, C, D, E, F, G, H, I, J); diff --git a/packages/common/udb-util/src/codes.rs b/packages/common/universaldb/src/utils/codes.rs similarity index 100% rename from packages/common/udb-util/src/codes.rs rename to packages/common/universaldb/src/utils/codes.rs diff --git a/packages/common/universaldb/src/utils/ext.rs b/packages/common/universaldb/src/utils/ext.rs new file mode 100644 index 0000000000..c18074c7d0 --- /dev/null +++ b/packages/common/universaldb/src/utils/ext.rs @@ -0,0 +1,58 @@ +use anyhow::{Context, Result}; + +use crate::{tuple::TupleUnpack, utils::FormalKey}; + +pub trait SliceExt { + fn read<'de, T: FormalKey + TupleUnpack<'de>>(&self, key: &'de T) -> Result; +} + +pub trait OptSliceExt { + fn read<'de, T: FormalKey + TupleUnpack<'de>>(&self, key: &'de T) -> Result; + fn read_opt<'de, T: FormalKey + TupleUnpack<'de>>( + &self, + key: &'de T, + ) -> Result>; +} + +impl SliceExt for crate::value::Slice { + fn read<'de, T: FormalKey + TupleUnpack<'de>>(&self, key: &'de T) -> Result { + key.deserialize(self).with_context(|| { + format!( + "failed deserializing key value of {}", + std::any::type_name::(), + ) + }) + } +} + +impl OptSliceExt for Option { + fn read<'de, T: FormalKey + TupleUnpack<'de>>(&self, key: &'de T) -> Result { + key.deserialize( + &self + .as_ref() + .with_context(|| format!("key should exist: {}", std::any::type_name::()))?, + ) + .with_context(|| { + format!( + "failed deserializing key value of {}", + std::any::type_name::(), + ) + }) + } + + fn read_opt<'de, T: FormalKey + TupleUnpack<'de>>( + &self, + key: &'de T, + ) -> Result> { + if let Some(data) = self { + key.deserialize(data).map(Some).with_context(|| { + format!( + "failed deserializing key value of {}", + std::any::type_name::(), + ) + }) + } else { + Ok(None) + } + } +} diff --git a/packages/common/udb-util/src/formal_key.rs b/packages/common/universaldb/src/utils/formal_key.rs similarity index 54% rename from packages/common/udb-util/src/formal_key.rs rename to packages/common/universaldb/src/utils/formal_key.rs index d3b00f36b0..39551d497b 100644 --- a/packages/common/udb-util/src/formal_key.rs +++ b/packages/common/universaldb/src/utils/formal_key.rs @@ -1,5 +1,6 @@ -use anyhow::*; -use universaldb::{self as udb, future::FdbValue}; +use anyhow::Result; + +use crate::value::Value; pub trait FormalKey { type Value; @@ -7,11 +8,6 @@ pub trait FormalKey { fn deserialize(&self, raw: &[u8]) -> Result; fn serialize(&self, value: Self::Value) -> Result>; - - fn read(&self, value: &[u8]) -> std::result::Result { - self.deserialize(value) - .map_err(|x| udb::FdbBindingError::CustomError(x.into())) - } } pub trait FormalChunkedKey { @@ -21,7 +17,7 @@ pub trait FormalChunkedKey { fn chunk(&self, chunk: usize) -> Self::ChunkKey; /// Assumes chunks are in order. - fn combine(&self, chunks: Vec) -> Result; + fn combine(&self, chunks: Vec) -> Result; fn split(&self, value: Self::Value) -> Result>>; } diff --git a/packages/common/udb-util/src/keys.rs b/packages/common/universaldb/src/utils/keys.rs similarity index 100% rename from packages/common/udb-util/src/keys.rs rename to packages/common/universaldb/src/utils/keys.rs diff --git a/packages/common/udb-util/src/lib.rs b/packages/common/universaldb/src/utils/mod.rs similarity index 69% rename from packages/common/udb-util/src/lib.rs rename to packages/common/universaldb/src/utils/mod.rs index c1036d82dd..8b6d673c87 100644 --- a/packages/common/udb-util/src/lib.rs +++ b/packages/common/universaldb/src/utils/mod.rs @@ -1,34 +1,38 @@ -use std::result::Result::Ok; - -use universaldb::tuple::{PackError, PackResult}; +use crate::tuple::{PackError, PackResult}; +mod cherry_pick; pub mod codes; mod ext; mod formal_key; pub mod keys; -mod metrics; mod subspace; +pub use cherry_pick::*; pub use ext::*; pub use formal_key::*; pub use subspace::Subspace; -/// Makes the code blatantly obvious if its using a snapshot read. -pub const SNAPSHOT: bool = true; -pub const SERIALIZABLE: bool = false; pub const CHUNK_SIZE: usize = 10_000; // 10 KB, not KiB, see https://apple.github.io/foundationdb/blob.html -pub mod prelude { - pub use universaldb::{ - FdbBindingError, KeySelector, RangeOption, - future::FdbValue, - options::StreamingMode, - tuple::{PackError, PackResult, TupleDepth, TuplePack, TupleUnpack, VersionstampOffset}, - }; +#[derive(Debug, Clone, Copy)] +pub enum IsolationLevel { + Serializable, + Snapshot, +} + +/// Indicates the transaction might have committed +#[derive(Debug, Clone, Copy)] +pub struct MaybeCommitted(pub bool); + +pub fn calculate_tx_retry_backoff(attempt: usize) -> u64 { + // TODO: Update this to mirror fdb 1:1: + // https://github.com/apple/foundationdb/blob/21407341d9b49e1d343514a7a5f395bd5f232079/fdbclient/NativeAPI.actor.cpp#L3162 + + let base_backoff_ms = 2_u64.pow((attempt as u32).min(10)) * 10; - pub use super::{FormalChunkedKey, FormalKey, OptSliceExt, SliceExt, TxnExt, keys::*}; + let jitter_ms = rand::random::() % 100; - pub use crate::{SERIALIZABLE, SNAPSHOT}; + base_backoff_ms + jitter_ms } /// When using `add_conflict_range` to add a conflict for a single key, you cannot set both the start and end diff --git a/packages/common/udb-util/src/subspace.rs b/packages/common/universaldb/src/utils/subspace.rs similarity index 88% rename from packages/common/udb-util/src/subspace.rs rename to packages/common/universaldb/src/utils/subspace.rs index 13a27f79b3..1c743d0fda 100644 --- a/packages/common/udb-util/src/subspace.rs +++ b/packages/common/universaldb/src/utils/subspace.rs @@ -1,10 +1,11 @@ use std::{borrow::Cow, ops::Deref}; -use rivet_metrics::KeyValue; -use universaldb::{ - KeySelector, RangeOption, +use crate::{ + key_selector::KeySelector, + range_option::RangeOption, tuple::{self, PackResult, TuplePack, TupleUnpack}, }; +use rivet_metrics::KeyValue; use crate::metrics; @@ -22,6 +23,12 @@ impl Subspace { } } + pub fn all() -> Self { + Self { + inner: tuple::Subspace::all(), + } + } + /// Returns a new Subspace whose prefix extends this Subspace with a given tuple encodable. pub fn subspace(&self, t: &T) -> Self { Self { @@ -63,6 +70,12 @@ impl Deref for Subspace { } } +impl From for Subspace { + fn from(value: tuple::Subspace) -> Self { + Subspace { inner: value } + } +} + impl<'a> From<&'a Subspace> for RangeOption<'static> { fn from(subspace: &Subspace) -> Self { let (begin, end) = subspace.range(); diff --git a/packages/common/universaldb/src/value.rs b/packages/common/universaldb/src/value.rs new file mode 100644 index 0000000000..8e556e7bb9 --- /dev/null +++ b/packages/common/universaldb/src/value.rs @@ -0,0 +1,159 @@ +use std::{ + ops::{Deref, DerefMut}, + pin::Pin, +}; + +use anyhow::Result; + +#[derive(Debug, PartialEq, Eq)] +pub struct Slice(Vec); + +impl Slice { + pub fn new() -> Self { + Slice(Vec::new()) + } +} + +impl Deref for Slice { + type Target = Vec; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for Slice { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +impl From> for Slice { + fn from(value: Vec) -> Self { + Slice(value) + } +} + +impl From for Vec { + fn from(value: Slice) -> Self { + value.0 + } +} + +#[derive(Debug, Clone)] +pub struct Value(KeyValue); + +impl Value { + pub fn new(key: Vec, value: Vec) -> Self { + Value(KeyValue::new(key, value)) + } + + pub fn from_keyvalue(kv: KeyValue) -> Self { + Value(kv) + } + + pub fn key(&self) -> &[u8] { + self.0.key() + } + + pub fn value(&self) -> &[u8] { + self.0.value() + } + + pub fn into_parts(self) -> (Vec, Vec) { + self.0.into_parts() + } +} + +// Values wraps a Vec to match FoundationDB API +#[derive(Debug, Clone)] +pub struct Values { + values: Vec, + more: bool, +} + +impl Values { + pub fn new(values: Vec) -> Self { + Values { + values, + more: false, + } + } + + pub fn with_more(values: Vec, more: bool) -> Self { + Values { values, more } + } + + pub fn more(&self) -> bool { + self.more + } + + pub fn into_vec(self) -> Vec { + self.values + } + + pub fn len(&self) -> usize { + self.values.len() + } + + pub fn is_empty(&self) -> bool { + self.values.is_empty() + } + + pub fn iter(&self) -> std::slice::Iter<'_, KeyValue> { + self.values.iter() + } + + pub fn into_iter(self) -> std::vec::IntoIter { + self.values.into_iter() + } +} + +// impl Deref for Values { +// type Target = [KeyValue]; +// fn deref(&self) -> &Self::Target { +// &self.values +// } +// } +// impl AsRef<[KeyValue]> for Values { +// fn as_ref(&self) -> &[KeyValue] { +// self.deref() +// } +// } + +// KeyValue type with key() and value() methods +#[derive(Debug, Clone)] +pub struct KeyValue { + key: Vec, + value: Vec, +} + +impl KeyValue { + pub fn new(key: Vec, value: Vec) -> Self { + KeyValue { key, value } + } + + pub fn key(&self) -> &[u8] { + &self.key + } + + pub fn value(&self) -> &[u8] { + &self.value + } + + pub fn into_parts(self) -> (Vec, Vec) { + (self.key, self.value) + } + + pub fn to_value(self) -> Value { + Value::from_keyvalue(self) + } + + pub fn value_ref(&self) -> Value { + Value::from_keyvalue(self.clone()) + } +} + +// Stream type for range queries - generic over item type +pub type Stream<'a, T = KeyValue> = + Pin> + Send + 'a>>; diff --git a/packages/common/universaldb/tests/integration.rs b/packages/common/universaldb/tests/integration.rs index a20f44f248..f066da4747 100644 --- a/packages/common/universaldb/tests/integration.rs +++ b/packages/common/universaldb/tests/integration.rs @@ -1,9 +1,13 @@ +use anyhow::Result; use rivet_test_deps_docker::TestDatabase; use std::{borrow::Cow, sync::Arc}; use universaldb::{ - Database, FdbBindingError, KeySelector, RangeOption, + Database, + key_selector::KeySelector, options::{ConflictRangeType, StreamingMode}, + range_option::RangeOption, tuple::{Element, Subspace, Versionstamp, pack_with_versionstamp}, + utils::IsolationLevel::*, versionstamp::generate_versionstamp, }; use uuid::Uuid; @@ -123,7 +127,7 @@ async fn run_all_tests(db: universaldb::Database) { async fn test_database_options(db: &Database) { use std::sync::Arc; use std::sync::atomic::{AtomicU32, Ordering}; - use universaldb::FdbError; + use universaldb::error::DatabaseError; use universaldb::options::DatabaseOption; // Test setting transaction retry limit @@ -135,7 +139,7 @@ async fn test_database_options(db: &Database) { let counter_clone = conflict_counter.clone(); let result = db - .run(|tx, _maybe_committed| { + .run(|tx| { let counter = counter_clone.clone(); async move { // Increment counter to track retry attempts @@ -143,7 +147,7 @@ async fn test_database_options(db: &Database) { // Force a retry on first few attempts by returning a retryable error if attempts < 3 { - return Err(FdbBindingError::from(FdbError::from_code(1020))); // not_committed + return Err(DatabaseError::NotCommitted.into()); } // Should succeed on the third attempt @@ -166,7 +170,7 @@ async fn test_database_options(db: &Database) { let counter_clone2 = conflict_counter2.clone(); let result = db - .run(|_tx, _maybe_committed| { + .run(|_tx| { let counter = counter_clone2.clone(); async move { // Increment counter to track retry attempts @@ -174,7 +178,7 @@ async fn test_database_options(db: &Database) { // Always force a retry if attempts < 10 { - return Err(FdbBindingError::from(FdbError::from_code(1020))); // not_committed + return Err(DatabaseError::NotCommitted.into()); } Ok(()) @@ -195,8 +199,8 @@ async fn test_database_options(db: &Database) { .unwrap(); } -async fn clear_test_namespace(db: &Database) -> Result<(), FdbBindingError> { - db.run(|tx, _maybe_committed| async move { +async fn clear_test_namespace(db: &Database) -> Result<()> { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let (begin, end) = test_subspace.range(); tx.clear_range(&begin, &end); @@ -207,7 +211,7 @@ async fn clear_test_namespace(db: &Database) -> Result<(), FdbBindingError> { async fn test_basic_operations(db: &Database) { // Test set and get using subspace and tuple syntax - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("key1",)); tx.set(&key, b"value1"); @@ -217,22 +221,22 @@ async fn test_basic_operations(db: &Database) { .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("key1",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - assert_eq!(value, Some(b"value1".to_vec())); + assert_eq!(value, Some(b"value1".to_vec().into())); // Test get non-existent key let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("nonexistent",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await @@ -240,7 +244,7 @@ async fn test_basic_operations(db: &Database) { assert_eq!(value, None); // Test clear - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("key1",)); tx.clear(&key); @@ -250,10 +254,10 @@ async fn test_basic_operations(db: &Database) { .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("key1",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await @@ -263,7 +267,7 @@ async fn test_basic_operations(db: &Database) { async fn test_range_operations(db: &Database) { // Setup test data using subspace keys - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key_a = test_subspace.pack(&("a",)); let key_b = test_subspace.pack(&("b",)); @@ -281,7 +285,7 @@ async fn test_range_operations(db: &Database) { // Test get_range using subspace range for keys "b" through "d" (exclusive) let results = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key_b = test_subspace.pack(&("b",)); let key_d = test_subspace.pack(&("d",)); @@ -296,7 +300,7 @@ async fn test_range_operations(db: &Database) { ..RangeOption::default() }; - let vals = tx.get_range(&range, 1, false).await?; + let vals = tx.get_range(&range, 1, Serializable).await?; Ok(vals) }) .await @@ -316,7 +320,7 @@ async fn test_range_operations(db: &Database) { assert_eq!(values[1].value(), b"3"); // Test clear_range using subspace keys - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key_b = test_subspace.pack(&("b",)); let key_d = test_subspace.pack(&("d",)); @@ -328,7 +332,7 @@ async fn test_range_operations(db: &Database) { // Verify range was cleared let results = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key_b = test_subspace.pack(&("b",)); let key_d = test_subspace.pack(&("d",)); @@ -343,7 +347,7 @@ async fn test_range_operations(db: &Database) { ..RangeOption::default() }; - let vals = tx.get_range(&range, 1, false).await?; + let vals = tx.get_range(&range, 1, Serializable).await?; Ok(vals) }) .await @@ -352,31 +356,31 @@ async fn test_range_operations(db: &Database) { // Verify keys outside range still exist let value_a = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("a",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - assert_eq!(value_a, Some(b"1".to_vec())); + assert_eq!(value_a, Some(b"1".to_vec().into())); let value_d = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("d",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - assert_eq!(value_d, Some(b"4".to_vec())); + assert_eq!(value_d, Some(b"4".to_vec().into())); } async fn test_transaction_isolation(db: &Database) { // Set initial value using subspace - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("counter",)); tx.set(&key, b"0"); @@ -387,18 +391,18 @@ async fn test_transaction_isolation(db: &Database) { // Test that each transaction sees consistent state let val1 = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("counter",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - assert_eq!(val1, Some(b"0".to_vec())); + assert_eq!(val1, Some(b"0".to_vec().into())); // Set value in one transaction - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("counter",)); tx.set(&key, b"1"); @@ -409,20 +413,20 @@ async fn test_transaction_isolation(db: &Database) { // Verify the change is visible in new transaction let val3 = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("counter",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - assert_eq!(val3, Some(b"1".to_vec())); + assert_eq!(val3, Some(b"1".to_vec().into())); } async fn test_conflict_ranges(db: &Database) { // Test 1: Basic conflict range with read type - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("conflict",)); tx.set(&key, b"initial"); @@ -435,7 +439,7 @@ async fn test_conflict_ranges(db: &Database) { .unwrap(); // Test 2: Conflict range with write type - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key1 = test_subspace.pack(&("range_test1",)); let key2 = test_subspace.pack(&("range_test2",)); @@ -452,7 +456,7 @@ async fn test_conflict_ranges(db: &Database) { .unwrap(); // Test 3: Multiple conflict ranges - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); // Add multiple conflict ranges @@ -472,7 +476,7 @@ async fn test_conflict_ranges(db: &Database) { async fn test_get_key(db: &Database) { // Setup test data using subspace keys - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key1 = test_subspace.pack(&("k1",)); let key2 = test_subspace.pack(&("k2",)); @@ -488,11 +492,11 @@ async fn test_get_key(db: &Database) { // Test first_greater_or_equal let key = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let search_key = test_subspace.pack(&("k2",)); let selector = KeySelector::first_greater_or_equal(Cow::Owned(search_key)); - let k = tx.get_key(&selector, false).await?; + let k = tx.get_key(&selector, Serializable).await?; Ok(k) }) .await @@ -500,27 +504,27 @@ async fn test_get_key(db: &Database) { let test_subspace = Subspace::from("test"); let expected_key = test_subspace.pack(&("k2",)); - assert_eq!(key, expected_key); + assert_eq!(key, expected_key.into()); // Test with first_greater_than let key = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let search_key = test_subspace.pack(&("k1",)); let selector = KeySelector::first_greater_than(Cow::Owned(search_key)); - let k = tx.get_key(&selector, false).await?; + let k = tx.get_key(&selector, Serializable).await?; Ok(k) }) .await .unwrap(); let expected_key = test_subspace.pack(&("k2",)); - assert_eq!(key, expected_key); + assert_eq!(key, expected_key.into()); } async fn test_range_options(db: &Database) { // Setup test data - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key_a = test_subspace.pack(&("range_a",)); let key_b = test_subspace.pack(&("range_b",)); @@ -540,7 +544,7 @@ async fn test_range_options(db: &Database) { // Test 1: first_greater_or_equal on both bounds (inclusive range [b, d)) let results = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key_b = test_subspace.pack(&("range_b",)); let key_d = test_subspace.pack(&("range_d",)); @@ -555,7 +559,7 @@ async fn test_range_options(db: &Database) { ..RangeOption::default() }; - let vals = tx.get_range(&range, 1, false).await?; + let vals = tx.get_range(&range, 1, Serializable).await?; Ok(vals.into_vec()) }) .await @@ -575,7 +579,7 @@ async fn test_range_options(db: &Database) { // Test 2: first_greater_than on lower, first_greater_or_equal on upper (b, d) // Note: Some drivers may not correctly implement first_greater_than and include the boundary key let results = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key_b = test_subspace.pack(&("range_b",)); let key_d = test_subspace.pack(&("range_d",)); @@ -590,7 +594,7 @@ async fn test_range_options(db: &Database) { ..RangeOption::default() }; - let vals = tx.get_range(&range, 1, false).await?; + let vals = tx.get_range(&range, 1, Serializable).await?; Ok(vals.into_vec()) }) .await @@ -617,7 +621,7 @@ async fn test_range_options(db: &Database) { // Test 3: first_greater_or_equal on lower, first_greater_than on upper [b, d] let results = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key_b = test_subspace.pack(&("range_b",)); let key_d = test_subspace.pack(&("range_d",)); @@ -632,7 +636,7 @@ async fn test_range_options(db: &Database) { ..RangeOption::default() }; - let vals = tx.get_range(&range, 1, false).await?; + let vals = tx.get_range(&range, 1, Serializable).await?; Ok(vals.into_vec()) }) .await @@ -652,7 +656,7 @@ async fn test_range_options(db: &Database) { // Test 4: first_greater_than on both bounds (b, e) let results = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key_b = test_subspace.pack(&("range_b",)); let key_e = test_subspace.pack(&("range_e",)); @@ -667,7 +671,7 @@ async fn test_range_options(db: &Database) { ..RangeOption::default() }; - let vals = tx.get_range(&range, 1, false).await?; + let vals = tx.get_range(&range, 1, Serializable).await?; Ok(vals.into_vec()) }) .await @@ -686,7 +690,7 @@ async fn test_range_options(db: &Database) { assert_eq!(results[2].value(), b"val_e"); // Clear test data - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let (begin, end) = test_subspace.range(); tx.clear_range(&begin, &end); @@ -698,7 +702,7 @@ async fn test_range_options(db: &Database) { async fn test_read_after_write(db: &Database) { // Test 1: Basic set and get within same transaction - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key1 = test_subspace.pack(&("raw_key1",)); @@ -706,12 +710,12 @@ async fn test_read_after_write(db: &Database) { tx.set(&key1, b"value1"); // Read it back immediately (read-after-write) - let value = tx.get(&key1, false).await?; - assert_eq!(value, Some(b"value1".to_vec())); + let value = tx.get(&key1, Serializable).await?; + assert_eq!(value, Some(b"value1".to_vec().into())); // Read a non-existent key let key2 = test_subspace.pack(&("raw_key2",)); - let value = tx.get(&key2, false).await?; + let value = tx.get(&key2, Serializable).await?; assert_eq!(value, None); Ok(()) @@ -720,19 +724,19 @@ async fn test_read_after_write(db: &Database) { .unwrap(); // Test 2: Clear and get - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key1 = test_subspace.pack(&("raw_key1",)); // First verify the key exists from previous test - let value = tx.get(&key1, false).await?; - assert_eq!(value, Some(b"value1".to_vec())); + let value = tx.get(&key1, Serializable).await?; + assert_eq!(value, Some(b"value1".to_vec().into())); // Clear it tx.clear(&key1); // Read should return None - let value = tx.get(&key1, false).await?; + let value = tx.get(&key1, Serializable).await?; assert_eq!(value, None); Ok(()) @@ -741,7 +745,7 @@ async fn test_read_after_write(db: &Database) { .unwrap(); // Test 3: Clear range and get - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key_a = test_subspace.pack(&("raw_a",)); let key_b = test_subspace.pack(&("raw_b",)); @@ -758,10 +762,16 @@ async fn test_read_after_write(db: &Database) { tx.clear_range(&key_b, &key_d); // Check values - assert_eq!(tx.get(&key_a, false).await?, Some(b"value_a".to_vec())); - assert_eq!(tx.get(&key_b, false).await?, None); - assert_eq!(tx.get(&key_c, false).await?, None); - assert_eq!(tx.get(&key_d, false).await?, Some(b"value_d".to_vec())); + assert_eq!( + tx.get(&key_a, Serializable).await?, + Some(b"value_a".to_vec().into()) + ); + assert_eq!(tx.get(&key_b, Serializable).await?, None); + assert_eq!(tx.get(&key_c, Serializable).await?, None); + assert_eq!( + tx.get(&key_d, Serializable).await?, + Some(b"value_d".to_vec().into()) + ); Ok(()) }) @@ -769,7 +779,7 @@ async fn test_read_after_write(db: &Database) { .unwrap(); // Test 4: Get range with local modifications - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let range_key1 = test_subspace.pack(&("range_key1",)); let range_key3 = test_subspace.pack(&("range_key3",)); @@ -786,7 +796,7 @@ async fn test_read_after_write(db: &Database) { end: KeySelector::first_greater_or_equal(Cow::Owned(end)), ..RangeOption::default() }; - let values = tx.get_range(&range_opt, 1, false).await?; + let values = tx.get_range(&range_opt, 1, Serializable).await?; let mut keys = Vec::new(); for kv in values.into_iter() { @@ -803,18 +813,27 @@ async fn test_read_after_write(db: &Database) { .unwrap(); // Test 5: Overwrite value multiple times - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("overwrite_key",)); tx.set(&key, b"value1"); - assert_eq!(tx.get(&key, false).await?, Some(b"value1".to_vec())); + assert_eq!( + tx.get(&key, Serializable).await?, + Some(b"value1".to_vec().into()) + ); tx.set(&key, b"value2"); - assert_eq!(tx.get(&key, false).await?, Some(b"value2".to_vec())); + assert_eq!( + tx.get(&key, Serializable).await?, + Some(b"value2".to_vec().into()) + ); tx.set(&key, b"value3"); - assert_eq!(tx.get(&key, false).await?, Some(b"value3".to_vec())); + assert_eq!( + tx.get(&key, Serializable).await?, + Some(b"value3".to_vec().into()) + ); Ok(()) }) @@ -824,7 +843,7 @@ async fn test_read_after_write(db: &Database) { async fn test_set_clear_set(db: &Database) { // Test the bug where set → clear → set sequence doesn't work correctly - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bug_key",)); @@ -838,10 +857,10 @@ async fn test_set_clear_set(db: &Database) { tx.set(&key, b"value2"); // This should return the latest value "value2", not None or Cleared - let value = tx.get(&key, false).await?; + let value = tx.get(&key, Serializable).await?; assert_eq!( value, - Some(b"value2".to_vec()), + Some(b"value2".to_vec().into()), "Expected to get the latest set value after set-clear-set sequence" ); @@ -853,7 +872,7 @@ async fn test_set_clear_set(db: &Database) { async fn test_get_key_with_local_writes(db: &Database) { // Setup: Store keys with values 2 and 10 in the database - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key2 = test_subspace.pack(&(2,)); let key10 = test_subspace.pack(&(10,)); @@ -866,7 +885,7 @@ async fn test_get_key_with_local_writes(db: &Database) { // Test: Write a key with value 5 in the transaction, then get_key with >= 3 let result_key = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); // Write key 5 in the transaction @@ -876,7 +895,7 @@ async fn test_get_key_with_local_writes(db: &Database) { // Use get_key with >= 3 selector let search_key = test_subspace.pack(&(3,)); let selector = KeySelector::first_greater_or_equal(Cow::Owned(search_key)); - let k = tx.get_key(&selector, false).await?; + let k = tx.get_key(&selector, Serializable).await?; Ok(k) }) .await @@ -886,13 +905,14 @@ async fn test_get_key_with_local_writes(db: &Database) { let test_subspace = Subspace::from("test"); let expected_key5 = test_subspace.pack(&(5,)); assert_eq!( - result_key, expected_key5, + result_key, + expected_key5.clone().into(), "get_key should return key 5 from local writes, not key 10 from database" ); // Test with first_greater_than let result_key = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); // Write key 5 in the transaction @@ -902,7 +922,7 @@ async fn test_get_key_with_local_writes(db: &Database) { // Use get_key with > 4 selector let search_key = test_subspace.pack(&(4,)); let selector = KeySelector::first_greater_than(Cow::Owned(search_key)); - let k = tx.get_key(&selector, false).await?; + let k = tx.get_key(&selector, Serializable).await?; Ok(k) }) .await @@ -910,14 +930,15 @@ async fn test_get_key_with_local_writes(db: &Database) { // Should return key5, not key10 assert_eq!( - result_key, expected_key5, + result_key, + expected_key5.into(), "get_key with > selector should return key 5 from local writes" ); } async fn test_snapshot_reads(db: &Database) { // Setup: Store initial data in the database - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key1 = test_subspace.pack(&("snap_key1",)); let key2 = test_subspace.pack(&("snap_key2",)); @@ -932,7 +953,7 @@ async fn test_snapshot_reads(db: &Database) { .unwrap(); // Test 1: Just snapshot reads - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key1 = test_subspace.pack(&("snap_key1",)); let key2 = test_subspace.pack(&("snap_key2",)); @@ -940,10 +961,10 @@ async fn test_snapshot_reads(db: &Database) { let key4 = test_subspace.pack(&("snap_key4",)); // Snapshot read should see database value - let snapshot_value = tx.get(&key1, true).await?; + let snapshot_value = tx.get(&key1, Snapshot).await?; assert_eq!( snapshot_value, - Some(b"db_value1".to_vec()), + Some(b"db_value1".to_vec().into()), "Snapshot read should see database value" ); @@ -957,7 +978,7 @@ async fn test_snapshot_reads(db: &Database) { }; // Snapshot range read - let snapshot_values = tx.get_range(&range_opt, 1, true).await?; + let snapshot_values = tx.get_range(&range_opt, 1, Snapshot).await?; let mut snapshot_keys = Vec::new(); for kv in snapshot_values.into_iter() { snapshot_keys.push((kv.key().to_vec(), kv.value().to_vec())); @@ -987,7 +1008,7 @@ async fn test_snapshot_reads(db: &Database) { .unwrap(); // Test 2: Snapshot read should skip local set operations within a transaction - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key1 = test_subspace.pack(&("snap_key1",)); @@ -995,18 +1016,18 @@ async fn test_snapshot_reads(db: &Database) { tx.set(&key1, b"local_value1"); // Non-snapshot read should see local value - let value = tx.get(&key1, false).await?; + let value = tx.get(&key1, Serializable).await?; assert_eq!( value, - Some(b"local_value1".to_vec()), + Some(b"local_value1".to_vec().into()), "Non-snapshot read should see local write" ); // Snapshot read should see local write - let snapshot_value = tx.get(&key1, true).await?; + let snapshot_value = tx.get(&key1, Snapshot).await?; assert_eq!( snapshot_value, - Some(b"local_value1".to_vec()), + Some(b"local_value1".to_vec().into()), "Snapshot read should see local write" ); @@ -1018,7 +1039,7 @@ async fn test_snapshot_reads(db: &Database) { // Reset state { clear_test_namespace(&db).await.unwrap(); - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key1 = test_subspace.pack(&("snap_key1",)); let key2 = test_subspace.pack(&("snap_key2",)); @@ -1034,7 +1055,7 @@ async fn test_snapshot_reads(db: &Database) { } // Test 3: Snapshot read should skip local clear operations - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key2 = test_subspace.pack(&("snap_key2",)); @@ -1042,11 +1063,11 @@ async fn test_snapshot_reads(db: &Database) { tx.clear(&key2); // Non-snapshot read should see None (cleared) - let value = tx.get(&key2, false).await?; + let value = tx.get(&key2, Serializable).await?; assert_eq!(value, None, "Non-snapshot read should see cleared value"); // Snapshot read should still see database value - let snapshot_value = tx.get(&key2, true).await?; + let snapshot_value = tx.get(&key2, Snapshot).await?; assert_eq!(snapshot_value, None, "Snapshot read should see local clear"); Ok(()) @@ -1057,7 +1078,7 @@ async fn test_snapshot_reads(db: &Database) { // Reset state { clear_test_namespace(&db).await.unwrap(); - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key1 = test_subspace.pack(&("snap_key1",)); let key2 = test_subspace.pack(&("snap_key2",)); @@ -1073,7 +1094,7 @@ async fn test_snapshot_reads(db: &Database) { } // Test 4: Snapshot get_range should skip local operations - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key1 = test_subspace.pack(&("snap_key1",)); let key2 = test_subspace.pack(&("snap_key2",)); @@ -1095,7 +1116,7 @@ async fn test_snapshot_reads(db: &Database) { }; // Non-snapshot range read - let values = tx.get_range(&range_opt, 1, false).await?; + let values = tx.get_range(&range_opt, 1, Serializable).await?; let mut non_snapshot_keys = Vec::new(); for kv in values.into_iter() { non_snapshot_keys.push((kv.key().to_vec(), kv.value().to_vec())); @@ -1121,7 +1142,7 @@ async fn test_snapshot_reads(db: &Database) { ); // Snapshot range read - let snapshot_values = tx.get_range(&range_opt, 1, true).await?; + let snapshot_values = tx.get_range(&range_opt, 1, Snapshot).await?; let mut snapshot_keys = Vec::new(); for kv in snapshot_values.into_iter() { snapshot_keys.push((kv.key().to_vec(), kv.value().to_vec())); @@ -1154,7 +1175,7 @@ async fn test_snapshot_reads(db: &Database) { // Reset state { clear_test_namespace(&db).await.unwrap(); - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key1 = test_subspace.pack(&("snap_key1",)); let key2 = test_subspace.pack(&("snap_key2",)); @@ -1170,7 +1191,7 @@ async fn test_snapshot_reads(db: &Database) { } // Test 5: Snapshot get_key should skip local operations - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); // Add a local key between existing database keys @@ -1180,13 +1201,18 @@ async fn test_snapshot_reads(db: &Database) { // Non-snapshot get_key >= "snap_key14" should find the local key15 let search_key = test_subspace.pack(&("snap_key14",)); let selector = KeySelector::first_greater_or_equal(Cow::Owned(search_key)); - let result = tx.get_key(&selector, false).await?; - assert_eq!(result, key15, "Non-snapshot get_key should find local key"); + let result = tx.get_key(&selector, Serializable).await?; + assert_eq!( + result, + key15.clone().into(), + "Non-snapshot get_key should find local key" + ); // Snapshot get_key >= "snap_key14" should find the local key15 - let snapshot_result = tx.get_key(&selector, true).await?; + let snapshot_result = tx.get_key(&selector, Snapshot).await?; assert_eq!( - snapshot_result, key15, + snapshot_result, + key15.into(), "Snapshot get_key should find local key" ); @@ -1234,12 +1260,13 @@ async fn test_atomic_add(db: &Database) { use universaldb::options::MutationType; // Test 1: Add to non-existent key (should treat as 0) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("add_key1",)); // Add 42 to non-existent key - tx.atomic_op(&key, &42i64.to_le_bytes(), MutationType::Add); + tx.informal() + .atomic_op(&key, &42i64.to_le_bytes(), MutationType::Add); Ok(()) }) .await @@ -1247,73 +1274,75 @@ async fn test_atomic_add(db: &Database) { // Verify the result let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("add_key1",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!( result, 42, "Add to non-existent key should equal the parameter" ); // Test 2: Add to existing value - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("add_key1",)); // Add 10 to existing value (42) - tx.atomic_op(&key, &10i64.to_le_bytes(), MutationType::Add); + tx.informal() + .atomic_op(&key, &10i64.to_le_bytes(), MutationType::Add); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("add_key1",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!(result, 52, "42 + 10 should equal 52"); // Test 3: Add negative number - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("add_key1",)); // Add -20 to existing value (52) - tx.atomic_op(&key, &(-20i64).to_le_bytes(), MutationType::Add); + tx.informal() + .atomic_op(&key, &(-20i64).to_le_bytes(), MutationType::Add); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("add_key1",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!(result, 32, "52 + (-20) should equal 32"); // Test 4: Test wrapping behavior with overflow - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("add_overflow",)); @@ -1324,28 +1353,29 @@ async fn test_atomic_add(db: &Database) { .await .unwrap(); - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("add_overflow",)); // Add 1 to max i64 (should wrap) - tx.atomic_op(&key, &1i64.to_le_bytes(), MutationType::Add); + tx.informal() + .atomic_op(&key, &1i64.to_le_bytes(), MutationType::Add); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("add_overflow",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!(result, i64::MIN, "Max i64 + 1 should wrap to min i64"); } @@ -1353,7 +1383,7 @@ async fn test_atomic_bitwise(db: &Database) { use universaldb::options::MutationType; // Test BitAnd operation - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bit_and",)); @@ -1364,22 +1394,23 @@ async fn test_atomic_bitwise(db: &Database) { .await .unwrap(); - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bit_and",)); // AND with 0b10101010 (170) - tx.atomic_op(&key, &[0b10101010], MutationType::BitAnd); + tx.informal() + .atomic_op(&key, &[0b10101010], MutationType::BitAnd); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bit_and",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await @@ -1392,7 +1423,7 @@ async fn test_atomic_bitwise(db: &Database) { ); // Test BitOr operation - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bit_or",)); @@ -1403,22 +1434,23 @@ async fn test_atomic_bitwise(db: &Database) { .await .unwrap(); - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bit_or",)); // OR with 0b00001111 (15) - tx.atomic_op(&key, &[0b00001111], MutationType::BitOr); + tx.informal() + .atomic_op(&key, &[0b00001111], MutationType::BitOr); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bit_or",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await @@ -1431,7 +1463,7 @@ async fn test_atomic_bitwise(db: &Database) { ); // Test BitXor operation - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bit_xor",)); @@ -1442,22 +1474,23 @@ async fn test_atomic_bitwise(db: &Database) { .await .unwrap(); - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bit_xor",)); // XOR with 0b10101010 (170) - tx.atomic_op(&key, &[0b10101010], MutationType::BitXor); + tx.informal() + .atomic_op(&key, &[0b10101010], MutationType::BitXor); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bit_xor",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await @@ -1470,7 +1503,7 @@ async fn test_atomic_bitwise(db: &Database) { ); // Test bitwise operations with different lengths - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bit_len",)); @@ -1481,22 +1514,23 @@ async fn test_atomic_bitwise(db: &Database) { .await .unwrap(); - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bit_len",)); // AND with 1-byte value (should extend current to match param length) - tx.atomic_op(&key, &[0b10101010], MutationType::BitAnd); + tx.informal() + .atomic_op(&key, &[0b10101010], MutationType::BitAnd); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("bit_len",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await @@ -1518,61 +1552,63 @@ async fn test_atomic_append_if_fits(db: &Database) { use universaldb::options::MutationType; // Test 1: Append to non-existent key - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("append_key1",)); - tx.atomic_op(&key, b"hello", MutationType::AppendIfFits); + tx.informal() + .atomic_op(&key, b"hello", MutationType::AppendIfFits); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("append_key1",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); assert_eq!( - value.unwrap(), + value.unwrap().as_slice(), b"hello", "Append to non-existent key should create the key" ); // Test 2: Append to existing key - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("append_key1",)); - tx.atomic_op(&key, b" world", MutationType::AppendIfFits); + tx.informal() + .atomic_op(&key, b" world", MutationType::AppendIfFits); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("append_key1",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); assert_eq!( - value.unwrap(), + value.unwrap().as_slice(), b"hello world", "Append should concatenate values" ); // Test 3: Append that would exceed size limit (should not append) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("append_large",)); @@ -1584,23 +1620,24 @@ async fn test_atomic_append_if_fits(db: &Database) { .await .unwrap(); - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("append_large",)); // Try to append 2KB more (should not fit) let append_value = vec![b'y'; 2000]; - tx.atomic_op(&key, &append_value, MutationType::AppendIfFits); + tx.informal() + .atomic_op(&key, &append_value, MutationType::AppendIfFits); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("append_large",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await @@ -1622,7 +1659,7 @@ async fn test_atomic_min_max(db: &Database) { use universaldb::options::MutationType; // Test Max operation - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("max_key",)); @@ -1634,55 +1671,57 @@ async fn test_atomic_min_max(db: &Database) { .unwrap(); // Max with larger value (should replace) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("max_key",)); - tx.atomic_op(&key, &20i64.to_le_bytes(), MutationType::Max); + tx.informal() + .atomic_op(&key, &20i64.to_le_bytes(), MutationType::Max); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("max_key",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!(result, 20, "Max should select the larger value"); // Max with smaller value (should not replace) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("max_key",)); - tx.atomic_op(&key, &15i64.to_le_bytes(), MutationType::Max); + tx.informal() + .atomic_op(&key, &15i64.to_le_bytes(), MutationType::Max); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("max_key",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!(result, 20, "Max should keep the larger value"); // Test Min operation - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("min_key",)); @@ -1694,75 +1733,78 @@ async fn test_atomic_min_max(db: &Database) { .unwrap(); // Min with smaller value (should replace) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("min_key",)); - tx.atomic_op(&key, &5i64.to_le_bytes(), MutationType::Min); + tx.informal() + .atomic_op(&key, &5i64.to_le_bytes(), MutationType::Min); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("min_key",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!(result, 5, "Min should select the smaller value"); // Min with larger value (should not replace) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("min_key",)); - tx.atomic_op(&key, &15i64.to_le_bytes(), MutationType::Min); + tx.informal() + .atomic_op(&key, &15i64.to_le_bytes(), MutationType::Min); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("min_key",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!(result, 5, "Min should keep the smaller value"); // Test Max/Min with non-existent key - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("max_nonexistent",)); - tx.atomic_op(&key, &42i64.to_le_bytes(), MutationType::Max); + tx.informal() + .atomic_op(&key, &42i64.to_le_bytes(), MutationType::Max); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("max_nonexistent",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!(result, 42, "Max on non-existent key should set the value"); } @@ -1770,7 +1812,7 @@ async fn test_atomic_byte_min_max(db: &Database) { use universaldb::options::MutationType; // Test ByteMax operation (lexicographic comparison) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("byte_max",)); @@ -1782,61 +1824,63 @@ async fn test_atomic_byte_min_max(db: &Database) { .unwrap(); // ByteMax with lexicographically larger string (should replace) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("byte_max",)); - tx.atomic_op(&key, b"cherry", MutationType::ByteMax); + tx.informal() + .atomic_op(&key, b"cherry", MutationType::ByteMax); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("byte_max",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); assert_eq!( - value.unwrap(), + value.unwrap().as_slice(), b"cherry", "ByteMax should select lexicographically larger value" ); // ByteMax with lexicographically smaller string (should not replace) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("byte_max",)); - tx.atomic_op(&key, b"apple", MutationType::ByteMax); + tx.informal() + .atomic_op(&key, b"apple", MutationType::ByteMax); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("byte_max",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); assert_eq!( - value.unwrap(), + value.unwrap().as_slice(), b"cherry", "ByteMax should keep lexicographically larger value" ); // Test ByteMin operation - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("byte_min",)); @@ -1848,82 +1892,85 @@ async fn test_atomic_byte_min_max(db: &Database) { .unwrap(); // ByteMin with lexicographically smaller string (should replace) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("byte_min",)); - tx.atomic_op(&key, b"apple", MutationType::ByteMin); + tx.informal() + .atomic_op(&key, b"apple", MutationType::ByteMin); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("byte_min",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); assert_eq!( - value.unwrap(), + value.unwrap().as_slice(), b"apple", "ByteMin should select lexicographically smaller value" ); // ByteMin with lexicographically larger string (should not replace) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("byte_min",)); - tx.atomic_op(&key, b"cherry", MutationType::ByteMin); + tx.informal() + .atomic_op(&key, b"cherry", MutationType::ByteMin); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("byte_min",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); assert_eq!( - value.unwrap(), + value.unwrap().as_slice(), b"apple", "ByteMin should keep lexicographically smaller value" ); // Test ByteMin/ByteMax with non-existent key - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("byte_nonexistent",)); - tx.atomic_op(&key, b"first", MutationType::ByteMin); + tx.informal() + .atomic_op(&key, b"first", MutationType::ByteMin); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("byte_nonexistent",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); assert_eq!( - value.unwrap(), + value.unwrap().as_slice(), b"first", "ByteMin on non-existent key should set the value" ); @@ -1933,7 +1980,7 @@ async fn test_atomic_compare_and_clear(db: &Database) { use universaldb::options::MutationType; // Test 1: Compare and clear with matching value - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("cac_key1",)); @@ -1944,22 +1991,23 @@ async fn test_atomic_compare_and_clear(db: &Database) { .await .unwrap(); - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("cac_key1",)); // Compare and clear with matching value - tx.atomic_op(&key, b"target_value", MutationType::CompareAndClear); + tx.informal() + .atomic_op(&key, b"target_value", MutationType::CompareAndClear); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("cac_key1",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await @@ -1968,7 +2016,7 @@ async fn test_atomic_compare_and_clear(db: &Database) { assert_eq!(value, None, "Key should be cleared when values match"); // Test 2: Compare and clear with non-matching value - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("cac_key2",)); @@ -1979,50 +2027,52 @@ async fn test_atomic_compare_and_clear(db: &Database) { .await .unwrap(); - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("cac_key2",)); // Compare and clear with non-matching value - tx.atomic_op(&key, b"different_value", MutationType::CompareAndClear); + tx.informal() + .atomic_op(&key, b"different_value", MutationType::CompareAndClear); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("cac_key2",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); assert_eq!( - value.unwrap(), + value.unwrap().as_slice(), b"keep_this", "Key should remain unchanged when values don't match" ); // Test 3: Compare and clear on non-existent key - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("cac_nonexistent",)); // Compare and clear on non-existent key (treated as empty value) - tx.atomic_op(&key, b"", MutationType::CompareAndClear); + tx.informal() + .atomic_op(&key, b"", MutationType::CompareAndClear); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("cac_nonexistent",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await @@ -2034,7 +2084,7 @@ async fn test_atomic_compare_and_clear(db: &Database) { ); // Test 4: Compare and clear with empty value on existing key - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("cac_empty",)); @@ -2045,22 +2095,23 @@ async fn test_atomic_compare_and_clear(db: &Database) { .await .unwrap(); - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("cac_empty",)); // Compare and clear with empty value - tx.atomic_op(&key, b"", MutationType::CompareAndClear); + tx.informal() + .atomic_op(&key, b"", MutationType::CompareAndClear); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("cac_empty",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await @@ -2076,7 +2127,7 @@ async fn test_atomic_transaction_isolation(db: &Database) { use universaldb::options::MutationType; // Test that atomic operations within a transaction are visible to subsequent reads - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("isolation_key",)); @@ -2084,11 +2135,12 @@ async fn test_atomic_transaction_isolation(db: &Database) { tx.set(&key, &10i64.to_le_bytes()); // Perform atomic add - tx.atomic_op(&key, &5i64.to_le_bytes(), MutationType::Add); + tx.informal() + .atomic_op(&key, &5i64.to_le_bytes(), MutationType::Add); // Read the value within the same transaction - let value = tx.get(&key, false).await?; - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let value = tx.get(&key, Serializable).await?; + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!( result, 15, @@ -2096,11 +2148,12 @@ async fn test_atomic_transaction_isolation(db: &Database) { ); // Perform another atomic operation - tx.atomic_op(&key, &3i64.to_le_bytes(), MutationType::Add); + tx.informal() + .atomic_op(&key, &3i64.to_le_bytes(), MutationType::Add); // Read again - let value = tx.get(&key, false).await?; - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let value = tx.get(&key, Serializable).await?; + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!( result, 18, @@ -2114,7 +2167,7 @@ async fn test_atomic_transaction_isolation(db: &Database) { // Test that atomic operations are isolated between transactions // Set initial value in one transaction - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("isolation_key2",)); @@ -2125,11 +2178,12 @@ async fn test_atomic_transaction_isolation(db: &Database) { .unwrap(); // Perform atomic operation in another transaction - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("isolation_key2",)); - tx.atomic_op(&key, &50i64.to_le_bytes(), MutationType::Add); + tx.informal() + .atomic_op(&key, &50i64.to_le_bytes(), MutationType::Add); Ok(()) }) .await @@ -2137,16 +2191,16 @@ async fn test_atomic_transaction_isolation(db: &Database) { // Verify the result in a third transaction let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("isolation_key2",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!( result, 150, "Atomic operation should be committed and visible in new transaction" @@ -2159,45 +2213,47 @@ async fn test_atomic_nonexistent_keys(db: &Database) { // Test atomic operations on non-existent keys behave correctly // Test Add (should treat as 0) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("nonexistent_add",)); - tx.atomic_op(&key, &42i64.to_le_bytes(), MutationType::Add); + tx.informal() + .atomic_op(&key, &42i64.to_le_bytes(), MutationType::Add); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("nonexistent_add",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!(result, 42, "Add on non-existent key should treat as 0"); // Test BitOr (should treat as 0) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("nonexistent_or",)); - tx.atomic_op(&key, &[0b11110000], MutationType::BitOr); + tx.informal() + .atomic_op(&key, &[0b11110000], MutationType::BitOr); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("nonexistent_or",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await @@ -2210,75 +2266,78 @@ async fn test_atomic_nonexistent_keys(db: &Database) { ); // Test Max (should set the parameter value) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("nonexistent_max",)); - tx.atomic_op(&key, &123i64.to_le_bytes(), MutationType::Max); + tx.informal() + .atomic_op(&key, &123i64.to_le_bytes(), MutationType::Max); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("nonexistent_max",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); - let result = i64::from_le_bytes(value.unwrap().try_into().unwrap()); + let result = i64::from_le_bytes(Vec::from(value.unwrap()).try_into().unwrap()); assert_eq!( result, 123, "Max on non-existent key should set the parameter value" ); // Test ByteMin (should set the parameter value) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("nonexistent_bytemin",)); - tx.atomic_op(&key, b"hello", MutationType::ByteMin); + tx.informal() + .atomic_op(&key, b"hello", MutationType::ByteMin); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("nonexistent_bytemin",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await .unwrap(); assert_eq!( - value.unwrap(), + value.unwrap().as_slice(), b"hello", "ByteMin on non-existent key should set the parameter value" ); // Test CompareAndClear with empty comparison (should clear since non-existent = empty) - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("nonexistent_cac",)); - tx.atomic_op(&key, b"", MutationType::CompareAndClear); + tx.informal() + .atomic_op(&key, b"", MutationType::CompareAndClear); Ok(()) }) .await .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test"); let key = test_subspace.pack(&("nonexistent_cac",)); - let val = tx.get(&key, false).await?; + let val = tx.get(&key, Serializable).await?; Ok(val) }) .await @@ -2292,7 +2351,7 @@ async fn test_atomic_nonexistent_keys(db: &Database) { async fn test_versionstamps(db: &Database) { // Test 1: Basic versionstamp insertion and ordering within a single transaction - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test_vs"); // Create multiple values with incomplete versionstamps in the same transaction @@ -2317,7 +2376,7 @@ async fn test_versionstamps(db: &Database) { // Verify that versionstamps were substituted and have the same transaction version // but different user versions (counter values) let results = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test_vs"); let (begin, end) = test_subspace.range(); @@ -2327,7 +2386,7 @@ async fn test_versionstamps(db: &Database) { ..RangeOption::default() }; - let values = tx.get_range(&range_opt, 1, false).await?; + let values = tx.get_range(&range_opt, 1, Serializable).await?; let mut results = Vec::new(); for kv in values.into_iter() { @@ -2375,7 +2434,7 @@ async fn test_versionstamps(db: &Database) { for i in 0..3 { let vs = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test_vs"); let incomplete = Versionstamp::from([0xff; 12]); @@ -2401,7 +2460,7 @@ async fn test_versionstamps(db: &Database) { // Read back and verify ordering let multi_tx_results = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test_vs"); let begin = test_subspace.pack(&("tx_",)); let end = test_subspace.pack(&("tx_z",)); @@ -2412,7 +2471,7 @@ async fn test_versionstamps(db: &Database) { ..RangeOption::default() }; - let values = tx.get_range(&range_opt, 1, false).await?; + let values = tx.get_range(&range_opt, 1, Serializable).await?; let mut results = Vec::new(); for kv in values.into_iter() { @@ -2444,7 +2503,7 @@ async fn test_versionstamps(db: &Database) { } // Test 3: Already complete versionstamps should not be modified - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test_vs"); // Create a complete versionstamp manually @@ -2468,10 +2527,10 @@ async fn test_versionstamps(db: &Database) { // Read back and verify the versionstamp remains unchanged let complete_result = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test_vs"); let key = test_subspace.pack(&("complete_entry",)); - let value = tx.get(&key, false).await?.unwrap(); + let value = tx.get(&key, Serializable).await?.unwrap(); let unpacked: Vec = universaldb::tuple::unpack(&value).unwrap(); if let Element::Versionstamp(vs) = &unpacked[1] { @@ -2488,7 +2547,7 @@ async fn test_versionstamps(db: &Database) { // Test 4: Verify correct count and order within a transaction // Insert 10 entries in one transaction and verify they have sequential counters - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test_vs"); for i in 0..10 { @@ -2510,7 +2569,7 @@ async fn test_versionstamps(db: &Database) { // Read back and verify count and ordering let count_results = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test_vs"); let begin = test_subspace.pack(&("count_",)); let end = test_subspace.pack(&("count_z",)); @@ -2521,7 +2580,7 @@ async fn test_versionstamps(db: &Database) { ..RangeOption::default() }; - let values = tx.get_range(&range_opt, 1, false).await?; + let values = tx.get_range(&range_opt, 1, Serializable).await?; let mut results = Vec::new(); for kv in values.into_iter() { @@ -2566,7 +2625,7 @@ async fn test_versionstamps(db: &Database) { } // Test 5: Mixed incomplete and complete versionstamps in same transaction - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { let test_subspace = Subspace::from("test_vs"); // Insert an incomplete versionstamp @@ -2598,7 +2657,7 @@ async fn test_versionstamps(db: &Database) { // Verify both were stored correctly let mixed_results = db - .run(|tx, _maybe_committed| async move { + .run(|tx| async move { let test_subspace = Subspace::from("test_vs"); let begin = test_subspace.pack(&("mixed_",)); let end = test_subspace.pack(&("mixed_z",)); @@ -2609,7 +2668,7 @@ async fn test_versionstamps(db: &Database) { ..RangeOption::default() }; - let values = tx.get_range(&range_opt, 1, false).await?; + let values = tx.get_range(&range_opt, 1, Serializable).await?; let mut results = Vec::new(); for kv in values.into_iter() { diff --git a/packages/common/universaldb/tests/integration_gas.rs b/packages/common/universaldb/tests/integration_gas.rs index c3b48b77aa..26aec44a9d 100644 --- a/packages/common/universaldb/tests/integration_gas.rs +++ b/packages/common/universaldb/tests/integration_gas.rs @@ -1,7 +1,9 @@ use futures_util::TryStreamExt; use rivet_test_deps_docker::TestDatabase; use std::sync::Arc; -use universaldb::{Database, RangeOption, options::StreamingMode, tuple::Subspace}; +use universaldb::{ + Database, RangeOption, options::StreamingMode, tuple::Subspace, utils::IsolationLevel::*, +}; use uuid::Uuid; #[tokio::test] @@ -51,7 +53,7 @@ pub async fn test_gasoline_operations(db: &Database) { println!("Running simple write/read test..."); // Simple test: write a single key and read it back - db.run(|tx, _maybe_committed| async move { + db.run(|tx| async move { tx.set(b"simple_test_key", b"simple_test_value"); Ok(()) }) @@ -59,8 +61,8 @@ pub async fn test_gasoline_operations(db: &Database) { .unwrap(); let value = db - .run(|tx, _maybe_committed| async move { - let val = tx.get(b"simple_test_key", false).await?; + .run(|tx| async move { + let val = tx.get(b"simple_test_key", Serializable).await?; println!( "Simple test read result: {:?}", val.as_ref() @@ -87,7 +89,7 @@ pub async fn test_gasoline_operations(db: &Database) { let workflow_id = Uuid::new_v4(); // Test 1: Write workflow data like gasoline does - db.run(|tx, _maybe_committed| { + db.run(|tx| { let workflow_subspace = workflow_subspace.clone(); async move { // Write create timestamp (similar to CreateTsKey) @@ -128,7 +130,7 @@ pub async fn test_gasoline_operations(db: &Database) { // Test 2: Read workflow data back like gasoline does let (input_found, state_found, wake_found) = db - .run(|tx, _maybe_committed| { + .run(|tx| { let workflow_subspace = workflow_subspace.clone(); async move { // Read input chunks using range query @@ -141,7 +143,7 @@ pub async fn test_gasoline_operations(db: &Database) { mode: StreamingMode::WantAll, ..(&input_subspace).into() }, - false, + Serializable, ) .try_collect::>() .await?; @@ -156,7 +158,7 @@ pub async fn test_gasoline_operations(db: &Database) { mode: StreamingMode::WantAll, ..(&state_subspace).into() }, - false, + Serializable, ) .try_collect::>() .await?; @@ -164,7 +166,7 @@ pub async fn test_gasoline_operations(db: &Database) { // Read wake condition let wake_condition_key = workflow_subspace.pack(&("workflow", "has_wake_condition", workflow_id)); - let wake_condition = tx.get(&wake_condition_key, false).await?; + let wake_condition = tx.get(&wake_condition_key, Serializable).await?; println!("Input chunks found: {}", input_chunks.len()); println!("State chunks found: {}", state_chunks.len()); @@ -185,7 +187,7 @@ pub async fn test_gasoline_operations(db: &Database) { assert!(wake_found, "Should find wake condition"); // Test 3: Test the exact pattern gasoline uses with subspace operations - db.run(|tx, _maybe_committed| { + db.run(|tx| { let workflow_subspace = workflow_subspace.clone(); async move { // Create a new workflow ID @@ -208,7 +210,7 @@ pub async fn test_gasoline_operations(db: &Database) { // Test 4: Verify the data was written correctly let workflow_id2 = db - .run(|tx, _maybe_committed| { + .run(|tx| { let workflow_subspace = workflow_subspace.clone(); async move { // Generate the same workflow_id2 again (for test purposes, we'll store it) @@ -223,10 +225,10 @@ pub async fn test_gasoline_operations(db: &Database) { tx.set(&chunk_key, b"test_data"); // Read it back in the same transaction - let value = tx.get(&chunk_key, false).await?; + let value = tx.get(&chunk_key, Serializable).await?; assert_eq!( value, - Some(b"test_data".to_vec()), + Some(b"test_data".to_vec().into()), "Should read back the same data" ); @@ -237,7 +239,7 @@ pub async fn test_gasoline_operations(db: &Database) { .unwrap(); // Test 5: Read in a separate transaction (like gasoline does) - db.run(|tx, _maybe_committed| { + db.run(|tx| { let workflow_subspace = workflow_subspace.clone(); async move { let input_key_base = workflow_subspace.pack(&("workflow", "input", workflow_id2)); @@ -245,7 +247,7 @@ pub async fn test_gasoline_operations(db: &Database) { // Read using range query like gasoline let input_chunks = tx - .get_ranges_keyvalues((&input_subspace).into(), false) + .get_ranges_keyvalues((&input_subspace).into(), Serializable) .try_collect::>() .await?; diff --git a/packages/common/util/id/Cargo.toml b/packages/common/util/id/Cargo.toml index 227be8b251..b2e5988d06 100644 --- a/packages/common/util/id/Cargo.toml +++ b/packages/common/util/id/Cargo.toml @@ -6,8 +6,8 @@ license.workspace = true edition.workspace = true [dependencies] -udb-util.workspace = true serde.workspace = true thiserror.workspace = true -uuid.workspace = true +universaldb.workspace = true utoipa.workspace = true +uuid.workspace = true diff --git a/packages/common/util/id/src/lib.rs b/packages/common/util/id/src/lib.rs index e10fef2781..c6c4c9338f 100644 --- a/packages/common/util/id/src/lib.rs +++ b/packages/common/util/id/src/lib.rs @@ -1,7 +1,7 @@ use std::{fmt, str::FromStr}; use thiserror::Error; -use udb_util::prelude::*; +use universaldb::prelude::*; use uuid::Uuid; #[derive(Debug, Error)] @@ -226,7 +226,7 @@ impl TuplePack for Id { ) -> std::io::Result { let mut offset = VersionstampOffset::None { size: 0 }; - w.write_all(&[udb_util::codes::ID])?; + w.write_all(&[universaldb::utils::codes::ID])?; let bytes = self.as_bytes(); @@ -242,14 +242,14 @@ impl TuplePack for Id { impl<'de> TupleUnpack<'de> for Id { fn unpack(input: &[u8], _tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let input = udb_util::parse_code(input, udb_util::codes::ID)?; - let (_, version) = udb_util::parse_byte(input)?; + let input = universaldb::utils::parse_code(input, universaldb::utils::codes::ID)?; + let (_, version) = universaldb::utils::parse_byte(input)?; let (input, slice) = if version == 1 { // Parse 19 bytes including version - udb_util::parse_bytes(input, 19)? + universaldb::utils::parse_bytes(input, 19)? } else { - udb_util::parse_bytes(input, 1)? + universaldb::utils::parse_bytes(input, 1)? }; let v = Id::from_slice(slice) diff --git a/packages/core/actor-kv/Cargo.toml b/packages/core/actor-kv/Cargo.toml index f641fc2d05..8c92c45a1f 100644 --- a/packages/core/actor-kv/Cargo.toml +++ b/packages/core/actor-kv/Cargo.toml @@ -8,6 +8,7 @@ edition.workspace = true [dependencies] anyhow.workspace = true futures-util.workspace = true +rivet-runner-protocol.workspace = true rivet-util-id.workspace = true serde_bare.workspace = true serde.workspace = true @@ -15,8 +16,6 @@ tokio.workspace = true tracing-logfmt.workspace = true tracing-subscriber.workspace = true tracing.workspace = true -udb-util.workspace = true universaldb.workspace = true -rivet-runner-protocol.workspace = true pegboard.workspace = true diff --git a/packages/core/actor-kv/src/entry.rs b/packages/core/actor-kv/src/entry.rs index 143348be2f..81a57d0d96 100644 --- a/packages/core/actor-kv/src/entry.rs +++ b/packages/core/actor-kv/src/entry.rs @@ -1,7 +1,7 @@ use std::result::Result::Ok; use anyhow::*; -use udb_util::prelude::*; +use universaldb::prelude::*; use rivet_runner_protocol as rp; diff --git a/packages/core/actor-kv/src/key.rs b/packages/core/actor-kv/src/key.rs index a186caed6f..a10bbd2e44 100644 --- a/packages/core/actor-kv/src/key.rs +++ b/packages/core/actor-kv/src/key.rs @@ -20,12 +20,12 @@ impl TuplePack for KeyWrapper { ) -> std::io::Result { let mut offset = VersionstampOffset::None { size: 0 }; - w.write_all(&[udb_util::codes::NESTED])?; + w.write_all(&[universaldb::utils::codes::NESTED])?; offset += 1; offset += self.0.pack(w, tuple_depth.increment())?; - w.write_all(&[udb_util::codes::NIL])?; + w.write_all(&[universaldb::utils::codes::NIL])?; offset += 1; Ok(offset) @@ -34,11 +34,11 @@ impl TuplePack for KeyWrapper { impl<'de> TupleUnpack<'de> for KeyWrapper { fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let input = udb_util::parse_code(input, udb_util::codes::NESTED)?; + let input = universaldb::utils::parse_code(input, universaldb::utils::codes::NESTED)?; let (input, inner) = Bytes::unpack(input, tuple_depth.increment())?; - let input = udb_util::parse_code(input, udb_util::codes::NIL)?; + let input = universaldb::utils::parse_code(input, universaldb::utils::codes::NIL)?; Ok((input, KeyWrapper(inner.into_owned()))) } @@ -55,7 +55,7 @@ impl TuplePack for ListKeyWrapper { ) -> std::io::Result { let mut offset = VersionstampOffset::None { size: 0 }; - w.write_all(&[udb_util::codes::NESTED])?; + w.write_all(&[universaldb::utils::codes::NESTED])?; offset += 1; offset += self.0.pack(w, tuple_depth.increment())?; diff --git a/packages/core/actor-kv/src/lib.rs b/packages/core/actor-kv/src/lib.rs index 253e36c458..f42eb14280 100644 --- a/packages/core/actor-kv/src/lib.rs +++ b/packages/core/actor-kv/src/lib.rs @@ -6,8 +6,8 @@ use futures_util::{StreamExt, TryStreamExt}; use key::{KeyWrapper, ListKeyWrapper}; use rivet_runner_protocol as rp; use rivet_util_id::Id; -use udb_util::prelude::*; -use universaldb::{self as udb, tuple::Subspace}; +use universaldb::prelude::*; +use universaldb::tuple::Subspace; use utils::{validate_entries, validate_keys}; mod entry; @@ -22,12 +22,12 @@ const MAX_PUT_PAYLOAD_SIZE: usize = 976 * 1024; const MAX_STORAGE_SIZE: usize = 1024 * 1024 * 1024; // 1 GiB const VALUE_CHUNK_SIZE: usize = 10_000; // 10 KB, not KiB, see https://apple.github.io/foundationdb/blob.html -fn subspace(actor_id: Id) -> udb_util::Subspace { +fn subspace(actor_id: Id) -> universaldb::utils::Subspace { pegboard::keys::actor_kv_subspace().subspace(&actor_id) } /// Returns estimated size of the given subspace. -pub async fn get_subspace_size(db: &udb::Database, subspace: &Subspace) -> Result { +pub async fn get_subspace_size(db: &universaldb::Database, subspace: &Subspace) -> Result { let (start, end) = subspace.range(); // This txn does not have to be committed because we are not modifying any data @@ -39,30 +39,30 @@ pub async fn get_subspace_size(db: &udb::Database, subspace: &Subspace) -> Resul /// Gets keys from the KV store. pub async fn get( - db: &udb::Database, + db: &universaldb::Database, actor_id: Id, keys: Vec, ) -> Result<(Vec, Vec, Vec)> { validate_keys(&keys)?; - db.run(|tx, _mc| { + db.run(|tx| { let keys = keys.clone(); async move { - let txs = tx.subspace(subspace(actor_id)); + let tx = tx.with_subspace(subspace(actor_id)); let size_estimate = keys.len().min(1024); let mut stream = futures_util::stream::iter(keys) .map(|key| { - let key_subspace = txs.subspace(&KeyWrapper(key)); + let key_subspace = subspace(actor_id).subspace(&KeyWrapper(key)); // Get all sub keys in the key subspace - txs.get_ranges_keyvalues( - udb::RangeOption { - mode: udb::options::StreamingMode::WantAll, + tx.get_ranges_keyvalues( + universaldb::RangeOption { + mode: universaldb::options::StreamingMode::WantAll, ..key_subspace.range().into() }, - false, + Serializable, ) }) // Should remain in order @@ -79,13 +79,12 @@ pub async fn get( break; }; - let key = txs.unpack::(&entry.key())?.key; + let key = tx.unpack::(&entry.key())?.key; let current_entry = if let Some(inner) = &mut current_entry { if inner.key != key { - let (key, value, meta) = std::mem::replace(inner, EntryBuilder::new(key)) - .build() - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let (key, value, meta) = + std::mem::replace(inner, EntryBuilder::new(key)).build()?; keys.push(key); values.push(value); @@ -99,25 +98,19 @@ pub async fn get( current_entry.as_mut().expect("must be set") }; - if let Ok(chunk_key) = txs.unpack::(&entry.key()) { + if let Ok(chunk_key) = tx.unpack::(&entry.key()) { current_entry.append_chunk(chunk_key.chunk, entry.value()); - } else if let Ok(metadata_key) = txs.unpack::(&entry.key()) { - let value = metadata_key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + } else if let Ok(metadata_key) = tx.unpack::(&entry.key()) { + let value = metadata_key.deserialize(entry.value())?; current_entry.append_metadata(value); } else { - return Err(udb::FdbBindingError::CustomError( - "unexpected sub key".into(), - )); + bail!("unexpected sub key"); } } if let Some(inner) = current_entry { - let (key, value, meta) = inner - .build() - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let (key, value, meta) = inner.build()?; keys.push(key); values.push(value); @@ -133,7 +126,7 @@ pub async fn get( /// Gets keys from the KV store. pub async fn list( - db: &udb::Database, + db: &universaldb::Database, actor_id: Id, query: rp::KvListQuery, reverse: bool, @@ -145,20 +138,20 @@ pub async fn list( let subspace = subspace(actor_id); let list_range = list_query_range(query, &subspace); - db.run(|tx, _mc| { + db.run(|tx| { let list_range = list_range.clone(); let subspace = subspace.clone(); async move { - let txs = tx.subspace(subspace); + let tx = tx.with_subspace(subspace); - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { - mode: udb::options::StreamingMode::Iterator, + let mut stream = tx.get_ranges_keyvalues( + universaldb::RangeOption { + mode: universaldb::options::StreamingMode::Iterator, reverse, ..list_range.into() }, - false, + Serializable, ); let mut keys = Vec::new(); @@ -171,13 +164,12 @@ pub async fn list( break; }; - let key = txs.unpack::(&entry.key())?.key; + let key = tx.unpack::(&entry.key())?.key; let curr = if let Some(inner) = &mut current_entry { if inner.key != key { - let (key, value, meta) = std::mem::replace(inner, EntryBuilder::new(key)) - .build() - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let (key, value, meta) = + std::mem::replace(inner, EntryBuilder::new(key)).build()?; keys.push(key); values.push(value); @@ -196,25 +188,19 @@ pub async fn list( current_entry.as_mut().expect("must be set") }; - if let Ok(chunk_key) = txs.unpack::(&entry.key()) { + if let Ok(chunk_key) = tx.unpack::(&entry.key()) { curr.append_chunk(chunk_key.chunk, entry.value()); - } else if let Ok(metadata_key) = txs.unpack::(&entry.key()) { - let value = metadata_key - .deserialize(entry.value()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + } else if let Ok(metadata_key) = tx.unpack::(&entry.key()) { + let value = metadata_key.deserialize(entry.value())?; curr.append_metadata(value); } else { - return Err(udb::FdbBindingError::CustomError( - "unexpected sub key".into(), - )); + bail!("unexpected sub key"); } } if let Some(inner) = current_entry { - let (key, value, meta) = inner - .build() - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let (key, value, meta) = inner.build()?; keys.push(key); values.push(value); @@ -230,7 +216,7 @@ pub async fn list( /// Puts keys into the KV store. pub async fn put( - db: &udb::Database, + db: &universaldb::Database, actor_id: Id, keys: Vec, values: Vec, @@ -240,27 +226,27 @@ pub async fn put( validate_entries(&keys, &values, total_size)?; - db.run(|tx, _mc| { + db.run(|tx| { // TODO: Costly clone let keys = keys.clone(); let values = values.clone(); let subspace = subspace.clone(); async move { - let txs = tx.subspace(subspace.clone()); + let tx = tx.with_subspace(subspace.clone()); futures_util::stream::iter(keys.into_iter().zip(values.into_iter())) .map(|(key, value)| { - let txs = txs.clone(); + let tx = tx.clone(); let key = KeyWrapper(key.clone()); let subspace = subspace.clone(); async move { // Clear previous key data before setting - txs.clear_subspace_range(&subspace.subspace(&key)); + tx.clear_subspace_range(&subspace.subspace(&key)); // Set metadata - txs.write( + tx.write( &EntryMetadataKey::new(key.clone()), rp::KvMetadata { version: VERSION.as_bytes().to_vec(), @@ -273,12 +259,9 @@ pub async fn put( let idx = start / VALUE_CHUNK_SIZE; let end = (start + VALUE_CHUNK_SIZE).min(value.len()); - txs.set( + tx.set( &subspace.pack(&EntryValueChunkKey::new(key.clone(), idx)), - &value - .get(start..end) - .context("bad slice") - .map_err(|err| udb::FdbBindingError::CustomError(err.into()))?, + &value.get(start..end).context("bad slice")?, ); } @@ -295,10 +278,10 @@ pub async fn put( } /// Deletes keys from the KV store. Cannot be undone. -pub async fn delete(db: &udb::Database, actor_id: Id, keys: Vec) -> Result<()> { +pub async fn delete(db: &universaldb::Database, actor_id: Id, keys: Vec) -> Result<()> { validate_keys(&keys)?; - db.run(|tx, _mc| { + db.run(|tx| { let keys = keys.clone(); async move { for key in keys { @@ -315,8 +298,8 @@ pub async fn delete(db: &udb::Database, actor_id: Id, keys: Vec) -> R } /// Deletes all keys from the KV store. Cannot be undone. -pub async fn delete_all(db: &udb::Database, actor_id: Id) -> Result<()> { - db.run(|tx, _mc| async move { +pub async fn delete_all(db: &universaldb::Database, actor_id: Id) -> Result<()> { + db.run(|tx| async move { tx.clear_subspace_range(&subspace(actor_id)); Ok(()) }) diff --git a/packages/core/guard/server/Cargo.toml b/packages/core/guard/server/Cargo.toml index ebf82c812f..0b9944e328 100644 --- a/packages/core/guard/server/Cargo.toml +++ b/packages/core/guard/server/Cargo.toml @@ -18,25 +18,22 @@ http-body.workspace = true http-body-util.workspace = true hyper-tungstenite.workspace = true tower.workspace = true -udb-util.workspace = true -universaldb.workspace = true -universalpubsub.workspace = true futures.workspace = true # TODO: Make this use workspace version hyper = "1.6.0" indoc.workspace = true once_cell.workspace = true -pegboard.workspace = true pegboard-gateway.workspace = true pegboard-tunnel.workspace = true +pegboard.workspace = true regex.workspace = true rivet-api-peer.workspace = true rivet-api-public.workspace = true rivet-cache.workspace = true rivet-config.workspace = true +rivet-data.workspace = true rivet-error.workspace = true rivet-guard-core.workspace = true -rivet-data.workspace = true rivet-logs.workspace = true rivet-metrics.workspace = true rivet-pools.workspace = true @@ -47,6 +44,8 @@ serde_json.workspace = true serde.workspace = true tokio.workspace = true tracing.workspace = true +universaldb.workspace = true +universalpubsub.workspace = true url.workspace = true uuid.workspace = true diff --git a/packages/core/guard/server/src/routing/pegboard_gateway.rs b/packages/core/guard/server/src/routing/pegboard_gateway.rs index 7015a8a686..4a142d7526 100644 --- a/packages/core/guard/server/src/routing/pegboard_gateway.rs +++ b/packages/core/guard/server/src/routing/pegboard_gateway.rs @@ -4,7 +4,7 @@ use anyhow::Result; use gas::prelude::*; use hyper::header::HeaderName; use rivet_guard_core::proxy_service::{RouteConfig, RouteTarget, RoutingOutput, RoutingTimeout}; -use udb_util::{SERIALIZABLE, TxnExt}; +use universaldb::utils::IsolationLevel::*; use crate::{errors, shared_state::SharedState}; @@ -98,17 +98,17 @@ async fn find_actor( let actor_res = tokio::time::timeout( Duration::from_secs(5), ctx.udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(pegboard::keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(pegboard::keys::subspace()); let workflow_id_key = pegboard::keys::actor::WorkflowIdKey::new(actor_id); let sleep_ts_key = pegboard::keys::actor::SleepTsKey::new(actor_id); let destroy_ts_key = pegboard::keys::actor::DestroyTsKey::new(actor_id); let (workflow_id_entry, sleeping, destroyed) = tokio::try_join!( - txs.read_opt(&workflow_id_key, SERIALIZABLE), - txs.exists(&sleep_ts_key, SERIALIZABLE), - txs.exists(&destroy_ts_key, SERIALIZABLE), + tx.read_opt(&workflow_id_key, Serializable), + tx.exists(&sleep_ts_key, Serializable), + tx.exists(&destroy_ts_key, Serializable), )?; let Some(workflow_id) = workflow_id_entry else { @@ -177,10 +177,13 @@ async fn find_actor( // Get runner key from runner_id let runner_key = ctx .udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(pegboard::keys::subspace()); - let key_key = pegboard::keys::runner::KeyKey::new(runner_id); - txs.read_opt(&key_key, SERIALIZABLE).await + .run(|tx| async move { + let tx = tx.with_subspace(pegboard::keys::subspace()); + tx.read_opt( + &pegboard::keys::runner::KeyKey::new(runner_id), + Serializable, + ) + .await }) .await? .context("runner key not found")?; diff --git a/packages/core/guard/server/src/shared_state.rs b/packages/core/guard/server/src/shared_state.rs index 71bc25939e..0462fad5c8 100644 --- a/packages/core/guard/server/src/shared_state.rs +++ b/packages/core/guard/server/src/shared_state.rs @@ -1,5 +1,4 @@ use anyhow::*; -use gas::prelude::*; use std::{ops::Deref, sync::Arc}; use universalpubsub::PubSub; diff --git a/packages/core/pegboard-gateway/Cargo.toml b/packages/core/pegboard-gateway/Cargo.toml index 6a3ddc446a..bf4eba0740 100644 --- a/packages/core/pegboard-gateway/Cargo.toml +++ b/packages/core/pegboard-gateway/Cargo.toml @@ -20,8 +20,8 @@ rivet-error.workspace = true rivet-guard-core.workspace = true rivet-tunnel-protocol.workspace = true rivet-util.workspace = true +thiserror.workspace = true tokio-tungstenite.workspace = true tokio.workspace = true universalpubsub.workspace = true versioned-data-util.workspace = true -thiserror.workspace = true diff --git a/packages/core/pegboard-serverless/Cargo.toml b/packages/core/pegboard-serverless/Cargo.toml index 78eaca5978..44b97cbea1 100644 --- a/packages/core/pegboard-serverless/Cargo.toml +++ b/packages/core/pegboard-serverless/Cargo.toml @@ -14,7 +14,6 @@ rivet-config.workspace = true rivet-runner-protocol.workspace = true rivet-types.workspace = true tracing.workspace = true -udb-util.workspace = true universaldb.workspace = true namespace.workspace = true diff --git a/packages/core/pegboard-serverless/src/lib.rs b/packages/core/pegboard-serverless/src/lib.rs index 9d6ff472bf..c069dc050e 100644 --- a/packages/core/pegboard-serverless/src/lib.rs +++ b/packages/core/pegboard-serverless/src/lib.rs @@ -14,8 +14,8 @@ use pegboard::keys; use reqwest_eventsource as sse; use rivet_runner_protocol::protocol; use tokio::{sync::oneshot, task::JoinHandle, time::Duration}; -use udb_util::{SNAPSHOT, TxnExt}; -use universaldb::{self as udb, options::StreamingMode}; +use universaldb::options::StreamingMode; +use universaldb::utils::IsolationLevel::*; struct OutboundConnection { handle: JoinHandle<()>, @@ -54,25 +54,25 @@ async fn tick( ) -> Result<()> { let serverless_data = ctx .udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); - let serverless_desired_subspace = txs.subspace( + let serverless_desired_subspace = keys::subspace().subspace( &rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::entire_subspace(), ); - txs.get_ranges_keyvalues( - udb::RangeOption { + tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&serverless_desired_subspace).into() }, // NOTE: This is a snapshot to prevent conflict with updates to this subspace - SNAPSHOT, + Snapshot, ) .map(|res| match res { Ok(entry) => { let (key, desired_slots) = - txs.read_entry::(&entry)?; + tx.read_entry::(&entry)?; Ok((key.namespace_id, key.runner_name, desired_slots)) } diff --git a/packages/infra/engine/Cargo.toml b/packages/infra/engine/Cargo.toml index f141806f6c..b654a96003 100644 --- a/packages/infra/engine/Cargo.toml +++ b/packages/infra/engine/Cargo.toml @@ -46,7 +46,6 @@ tempfile.workspace = true thiserror.workspace = true tokio.workspace = true tracing.workspace = true -udb-util.workspace = true universaldb.workspace = true url.workspace = true uuid.workspace = true diff --git a/packages/infra/engine/src/commands/udb/cli.rs b/packages/infra/engine/src/commands/udb/cli.rs index 0149ea5678..b4926f9275 100644 --- a/packages/infra/engine/src/commands/udb/cli.rs +++ b/packages/infra/engine/src/commands/udb/cli.rs @@ -7,7 +7,7 @@ use clap::{Parser, ValueEnum}; use futures_util::TryStreamExt; use rivet_pools::UdbPool; use rivet_term::console::style; -use universaldb::{self as udb, options::StreamingMode}; +use universaldb::{options::StreamingMode, utils::IsolationLevel::*}; use crate::util::{ format::indent_string, @@ -123,11 +123,11 @@ impl SubCommand { return CommandResult::Error; } - let fut = pool.run(|tx, _mc| { + let fut = pool.run(|tx| { let current_tuple = current_tuple.clone(); async move { - let key = udb::tuple::pack(¤t_tuple); - let entry = tx.get(&key, true).await?; + let key = universaldb::tuple::pack(¤t_tuple); + let entry = tx.get(&key, Snapshot).await?; Ok(entry) } }); @@ -166,18 +166,18 @@ impl SubCommand { return CommandResult::Error; } - let subspace = udb::tuple::Subspace::all().subspace(¤t_tuple); + let subspace = universaldb::tuple::Subspace::all().subspace(¤t_tuple); - let fut = pool.run(|tx, _mc| { + let fut = pool.run(|tx| { let subspace = subspace.clone(); async move { let entries = tx .get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&subspace).into() }, - true, + Snapshot, ) .try_collect::>() .await?; @@ -355,21 +355,23 @@ impl SubCommand { return CommandResult::Error; } - let fut = pool.run(|tx, _mc| { + let fut = pool.run(|tx| { let old_tuple = old_tuple.clone(); let new_tuple = new_tuple.clone(); async move { if recursive { - let old_subspace = udb::tuple::Subspace::all().subspace(&old_tuple); - let new_subspace = udb::tuple::Subspace::all().subspace(&new_tuple); + let old_subspace = + universaldb::tuple::Subspace::all().subspace(&old_tuple); + let new_subspace = + universaldb::tuple::Subspace::all().subspace(&new_tuple); // Get all key-value pairs from the old subspace let mut stream = tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&old_subspace).into() }, - true, + Snapshot, ); let mut keys_moved = 0; @@ -393,10 +395,10 @@ impl SubCommand { Ok(keys_moved) } else { - let old_key = udb::tuple::pack(&old_tuple); - let new_key = udb::tuple::pack(&new_tuple); + let old_key = universaldb::tuple::pack(&old_tuple); + let new_key = universaldb::tuple::pack(&new_tuple); - let Some(value) = tx.get(&old_key, true).await? else { + let Some(value) = tx.get(&old_key, Snapshot).await? else { return Ok(0); }; @@ -451,14 +453,12 @@ impl SubCommand { } }; - let fut = pool.run(|tx, _mc| { + let fut = pool.run(|tx| { let current_tuple = current_tuple.clone(); let value = parsed_value.clone(); async move { - let key = udb::tuple::pack(¤t_tuple); - let value = value - .serialize() - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let key = universaldb::tuple::pack(¤t_tuple); + let value = value.serialize()?; tx.set(&key, &value); Ok(()) @@ -495,14 +495,15 @@ impl SubCommand { } } - let fut = pool.run(|tx, _mc| { + let fut = pool.run(|tx| { let current_tuple = current_tuple.clone(); async move { if clear_range { - let subspace = udb::tuple::Subspace::all().subspace(¤t_tuple); + let subspace = + universaldb::utils::Subspace::all().subspace(¤t_tuple); tx.clear_subspace_range(&subspace); } else { - let key = udb::tuple::pack(¤t_tuple); + let key = universaldb::tuple::pack(¤t_tuple); tx.clear(&key); } diff --git a/packages/infra/engine/src/util/udb.rs b/packages/infra/engine/src/util/udb.rs index 7165a891e5..0f1c3efd1b 100644 --- a/packages/infra/engine/src/util/udb.rs +++ b/packages/infra/engine/src/util/udb.rs @@ -43,7 +43,9 @@ impl SimpleTupleValue { SimpleTupleValue::Uuid(v) } else if let Ok(v) = rivet_util::Id::from_str(value) { SimpleTupleValue::Id(v) - } else if let (true, Some(v)) = (convert_keys, udb_util::prelude::key_from_str(value)) { + } else if let (true, Some(v)) = + (convert_keys, universaldb::utils::keys::key_from_str(value)) + { SimpleTupleValue::U64(v as u64) } else if nested && value.trim().starts_with('[') && value.trim().ends_with(']') { let mut items = Vec::new(); @@ -227,7 +229,9 @@ impl SimpleTupleValue { match &self { SimpleTupleValue::U64(v) => { if let Ok(v) = (*v).try_into() { - if let (true, Some(key)) = (convert_keys, udb_util::prelude::str_from_key(v)) { + if let (true, Some(key)) = + (convert_keys, universaldb::utils::keys::str_from_key(v)) + { write!( f, "{} {}", diff --git a/packages/services/epoxy/Cargo.toml b/packages/services/epoxy/Cargo.toml index c3c5564eeb..b20f72c94d 100644 --- a/packages/services/epoxy/Cargo.toml +++ b/packages/services/epoxy/Cargo.toml @@ -31,7 +31,6 @@ slog.workspace = true tokio.workspace = true tracing-slog.workspace = true tracing.workspace = true -udb-util.workspace = true universaldb.workspace = true url.workspace = true uuid.workspace = true diff --git a/packages/services/epoxy/src/keys/keys.rs b/packages/services/epoxy/src/keys/keys.rs index 0ad3cf0c38..bbe22ba2d3 100644 --- a/packages/services/epoxy/src/keys/keys.rs +++ b/packages/services/epoxy/src/keys/keys.rs @@ -1,7 +1,7 @@ use anyhow::*; use epoxy_protocol::protocol::ReplicaId; use std::result::Result::Ok; -use udb_util::prelude::*; +use universaldb::prelude::*; #[derive(Debug, Clone)] pub struct KvValueKey { diff --git a/packages/services/epoxy/src/keys/mod.rs b/packages/services/epoxy/src/keys/mod.rs index 8a582bbd83..68c8010140 100644 --- a/packages/services/epoxy/src/keys/mod.rs +++ b/packages/services/epoxy/src/keys/mod.rs @@ -1,9 +1,9 @@ use epoxy_protocol::protocol::ReplicaId; -use udb_util::prelude::*; +use universaldb::prelude::*; pub mod keys; pub mod replica; -pub fn subspace(replica_id: ReplicaId) -> udb_util::Subspace { - udb_util::Subspace::new(&(RIVET, EPOXY, REPLICA, replica_id)) +pub fn subspace(replica_id: ReplicaId) -> universaldb::utils::Subspace { + universaldb::utils::Subspace::new(&(RIVET, EPOXY, REPLICA, replica_id)) } diff --git a/packages/services/epoxy/src/keys/replica.rs b/packages/services/epoxy/src/keys/replica.rs index 7a8f83eaea..809abb5fa4 100644 --- a/packages/services/epoxy/src/keys/replica.rs +++ b/packages/services/epoxy/src/keys/replica.rs @@ -1,7 +1,7 @@ use anyhow::*; use epoxy_protocol::protocol::{ReplicaId, SlotId}; -use udb_util::prelude::*; -use versioned_data_util::OwnedVersionedData as _; +use universaldb::prelude::*; +use versioned_data_util::OwnedVersionedData; #[derive(Debug)] pub struct InstanceNumberKey; diff --git a/packages/services/epoxy/src/ops/explicit_prepare.rs b/packages/services/epoxy/src/ops/explicit_prepare.rs index 97a2d85133..9fef6278be 100644 --- a/packages/services/epoxy/src/ops/explicit_prepare.rs +++ b/packages/services/epoxy/src/ops/explicit_prepare.rs @@ -1,7 +1,6 @@ use anyhow::*; use epoxy_protocol::protocol::{self, ReplicaId}; use gas::prelude::*; -use universaldb::FdbBindingError; use crate::{http_client, replica, types, utils}; @@ -33,21 +32,13 @@ pub async fn explicit_prepare(ctx: &OperationCtx, input: &Input) -> Result Result { let value = ctx .udb()? - .run(|tx, _| { + .run(|tx| { let packed_key = packed_key.clone(); let kv_key = kv_key.clone(); async move { (async move { - let value = tx.get(&packed_key, false).await?; + let value = tx.get(&packed_key, Serializable).await?; if let Some(v) = value { Ok(Some(kv_key.deserialize(&v)?)) } else { @@ -39,7 +39,6 @@ pub async fn get_local(ctx: &OperationCtx, input: &Input) -> Result { } }) .await - .map_err(|e: anyhow::Error| universaldb::FdbBindingError::CustomError(e.into())) } }) .await?; diff --git a/packages/services/epoxy/src/ops/kv/get_optimistic.rs b/packages/services/epoxy/src/ops/kv/get_optimistic.rs index 4a00dadef2..1f9802e68a 100644 --- a/packages/services/epoxy/src/ops/kv/get_optimistic.rs +++ b/packages/services/epoxy/src/ops/kv/get_optimistic.rs @@ -1,7 +1,7 @@ use anyhow::*; use epoxy_protocol::protocol::{self, ReplicaId}; use gas::prelude::*; -use udb_util::FormalKey; +use universaldb::utils::{FormalKey, IsolationLevel::*}; use crate::{http_client, keys, utils}; @@ -45,7 +45,7 @@ pub async fn get_optimistic(ctx: &OperationCtx, input: &Input) -> Result let value = ctx .udb()? - .run(|tx, _| { + .run(|tx| { let packed_key = packed_key.clone(); let packed_cache_key = packed_cache_key.clone(); let kv_key = kv_key.clone(); @@ -54,7 +54,7 @@ pub async fn get_optimistic(ctx: &OperationCtx, input: &Input) -> Result (async move { let (value, cache_value) = tokio::try_join!( async { - let v = tx.get(&packed_key, false).await?; + let v = tx.get(&packed_key, Serializable).await?; if let Some(ref bytes) = v { Ok(Some(kv_key.deserialize(bytes)?)) } else { @@ -62,7 +62,7 @@ pub async fn get_optimistic(ctx: &OperationCtx, input: &Input) -> Result } }, async { - let v = tx.get(&packed_cache_key, false).await?; + let v = tx.get(&packed_cache_key, Serializable).await?; if let Some(ref bytes) = v { Ok(Some(cache_key.deserialize(bytes)?)) } else { @@ -74,7 +74,6 @@ pub async fn get_optimistic(ctx: &OperationCtx, input: &Input) -> Result Ok(value.or(cache_value)) }) .await - .map_err(|e: anyhow::Error| universaldb::FdbBindingError::CustomError(e.into())) } }) .await?; @@ -129,7 +128,7 @@ pub async fn get_optimistic(ctx: &OperationCtx, input: &Input) -> Result if let Some(value) = response { // Cache value ctx.udb()? - .run(|tx, _| { + .run(|tx| { let packed_cache_key = packed_cache_key.clone(); let cache_key = cache_key.clone(); let value_to_cache = value.clone(); @@ -140,9 +139,6 @@ pub async fn get_optimistic(ctx: &OperationCtx, input: &Input) -> Result Ok(()) }) .await - .map_err(|e: anyhow::Error| { - universaldb::FdbBindingError::CustomError(e.into()) - }) } }) .await?; diff --git a/packages/services/epoxy/src/ops/propose.rs b/packages/services/epoxy/src/ops/propose.rs index dafacc6d32..6f9b05f7dc 100644 --- a/packages/services/epoxy/src/ops/propose.rs +++ b/packages/services/epoxy/src/ops/propose.rs @@ -3,7 +3,6 @@ use epoxy_protocol::protocol::{self, Path, Payload, ReplicaId}; use gas::prelude::*; use rivet_api_builder::prelude::*; use rivet_config::Config; -use universaldb::FdbBindingError; use crate::{http_client, replica, utils}; @@ -34,23 +33,15 @@ pub async fn propose(ctx: &OperationCtx, input: &Input) -> Result Result Result { let config = ctx .udb()? - .run(|tx, _| { + .run(|tx| { let replica_id = input.replica_id; - async move { - utils::read_config(&tx, replica_id) - .await - .map_err(|e: anyhow::Error| universaldb::FdbBindingError::CustomError(e.into())) - } + async move { utils::read_config(&tx, replica_id).await } }) .await?; diff --git a/packages/services/epoxy/src/replica/ballot.rs b/packages/services/epoxy/src/replica/ballot.rs index f81aa92d6c..ae52749299 100644 --- a/packages/services/epoxy/src/replica/ballot.rs +++ b/packages/services/epoxy/src/replica/ballot.rs @@ -1,6 +1,7 @@ +use anyhow::Result; use epoxy_protocol::protocol; -use udb_util::FormalKey; -use universaldb::{FdbBindingError, Transaction}; +use universaldb::Transaction; +use universaldb::utils::{FormalKey, IsolationLevel::*}; use crate::keys; @@ -8,16 +9,14 @@ use crate::keys; pub async fn get_ballot( tx: &Transaction, replica_id: protocol::ReplicaId, -) -> Result { +) -> Result { let ballot_key = keys::replica::CurrentBallotKey; let subspace = keys::subspace(replica_id); let packed_key = subspace.pack(&ballot_key); - match tx.get(&packed_key, false).await? { + match tx.get(&packed_key, Serializable).await? { Some(bytes) => { - let ballot = ballot_key - .deserialize(&bytes) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + let ballot = ballot_key.deserialize(&bytes)?; Ok(ballot) } None => { @@ -35,7 +34,7 @@ pub async fn get_ballot( pub async fn increment_ballot( tx: &Transaction, replica_id: protocol::ReplicaId, -) -> Result { +) -> Result { let mut current_ballot = get_ballot(tx, replica_id).await?; // Increment ballot number @@ -45,9 +44,7 @@ pub async fn increment_ballot( let ballot_key = keys::replica::CurrentBallotKey; let subspace = keys::subspace(replica_id); let packed_key = subspace.pack(&ballot_key); - let serialized = ballot_key - .serialize(current_ballot.clone()) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + let serialized = ballot_key.serialize(current_ballot.clone())?; tx.set(&packed_key, &serialized); @@ -75,18 +72,16 @@ pub async fn validate_and_update_ballot_for_instance( replica_id: protocol::ReplicaId, ballot: &protocol::Ballot, instance: &protocol::Instance, -) -> Result { +) -> Result { let instance_ballot_key = keys::replica::InstanceBallotKey::new(instance.replica_id, instance.slot_id); let subspace = keys::subspace(replica_id); let packed_key = subspace.pack(&instance_ballot_key); // Get the highest ballot seen for this instance - let highest_ballot = match tx.get(&packed_key, false).await? { + let highest_ballot = match tx.get(&packed_key, Serializable).await? { Some(bytes) => { - let stored_ballot = instance_ballot_key - .deserialize(&bytes) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + let stored_ballot = instance_ballot_key.deserialize(&bytes)?; stored_ballot } None => { @@ -107,9 +102,7 @@ pub async fn validate_and_update_ballot_for_instance( // If the incoming ballot is higher, update our stored highest if compare_ballots(ballot, &highest_ballot) == std::cmp::Ordering::Greater { - let serialized = instance_ballot_key - .serialize(ballot.clone()) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + let serialized = instance_ballot_key.serialize(ballot.clone())?; tx.set(&packed_key, &serialized); tracing::debug!(?ballot, ?instance, "updated highest ballot for instance"); diff --git a/packages/services/epoxy/src/replica/commit_kv.rs b/packages/services/epoxy/src/replica/commit_kv.rs index d44315a6e0..7a9fcfa450 100644 --- a/packages/services/epoxy/src/replica/commit_kv.rs +++ b/packages/services/epoxy/src/replica/commit_kv.rs @@ -1,7 +1,7 @@ use anyhow::*; use epoxy_protocol::protocol::{self, ReplicaId}; -use udb_util::prelude::*; -use universaldb::{FdbBindingError, Transaction}; +use universaldb::Transaction; +use universaldb::prelude::*; use crate::{keys, ops::propose::CommandError, replica::utils}; @@ -10,7 +10,7 @@ pub async fn commit_kv( tx: &Transaction, replica_id: ReplicaId, commands: &[protocol::Command], -) -> Result, FdbBindingError> { +) -> Result> { let subspace = keys::subspace(replica_id); for command in commands.iter() { @@ -23,12 +23,8 @@ pub async fn commit_kv( // Read current value let kv_key = keys::keys::KvValueKey::new(cmd.key.clone()); let packed_key = subspace.pack(&kv_key); - let current_value = if let Some(bytes) = tx.get(&packed_key, false).await? { - Some( - kv_key - .deserialize(&bytes) - .map_err(|x| FdbBindingError::CustomError(x.into()))?, - ) + let current_value = if let Some(bytes) = tx.get(&packed_key, Serializable).await? { + Some(kv_key.deserialize(&bytes)?) } else { None }; @@ -67,9 +63,7 @@ pub async fn commit_kv( // Update the value if let Some(value) = new_value { - let serialized = kv_key - .serialize(value.clone()) - .map_err(|x| FdbBindingError::CustomError(x.into()))?; + let serialized = kv_key.serialize(value.clone())?; tx.set(&packed_key, &serialized); } else { tx.clear(&packed_key); diff --git a/packages/services/epoxy/src/replica/decide_path.rs b/packages/services/epoxy/src/replica/decide_path.rs index af1345fa38..1b7022728f 100644 --- a/packages/services/epoxy/src/replica/decide_path.rs +++ b/packages/services/epoxy/src/replica/decide_path.rs @@ -1,5 +1,6 @@ +use anyhow::Result; use epoxy_protocol::protocol; -use universaldb::{FdbBindingError, Transaction}; +use universaldb::Transaction; use crate::replica::utils; @@ -7,7 +8,7 @@ pub fn decide_path( _tx: &Transaction, pre_accept_oks: Vec, payload: &protocol::Payload, -) -> Result { +) -> Result { tracing::info!(instance=?payload.instance, "deciding path"); let mut new_payload = payload.clone(); diff --git a/packages/services/epoxy/src/replica/lead_consensus.rs b/packages/services/epoxy/src/replica/lead_consensus.rs index 1549c0dc60..96de6e75a3 100644 --- a/packages/services/epoxy/src/replica/lead_consensus.rs +++ b/packages/services/epoxy/src/replica/lead_consensus.rs @@ -1,6 +1,7 @@ +use anyhow::Result; use epoxy_protocol::protocol; -use udb_util::FormalKey as _; -use universaldb::{FdbBindingError, Transaction}; +use universaldb::Transaction; +use universaldb::utils::{FormalKey, IsolationLevel::*}; use crate::keys; use crate::replica::{ballot, messages, utils}; @@ -9,18 +10,16 @@ pub async fn lead_consensus( tx: &Transaction, replica_id: protocol::ReplicaId, proposal: protocol::Proposal, -) -> Result { +) -> Result { tracing::info!(?replica_id, "leading consensus"); // EPaxos Step 1 let instance_num_key = keys::replica::InstanceNumberKey; let packed_key = keys::subspace(replica_id).pack(&instance_num_key); - let value = tx.get(&packed_key, false).await?; + let value = tx.get(&packed_key, Serializable).await?; let current_slot = if let Some(ref bytes) = value { - let current = instance_num_key - .deserialize(bytes) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + let current = instance_num_key.deserialize(bytes)?; current } else { 0 @@ -28,12 +27,7 @@ pub async fn lead_consensus( // Increment and store the new instance number let slot_id = current_slot + 1; - tx.set( - &packed_key, - &instance_num_key - .serialize(slot_id) - .map_err(|e| FdbBindingError::CustomError(e.into()))?, - ); + tx.set(&packed_key, &instance_num_key.serialize(slot_id)?); // Find interference for this key let interf = utils::find_interference(tx, replica_id, &proposal.commands).await?; diff --git a/packages/services/epoxy/src/replica/log.rs b/packages/services/epoxy/src/replica/log.rs index a2f4dd2da8..ab9ba95d13 100644 --- a/packages/services/epoxy/src/replica/log.rs +++ b/packages/services/epoxy/src/replica/log.rs @@ -1,6 +1,7 @@ +use anyhow::{Result, ensure}; use epoxy_protocol::protocol::{self, ReplicaId}; -use udb_util::FormalKey; -use universaldb::{FdbBindingError, Transaction}; +use universaldb::Transaction; +use universaldb::utils::{FormalKey, IsolationLevel::*}; use crate::{keys, replica::utils}; @@ -19,7 +20,7 @@ pub async fn update_log( replica_id: ReplicaId, log_entry: protocol::LogEntry, instance: &protocol::Instance, -) -> Result<(), FdbBindingError> { +) -> Result<()> { tracing::debug!(?replica_id, ?instance, ?log_entry.state, "updating log"); let subspace = keys::subspace(replica_id); @@ -27,12 +28,8 @@ pub async fn update_log( let packed_key = subspace.pack(&log_key); // Read existing log entry to validate state progression - let current_entry = match tx.get(&packed_key, false).await? { - Some(bytes) => Some( - log_key - .deserialize(&bytes) - .map_err(|e| FdbBindingError::CustomError(e.into()))?, - ), + let current_entry = match tx.get(&packed_key, Serializable).await? { + Some(bytes) => Some(log_key.deserialize(&bytes)?), None => None, }; @@ -41,18 +38,14 @@ pub async fn update_log( let current_order = state_order(¤t.state); let new_order = state_order(&log_entry.state); - if new_order <= current_order { - return Err(FdbBindingError::CustomError( - anyhow::anyhow!( - "invalid state transition: cannot transition from {:?} to {:?} (order {} to {})", - current.state, - log_entry.state, - current_order, - new_order - ) - .into(), - )); - } + ensure!( + new_order > current_order, + "invalid state transition: cannot transition from {:?} to {:?} (order {} to {})", + current.state, + log_entry.state, + current_order, + new_order, + ); tracing::debug!( ?current.state, @@ -66,9 +59,7 @@ pub async fn update_log( } // Store log entry in UDB - let value = log_key - .serialize(log_entry.clone()) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + let value = log_key.serialize(log_entry.clone())?; tx.set(&packed_key, &value); // Store in keys for interference diff --git a/packages/services/epoxy/src/replica/message_request.rs b/packages/services/epoxy/src/replica/message_request.rs index fe81278c70..9dd0a3655d 100644 --- a/packages/services/epoxy/src/replica/message_request.rs +++ b/packages/services/epoxy/src/replica/message_request.rs @@ -20,12 +20,9 @@ pub async fn message_request( // Store the configuration ctx.udb()? - .run(move |tx, _| { + .run(move |tx| { let req = req.clone(); - async move { - replica::update_config::update_config(&*tx, replica_id, req) - .map_err(|e| universaldb::FdbBindingError::CustomError(e.into())) - } + async move { replica::update_config::update_config(&*tx, replica_id, req) } }) .await?; @@ -34,13 +31,9 @@ pub async fn message_request( protocol::RequestKind::PreAcceptRequest(req) => { let response = ctx .udb()? - .run(move |tx, _| { + .run(move |tx| { let req = req.clone(); - async move { - replica::messages::pre_accept(&*tx, replica_id, req) - .await - .map_err(|e| universaldb::FdbBindingError::CustomError(e.into())) - } + async move { replica::messages::pre_accept(&*tx, replica_id, req).await } }) .await?; protocol::ResponseKind::PreAcceptResponse(response) @@ -48,13 +41,9 @@ pub async fn message_request( protocol::RequestKind::AcceptRequest(req) => { let response = ctx .udb()? - .run(move |tx, _| { + .run(move |tx| { let req = req.clone(); - async move { - replica::messages::accept(&*tx, replica_id, req) - .await - .map_err(|e| universaldb::FdbBindingError::CustomError(e.into())) - } + async move { replica::messages::accept(&*tx, replica_id, req).await } }) .await?; protocol::ResponseKind::AcceptResponse(response) @@ -62,13 +51,10 @@ pub async fn message_request( protocol::RequestKind::CommitRequest(req) => { // Commit and update KV store ctx.udb()? - .run(move |tx, _| { + .run(move |tx| { let req = req.clone(); async move { - replica::messages::commit(&*tx, replica_id, req, true) - .await - .map_err(|e| universaldb::FdbBindingError::CustomError(e.into()))?; - + replica::messages::commit(&*tx, replica_id, req, true).await?; Result::Ok(()) } }) @@ -79,28 +65,20 @@ pub async fn message_request( protocol::RequestKind::PrepareRequest(req) => { let response = ctx .udb()? - .run(move |tx, _| { + .run(move |tx| { let req = req.clone(); - async move { - replica::messages::prepare(&*tx, replica_id, req) - .await - .map_err(|e| universaldb::FdbBindingError::CustomError(e.into())) - } + async move { replica::messages::prepare(&*tx, replica_id, req).await } }) .await?; protocol::ResponseKind::PrepareResponse(response) } protocol::RequestKind::DownloadInstancesRequest(req) => { - // Handle download instances request - read from FDB and return instances + // Handle download instances request - read from UDB and return instances let instances = ctx .udb()? - .run(move |tx, _| { + .run(move |tx| { let req = req.clone(); - async move { - replica::messages::download_instances(&*tx, replica_id, req) - .await - .map_err(|e| universaldb::FdbBindingError::CustomError(e.into())) - } + async move { replica::messages::download_instances(&*tx, replica_id, req).await } }) .await?; diff --git a/packages/services/epoxy/src/replica/messages/accept.rs b/packages/services/epoxy/src/replica/messages/accept.rs index 63295464b1..0f570e892e 100644 --- a/packages/services/epoxy/src/replica/messages/accept.rs +++ b/packages/services/epoxy/src/replica/messages/accept.rs @@ -1,5 +1,6 @@ +use anyhow::{Result, ensure}; use epoxy_protocol::protocol; -use universaldb::{FdbBindingError, Transaction}; +use universaldb::Transaction; use crate::replica::{ballot, messages}; @@ -7,7 +8,7 @@ pub async fn accept( tx: &Transaction, replica_id: protocol::ReplicaId, accept_req: protocol::AcceptRequest, -) -> Result { +) -> Result { let protocol::Payload { proposal, seq, @@ -22,11 +23,7 @@ pub async fn accept( let is_valid = ballot::validate_and_update_ballot_for_instance(tx, replica_id, ¤t_ballot, &instance) .await?; - if !is_valid { - return Err(FdbBindingError::CustomError( - anyhow::anyhow!("ballot validation failed for pre_accept").into(), - )); - } + ensure!(is_valid, "ballot validation failed for pre_accept"); // EPaxos Step 18 let log_entry = protocol::LogEntry { diff --git a/packages/services/epoxy/src/replica/messages/accepted.rs b/packages/services/epoxy/src/replica/messages/accepted.rs index ae765fbaa0..5314ef88c2 100644 --- a/packages/services/epoxy/src/replica/messages/accepted.rs +++ b/packages/services/epoxy/src/replica/messages/accepted.rs @@ -1,5 +1,6 @@ +use anyhow::Result; use epoxy_protocol::protocol; -use universaldb::{FdbBindingError, Transaction}; +use universaldb::Transaction; use crate::replica::{ballot, messages, utils}; @@ -8,7 +9,7 @@ pub async fn accepted( tx: &Transaction, replica_id: protocol::ReplicaId, payload: protocol::Payload, -) -> Result<(), FdbBindingError> { +) -> Result<()> { let protocol::Payload { proposal, seq, diff --git a/packages/services/epoxy/src/replica/messages/commit.rs b/packages/services/epoxy/src/replica/messages/commit.rs index 60197c86ff..fd52d64833 100644 --- a/packages/services/epoxy/src/replica/messages/commit.rs +++ b/packages/services/epoxy/src/replica/messages/commit.rs @@ -1,5 +1,6 @@ +use anyhow::Result; use epoxy_protocol::protocol; -use universaldb::{FdbBindingError, Transaction}; +use universaldb::Transaction; use crate::replica::ballot; @@ -9,7 +10,7 @@ pub async fn commit( replica_id: protocol::ReplicaId, commit_req: protocol::CommitRequest, commit_to_kv: bool, -) -> Result<(), FdbBindingError> { +) -> Result<()> { let protocol::Payload { proposal, seq, diff --git a/packages/services/epoxy/src/replica/messages/committed.rs b/packages/services/epoxy/src/replica/messages/committed.rs index 0694a9c584..444c7681d9 100644 --- a/packages/services/epoxy/src/replica/messages/committed.rs +++ b/packages/services/epoxy/src/replica/messages/committed.rs @@ -1,5 +1,6 @@ +use anyhow::Result; use epoxy_protocol::protocol; -use universaldb::{FdbBindingError, Transaction}; +use universaldb::Transaction; use crate::replica::ballot; @@ -8,7 +9,7 @@ pub async fn committed( tx: &Transaction, replica_id: protocol::ReplicaId, payload: &protocol::Payload, -) -> Result, FdbBindingError> { +) -> Result> { let protocol::Payload { proposal, seq, diff --git a/packages/services/epoxy/src/replica/messages/download_instances.rs b/packages/services/epoxy/src/replica/messages/download_instances.rs index bb92103d78..aa997f1d04 100644 --- a/packages/services/epoxy/src/replica/messages/download_instances.rs +++ b/packages/services/epoxy/src/replica/messages/download_instances.rs @@ -1,7 +1,8 @@ +use anyhow::Result; use epoxy_protocol::protocol::{self, ReplicaId}; use futures_util::TryStreamExt; -use udb_util::prelude::*; -use universaldb::{FdbBindingError, KeySelector, RangeOption, Transaction, options::StreamingMode}; +use universaldb::prelude::*; +use universaldb::{KeySelector, RangeOption, Transaction, options::StreamingMode}; use crate::keys; @@ -9,7 +10,7 @@ pub async fn download_instances( tx: &Transaction, replica_id: ReplicaId, req: protocol::DownloadInstancesRequest, -) -> Result, FdbBindingError> { +) -> Result> { tracing::info!(?replica_id, "handling download instances message"); let mut entries = Vec::new(); @@ -25,13 +26,13 @@ pub async fn download_instances( } else { // TODO: Use ::subspace() // Start from the beginning of the log - let prefix = subspace.pack(&(udb_util::keys::LOG,)); + let prefix = subspace.pack(&(universaldb::utils::keys::LOG,)); KeySelector::first_greater_or_equal(prefix) }; // TODO: Is there a cleaner way to do this // End key is after all log entries - let end_prefix = subspace.pack(&(udb_util::keys::LOG + 1,)); + let end_prefix = subspace.pack(&(universaldb::utils::keys::LOG + 1,)); let end_key = KeySelector::first_greater_or_equal(end_prefix); let range = RangeOption { @@ -43,20 +44,14 @@ pub async fn download_instances( }; // Query the range - let mut stream = tx.get_ranges_keyvalues(range, SERIALIZABLE); + let mut stream = tx.get_ranges_keyvalues(range, Serializable); while let Some(kv) = stream.try_next().await? { // Parse the key to extract instance info let key_bytes = kv.key(); - let log_key = subspace - .unpack::(key_bytes) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; - + let log_key = subspace.unpack::(key_bytes)?; // Deserialize the log entry - let log_entry = log_key - .deserialize(kv.value()) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; - + let log_entry = log_key.deserialize(kv.value())?; // Create the instance from the key let instance = protocol::Instance { replica_id: log_key.instance_replica_id, diff --git a/packages/services/epoxy/src/replica/messages/pre_accept.rs b/packages/services/epoxy/src/replica/messages/pre_accept.rs index 3801bfc439..23f33755a6 100644 --- a/packages/services/epoxy/src/replica/messages/pre_accept.rs +++ b/packages/services/epoxy/src/replica/messages/pre_accept.rs @@ -1,6 +1,7 @@ +use anyhow::{Result, ensure}; use epoxy_protocol::protocol; use std::cmp; -use universaldb::{FdbBindingError, Transaction}; +use universaldb::Transaction; use crate::replica::{ballot, messages, utils}; @@ -8,7 +9,7 @@ pub async fn pre_accept( tx: &Transaction, replica_id: protocol::ReplicaId, pre_accept_req: protocol::PreAcceptRequest, -) -> Result { +) -> Result { tracing::info!(?replica_id, "handling pre-accept message"); let protocol::Payload { @@ -25,11 +26,7 @@ pub async fn pre_accept( let is_valid = ballot::validate_and_update_ballot_for_instance(tx, replica_id, ¤t_ballot, &instance) .await?; - if !is_valid { - return Err(FdbBindingError::CustomError( - anyhow::anyhow!("ballot validation failed for pre_accept").into(), - )); - } + ensure!(is_valid, "ballot validation failed for pre_accept"); // Find interference for this key let interf = utils::find_interference(tx, replica_id, &proposal.commands).await?; diff --git a/packages/services/epoxy/src/replica/messages/prepare.rs b/packages/services/epoxy/src/replica/messages/prepare.rs index 296f3c4d68..df3581eb3b 100644 --- a/packages/services/epoxy/src/replica/messages/prepare.rs +++ b/packages/services/epoxy/src/replica/messages/prepare.rs @@ -1,6 +1,7 @@ +use anyhow::Result; use epoxy_protocol::protocol; -use udb_util::FormalKey; -use universaldb::{FdbBindingError, Transaction}; +use universaldb::Transaction; +use universaldb::utils::{FormalKey, IsolationLevel::*}; use crate::{keys, replica::ballot}; @@ -8,7 +9,7 @@ pub async fn prepare( tx: &Transaction, replica_id: protocol::ReplicaId, prepare_req: protocol::PrepareRequest, -) -> Result { +) -> Result { tracing::info!(?replica_id, "handling prepare message"); let protocol::PrepareRequest { ballot, instance } = prepare_req; @@ -19,14 +20,10 @@ pub async fn prepare( let subspace = keys::subspace(replica_id); let log_key = keys::replica::LogEntryKey::new(instance.replica_id, instance.slot_id); - let current_entry = match tx.get(&subspace.pack(&log_key), false).await? { + let current_entry = match tx.get(&subspace.pack(&log_key), Serializable).await? { Some(bytes) => { // Deserialize the existing log entry - Some( - log_key - .deserialize(&bytes) - .map_err(|e| FdbBindingError::CustomError(e.into()))?, - ) + Some(log_key.deserialize(&bytes)?) } None => None, }; @@ -72,10 +69,8 @@ pub async fn prepare( let subspace = keys::subspace(replica_id); let packed_key = subspace.pack(&instance_ballot_key); - let highest_ballot = match tx.get(&packed_key, false).await? { - Some(bytes) => instance_ballot_key - .deserialize(&bytes) - .map_err(|e| FdbBindingError::CustomError(e.into()))?, + let highest_ballot = match tx.get(&packed_key, Serializable).await? { + Some(bytes) => instance_ballot_key.deserialize(&bytes)?, None => { // Default ballot for the original replica protocol::Ballot { diff --git a/packages/services/epoxy/src/replica/update_config.rs b/packages/services/epoxy/src/replica/update_config.rs index b44987a4f6..c54a26ee2e 100644 --- a/packages/services/epoxy/src/replica/update_config.rs +++ b/packages/services/epoxy/src/replica/update_config.rs @@ -1,6 +1,7 @@ +use anyhow::Result; use epoxy_protocol::protocol::{self, ReplicaId}; -use udb_util::FormalKey; -use universaldb::{FdbBindingError, Transaction}; +use universaldb::Transaction; +use universaldb::utils::FormalKey; use crate::keys; @@ -8,16 +9,14 @@ pub fn update_config( tx: &Transaction, replica_id: ReplicaId, update_config_req: protocol::UpdateConfigRequest, -) -> Result<(), FdbBindingError> { +) -> Result<()> { tracing::debug!("updating config"); // Store config in UDB let config_key = keys::replica::ConfigKey; let subspace = keys::subspace(replica_id); let packed_key = subspace.pack(&config_key); - let value = config_key - .serialize(update_config_req.config) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + let value = config_key.serialize(update_config_req.config)?; tx.set(&packed_key, &value); diff --git a/packages/services/epoxy/src/replica/utils.rs b/packages/services/epoxy/src/replica/utils.rs index 2efa1bea3d..870d3db7ef 100644 --- a/packages/services/epoxy/src/replica/utils.rs +++ b/packages/services/epoxy/src/replica/utils.rs @@ -1,8 +1,9 @@ +use anyhow::Result; use epoxy_protocol::protocol::{self, ReplicaId}; use futures_util::TryStreamExt; use std::{cmp::Ordering, collections::HashSet}; -use udb_util::prelude::*; -use universaldb::{FdbBindingError, KeySelector, RangeOption, Transaction, options::StreamingMode}; +use universaldb::prelude::*; +use universaldb::{KeySelector, RangeOption, Transaction, options::StreamingMode}; use crate::keys; @@ -11,7 +12,7 @@ pub async fn find_interference( tx: &Transaction, replica_id: ReplicaId, commands: &Vec, -) -> Result, FdbBindingError> { +) -> Result> { let mut interf = Vec::new(); // Get deduplicated keys @@ -34,14 +35,12 @@ pub async fn find_interference( ..Default::default() }; - let mut stream = tx.get_ranges_keyvalues(range, SERIALIZABLE); + let mut stream = tx.get_ranges_keyvalues(range, Serializable); while let Some(kv) = stream.try_next().await? { // Parse the key to extract replica_id and slot_id let key_bytes = kv.key(); - let key = subspace - .unpack::(key_bytes) - .map_err(|x| FdbBindingError::CustomError(x.into()))?; + let key = subspace.unpack::(key_bytes)?; interf.push(protocol::Instance { replica_id: key.instance_replica_id, @@ -59,7 +58,7 @@ pub async fn find_max_seq( tx: &Transaction, replica_id: protocol::ReplicaId, interf: &Vec, -) -> Result { +) -> Result { let mut seq = 0; for instance in interf { @@ -67,11 +66,9 @@ pub async fn find_max_seq( let key = keys::replica::LogEntryKey::new(instance.replica_id, instance.slot_id); let subspace = keys::subspace(replica_id); - let value = tx.get(&subspace.pack(&key), false).await?; + let value = tx.get(&subspace.pack(&key), Serializable).await?; if let Some(ref bytes) = value { - let log_entry: protocol::LogEntry = key - .deserialize(bytes) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + let log_entry: protocol::LogEntry = key.deserialize(bytes)?; if log_entry.seq > seq { seq = log_entry.seq; } diff --git a/packages/services/epoxy/src/utils.rs b/packages/services/epoxy/src/utils.rs index b6896f1a08..88e6fbee1d 100644 --- a/packages/services/epoxy/src/utils.rs +++ b/packages/services/epoxy/src/utils.rs @@ -1,6 +1,6 @@ use anyhow::*; use epoxy_protocol::protocol::{self, ReplicaId}; -use universaldb::Transaction; +use universaldb::{Transaction, utils::IsolationLevel::*}; #[derive(Clone, Copy, Debug)] pub enum QuorumType { @@ -44,13 +44,13 @@ pub async fn read_config( tx: &Transaction, replica_id: ReplicaId, ) -> Result { - use udb_util::FormalKey; + use universaldb::utils::FormalKey; let config_key = crate::keys::replica::ConfigKey; let subspace = crate::keys::subspace(replica_id); let packed_key = subspace.pack(&config_key); - match tx.get(&packed_key, false).await? { + match tx.get(&packed_key, Serializable).await? { Some(value) => { let config = config_key.deserialize(&value)?; Ok(config) diff --git a/packages/services/epoxy/src/workflows/replica/setup.rs b/packages/services/epoxy/src/workflows/replica/setup.rs index 50ff1929ac..278c77269d 100644 --- a/packages/services/epoxy/src/workflows/replica/setup.rs +++ b/packages/services/epoxy/src/workflows/replica/setup.rs @@ -4,7 +4,7 @@ use futures_util::{FutureExt, TryStreamExt}; use gas::prelude::*; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, VecDeque}; -use udb_util::prelude::*; +use universaldb::prelude::*; use universaldb::{KeySelector, RangeOption, options::StreamingMode}; use crate::types; @@ -245,7 +245,7 @@ async fn apply_log_entry( // Replay the log entry ctx.udb()? - .run(move |tx, _| { + .run(move |tx| { let log_entry = log_entry.clone(); let instance = instance.clone(); @@ -256,10 +256,8 @@ async fn apply_log_entry( let packed_key = subspace.pack(&log_key); // Read existing entry to determine if we need to replay this log entry - if let Some(bytes) = tx.get(&packed_key, SERIALIZABLE).await? { - let existing = log_key - .deserialize(&bytes) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + if let Some(bytes) = tx.get(&packed_key, Serializable).await? { + let existing = log_key.deserialize(&bytes)?; let existing_order = crate::replica::log::state_order(&existing.state); let new_order = crate::replica::log::state_order(&log_entry.state); @@ -286,21 +284,15 @@ async fn apply_log_entry( match log_entry.state { protocol::State::PreAccepted => { let request = protocol::PreAcceptRequest { payload }; - crate::replica::messages::pre_accept(&*tx, replica_id, request) - .await - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + crate::replica::messages::pre_accept(&*tx, replica_id, request).await?; } protocol::State::Accepted => { let request = protocol::AcceptRequest { payload }; - crate::replica::messages::accept(&*tx, replica_id, request) - .await - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + crate::replica::messages::accept(&*tx, replica_id, request).await?; } protocol::State::Committed => { let request = protocol::CommitRequest { payload }; - crate::replica::messages::commit(&*tx, replica_id, request, false) - .await - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + crate::replica::messages::commit(&*tx, replica_id, request, false).await?; } } @@ -370,7 +362,7 @@ pub async fn recover_keys_chunk( let (last_key, recovered_count) = ctx .udb()? - .run(move |tx, _| { + .run(move |tx| { let after_key = input.after_key.clone(); let count = input.count; @@ -409,7 +401,7 @@ pub async fn recover_keys_chunk( ..Default::default() }; - let mut stream = tx.get_ranges_keyvalues(range_option, SERIALIZABLE); + let mut stream = tx.get_ranges_keyvalues(range_option, Serializable); // Iterate over stream and aggregate data for each key let mut current_key: Option> = None; @@ -423,10 +415,8 @@ pub async fn recover_keys_chunk( scanned_count += 1; // Parse the key instance entry to extract the key and instance info - let key_instance = subspace - .unpack::(kv.key()) - .map_err(|e| universaldb::FdbBindingError::CustomError(e.into()))?; - + let key_instance = + subspace.unpack::(kv.key())?; let key = key_instance.key; let instance = ( key_instance.instance_replica_id, @@ -483,13 +473,9 @@ pub async fn recover_keys_chunk( // it means a single key has too many instances (i.e. larger than // the range limit) if recovered_count == 0 && scanned_count >= count { - return Err(universaldb::FdbBindingError::CustomError( - anyhow!( - "single key has more than {} instances, cannot process in one chunk", - count - ) - .into(), - )); + bail!( + "single key has more than {count} instances, cannot process in one chunk", + ); } tracing::info!( @@ -640,7 +626,7 @@ async fn recover_key_value_with_instances( replica_id: protocol::ReplicaId, key: &[u8], instances: &[(protocol::ReplicaId, protocol::SlotId)], -) -> Result<(), universaldb::FdbBindingError> { +) -> Result<()> { let subspace = crate::keys::subspace(replica_id); tracing::debug!( @@ -659,7 +645,7 @@ async fn recover_key_value_with_instances( let log_key = crate::keys::replica::LogEntryKey::new(instance_replica_id, instance_slot_id); let packed_key = subspace.pack(&log_key); - futures.push(tx.get(&packed_key, SERIALIZABLE)); + futures.push(tx.get(&packed_key, Serializable)); batch_keys.push((packed_key, log_key, instance_replica_id, instance_slot_id)); } let batch_results = futures_util::future::try_join_all(futures).await?; @@ -669,21 +655,15 @@ async fn recover_key_value_with_instances( batch_results.into_iter().zip(batch_keys.iter()) { // Missing log entry indicates data corruption - let bytes = bytes.ok_or_else(|| { - universaldb::FdbBindingError::CustomError( - anyhow!( - "missing log entry for instance ({}, {}), data corruption detected", - instance_replica_id, - instance_slot_id - ) - .into(), + let bytes = bytes.with_context(|| { + format!( + "missing log entry for instance ({}, {}), data corruption detected", + instance_replica_id, instance_slot_id ) })?; // Collect committed entries - let entry = log_key - .deserialize(&bytes) - .map_err(|e| universaldb::FdbBindingError::CustomError(e.into()))?; + let entry = log_key.deserialize(&bytes)?; if matches!(entry.state, protocol::State::Committed) { committed_entries.push(CommittedEntry { instance: (*instance_replica_id, *instance_slot_id), @@ -702,9 +682,7 @@ async fn recover_key_value_with_instances( // Sort entries topologically to respect dependencies // This ensures that operations are applied in the correct order, // particularly important for dependent operations like check-and-set - let sorted_entries = topological_sort_entries(&committed_entries) - .map_err(|e| universaldb::FdbBindingError::CustomError(e.into()))?; - + let sorted_entries = topological_sort_entries(&committed_entries)?; tracing::debug!( key_len = key.len(), sorted_count = sorted_entries.len(), diff --git a/packages/services/epoxy/tests/reconfigure.rs b/packages/services/epoxy/tests/reconfigure.rs index b54ef7464f..948e358b29 100644 --- a/packages/services/epoxy/tests/reconfigure.rs +++ b/packages/services/epoxy/tests/reconfigure.rs @@ -7,8 +7,8 @@ use futures_util::TryStreamExt; use gas::prelude::*; use serde_json::json; use std::collections::HashSet; -use udb_util::prelude::*; -use universaldb::{FdbBindingError, KeySelector, RangeOption, options::StreamingMode}; +use universaldb::prelude::*; +use universaldb::{KeySelector, RangeOption, options::StreamingMode}; mod common; @@ -456,7 +456,7 @@ async fn verify_log_entries_match( // Read log entries from replica 1 let log_entries_1 = ctx_1 .udb()? - .run(move |tx, _| async move { + .run(move |tx| async move { let subspace = epoxy::keys::subspace(replica_1_id); // Range scan to get all log entries for this replica @@ -467,7 +467,7 @@ async fn verify_log_entries_match( ..Default::default() }; - let mut stream = tx.get_ranges_keyvalues(range, SERIALIZABLE); + let mut stream = tx.get_ranges_keyvalues(range, Serializable); let mut log_entries = Vec::new(); while let Some(kv) = stream.try_next().await? { @@ -475,29 +475,26 @@ async fn verify_log_entries_match( let value_bytes = kv.value(); // Parse the key to get replica_id and slot_id - let key = subspace - .unpack::(key_bytes) - .map_err(|x| FdbBindingError::CustomError(x.into()))?; + let key = subspace.unpack::(key_bytes)?; // Deserialize the log entry let log_entry = epoxy::keys::replica::LogEntryKey::new( key.instance_replica_id, key.instance_slot_id, ) - .deserialize(value_bytes) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + .deserialize(value_bytes)?; log_entries.push((key.instance_replica_id, key.instance_slot_id, log_entry)); } - Result::<_, FdbBindingError>::Ok(log_entries) + Ok(log_entries) }) .await?; // Read log entries from replica 2 let log_entries_2 = ctx_2 .udb()? - .run(move |tx, _| async move { + .run(move |tx| async move { let subspace = epoxy::keys::subspace(replica_2_id); // Range scan to get all log entries for this replica @@ -508,7 +505,7 @@ async fn verify_log_entries_match( ..Default::default() }; - let mut stream = tx.get_ranges_keyvalues(range, SERIALIZABLE); + let mut stream = tx.get_ranges_keyvalues(range, Serializable); let mut log_entries = Vec::new(); while let Some(kv) = stream.try_next().await? { @@ -516,22 +513,19 @@ async fn verify_log_entries_match( let value_bytes = kv.value(); // Parse the key to get replica_id and slot_id - let key = subspace - .unpack::(key_bytes) - .map_err(|x| FdbBindingError::CustomError(x.into()))?; + let key = subspace.unpack::(key_bytes)?; // Deserialize the log entry let log_entry = epoxy::keys::replica::LogEntryKey::new( key.instance_replica_id, key.instance_slot_id, ) - .deserialize(value_bytes) - .map_err(|e| FdbBindingError::CustomError(e.into()))?; + .deserialize(value_bytes)?; log_entries.push((key.instance_replica_id, key.instance_slot_id, log_entry)); } - Result::<_, FdbBindingError>::Ok(log_entries) + Ok(log_entries) }) .await?; @@ -593,15 +587,15 @@ async fn verify_kv_replication( for (i, (key, expected_value)) in expected_keys.iter().enumerate() { // Read the KV value from the replica's UDB let actual_value = udb - .run(move |tx, _| { + .run(move |tx| { let key_clone = key.clone(); async move { let subspace = epoxy::keys::subspace(replica_id); let kv_key = epoxy::keys::keys::KvValueKey::new(key_clone); - let result = tx.get(&subspace.pack(&kv_key), SERIALIZABLE).await?; + let result = tx.get(&subspace.pack(&kv_key), Serializable).await?; // KvValueKey stores Vec directly, so we can return it as is - Result::<_, FdbBindingError>::Ok(result) + Ok(result) } }) .await?; diff --git a/packages/services/namespace/Cargo.toml b/packages/services/namespace/Cargo.toml index 46ec60f082..efd8da845f 100644 --- a/packages/services/namespace/Cargo.toml +++ b/packages/services/namespace/Cargo.toml @@ -18,7 +18,6 @@ rivet-util.workspace = true serde.workspace = true strum.workspace = true tracing.workspace = true -udb-util.workspace = true universaldb.workspace = true url.workspace = true utoipa.workspace = true diff --git a/packages/services/namespace/src/keys.rs b/packages/services/namespace/src/keys.rs index 93e133de25..caa49d1005 100644 --- a/packages/services/namespace/src/keys.rs +++ b/packages/services/namespace/src/keys.rs @@ -3,12 +3,12 @@ use std::result::Result::Ok; use anyhow::*; use gas::prelude::*; use serde::Serialize; -use udb_util::prelude::*; +use universaldb::prelude::*; use utoipa::ToSchema; use versioned_data_util::OwnedVersionedData; -pub fn subspace() -> udb_util::Subspace { - udb_util::Subspace::new(&(RIVET, NAMESPACE)) +pub fn subspace() -> universaldb::utils::Subspace { + universaldb::utils::Subspace::new(&(RIVET, NAMESPACE)) } #[derive(Debug)] diff --git a/packages/services/namespace/src/ops/get_local.rs b/packages/services/namespace/src/ops/get_local.rs index 3f1b5f4474..156632f384 100644 --- a/packages/services/namespace/src/ops/get_local.rs +++ b/packages/services/namespace/src/ops/get_local.rs @@ -1,7 +1,6 @@ use futures_util::{StreamExt, TryStreamExt}; use gas::prelude::*; -use udb_util::{SERIALIZABLE, TxnExt}; -use universaldb as udb; +use universaldb::utils::IsolationLevel::*; use crate::{errors, keys, types::Namespace}; @@ -18,7 +17,7 @@ pub async fn namespace_get_local(ctx: &OperationCtx, input: &Input) -> Result Result std::result::Result, udb::FdbBindingError> { - let txs = tx.subspace(keys::subspace()); + tx: &universaldb::Transaction, +) -> Result> { + let tx = tx.with_subspace(keys::subspace()); let name_key = keys::NameKey::new(namespace_id); let display_name_key = keys::DisplayNameKey::new(namespace_id); let create_ts_key = keys::CreateTsKey::new(namespace_id); let (name, display_name, create_ts) = tokio::try_join!( - txs.read_opt(&name_key, SERIALIZABLE), - txs.read_opt(&display_name_key, SERIALIZABLE), - txs.read_opt(&create_ts_key, SERIALIZABLE), + tx.read_opt(&name_key, Serializable), + tx.read_opt(&display_name_key, Serializable), + tx.read_opt(&create_ts_key, Serializable), )?; // Namespace not found @@ -57,12 +56,8 @@ pub(crate) async fn get_inner( return Ok(None); }; - let display_name = display_name.ok_or(udb::FdbBindingError::CustomError( - format!("key should exist: {display_name_key:?}").into(), - ))?; - let create_ts = create_ts.ok_or(udb::FdbBindingError::CustomError( - format!("key should exist: {create_ts_key:?}").into(), - ))?; + let display_name = display_name.context("key should exist")?; + let create_ts = create_ts.context("key should exist")?; Ok(Some(Namespace { namespace_id, diff --git a/packages/services/namespace/src/ops/list.rs b/packages/services/namespace/src/ops/list.rs index abac196775..57955718b7 100644 --- a/packages/services/namespace/src/ops/list.rs +++ b/packages/services/namespace/src/ops/list.rs @@ -1,8 +1,8 @@ use anyhow::Result; use futures_util::TryStreamExt; use gas::prelude::*; -use udb_util::SNAPSHOT; -use universaldb::{self as udb, options::StreamingMode}; +use universaldb::options::StreamingMode; +use universaldb::utils::IsolationLevel::*; use crate::{errors, keys, types::Namespace}; @@ -24,16 +24,16 @@ pub async fn namespace_list(ctx: &OperationCtx, input: &Input) -> Result let namespaces = ctx .udb()? - .run(|tx, _mc| async move { + .run(|tx| async move { let mut namespaces = Vec::new(); let limit = input.limit.unwrap_or(1000); // Default limit to 1000 let mut stream = tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::Iterator, ..(&keys::subspace()).into() }, - SNAPSHOT, + Snapshot, ); let mut seen_namespaces = std::collections::HashSet::new(); @@ -59,7 +59,7 @@ pub async fn namespace_list(ctx: &OperationCtx, input: &Input) -> Result } } - Result::<_, udb::FdbBindingError>::Ok(namespaces) + Ok(namespaces) }) .custom_instrument(tracing::info_span!("namespace_list_tx")) .await?; diff --git a/packages/services/namespace/src/ops/resolve_for_name_local.rs b/packages/services/namespace/src/ops/resolve_for_name_local.rs index d70e382000..b02617a35d 100644 --- a/packages/services/namespace/src/ops/resolve_for_name_local.rs +++ b/packages/services/namespace/src/ops/resolve_for_name_local.rs @@ -1,5 +1,5 @@ use gas::prelude::*; -use udb_util::{SERIALIZABLE, TxnExt}; +use universaldb::utils::IsolationLevel::*; use crate::{errors, keys, ops::get_local::get_inner, types::Namespace}; @@ -18,13 +18,13 @@ pub async fn namespace_resolve_for_name_local( } ctx.udb()? - .run(|tx, _mc| { + .run(|tx| { let name = input.name.clone(); async move { - let txs = tx.subspace(keys::subspace()); + let tx = tx.with_subspace(keys::subspace()); - let Some(namespace_id) = txs - .read_opt(&keys::ByNameKey::new(name.clone()), SERIALIZABLE) + let Some(namespace_id) = tx + .read_opt(&keys::ByNameKey::new(name.clone()), Serializable) .await? else { // Namespace not found diff --git a/packages/services/namespace/src/ops/runner_config/delete.rs b/packages/services/namespace/src/ops/runner_config/delete.rs index 09ba689150..1b6f68cc3a 100644 --- a/packages/services/namespace/src/ops/runner_config/delete.rs +++ b/packages/services/namespace/src/ops/runner_config/delete.rs @@ -1,6 +1,6 @@ use gas::prelude::*; use rivet_cache::CacheKey; -use udb_util::{SERIALIZABLE, TxnExt}; +use universaldb::utils::IsolationLevel::*; use crate::{errors, keys}; @@ -17,18 +17,18 @@ pub async fn namespace_runner_config_delete(ctx: &OperationCtx, input: &Input) - } ctx.udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); // Read existing config to determine variant let runner_config_key = keys::RunnerConfigKey::new(input.namespace_id, input.name.clone()); - if let Some(config) = txs.read_opt(&runner_config_key, SERIALIZABLE).await? { - txs.delete(&runner_config_key); + if let Some(config) = tx.read_opt(&runner_config_key, Serializable).await? { + tx.delete(&runner_config_key); // Clear secondary idx - txs.delete(&keys::RunnerConfigByVariantKey::new( + tx.delete(&keys::RunnerConfigByVariantKey::new( input.namespace_id, config.variant(), input.name.clone(), diff --git a/packages/services/namespace/src/ops/runner_config/get_local.rs b/packages/services/namespace/src/ops/runner_config/get_local.rs index fd23ad9562..35a00aca94 100644 --- a/packages/services/namespace/src/ops/runner_config/get_local.rs +++ b/packages/services/namespace/src/ops/runner_config/get_local.rs @@ -1,7 +1,7 @@ use futures_util::{StreamExt, TryStreamExt}; use gas::prelude::*; use serde::{Deserialize, Serialize}; -use udb_util::{SERIALIZABLE, TxnExt}; +use universaldb::utils::IsolationLevel::*; use crate::{errors, keys}; @@ -28,20 +28,20 @@ pub async fn namespace_runner_config_get_local( let runner_configs = ctx .udb()? - .run(|tx, _mc| async move { + .run(|tx| async move { futures_util::stream::iter(input.runners.clone()) .map(|(namespace_id, runner_name)| { let tx = tx.clone(); async move { - let txs = tx.subspace(keys::subspace()); + let tx = tx.with_subspace(keys::subspace()); let runner_config_key = keys::RunnerConfigKey::new(namespace_id, runner_name.clone()); // Runner config not found let Some(runner_config) = - txs.read_opt(&runner_config_key, SERIALIZABLE).await? + tx.read_opt(&runner_config_key, Serializable).await? else { return Ok(None); }; diff --git a/packages/services/namespace/src/ops/runner_config/list.rs b/packages/services/namespace/src/ops/runner_config/list.rs index 1b15b21fb2..45f414304d 100644 --- a/packages/services/namespace/src/ops/runner_config/list.rs +++ b/packages/services/namespace/src/ops/runner_config/list.rs @@ -1,7 +1,7 @@ use futures_util::{StreamExt, TryStreamExt}; use gas::prelude::*; -use udb_util::{SERIALIZABLE, TxnExt}; -use universaldb::{self as udb, options::StreamingMode}; +use universaldb::options::StreamingMode; +use universaldb::utils::IsolationLevel::*; use crate::{errors, keys, types::RunnerConfig}; @@ -24,11 +24,11 @@ pub async fn namespace_runner_config_list( let runner_configs = ctx .udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); let (start, end) = if let Some(variant) = input.variant { - let (start, end) = txs + let (start, end) = keys::subspace() .subspace(&keys::RunnerConfigByVariantKey::subspace_with_variant( input.namespace_id, variant, @@ -36,7 +36,7 @@ pub async fn namespace_runner_config_list( .range(); let start = if let Some(name) = &input.after_name { - txs.pack(&keys::RunnerConfigByVariantKey::new( + tx.pack(&keys::RunnerConfigByVariantKey::new( input.namespace_id, variant, name.clone(), @@ -47,12 +47,12 @@ pub async fn namespace_runner_config_list( (start, end) } else { - let (start, end) = txs + let (start, end) = keys::subspace() .subspace(&keys::RunnerConfigKey::subspace(input.namespace_id)) .range(); let start = if let Some(name) = &input.after_name { - txs.pack(&keys::RunnerConfigKey::new( + tx.pack(&keys::RunnerConfigKey::new( input.namespace_id, name.clone(), )) @@ -63,22 +63,22 @@ pub async fn namespace_runner_config_list( (start, end) }; - txs.get_ranges_keyvalues( - udb::RangeOption { + tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::WantAll, limit: Some(input.limit), ..(start, end).into() }, - SERIALIZABLE, + Serializable, ) .map(|res| match res { Ok(entry) => { if input.variant.is_some() { let (key, config) = - txs.read_entry::(&entry)?; + tx.read_entry::(&entry)?; Ok((key.name, config)) } else { - let (key, config) = txs.read_entry::(&entry)?; + let (key, config) = tx.read_entry::(&entry)?; Ok((key.name, config)) } } diff --git a/packages/services/namespace/src/ops/runner_config/upsert.rs b/packages/services/namespace/src/ops/runner_config/upsert.rs index e530f8e3ce..5c88b74f62 100644 --- a/packages/services/namespace/src/ops/runner_config/upsert.rs +++ b/packages/services/namespace/src/ops/runner_config/upsert.rs @@ -1,6 +1,5 @@ use gas::prelude::*; use rivet_cache::CacheKey; -use udb_util::TxnExt; use universaldb::options::MutationType; use crate::{errors, keys, types::RunnerConfig}; @@ -19,17 +18,17 @@ pub async fn namespace_runner_config_upsert(ctx: &OperationCtx, input: &Input) - } ctx.udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); // TODO: Once other types of configs get added, delete previous config before writing - txs.write( + tx.write( &keys::RunnerConfigKey::new(input.namespace_id, input.name.clone()), input.config.clone(), )?; // Write to secondary idx - txs.write( + tx.write( &keys::RunnerConfigByVariantKey::new( input.namespace_id, input.config.variant(), @@ -59,8 +58,8 @@ pub async fn namespace_runner_config_upsert(ctx: &OperationCtx, input: &Input) - } // Sets desired count to 0 if it doesn't exist - let txs = tx.subspace(rivet_types::keys::pegboard::subspace()); - txs.atomic_op( + let tx = tx.with_subspace(rivet_types::keys::pegboard::subspace()); + tx.atomic_op( &rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new( input.namespace_id, input.name.clone(), diff --git a/packages/services/namespace/src/workflows/namespace.rs b/packages/services/namespace/src/workflows/namespace.rs index 4feac7a0d2..525aee56d6 100644 --- a/packages/services/namespace/src/workflows/namespace.rs +++ b/packages/services/namespace/src/workflows/namespace.rs @@ -1,7 +1,7 @@ use futures_util::FutureExt; use gas::prelude::*; use serde::{Deserialize, Serialize}; -use udb_util::{SERIALIZABLE, TxnExt}; +use universaldb::utils::IsolationLevel::*; use crate::{errors, keys}; @@ -32,7 +32,7 @@ pub async fn namespace(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> { } let insert_res = ctx - .activity(InsertFdbInput { + .activity(InsertDbInput { namespace_id: input.namespace_id, name: input.name.clone(), display_name: input.display_name.clone(), @@ -135,39 +135,39 @@ pub async fn validate( } #[derive(Debug, Clone, Serialize, Deserialize, Hash)] -struct InsertFdbInput { +struct InsertDbInput { namespace_id: Id, name: String, display_name: String, create_ts: i64, } -#[activity(InsertFdb)] -async fn insert_fdb( +#[activity(InsertDb)] +async fn insert_db( ctx: &ActivityCtx, - input: &InsertFdbInput, + input: &InsertDbInput, ) -> Result> { ctx.udb()? - .run(|tx, _mc| { + .run(|tx| { let namespace_id = input.namespace_id; let name = input.name.clone(); let display_name = input.display_name.clone(); async move { - let txs = tx.subspace(keys::subspace()); + let tx = tx.with_subspace(keys::subspace()); let name_idx_key = keys::ByNameKey::new(name.clone()); - if txs.exists(&name_idx_key, SERIALIZABLE).await? { + if tx.exists(&name_idx_key, Serializable).await? { return Ok(Err(errors::Namespace::NameNotUnique)); } - txs.write(&keys::NameKey::new(namespace_id), name)?; - txs.write(&keys::DisplayNameKey::new(namespace_id), display_name)?; - txs.write(&keys::CreateTsKey::new(namespace_id), input.create_ts)?; + tx.write(&keys::NameKey::new(namespace_id), name)?; + tx.write(&keys::DisplayNameKey::new(namespace_id), display_name)?; + tx.write(&keys::CreateTsKey::new(namespace_id), input.create_ts)?; // Insert idx - txs.write(&name_idx_key, namespace_id)?; + tx.write(&name_idx_key, namespace_id)?; Ok(Ok(())) } diff --git a/packages/services/pegboard/Cargo.toml b/packages/services/pegboard/Cargo.toml index f9569fa571..cc08e962cf 100644 --- a/packages/services/pegboard/Cargo.toml +++ b/packages/services/pegboard/Cargo.toml @@ -25,7 +25,6 @@ serde_json.workspace = true serde.workspace = true strum.workspace = true tracing.workspace = true -udb-util.workspace = true universaldb.workspace = true utoipa.workspace = true versioned-data-util.workspace = true diff --git a/packages/services/pegboard/src/keys/actor.rs b/packages/services/pegboard/src/keys/actor.rs index 022f047c93..e85d5fc07a 100644 --- a/packages/services/pegboard/src/keys/actor.rs +++ b/packages/services/pegboard/src/keys/actor.rs @@ -2,7 +2,7 @@ use std::result::Result::Ok; use anyhow::*; use gas::prelude::*; -use udb_util::prelude::*; +use universaldb::prelude::*; #[derive(Debug)] pub struct CreateTsKey { diff --git a/packages/services/pegboard/src/keys/epoxy/ns.rs b/packages/services/pegboard/src/keys/epoxy/ns.rs index e2807aac9e..6ca366be8c 100644 --- a/packages/services/pegboard/src/keys/epoxy/ns.rs +++ b/packages/services/pegboard/src/keys/epoxy/ns.rs @@ -2,7 +2,7 @@ use std::result::Result::Ok; use anyhow::*; use gas::prelude::*; -use udb_util::prelude::*; +use universaldb::prelude::*; #[derive(Debug)] pub struct ReservationByKeyKey { diff --git a/packages/services/pegboard/src/keys/mod.rs b/packages/services/pegboard/src/keys/mod.rs index 3cb17c5bbb..253fdcb409 100644 --- a/packages/services/pegboard/src/keys/mod.rs +++ b/packages/services/pegboard/src/keys/mod.rs @@ -1,14 +1,14 @@ -use udb_util::prelude::*; +use universaldb::prelude::*; pub mod actor; pub mod epoxy; pub mod ns; pub mod runner; -pub fn subspace() -> udb_util::Subspace { +pub fn subspace() -> universaldb::utils::Subspace { rivet_types::keys::pegboard::subspace() } -pub fn actor_kv_subspace() -> udb_util::Subspace { - udb_util::Subspace::new(&(RIVET, PEGBOARD, ACTOR_KV)) +pub fn actor_kv_subspace() -> universaldb::utils::Subspace { + universaldb::utils::Subspace::new(&(RIVET, PEGBOARD, ACTOR_KV)) } diff --git a/packages/services/pegboard/src/keys/ns.rs b/packages/services/pegboard/src/keys/ns.rs index 5ccf65fcb5..a30a82d8c6 100644 --- a/packages/services/pegboard/src/keys/ns.rs +++ b/packages/services/pegboard/src/keys/ns.rs @@ -2,7 +2,7 @@ use std::result::Result::Ok; use anyhow::*; use gas::prelude::*; -use udb_util::prelude::*; +use universaldb::prelude::*; use versioned_data_util::OwnedVersionedData; #[derive(Debug)] diff --git a/packages/services/pegboard/src/keys/runner.rs b/packages/services/pegboard/src/keys/runner.rs index 9c20248557..fcdd520ee1 100644 --- a/packages/services/pegboard/src/keys/runner.rs +++ b/packages/services/pegboard/src/keys/runner.rs @@ -2,7 +2,7 @@ use std::result::Result::Ok; use anyhow::*; use gas::prelude::*; -use udb_util::prelude::*; +use universaldb::prelude::*; use versioned_data_util::OwnedVersionedData; #[derive(Debug)] @@ -753,7 +753,7 @@ impl FormalChunkedKey for MetadataKey { } } - fn combine(&self, chunks: Vec) -> Result { + fn combine(&self, chunks: Vec) -> Result { rivet_data::versioned::MetadataKeyData::deserialize_with_embedded_version( &chunks .iter() @@ -768,7 +768,7 @@ impl FormalChunkedKey for MetadataKey { Ok( rivet_data::versioned::MetadataKeyData::latest(value.try_into()?) .serialize_with_embedded_version(rivet_data::PEGBOARD_RUNNER_METADATA_VERSION)? - .chunks(udb_util::CHUNK_SIZE) + .chunks(universaldb::utils::CHUNK_SIZE) .map(|x| x.to_vec()) .collect(), ) diff --git a/packages/services/pegboard/src/ops/actor/get.rs b/packages/services/pegboard/src/ops/actor/get.rs index 8c7ba0db56..22c9b55f17 100644 --- a/packages/services/pegboard/src/ops/actor/get.rs +++ b/packages/services/pegboard/src/ops/actor/get.rs @@ -1,8 +1,7 @@ use futures_util::{StreamExt, TryStreamExt}; use gas::prelude::*; use rivet_types::actors::Actor; -use udb_util::{FormalKey, SERIALIZABLE}; -use universaldb as udb; +use universaldb::utils::{FormalKey, IsolationLevel::*}; use crate::keys; @@ -20,7 +19,7 @@ pub struct Output { pub async fn pegboard_actor_get(ctx: &OperationCtx, input: &Input) -> Result { let actors_with_wf_ids = ctx .udb()? - .run(|tx, _mc| async move { + .run(|tx| async move { futures_util::stream::iter(input.actor_ids.clone()) .map(|actor_id| { let tx = tx.clone(); @@ -28,16 +27,14 @@ pub async fn pegboard_actor_get(ctx: &OperationCtx, input: &Input) -> Result Result { let actors = ctx .udb()? - .run(|tx, _mc| async move { + .run(|tx| async move { futures_util::stream::iter(input.actor_ids.clone()) .map(|actor_id| { let tx = tx.clone(); @@ -36,17 +35,15 @@ pub async fn pegboard_actor_get_runner(ctx: &OperationCtx, input: &Input) -> Res let connectable_key = keys::actor::ConnectableKey::new(actor_id); let (runner_id_entry, connectable_entry) = tokio::try_join!( - tx.get(&keys::subspace().pack(&runner_id_key), SERIALIZABLE), - tx.get(&keys::subspace().pack(&connectable_key), SERIALIZABLE), + tx.get(&keys::subspace().pack(&runner_id_key), Serializable), + tx.get(&keys::subspace().pack(&connectable_key), Serializable), )?; let Some(runner_id_entry) = runner_id_entry else { return Ok(None); }; - let runner_id = runner_id_key - .deserialize(&runner_id_entry) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?; + let runner_id = runner_id_key.deserialize(&runner_id_entry)?; Ok(Some(Actor { actor_id, diff --git a/packages/services/pegboard/src/ops/actor/list_for_ns.rs b/packages/services/pegboard/src/ops/actor/list_for_ns.rs index 864359ee4e..fdf4373c8a 100644 --- a/packages/services/pegboard/src/ops/actor/list_for_ns.rs +++ b/packages/services/pegboard/src/ops/actor/list_for_ns.rs @@ -1,8 +1,8 @@ use futures_util::TryStreamExt; use gas::prelude::*; use rivet_types::actors::Actor; -use udb_util::{SNAPSHOT, TxnExt}; -use universaldb::{self as udb, options::StreamingMode}; +use universaldb::options::StreamingMode; +use universaldb::utils::IsolationLevel::*; use crate::keys; @@ -25,12 +25,12 @@ pub struct Output { pub async fn pegboard_actor_list_for_ns(ctx: &OperationCtx, input: &Input) -> Result { let actors_with_wf_ids = ctx .udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); let mut results = Vec::new(); if let Some(key) = &input.key { - let actor_subspace = txs.subspace(&keys::ns::ActorByKeyKey::subspace( + let actor_subspace = keys::subspace().subspace(&keys::ns::ActorByKeyKey::subspace( input.namespace_id, input.name.clone(), key.clone(), @@ -38,7 +38,7 @@ pub async fn pegboard_actor_list_for_ns(ctx: &OperationCtx, input: &Input) -> Re let (start, end) = actor_subspace.range(); let end = if let Some(created_before) = input.created_before { - udb_util::end_of_key_range(&txs.pack( + universaldb::utils::end_of_key_range(&tx.pack( &keys::ns::ActorByKeyKey::subspace_with_create_ts( input.namespace_id, input.name.clone(), @@ -50,18 +50,18 @@ pub async fn pegboard_actor_list_for_ns(ctx: &OperationCtx, input: &Input) -> Re end }; - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { + let mut stream = tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::Iterator, reverse: true, ..(start, end).into() }, // NOTE: Does not have to be serializable because we are listing, stale data does not matter - SNAPSHOT, + Snapshot, ); while let Some(entry) = stream.try_next().await? { - let (idx_key, data) = txs.read_entry::(&entry)?; + let (idx_key, data) = tx.read_entry::(&entry)?; if !data.is_destroyed || input.include_destroyed { results.push((idx_key.actor_id, data.workflow_id)); @@ -72,14 +72,14 @@ pub async fn pegboard_actor_list_for_ns(ctx: &OperationCtx, input: &Input) -> Re } } } else if input.include_destroyed { - let actor_subspace = txs.subspace(&keys::ns::AllActorKey::subspace( + let actor_subspace = keys::subspace().subspace(&keys::ns::AllActorKey::subspace( input.namespace_id, input.name.clone(), )); let (start, end) = actor_subspace.range(); let end = if let Some(created_before) = input.created_before { - udb_util::end_of_key_range(&txs.pack( + universaldb::utils::end_of_key_range(&tx.pack( &keys::ns::AllActorKey::subspace_with_create_ts( input.namespace_id, input.name.clone(), @@ -90,18 +90,18 @@ pub async fn pegboard_actor_list_for_ns(ctx: &OperationCtx, input: &Input) -> Re end }; - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { + let mut stream = tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::Iterator, reverse: true, ..(start, end).into() }, // NOTE: Does not have to be serializable because we are listing, stale data does not matter - SNAPSHOT, + Snapshot, ); while let Some(entry) = stream.try_next().await? { - let (idx_key, workflow_id) = txs.read_entry::(&entry)?; + let (idx_key, workflow_id) = tx.read_entry::(&entry)?; results.push((idx_key.actor_id, workflow_id)); @@ -110,14 +110,13 @@ pub async fn pegboard_actor_list_for_ns(ctx: &OperationCtx, input: &Input) -> Re } } } else { - let actor_subspace = txs.subspace(&keys::ns::ActiveActorKey::subspace( - input.namespace_id, - input.name.clone(), - )); + let actor_subspace = keys::subspace().subspace( + &keys::ns::ActiveActorKey::subspace(input.namespace_id, input.name.clone()), + ); let (start, end) = actor_subspace.range(); let end = if let Some(created_before) = input.created_before { - udb_util::end_of_key_range(&txs.pack( + universaldb::utils::end_of_key_range(&tx.pack( &keys::ns::ActiveActorKey::subspace_with_create_ts( input.namespace_id, input.name.clone(), @@ -128,19 +127,19 @@ pub async fn pegboard_actor_list_for_ns(ctx: &OperationCtx, input: &Input) -> Re end }; - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { + let mut stream = tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::Iterator, reverse: true, ..(start, end).into() }, // NOTE: Does not have to be serializable because we are listing, stale data does not matter - SNAPSHOT, + Snapshot, ); while let Some(entry) = stream.try_next().await? { let (idx_key, workflow_id) = - txs.read_entry::(&entry)?; + tx.read_entry::(&entry)?; results.push((idx_key.actor_id, workflow_id)); diff --git a/packages/services/pegboard/src/ops/actor/list_names.rs b/packages/services/pegboard/src/ops/actor/list_names.rs index 0cd3752ce3..ac15a99f6c 100644 --- a/packages/services/pegboard/src/ops/actor/list_names.rs +++ b/packages/services/pegboard/src/ops/actor/list_names.rs @@ -1,8 +1,8 @@ use futures_util::{StreamExt, TryStreamExt}; use gas::prelude::*; use rivet_data::converted::ActorNameKeyData; -use udb_util::{SNAPSHOT, TxnExt}; -use universaldb::{self as udb, options::StreamingMode}; +use universaldb::options::StreamingMode; +use universaldb::utils::IsolationLevel::*; use crate::keys; @@ -22,15 +22,15 @@ pub struct Output { pub async fn pegboard_actor_list_names(ctx: &OperationCtx, input: &Input) -> Result { let names = ctx .udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); let actor_name_subspace = - txs.subspace(&keys::ns::ActorNameKey::subspace(input.namespace_id)); + keys::subspace().subspace(&keys::ns::ActorNameKey::subspace(input.namespace_id)); let (start, end) = actor_name_subspace.range(); let start = if let Some(name) = &input.after_name { - txs.pack(&keys::ns::ActorNameKey::new( + tx.pack(&keys::ns::ActorNameKey::new( input.namespace_id, name.clone(), )) @@ -38,22 +38,19 @@ pub async fn pegboard_actor_list_names(ctx: &OperationCtx, input: &Input) -> Res start }; - txs.get_ranges_keyvalues( - udb::RangeOption { + tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::WantAll, limit: Some(input.limit), ..(start, end).into() }, - // NOTE: This is not SERIALIZABLE to prevent contention with inserting new names - SNAPSHOT, + // NOTE: This is not Serializable to prevent contention with inserting new names + Snapshot, ) - .map(|res| match res { - Ok(entry) => { - let (key, metadata) = txs.read_entry::(&entry)?; + .map(|res| { + let (key, metadata) = tx.read_entry::(&res?)?; - Ok((key.name, metadata)) - } - Err(err) => Err(Into::::into(err)), + Ok((key.name, metadata)) }) .try_collect::>() .await diff --git a/packages/services/pegboard/src/ops/runner/get.rs b/packages/services/pegboard/src/ops/runner/get.rs index 484eb8e0de..3c588170fb 100644 --- a/packages/services/pegboard/src/ops/runner/get.rs +++ b/packages/services/pegboard/src/ops/runner/get.rs @@ -2,8 +2,8 @@ use anyhow::Result; use futures_util::TryStreamExt; use gas::prelude::*; use rivet_types::runners::Runner; -use udb_util::{FormalChunkedKey, SERIALIZABLE, SNAPSHOT, TxnExt}; -use universaldb::{self as udb, options::StreamingMode}; +use universaldb::options::StreamingMode; +use universaldb::utils::{FormalChunkedKey, IsolationLevel::*}; use crate::keys; @@ -23,7 +23,7 @@ pub async fn pegboard_runner_get(ctx: &OperationCtx, input: &Input) -> Result Result::Ok(runners) + Ok(runners) } }) .await?; @@ -44,18 +44,18 @@ pub async fn pegboard_runner_get(ctx: &OperationCtx, input: &Input) -> Result std::result::Result, udb::FdbBindingError> { - let txs = tx.subspace(keys::subspace()); +) -> Result> { + let tx = tx.with_subspace(keys::subspace()); // TODO: Make this part of the below try join to reduce round trip count // Check if runner exists by looking for workflow ID - if !txs - .exists(&keys::runner::WorkflowIdKey::new(runner_id), SERIALIZABLE) + if !tx + .exists(&keys::runner::WorkflowIdKey::new(runner_id), Serializable) .await? { - return std::result::Result::Ok(None); + return Ok(None); } let namespace_id_key = keys::runner::NamespaceIdKey::new(runner_id); @@ -71,7 +71,7 @@ pub(crate) async fn get_inner( let last_ping_ts_key = keys::runner::LastPingTsKey::new(runner_id); let last_rtt_key = keys::runner::LastRttKey::new(runner_id); let metadata_key = keys::runner::MetadataKey::new(runner_id); - let metadata_subspace = txs.subspace(&metadata_key); + let metadata_subspace = keys::subspace().subspace(&metadata_key); let ( namespace_id, @@ -88,27 +88,27 @@ pub(crate) async fn get_inner( last_rtt, metadata_chunks, ) = tokio::try_join!( - // NOTE: These are not SERIALIZABLE because this op is meant for basic information (i.e. data for the + // NOTE: These are not Serializable because this op is meant for basic information (i.e. data for the // API) - txs.read(&namespace_id_key, SNAPSHOT), - txs.read(&name_key, SNAPSHOT), - txs.read(&key_key, SNAPSHOT), - txs.read(&version_key, SNAPSHOT), - txs.read(&total_slots_key, SNAPSHOT), - txs.read(&remaining_slots_key, SNAPSHOT), - txs.read(&create_ts_key, SNAPSHOT), - txs.read_opt(&connected_ts_key, SNAPSHOT), - txs.read_opt(&drain_ts_key, SNAPSHOT), - txs.read_opt(&stop_ts_key, SNAPSHOT), - txs.read_opt(&last_ping_ts_key, SNAPSHOT), - txs.read_opt(&last_rtt_key, SNAPSHOT), + tx.read(&namespace_id_key, Snapshot), + tx.read(&name_key, Snapshot), + tx.read(&key_key, Snapshot), + tx.read(&version_key, Snapshot), + tx.read(&total_slots_key, Snapshot), + tx.read(&remaining_slots_key, Snapshot), + tx.read(&create_ts_key, Snapshot), + tx.read_opt(&connected_ts_key, Snapshot), + tx.read_opt(&drain_ts_key, Snapshot), + tx.read_opt(&stop_ts_key, Snapshot), + tx.read_opt(&last_ping_ts_key, Snapshot), + tx.read_opt(&last_rtt_key, Snapshot), async { - txs.get_ranges_keyvalues( - udb::RangeOption { + tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&metadata_subspace).into() }, - SNAPSHOT, + Snapshot, ) .try_collect::>() .await @@ -119,12 +119,7 @@ pub(crate) async fn get_inner( let metadata = if metadata_chunks.is_empty() { None } else { - Some( - metadata_key - .combine(metadata_chunks) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))? - .metadata, - ) + Some(metadata_key.combine(metadata_chunks)?.metadata) }; std::result::Result::Ok(Some(Runner { diff --git a/packages/services/pegboard/src/ops/runner/get_by_key.rs b/packages/services/pegboard/src/ops/runner/get_by_key.rs index 5a6ea63bf1..ab1a4919be 100644 --- a/packages/services/pegboard/src/ops/runner/get_by_key.rs +++ b/packages/services/pegboard/src/ops/runner/get_by_key.rs @@ -1,8 +1,7 @@ -use anyhow::*; +use anyhow::Result; use gas::prelude::*; use rivet_types::runners::Runner; -use udb_util::{SERIALIZABLE, TxnExt}; -use universaldb as udb; +use universaldb::utils::IsolationLevel::*; use crate::keys; @@ -24,24 +23,24 @@ pub async fn pegboard_runner_get_by_key(ctx: &OperationCtx, input: &Input) -> Re let runner = ctx .udb()? - .run(|tx, _mc| { + .run(|tx| { let dc_name = dc_name.to_string(); let input = input.clone(); async move { - let txs = tx.subspace(keys::subspace()); + let tx = tx.with_subspace(keys::subspace()); // Look up runner by key let runner_by_key_key = keys::ns::RunnerByKeyKey::new(input.namespace_id, input.name, input.key); - let runner_data = txs.read_opt(&runner_by_key_key, SERIALIZABLE).await?; + let runner_data = tx.read_opt(&runner_by_key_key, Serializable).await?; if let Some(data) = runner_data { // Get full runner details using the runner_id let runner = super::get::get_inner(&dc_name, &tx, data.runner_id).await?; - std::result::Result::<_, udb::FdbBindingError>::Ok(runner) + Ok(runner) } else { - std::result::Result::<_, udb::FdbBindingError>::Ok(None) + Ok(None) } } }) diff --git a/packages/services/pegboard/src/ops/runner/list_for_ns.rs b/packages/services/pegboard/src/ops/runner/list_for_ns.rs index 58bb974b1e..0947a8d1c6 100644 --- a/packages/services/pegboard/src/ops/runner/list_for_ns.rs +++ b/packages/services/pegboard/src/ops/runner/list_for_ns.rs @@ -2,8 +2,8 @@ use anyhow::Result; use futures_util::{StreamExt, TryStreamExt}; use gas::prelude::*; use rivet_types::runners::Runner; -use udb_util::{SNAPSHOT, TxnExt}; -use universaldb::{self as udb, options::StreamingMode}; +use universaldb::options::StreamingMode; +use universaldb::utils::IsolationLevel::*; use crate::keys; @@ -27,24 +27,24 @@ pub async fn pegboard_runner_list_for_ns(ctx: &OperationCtx, input: &Input) -> R let runners = ctx .udb()? - .run(|tx, _mc| { + .run(|tx| { let dc_name = dc_name.to_string(); async move { - let txs = tx.subspace(keys::subspace()); + let tx = tx.with_subspace(keys::subspace()); let mut results = Vec::new(); // TODO: Lots of duplicate code if let Some(name) = &input.name { if input.include_stopped { let runner_subspace = - txs.subspace(&keys::ns::AllRunnerByNameKey::subspace( + keys::subspace().subspace(&keys::ns::AllRunnerByNameKey::subspace( input.namespace_id, name.clone(), )); let (start, end) = runner_subspace.range(); let end = if let Some(created_before) = input.created_before { - udb_util::end_of_key_range(&txs.pack( + universaldb::utils::end_of_key_range(&tx.pack( &keys::ns::AllRunnerByNameKey::subspace_with_create_ts( input.namespace_id, name.clone(), @@ -55,19 +55,18 @@ pub async fn pegboard_runner_list_for_ns(ctx: &OperationCtx, input: &Input) -> R end }; - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { + let mut stream = tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::Iterator, reverse: true, ..(start, end).into() }, // NOTE: Does not have to be serializable because we are listing, stale data does not matter - SNAPSHOT, + Snapshot, ); while let Some(entry) = stream.try_next().await? { - let idx_key = - txs.unpack::(entry.key())?; + let idx_key = tx.unpack::(entry.key())?; results.push(idx_key.runner_id); @@ -77,14 +76,14 @@ pub async fn pegboard_runner_list_for_ns(ctx: &OperationCtx, input: &Input) -> R } } else { let runner_subspace = - txs.subspace(&keys::ns::ActiveRunnerByNameKey::subspace( + keys::subspace().subspace(&keys::ns::ActiveRunnerByNameKey::subspace( input.namespace_id, name.clone(), )); let (start, end) = runner_subspace.range(); let end = if let Some(created_before) = input.created_before { - udb_util::end_of_key_range(&txs.pack( + universaldb::utils::end_of_key_range(&tx.pack( &keys::ns::ActiveRunnerByNameKey::subspace_with_create_ts( input.namespace_id, name.clone(), @@ -95,19 +94,19 @@ pub async fn pegboard_runner_list_for_ns(ctx: &OperationCtx, input: &Input) -> R end }; - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { + let mut stream = tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::Iterator, reverse: true, ..(start, end).into() }, // NOTE: Does not have to be serializable because we are listing, stale data does not matter - SNAPSHOT, + Snapshot, ); while let Some(entry) = stream.try_next().await? { let idx_key = - txs.unpack::(entry.key())?; + tx.unpack::(entry.key())?; results.push(idx_key.runner_id); @@ -118,12 +117,12 @@ pub async fn pegboard_runner_list_for_ns(ctx: &OperationCtx, input: &Input) -> R } } else { if input.include_stopped { - let runner_subspace = - txs.subspace(&keys::ns::AllRunnerKey::subspace(input.namespace_id)); + let runner_subspace = keys::subspace() + .subspace(&keys::ns::AllRunnerKey::subspace(input.namespace_id)); let (start, end) = runner_subspace.range(); let end = if let Some(created_before) = input.created_before { - udb_util::end_of_key_range(&txs.pack( + universaldb::utils::end_of_key_range(&tx.pack( &keys::ns::AllRunnerKey::subspace_with_create_ts( input.namespace_id, created_before, @@ -133,18 +132,18 @@ pub async fn pegboard_runner_list_for_ns(ctx: &OperationCtx, input: &Input) -> R end }; - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { + let mut stream = tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::Iterator, reverse: true, ..(start, end).into() }, // NOTE: Does not have to be serializable because we are listing, stale data does not matter - SNAPSHOT, + Snapshot, ); while let Some(entry) = stream.try_next().await? { - let idx_key = txs.unpack::(entry.key())?; + let idx_key = tx.unpack::(entry.key())?; results.push(idx_key.runner_id); @@ -153,12 +152,12 @@ pub async fn pegboard_runner_list_for_ns(ctx: &OperationCtx, input: &Input) -> R } } } else { - let runner_subspace = - txs.subspace(&keys::ns::ActiveRunnerKey::subspace(input.namespace_id)); + let runner_subspace = keys::subspace() + .subspace(&keys::ns::ActiveRunnerKey::subspace(input.namespace_id)); let (start, end) = runner_subspace.range(); let end = if let Some(created_before) = input.created_before { - udb_util::end_of_key_range(&txs.pack( + universaldb::utils::end_of_key_range(&tx.pack( &keys::ns::ActiveRunnerKey::subspace_with_create_ts( input.namespace_id, created_before, @@ -168,18 +167,18 @@ pub async fn pegboard_runner_list_for_ns(ctx: &OperationCtx, input: &Input) -> R end }; - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { + let mut stream = tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::Iterator, reverse: true, ..(start, end).into() }, // NOTE: Does not have to be serializable because we are listing, stale data does not matter - SNAPSHOT, + Snapshot, ); while let Some(entry) = stream.try_next().await? { - let idx_key = txs.unpack::(entry.key())?; + let idx_key = tx.unpack::(entry.key())?; results.push(idx_key.runner_id); diff --git a/packages/services/pegboard/src/ops/runner/list_names.rs b/packages/services/pegboard/src/ops/runner/list_names.rs index cde5e8ea10..518971521e 100644 --- a/packages/services/pegboard/src/ops/runner/list_names.rs +++ b/packages/services/pegboard/src/ops/runner/list_names.rs @@ -1,7 +1,7 @@ use futures_util::{StreamExt, TryStreamExt}; use gas::prelude::*; -use udb_util::{SNAPSHOT, TxnExt}; -use universaldb::{self as udb, options::StreamingMode}; +use universaldb::options::StreamingMode; +use universaldb::utils::IsolationLevel::*; use crate::keys; @@ -21,15 +21,15 @@ pub struct Output { pub async fn pegboard_runner_list_names(ctx: &OperationCtx, input: &Input) -> Result { let names = ctx .udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); let runner_name_subspace = - txs.subspace(&keys::ns::RunnerNameKey::subspace(input.namespace_id)); + keys::subspace().subspace(&keys::ns::RunnerNameKey::subspace(input.namespace_id)); let (start, end) = runner_name_subspace.range(); let start = if let Some(name) = &input.after_name { - txs.pack(&keys::ns::RunnerNameKey::new( + tx.pack(&keys::ns::RunnerNameKey::new( input.namespace_id, name.clone(), )) @@ -37,22 +37,18 @@ pub async fn pegboard_runner_list_names(ctx: &OperationCtx, input: &Input) -> Re start }; - txs.get_ranges_keyvalues( - udb::RangeOption { + tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::WantAll, limit: Some(input.limit), ..(start, end).into() }, - // NOTE: This is not SERIALIZABLE to prevent contention with inserting new names - SNAPSHOT, + // NOTE: This is not Serializable to prevent contention with inserting new names + Snapshot, ) - .map(|res| match res { - Ok(entry) => { - let key = txs.unpack::(entry.key())?; - - Ok(key.name) - } - Err(err) => Err(Into::::into(err)), + .map(|res| { + let key = tx.unpack::(res?.key())?; + Ok(key.name) }) .try_collect::>() .await diff --git a/packages/services/pegboard/src/ops/runner/update_alloc_idx.rs b/packages/services/pegboard/src/ops/runner/update_alloc_idx.rs index 39519590df..5aa16e7572 100644 --- a/packages/services/pegboard/src/ops/runner/update_alloc_idx.rs +++ b/packages/services/pegboard/src/ops/runner/update_alloc_idx.rs @@ -1,6 +1,6 @@ use gas::prelude::*; -use udb_util::{SERIALIZABLE, TxnExt}; use universaldb::options::ConflictRangeType; +use universaldb::utils::IsolationLevel::*; use crate::{keys, workflows::runner::RUNNER_ELIGIBLE_THRESHOLD_MS}; @@ -47,11 +47,11 @@ pub enum RunnerEligibility { pub async fn pegboard_runner_update_alloc_idx(ctx: &OperationCtx, input: &Input) -> Result { let notifications = ctx .udb()? - .run(|tx, _mc| { + .run(|tx| { let runners = input.runners.clone(); async move { - let txs = tx.subspace(keys::subspace()); + let tx = tx.with_subspace(keys::subspace()); let mut notifications = Vec::new(); // TODO: Parallelize @@ -76,14 +76,14 @@ pub async fn pegboard_runner_update_alloc_idx(ctx: &OperationCtx, input: &Input) last_ping_ts_entry, expired_ts_entry, ) = tokio::try_join!( - txs.read_opt(&workflow_id_key, SERIALIZABLE), - txs.read_opt(&namespace_id_key, SERIALIZABLE), - txs.read_opt(&name_key, SERIALIZABLE), - txs.read_opt(&version_key, SERIALIZABLE), - txs.read_opt(&remaining_slots_key, SERIALIZABLE), - txs.read_opt(&total_slots_key, SERIALIZABLE), - txs.read_opt(&last_ping_ts_key, SERIALIZABLE), - txs.read_opt(&expired_ts_key, SERIALIZABLE), + tx.read_opt(&workflow_id_key, Serializable), + tx.read_opt(&namespace_id_key, Serializable), + tx.read_opt(&name_key, Serializable), + tx.read_opt(&version_key, Serializable), + tx.read_opt(&remaining_slots_key, Serializable), + tx.read_opt(&total_slots_key, Serializable), + tx.read_opt(&last_ping_ts_key, Serializable), + tx.read_opt(&expired_ts_key, Serializable), )?; let ( @@ -131,14 +131,14 @@ pub async fn pegboard_runner_update_alloc_idx(ctx: &OperationCtx, input: &Input) ); // Add read conflict - txs.add_conflict_key(&old_alloc_key, ConflictRangeType::Read)?; + tx.add_conflict_key(&old_alloc_key, ConflictRangeType::Read)?; match runner.action { Action::ClearIdx => { - txs.delete(&old_alloc_key); + tx.delete(&old_alloc_key); } Action::AddIdx => { - txs.write( + tx.write( &old_alloc_key, rivet_data::converted::RunnerAllocIdxKeyData { workflow_id, @@ -151,17 +151,17 @@ pub async fn pegboard_runner_update_alloc_idx(ctx: &OperationCtx, input: &Input) let last_ping_ts = util::timestamp::now(); // Write new ping - txs.write(&last_ping_ts_key, last_ping_ts)?; + tx.write(&last_ping_ts_key, last_ping_ts)?; let last_rtt_key = keys::runner::LastRttKey::new(runner.runner_id); - txs.write(&last_rtt_key, rtt)?; + tx.write(&last_rtt_key, rtt)?; // Only update allocation idx if it existed before - if txs.exists(&old_alloc_key, SERIALIZABLE).await? { + if tx.exists(&old_alloc_key, Serializable).await? { // Clear old key - txs.delete(&old_alloc_key); + tx.delete(&old_alloc_key); - txs.write( + tx.write( &keys::ns::RunnerAllocIdxKey::new( namespace_id, name.clone(), diff --git a/packages/services/pegboard/src/workflows/actor/actor_keys.rs b/packages/services/pegboard/src/workflows/actor/actor_keys.rs index e5cc89a10c..7bb8bcf851 100644 --- a/packages/services/pegboard/src/workflows/actor/actor_keys.rs +++ b/packages/services/pegboard/src/workflows/actor/actor_keys.rs @@ -5,8 +5,8 @@ use epoxy::{ use futures_util::TryStreamExt; use gas::prelude::*; use rivet_data::converted::ActorByKeyKeyData; -use udb_util::prelude::*; -use universaldb::{self as udb, FdbBindingError, options::StreamingMode}; +use universaldb::options::StreamingMode; +use universaldb::prelude::*; use crate::keys; @@ -231,27 +231,27 @@ pub async fn reserve_actor_key( ) -> Result { let res = ctx .udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); // Check if there are any actors that share the same key that are not destroyed - let actor_key_subspace = txs.subspace(&keys::ns::ActorByKeyKey::subspace( + let actor_key_subspace = keys::subspace().subspace(&keys::ns::ActorByKeyKey::subspace( input.namespace_id, input.name.clone(), input.key.clone(), )); let (start, end) = actor_key_subspace.range(); - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { + let mut stream = tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::Iterator, ..(start, end).into() }, - SERIALIZABLE, + Serializable, ); while let Some(entry) = stream.try_next().await? { - let (_idx_key, data) = txs.read_entry::(&entry)?; + let (_idx_key, data) = tx.read_entry::(&entry)?; if !data.is_destroyed { return Ok(ReserveActorKeyOutput::ExistingActor { existing_actor_id: _idx_key.actor_id, @@ -260,7 +260,7 @@ pub async fn reserve_actor_key( } // Write key - txs.write( + tx.write( &keys::ns::ActorByKeyKey::new( input.namespace_id, input.name.clone(), @@ -274,7 +274,7 @@ pub async fn reserve_actor_key( }, )?; - Result::<_, FdbBindingError>::Ok(ReserveActorKeyOutput::Success) + Ok(ReserveActorKeyOutput::Success) }) .await?; diff --git a/packages/services/pegboard/src/workflows/actor/destroy.rs b/packages/services/pegboard/src/workflows/actor/destroy.rs index 7ba691c760..f267328831 100644 --- a/packages/services/pegboard/src/workflows/actor/destroy.rs +++ b/packages/services/pegboard/src/workflows/actor/destroy.rs @@ -1,8 +1,8 @@ use gas::prelude::*; use rivet_data::converted::ActorByKeyKeyData; use rivet_runner_protocol::protocol; -use udb_util::{SERIALIZABLE, TxnExt}; -use universaldb::{self as udb, options::MutationType}; +use universaldb::options::MutationType; +use universaldb::utils::IsolationLevel::*; use super::{DestroyComplete, DestroyStarted, State}; @@ -28,7 +28,7 @@ pub(crate) async fn pegboard_actor_destroy(ctx: &mut WorkflowCtx, input: &Input) .await?; let res = ctx - .activity(UpdateStateAndFdbInput { + .activity(UpdateStateAndDbInput { actor_id: input.actor_id, }) .await?; @@ -53,31 +53,31 @@ pub(crate) async fn pegboard_actor_destroy(ctx: &mut WorkflowCtx, input: &Input) } #[derive(Debug, Serialize, Deserialize, Hash)] -struct UpdateStateAndFdbInput { +struct UpdateStateAndDbInput { actor_id: Id, } #[derive(Debug, Serialize, Deserialize, Hash)] -struct UpdateStateAndFdbOutput { +struct UpdateStateAndDbOutput { runner_workflow_id: Option, } -#[activity(UpdateStateAndFdb)] -async fn update_state_and_fdb( +#[activity(UpdateStateAndDb)] +async fn update_state_and_db( ctx: &ActivityCtx, - input: &UpdateStateAndFdbInput, -) -> Result { + input: &UpdateStateAndDbInput, +) -> Result { let mut state = ctx.state::()?; let destroy_ts = util::timestamp::now(); ctx.udb()? - .run(|tx, _mc| { + .run(|tx| { let state = (*state).clone(); async move { - let txs = tx.subspace(keys::subspace()); + let tx = tx.with_subspace(keys::subspace()); - txs.write(&keys::actor::DestroyTsKey::new(input.actor_id), destroy_ts)?; + tx.write(&keys::actor::DestroyTsKey::new(input.actor_id), destroy_ts)?; if let Some(runner_id) = state.runner_id { clear_slot( @@ -92,7 +92,7 @@ async fn update_state_and_fdb( } // Update namespace indexes - txs.delete(&keys::ns::ActiveActorKey::new( + tx.delete(&keys::ns::ActiveActorKey::new( state.namespace_id, state.name.clone(), state.create_ts, @@ -100,7 +100,7 @@ async fn update_state_and_fdb( )); if let Some(k) = &state.key { - txs.write( + tx.write( &keys::ns::ActorByKeyKey::new( state.namespace_id, state.name.clone(), @@ -125,7 +125,7 @@ async fn update_state_and_fdb( state.runner_id = None; let runner_workflow_id = state.runner_workflow_id.take(); - Ok(UpdateStateAndFdbOutput { runner_workflow_id }) + Ok(UpdateStateAndDbOutput { runner_workflow_id }) } #[derive(Debug, Serialize, Deserialize, Hash)] @@ -143,7 +143,7 @@ async fn clear_kv(ctx: &ActivityCtx, input: &ClearKvInput) -> Result Result<(), udb::FdbBindingError> { - let txs = tx.subspace(keys::subspace()); + tx: &universaldb::Transaction, +) -> Result<()> { + let tx = tx.with_subspace(keys::subspace()); - txs.delete(&keys::actor::RunnerIdKey::new(actor_id)); + tx.delete(&keys::actor::RunnerIdKey::new(actor_id)); // This is cleared when the state changes as well as when the actor is destroyed to ensure // consistency during rescheduling and forced deletion. - txs.delete(&keys::runner::ActorKey::new(runner_id, actor_id)); + tx.delete(&keys::runner::ActorKey::new(runner_id, actor_id)); let runner_workflow_id_key = keys::runner::WorkflowIdKey::new(runner_id); let runner_version_key = keys::runner::VersionKey::new(runner_id); @@ -187,18 +187,18 @@ pub(crate) async fn clear_slot( runner_total_slots, runner_last_ping_ts, ) = tokio::try_join!( - txs.read(&runner_workflow_id_key, SERIALIZABLE), - txs.read(&runner_version_key, SERIALIZABLE), - txs.read(&runner_remaining_slots_key, SERIALIZABLE), - txs.read(&runner_total_slots_key, SERIALIZABLE), - txs.read(&runner_last_ping_ts_key, SERIALIZABLE), + tx.read(&runner_workflow_id_key, Serializable), + tx.read(&runner_version_key, Serializable), + tx.read(&runner_remaining_slots_key, Serializable), + tx.read(&runner_total_slots_key, Serializable), + tx.read(&runner_last_ping_ts_key, Serializable), )?; let old_runner_remaining_millislots = (runner_remaining_slots * 1000) / runner_total_slots; let new_runner_remaining_slots = runner_remaining_slots + 1; // Write new remaining slots - txs.write(&runner_remaining_slots_key, new_runner_remaining_slots)?; + tx.write(&runner_remaining_slots_key, new_runner_remaining_slots)?; let old_runner_alloc_key = keys::ns::RunnerAllocIdxKey::new( namespace_id, @@ -210,9 +210,9 @@ pub(crate) async fn clear_slot( ); // Only update allocation idx if it existed before - if txs.exists(&old_runner_alloc_key, SERIALIZABLE).await? { + if tx.exists(&old_runner_alloc_key, Serializable).await? { // Clear old key - txs.delete(&old_runner_alloc_key); + tx.delete(&old_runner_alloc_key); let new_remaining_millislots = (new_runner_remaining_slots * 1000) / runner_total_slots; let new_runner_alloc_key = keys::ns::RunnerAllocIdxKey::new( @@ -224,7 +224,7 @@ pub(crate) async fn clear_slot( runner_id, ); - txs.write( + tx.write( &new_runner_alloc_key, rivet_data::converted::RunnerAllocIdxKeyData { workflow_id: runner_workflow_id, @@ -235,7 +235,7 @@ pub(crate) async fn clear_slot( } if for_serverless { - txs.atomic_op( + tx.atomic_op( &rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new( namespace_id, runner_name_selector.to_string(), diff --git a/packages/services/pegboard/src/workflows/actor/runtime.rs b/packages/services/pegboard/src/workflows/actor/runtime.rs index 9172583612..d01755d520 100644 --- a/packages/services/pegboard/src/workflows/actor/runtime.rs +++ b/packages/services/pegboard/src/workflows/actor/runtime.rs @@ -5,11 +5,8 @@ use futures_util::{FutureExt, TryStreamExt}; use gas::prelude::*; use rivet_metrics::KeyValue; use rivet_runner_protocol::protocol; -use udb_util::{FormalKey, SERIALIZABLE, SNAPSHOT, TxnExt}; -use universaldb::{ - self as udb, - options::{ConflictRangeType, MutationType, StreamingMode}, -}; +use universaldb::options::{ConflictRangeType, MutationType, StreamingMode}; +use universaldb::utils::{FormalKey, IsolationLevel::*}; use crate::{keys, metrics, workflows::runner::RUNNER_ELIGIBLE_THRESHOLD_MS}; @@ -107,24 +104,24 @@ async fn allocate_actor( // client wf let (for_serverless, res) = ctx .udb()? - .run(|tx, _mc| async move { + .run(|tx| async move { let ping_threshold_ts = util::timestamp::now() - RUNNER_ELIGIBLE_THRESHOLD_MS; - let txs = tx.subspace(keys::subspace()); + let tx = tx.with_subspace(keys::subspace()); // Check if runner is an serverless runner - let for_serverless = txs + let for_serverless = tx .exists( &namespace::keys::RunnerConfigByVariantKey::new( namespace_id, namespace::keys::RunnerConfigVariant::Serverless, input.runner_name_selector.clone(), ), - SERIALIZABLE, + Serializable, ) .await?; if for_serverless { - txs.atomic_op( + tx.atomic_op( &rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new( namespace_id, input.runner_name_selector.clone(), @@ -135,40 +132,42 @@ async fn allocate_actor( } // Check if a queue exists - let pending_actor_subspace = - txs.subspace(&keys::ns::PendingActorByRunnerNameSelectorKey::subspace( + let pending_actor_subspace = keys::subspace().subspace( + &keys::ns::PendingActorByRunnerNameSelectorKey::subspace( namespace_id, input.runner_name_selector.clone(), - )); - let queue_exists = txs + ), + ); + let queue_exists = tx .get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::Exact, limit: Some(1), ..(&pending_actor_subspace).into() }, - // NOTE: This is not SERIALIZABLE because we don't want to conflict with other + // NOTE: This is not Serializable because we don't want to conflict with other // inserts/clears to this range - SNAPSHOT, + Snapshot, ) .next() .await .is_some(); if !queue_exists { - let runner_alloc_subspace = txs.subspace(&keys::ns::RunnerAllocIdxKey::subspace( - namespace_id, - input.runner_name_selector.clone(), - )); + let runner_alloc_subspace = + keys::subspace().subspace(&keys::ns::RunnerAllocIdxKey::subspace( + namespace_id, + input.runner_name_selector.clone(), + )); - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { + let mut stream = tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::Iterator, ..(&runner_alloc_subspace).into() }, - // NOTE: This is not SERIALIZABLE because we don't want to conflict with all of the + // NOTE: This is not Serializable because we don't want to conflict with all of the // keys, just the one we choose - SNAPSHOT, + Snapshot, ); let mut highest_version = None; @@ -179,7 +178,7 @@ async fn allocate_actor( }; let (old_runner_alloc_key, old_runner_alloc_key_data) = - txs.read_entry::(&entry)?; + tx.read_entry::(&entry)?; if let Some(highest_version) = highest_version { // We have passed all of the runners with the highest version. This is reachable if @@ -202,10 +201,10 @@ async fn allocate_actor( } // Add read conflict only for this key - txs.add_conflict_key(&old_runner_alloc_key, ConflictRangeType::Read)?; + tx.add_conflict_key(&old_runner_alloc_key, ConflictRangeType::Read)?; // Clear old entry - txs.delete(&old_runner_alloc_key); + tx.delete(&old_runner_alloc_key); let new_remaining_slots = old_runner_alloc_key_data.remaining_slots.saturating_sub(1); @@ -213,7 +212,7 @@ async fn allocate_actor( (new_remaining_slots * 1000) / old_runner_alloc_key_data.total_slots; // Write new allocation key with 1 less slot - txs.write( + tx.write( &keys::ns::RunnerAllocIdxKey::new( namespace_id, input.runner_name_selector.clone(), @@ -230,19 +229,19 @@ async fn allocate_actor( )?; // Update runner record - txs.write( + tx.write( &keys::runner::RemainingSlotsKey::new(old_runner_alloc_key.runner_id), new_remaining_slots, )?; // Set runner id of actor - txs.write( + tx.write( &keys::actor::RunnerIdKey::new(input.actor_id), old_runner_alloc_key.runner_id, )?; // Insert actor index key - txs.write( + tx.write( &keys::runner::ActorKey::new( old_runner_alloc_key.runner_id, input.actor_id, @@ -251,7 +250,7 @@ async fn allocate_actor( )?; // Set actor as not sleeping - txs.delete(&keys::actor::SleepTsKey::new(input.actor_id)); + tx.delete(&keys::actor::SleepTsKey::new(input.actor_id)); return Ok(( for_serverless, @@ -270,7 +269,7 @@ async fn allocate_actor( // NOTE: This will conflict with serializable reads to the alloc queue, which is the behavior we // want. If a runner reads from the queue while this is being inserted, one of the two txns will // retry and we ensure the actor does not end up in queue limbo. - txs.write( + tx.write( &keys::ns::PendingActorByRunnerNameSelectorKey::new( namespace_id, input.runner_name_selector.clone(), @@ -316,7 +315,7 @@ pub async fn set_not_connectable(ctx: &ActivityCtx, input: &SetNotConnectableInp let mut state = ctx.state::()?; ctx.udb()? - .run(|tx, _mc| async move { + .run(|tx| async move { let connectable_key = keys::actor::ConnectableKey::new(input.actor_id); tx.clear(&keys::subspace().pack(&connectable_key)); @@ -344,10 +343,10 @@ pub async fn deallocate(ctx: &ActivityCtx, input: &DeallocateInput) -> Result<() let for_serverless = state.for_serverless; ctx.udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); - txs.delete(&keys::actor::ConnectableKey::new(input.actor_id)); + tx.delete(&keys::actor::ConnectableKey::new(input.actor_id)); if let Some(runner_id) = runner_id { destroy::clear_slot( @@ -360,7 +359,7 @@ pub async fn deallocate(ctx: &ActivityCtx, input: &DeallocateInput) -> Result<() ) .await?; } else if for_serverless { - txs.atomic_op( + tx.atomic_op( &rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new( namespace_id, runner_name_selector.clone(), @@ -574,7 +573,7 @@ pub async fn clear_pending_allocation( // Clear self from alloc queue let cleared = ctx .udb()? - .run(|tx, _mc| async move { + .run(|tx| async move { let pending_alloc_key = keys::subspace().pack(&keys::ns::PendingActorByRunnerNameSelectorKey::new( input.namespace_id, @@ -583,7 +582,7 @@ pub async fn clear_pending_allocation( input.actor_id, )); - let exists = tx.get(&pending_alloc_key, SERIALIZABLE).await?.is_some(); + let exists = tx.get(&pending_alloc_key, Serializable).await?.is_some(); tx.clear(&pending_alloc_key); @@ -620,13 +619,11 @@ pub async fn set_started(ctx: &ActivityCtx, input: &SetStartedInput) -> Result<( state.connectable_ts = Some(util::timestamp::now()); ctx.udb()? - .run(|tx, _mc| async move { + .run(|tx| async move { let connectable_key = keys::actor::ConnectableKey::new(input.actor_id); tx.set( &keys::subspace().pack(&connectable_key), - &connectable_key - .serialize(()) - .map_err(|x| udb::FdbBindingError::CustomError(x.into()))?, + &connectable_key.serialize(())?, ); Ok(()) @@ -650,13 +647,13 @@ pub async fn set_sleeping(ctx: &ActivityCtx, input: &SetSleepingInput) -> Result state.connectable_ts = None; ctx.udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); // Make not connectable - txs.delete(&keys::actor::ConnectableKey::new(input.actor_id)); + tx.delete(&keys::actor::ConnectableKey::new(input.actor_id)); - txs.write(&keys::actor::SleepTsKey::new(input.actor_id), sleep_ts)?; + tx.write(&keys::actor::SleepTsKey::new(input.actor_id), sleep_ts)?; Ok(()) }) diff --git a/packages/services/pegboard/src/workflows/actor/setup.rs b/packages/services/pegboard/src/workflows/actor/setup.rs index 068cc562f0..9ac87beeb2 100644 --- a/packages/services/pegboard/src/workflows/actor/setup.rs +++ b/packages/services/pegboard/src/workflows/actor/setup.rs @@ -1,7 +1,7 @@ use gas::prelude::*; use rivet_data::converted::ActorNameKeyData; use rivet_types::actors::CrashPolicy; -use udb_util::{SERIALIZABLE, TxnExt}; +use universaldb::utils::IsolationLevel::*; use super::State; @@ -69,8 +69,8 @@ pub struct InitStateAndUdbInput { pub create_ts: i64, } -#[activity(InitStateAndFdb)] -pub async fn insert_state_and_fdb(ctx: &ActivityCtx, input: &InitStateAndUdbInput) -> Result<()> { +#[activity(InitStateAndDb)] +pub async fn insert_state_and_db(ctx: &ActivityCtx, input: &InitStateAndUdbInput) -> Result<()> { let mut state = ctx.state::>()?; *state = Some(State::new( @@ -83,14 +83,14 @@ pub async fn insert_state_and_fdb(ctx: &ActivityCtx, input: &InitStateAndUdbInpu )); ctx.udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); - txs.write( + tx.write( &keys::actor::CreateTsKey::new(input.actor_id), input.create_ts, )?; - txs.write( + tx.write( &keys::actor::WorkflowIdKey::new(input.actor_id), ctx.workflow_id(), )?; @@ -120,15 +120,15 @@ pub async fn add_indexes_and_set_create_complete( // Populate indexes ctx.udb()? - .run(|tx, _mc| { + .run(|tx| { let namespace_id = state.namespace_id; let name = state.name.clone(); let create_ts = state.create_ts; async move { - let txs = tx.subspace(keys::subspace()); + let tx = tx.with_subspace(keys::subspace()); // Populate indexes - txs.write( + tx.write( &keys::ns::ActiveActorKey::new( namespace_id, name.clone(), @@ -138,7 +138,7 @@ pub async fn add_indexes_and_set_create_complete( ctx.workflow_id(), )?; - txs.write( + tx.write( &keys::ns::AllActorKey::new( namespace_id, name.clone(), @@ -150,8 +150,8 @@ pub async fn add_indexes_and_set_create_complete( // Write name into namespace actor names list with empty metadata (if it doesn't already exist) let name_key = keys::ns::ActorNameKey::new(namespace_id, name.clone()); - if !txs.exists(&name_key, SERIALIZABLE).await? { - txs.write( + if !tx.exists(&name_key, Serializable).await? { + tx.write( &name_key, ActorNameKeyData { metadata: serde_json::Map::new(), diff --git a/packages/services/pegboard/src/workflows/runner.rs b/packages/services/pegboard/src/workflows/runner.rs index b88cc019c2..7c2dd7ae93 100644 --- a/packages/services/pegboard/src/workflows/runner.rs +++ b/packages/services/pegboard/src/workflows/runner.rs @@ -2,11 +2,8 @@ use futures_util::{FutureExt, StreamExt, TryStreamExt}; use gas::prelude::*; use rivet_data::converted::{ActorNameKeyData, MetadataKeyData, RunnerByKeyKeyData}; use rivet_runner_protocol::protocol; -use udb_util::{FormalChunkedKey, SERIALIZABLE, SNAPSHOT, TxnExt}; -use universaldb::{ - self as udb, - options::{ConflictRangeType, StreamingMode}, -}; +use universaldb::options::{ConflictRangeType, StreamingMode}; +use universaldb::utils::{FormalChunkedKey, IsolationLevel::*}; use crate::{keys, workflows::actor::Allocate}; @@ -127,7 +124,7 @@ pub async fn pegboard_runner(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> } if !state.draining { - ctx.activity(InsertFdbInput { + ctx.activity(InsertDbInput { runner_id: input.runner_id, namespace_id: input.namespace_id, name: input.name.clone(), @@ -207,7 +204,7 @@ pub async fn pegboard_runner(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> state.draining = true; // Can't parallelize these two, requires reading from state - ctx.activity(ClearFdbInput { + ctx.activity(ClearDbInput { runner_id: input.runner_id, name: input.name.clone(), key: input.key.clone(), @@ -358,7 +355,7 @@ pub async fn pegboard_runner(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> }) .await?; - ctx.activity(ClearFdbInput { + ctx.activity(ClearDbInput { runner_id: input.runner_id, name: input.name.clone(), key: input.key.clone(), @@ -443,8 +440,8 @@ async fn init(ctx: &ActivityCtx, input: &InitInput) -> Result { let evict_workflow_id = ctx .udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); let runner_by_key_key = keys::ns::RunnerByKeyKey::new( input.namespace_id, @@ -453,13 +450,13 @@ async fn init(ctx: &ActivityCtx, input: &InitInput) -> Result { ); // Read existing runner by key slot - let evict_workflow_id = txs - .read_opt(&runner_by_key_key, SERIALIZABLE) + let evict_workflow_id = tx + .read_opt(&runner_by_key_key, Serializable) .await? .map(|x| x.workflow_id); // Allocate self - txs.write( + tx.write( &runner_by_key_key, RunnerByKeyKeyData { runner_id: input.runner_id, @@ -475,7 +472,7 @@ async fn init(ctx: &ActivityCtx, input: &InitInput) -> Result { } #[derive(Debug, Serialize, Deserialize, Hash)] -struct InsertFdbInput { +struct InsertDbInput { runner_id: Id, namespace_id: Id, name: String, @@ -485,19 +482,19 @@ struct InsertFdbInput { create_ts: i64, } -#[activity(InsertFdb)] -async fn insert_fdb(ctx: &ActivityCtx, input: &InsertFdbInput) -> Result<()> { +#[activity(InsertDb)] +async fn insert_db(ctx: &ActivityCtx, input: &InsertDbInput) -> Result<()> { ctx.udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); let remaining_slots_key = keys::runner::RemainingSlotsKey::new(input.runner_id); let last_ping_ts_key = keys::runner::LastPingTsKey::new(input.runner_id); let workflow_id_key = keys::runner::WorkflowIdKey::new(input.runner_id); let (remaining_slots_entry, last_ping_ts_entry) = tokio::try_join!( - txs.read_opt(&remaining_slots_key, SERIALIZABLE), - txs.read_opt(&last_ping_ts_key, SERIALIZABLE), + tx.read_opt(&remaining_slots_key, Serializable), + tx.read_opt(&last_ping_ts_key, Serializable), )?; let now = util::timestamp::now(); @@ -516,44 +513,44 @@ async fn insert_fdb(ctx: &ActivityCtx, input: &InsertFdbInput) -> Result<()> { } // NOTE: These properties are only inserted once else { - txs.write(&workflow_id_key, ctx.workflow_id())?; + tx.write(&workflow_id_key, ctx.workflow_id())?; - txs.write( + tx.write( &keys::runner::NamespaceIdKey::new(input.runner_id), input.namespace_id, )?; - txs.write( + tx.write( &keys::runner::NameKey::new(input.runner_id), input.name.clone(), )?; - txs.write( + tx.write( &keys::runner::KeyKey::new(input.runner_id), input.key.clone(), )?; - txs.write( + tx.write( &keys::runner::VersionKey::new(input.runner_id), input.version, )?; - txs.write(&remaining_slots_key, input.total_slots)?; + tx.write(&remaining_slots_key, input.total_slots)?; - txs.write( + tx.write( &keys::runner::TotalSlotsKey::new(input.runner_id), input.total_slots, )?; - txs.write( + tx.write( &keys::runner::CreateTsKey::new(input.runner_id), input.create_ts, )?; - txs.write(&last_ping_ts_key, now)?; + tx.write(&last_ping_ts_key, now)?; // Populate ns indexes - txs.write( + tx.write( &keys::ns::ActiveRunnerKey::new( input.namespace_id, input.create_ts, @@ -561,7 +558,7 @@ async fn insert_fdb(ctx: &ActivityCtx, input: &InsertFdbInput) -> Result<()> { ), ctx.workflow_id(), )?; - txs.write( + tx.write( &keys::ns::ActiveRunnerByNameKey::new( input.namespace_id, input.name.clone(), @@ -570,7 +567,7 @@ async fn insert_fdb(ctx: &ActivityCtx, input: &InsertFdbInput) -> Result<()> { ), ctx.workflow_id(), )?; - txs.write( + tx.write( &keys::ns::AllRunnerKey::new( input.namespace_id, input.create_ts, @@ -578,7 +575,7 @@ async fn insert_fdb(ctx: &ActivityCtx, input: &InsertFdbInput) -> Result<()> { ), ctx.workflow_id(), )?; - txs.write( + tx.write( &keys::ns::AllRunnerByNameKey::new( input.namespace_id, input.name.clone(), @@ -589,7 +586,7 @@ async fn insert_fdb(ctx: &ActivityCtx, input: &InsertFdbInput) -> Result<()> { )?; // Write name into namespace runner names list - txs.write( + tx.write( &keys::ns::RunnerNameKey::new(input.namespace_id, input.name.clone()), (), )?; @@ -598,12 +595,12 @@ async fn insert_fdb(ctx: &ActivityCtx, input: &InsertFdbInput) -> Result<()> { }; // Set last connect ts - txs.write(&keys::runner::ConnectedTsKey::new(input.runner_id), now)?; + tx.write(&keys::runner::ConnectedTsKey::new(input.runner_id), now)?; let remaining_millislots = (remaining_slots * 1000) / input.total_slots; // Insert into index (same as the `update_alloc_idx` op with `AddIdx`) - txs.write( + tx.write( &keys::ns::RunnerAllocIdxKey::new( input.namespace_id, input.name.clone(), @@ -628,7 +625,7 @@ async fn insert_fdb(ctx: &ActivityCtx, input: &InsertFdbInput) -> Result<()> { } #[derive(Debug, Serialize, Deserialize, Hash)] -struct ClearFdbInput { +struct ClearDbInput { runner_id: Id, name: String, key: String, @@ -641,44 +638,44 @@ enum RunnerState { Stopped, } -#[activity(ClearFdb)] -async fn clear_fdb(ctx: &ActivityCtx, input: &ClearFdbInput) -> Result<()> { +#[activity(ClearDb)] +async fn clear_Db(ctx: &ActivityCtx, input: &ClearDbInput) -> Result<()> { let state = ctx.state::()?; let namespace_id = state.namespace_id; let create_ts = state.create_ts; - // TODO: Combine into a single fdb txn + // TODO: Combine into a single udb txn ctx.udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); let now = util::timestamp::now(); // Clear runner by key idx if its still the current runner let runner_by_key_key = keys::ns::RunnerByKeyKey::new(namespace_id, input.name.clone(), input.key.clone()); - let runner_id = txs - .read_opt(&runner_by_key_key, SERIALIZABLE) + let runner_id = tx + .read_opt(&runner_by_key_key, Serializable) .await? .map(|x| x.runner_id); if runner_id == Some(input.runner_id) { - txs.delete(&runner_by_key_key); + tx.delete(&runner_by_key_key); } match input.update_state { RunnerState::Draining => { - txs.write(&keys::runner::DrainTsKey::new(input.runner_id), now)?; - txs.write(&keys::runner::ExpiredTsKey::new(input.runner_id), now)?; + tx.write(&keys::runner::DrainTsKey::new(input.runner_id), now)?; + tx.write(&keys::runner::ExpiredTsKey::new(input.runner_id), now)?; } RunnerState::Stopped => { - txs.write(&keys::runner::StopTsKey::new(input.runner_id), now)?; + tx.write(&keys::runner::StopTsKey::new(input.runner_id), now)?; // Update namespace indexes - txs.delete(&keys::ns::ActiveRunnerKey::new( + tx.delete(&keys::ns::ActiveRunnerKey::new( namespace_id, create_ts, input.runner_id, )); - txs.delete(&keys::ns::ActiveRunnerByNameKey::new( + tx.delete(&keys::ns::ActiveRunnerByNameKey::new( namespace_id, input.name.clone(), create_ts, @@ -723,8 +720,8 @@ async fn process_init(ctx: &ActivityCtx, input: &ProcessInitInput) -> Result()?; ctx.udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); // Populate actor names if provided if let Some(actor_names) = &input.prepopulate_actor_names { @@ -736,7 +733,7 @@ async fn process_init(ctx: &ActivityCtx, input: &ProcessInitInput) -> Result Result Result> { let actors = ctx .udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); let actor_subspace = keys::subspace().subspace(&keys::runner::ActorKey::subspace(input.runner_id)); tx.get_ranges_keyvalues( - udb::RangeOption { + universaldb::RangeOption { mode: StreamingMode::WantAll, ..(&actor_subspace).into() }, - SERIALIZABLE, + Serializable, ) - .map(|res| match res { - Ok(entry) => { - let (key, generation) = txs.read_entry::(&entry)?; + .map(|res| { + let (key, generation) = tx.read_entry::(&res?)?; - Ok((key.actor_id.into(), generation)) - } - Err(err) => Err(Into::::into(err)), + Ok((key.actor_id.into(), generation)) }) .try_collect::>() .await @@ -910,13 +899,13 @@ struct CheckExpiredInput { #[activity(CheckExpired)] async fn check_expired(ctx: &ActivityCtx, input: &CheckExpiredInput) -> Result { ctx.udb()? - .run(|tx, _mc| async move { - let txs = tx.subspace(keys::subspace()); + .run(|tx| async move { + let tx = tx.with_subspace(keys::subspace()); - let last_ping_ts = txs + let last_ping_ts = tx .read( &keys::runner::LastPingTsKey::new(input.runner_id), - SERIALIZABLE, + Serializable, ) .await?; @@ -924,7 +913,7 @@ async fn check_expired(ctx: &ActivityCtx, input: &CheckExpiredInput) -> Result(&queue_entry)?; + tx.read_entry::(&queue_entry)?; - let runner_alloc_subspace = txs.subspace(&keys::ns::RunnerAllocIdxKey::subspace( - input.namespace_id, - input.name.clone(), - )); + let runner_alloc_subspace = keys::subspace().subspace( + &keys::ns::RunnerAllocIdxKey::subspace(input.namespace_id, input.name.clone()), + ); - let mut stream = txs.get_ranges_keyvalues( - udb::RangeOption { + let mut stream = tx.get_ranges_keyvalues( + universaldb::RangeOption { mode: StreamingMode::Iterator, // Containers bin pack so we reverse the order reverse: true, ..(&runner_alloc_subspace).into() }, - // NOTE: This is not SERIALIZABLE because we don't want to conflict with all of the + // NOTE: This is not Serializable because we don't want to conflict with all of the // keys, just the one we choose - SNAPSHOT, + Snapshot, ); let mut highest_version = None; @@ -1012,7 +1001,7 @@ pub(crate) async fn allocate_pending_actors( }; let (old_runner_alloc_key, old_runner_alloc_key_data) = - txs.read_entry::(&entry)?; + tx.read_entry::(&entry)?; if let Some(highest_version) = highest_version { // We have passed all of the runners with the highest version. This is reachable if @@ -1035,12 +1024,12 @@ pub(crate) async fn allocate_pending_actors( } // Add read conflict only for this runner key - txs.add_conflict_key(&old_runner_alloc_key, ConflictRangeType::Read)?; - txs.delete(&old_runner_alloc_key); + tx.add_conflict_key(&old_runner_alloc_key, ConflictRangeType::Read)?; + tx.delete(&old_runner_alloc_key); // Add read conflict for the queue key - txs.add_conflict_key(&queue_key, ConflictRangeType::Read)?; - txs.delete(&queue_key); + tx.add_conflict_key(&queue_key, ConflictRangeType::Read)?; + tx.delete(&queue_key); let new_remaining_slots = old_runner_alloc_key_data.remaining_slots.saturating_sub(1); @@ -1048,7 +1037,7 @@ pub(crate) async fn allocate_pending_actors( (new_remaining_slots * 1000) / old_runner_alloc_key_data.total_slots; // Write new allocation key with 1 less slot - txs.write( + tx.write( &keys::ns::RunnerAllocIdxKey::new( input.namespace_id, input.name.clone(), @@ -1065,19 +1054,19 @@ pub(crate) async fn allocate_pending_actors( )?; // Update runner record - txs.write( + tx.write( &keys::runner::RemainingSlotsKey::new(old_runner_alloc_key.runner_id), new_remaining_slots, )?; // Set runner id of actor - txs.write( + tx.write( &keys::actor::RunnerIdKey::new(queue_key.actor_id), old_runner_alloc_key.runner_id, )?; // Insert actor index key - txs.write( + tx.write( &keys::runner::ActorKey::new( old_runner_alloc_key.runner_id, queue_key.actor_id, From be514ef5d27c4410db7e1b2576dc8a9070040614 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Thu, 11 Sep 2025 14:32:36 -0700 Subject: [PATCH 13/17] chore: cleanup, move ee changes --- frontend/packages/components/src/lib/config.ts | 2 ++ frontend/src/components/lib/config.ts | 2 ++ packages/common/universaldb/src/error.rs | 2 -- packages/core/bootstrap/src/lib.rs | 2 +- .../services/epoxy/src/ops/explicit_prepare.rs | 5 ++++- packages/services/epoxy/src/ops/kv/get_local.rs | 2 +- .../services/epoxy/src/ops/kv/get_optimistic.rs | 2 +- packages/services/epoxy/src/ops/propose.rs | 2 +- .../services/epoxy/src/ops/read_cluster_config.rs | 2 +- .../services/epoxy/src/replica/message_request.rs | 4 ++-- .../epoxy/src/workflows/coordinator/mod.rs | 14 +++++++------- .../workflows/coordinator/replica_status_change.rs | 2 +- .../services/epoxy/src/workflows/replica/mod.rs | 10 +++++----- .../services/epoxy/src/workflows/replica/setup.rs | 14 +++++++------- packages/services/epoxy/tests/common/mod.rs | 2 +- packages/services/epoxy/tests/reconfigure.rs | 2 +- sdks/typescript/test-runner/src/main.ts | 10 ++++++---- 17 files changed, 43 insertions(+), 36 deletions(-) diff --git a/frontend/packages/components/src/lib/config.ts b/frontend/packages/components/src/lib/config.ts index 91b9537457..2131a0d0fc 100644 --- a/frontend/packages/components/src/lib/config.ts +++ b/frontend/packages/components/src/lib/config.ts @@ -31,6 +31,8 @@ const getApiEndpoint = (apiEndpoint: string) => { } // Default to staging servers for all other endpoints return "https://api.staging2.gameinc.io"; + } else if (apiEndpoint === "__SAME__") { + return location.origin; } return apiEndpoint; }; diff --git a/frontend/src/components/lib/config.ts b/frontend/src/components/lib/config.ts index 91b9537457..2131a0d0fc 100644 --- a/frontend/src/components/lib/config.ts +++ b/frontend/src/components/lib/config.ts @@ -31,6 +31,8 @@ const getApiEndpoint = (apiEndpoint: string) => { } // Default to staging servers for all other endpoints return "https://api.staging2.gameinc.io"; + } else if (apiEndpoint === "__SAME__") { + return location.origin; } return apiEndpoint; }; diff --git a/packages/common/universaldb/src/error.rs b/packages/common/universaldb/src/error.rs index 04451239e5..f4b0ccd963 100644 --- a/packages/common/universaldb/src/error.rs +++ b/packages/common/universaldb/src/error.rs @@ -12,8 +12,6 @@ pub enum DatabaseError { #[error("operation issued while a commit was outstanding")] UsedDuringCommit, - // #[error(transparent)] - // Custom(Box), } impl DatabaseError { diff --git a/packages/core/bootstrap/src/lib.rs b/packages/core/bootstrap/src/lib.rs index 304ccafbee..f2684557b6 100644 --- a/packages/core/bootstrap/src/lib.rs +++ b/packages/core/bootstrap/src/lib.rs @@ -42,7 +42,7 @@ async fn setup_epoxy_coordinator(ctx: &StandaloneCtx) -> Result<()> { // // This does not guarantee the config will change immediately since we can't guarantee that the // coordinator workflow is running on a node with the newest version of the config. - ctx.signal(epoxy::workflows::coordinator::ReconfigureSignal {}) + ctx.signal(epoxy::workflows::coordinator::Reconfigure {}) .to_workflow_id(workflow_id) .send() .await?; diff --git a/packages/services/epoxy/src/ops/explicit_prepare.rs b/packages/services/epoxy/src/ops/explicit_prepare.rs index 9fef6278be..14ca28852a 100644 --- a/packages/services/epoxy/src/ops/explicit_prepare.rs +++ b/packages/services/epoxy/src/ops/explicit_prepare.rs @@ -20,7 +20,10 @@ pub enum ExplicitPrepareResult { } #[operation] -pub async fn explicit_prepare(ctx: &OperationCtx, input: &Input) -> Result { +pub async fn epoxy_explicit_prepare( + ctx: &OperationCtx, + input: &Input, +) -> Result { let replica_id = ctx.config().epoxy_replica_id(); let instance = &input.instance; diff --git a/packages/services/epoxy/src/ops/kv/get_local.rs b/packages/services/epoxy/src/ops/kv/get_local.rs index 74636d9d04..ef314689a3 100644 --- a/packages/services/epoxy/src/ops/kv/get_local.rs +++ b/packages/services/epoxy/src/ops/kv/get_local.rs @@ -18,7 +18,7 @@ pub struct Output { } #[operation] -pub async fn get_local(ctx: &OperationCtx, input: &Input) -> Result { +pub async fn epoxy_kv_get_local(ctx: &OperationCtx, input: &Input) -> Result { // Read from local KV store only let kv_key = keys::keys::KvValueKey::new(input.key.clone()); let subspace = keys::subspace(input.replica_id); diff --git a/packages/services/epoxy/src/ops/kv/get_optimistic.rs b/packages/services/epoxy/src/ops/kv/get_optimistic.rs index 1f9802e68a..6335cbbb90 100644 --- a/packages/services/epoxy/src/ops/kv/get_optimistic.rs +++ b/packages/services/epoxy/src/ops/kv/get_optimistic.rs @@ -35,7 +35,7 @@ pub struct Output { /// /// We cannot use quorum reads for the fanout read because of the constraints of Epaxos. #[operation] -pub async fn get_optimistic(ctx: &OperationCtx, input: &Input) -> Result { +pub async fn epoxy_kv_get_optimistic(ctx: &OperationCtx, input: &Input) -> Result { // Try to read locally let kv_key = keys::keys::KvValueKey::new(input.key.clone()); let cache_key = keys::keys::KvOptimisticCacheKey::new(input.key.clone()); diff --git a/packages/services/epoxy/src/ops/propose.rs b/packages/services/epoxy/src/ops/propose.rs index 6f9b05f7dc..8c42747036 100644 --- a/packages/services/epoxy/src/ops/propose.rs +++ b/packages/services/epoxy/src/ops/propose.rs @@ -27,7 +27,7 @@ pub struct Input { } #[operation] -pub async fn propose(ctx: &OperationCtx, input: &Input) -> Result { +pub async fn epoxy_propose(ctx: &OperationCtx, input: &Input) -> Result { let replica_id = ctx.config().epoxy_replica_id(); // Read config diff --git a/packages/services/epoxy/src/ops/read_cluster_config.rs b/packages/services/epoxy/src/ops/read_cluster_config.rs index 859f687aa0..ef4dc2d85c 100644 --- a/packages/services/epoxy/src/ops/read_cluster_config.rs +++ b/packages/services/epoxy/src/ops/read_cluster_config.rs @@ -15,7 +15,7 @@ pub struct Output { } #[operation] -pub async fn read_config(ctx: &OperationCtx, input: &Input) -> Result { +pub async fn epoxy_read_cluster_config(ctx: &OperationCtx, input: &Input) -> Result { let config = ctx .udb()? .run(|tx| { diff --git a/packages/services/epoxy/src/replica/message_request.rs b/packages/services/epoxy/src/replica/message_request.rs index 9dd0a3655d..44ec707be1 100644 --- a/packages/services/epoxy/src/replica/message_request.rs +++ b/packages/services/epoxy/src/replica/message_request.rs @@ -100,7 +100,7 @@ pub async fn message_request( "received coordinator update replica status request" ); - ctx.signal(crate::workflows::coordinator::ReplicaStatusChangeSignal { + ctx.signal(crate::workflows::coordinator::ReplicaStatusChange { replica_id: req.replica_id, status: req.status.into(), }) @@ -118,7 +118,7 @@ pub async fn message_request( "received begin learning request" ); - ctx.signal(crate::workflows::replica::BeginLearningSignal { + ctx.signal(crate::workflows::replica::BeginLearning { config: req.config.clone().into(), }) .to_workflow::() diff --git a/packages/services/epoxy/src/workflows/coordinator/mod.rs b/packages/services/epoxy/src/workflows/coordinator/mod.rs index 8c86d913e4..e5afb33f26 100644 --- a/packages/services/epoxy/src/workflows/coordinator/mod.rs +++ b/packages/services/epoxy/src/workflows/coordinator/mod.rs @@ -25,16 +25,16 @@ pub struct ReplicaState { } #[workflow] -pub async fn coordinator(ctx: &mut WorkflowCtx, _input: &Input) -> Result<()> { +pub async fn epoxy_coordinator(ctx: &mut WorkflowCtx, _input: &Input) -> Result<()> { ctx.activity(InitInput {}).await?; ctx.repeat(|ctx| { async move { match ctx.listen::
().await? { - Main::ReconfigureSignal(_) => { + Main::Reconfigure(_) => { reconfigure::reconfigure(ctx).await?; } - Main::ReplicaStatusChangeSignal(sig) => { + Main::ReplicaStatusChange(sig) => { replica_status_change::replica_status_change(ctx, sig).await?; } } @@ -73,15 +73,15 @@ pub struct ConfigChangeMessage { /// /// This gets called any time an engine node starts. #[signal("epoxy_coordinator_reconfigure")] -pub struct ReconfigureSignal {} +pub struct Reconfigure {} #[signal("epoxy_coordinator_replica_status_change")] -pub struct ReplicaStatusChangeSignal { +pub struct ReplicaStatusChange { pub replica_id: protocol::ReplicaId, pub status: types::ReplicaStatus, } join_signal!(Main { - ReconfigureSignal, - ReplicaStatusChangeSignal, + Reconfigure, + ReplicaStatusChange, }); diff --git a/packages/services/epoxy/src/workflows/coordinator/replica_status_change.rs b/packages/services/epoxy/src/workflows/coordinator/replica_status_change.rs index 2c92745cab..41da43f595 100644 --- a/packages/services/epoxy/src/workflows/coordinator/replica_status_change.rs +++ b/packages/services/epoxy/src/workflows/coordinator/replica_status_change.rs @@ -8,7 +8,7 @@ use crate::types; pub async fn replica_status_change( ctx: &mut WorkflowCtx, - signal: super::ReplicaStatusChangeSignal, + signal: super::ReplicaStatusChange, ) -> Result<()> { // Update replica status let should_increment_epoch = ctx diff --git a/packages/services/epoxy/src/workflows/replica/mod.rs b/packages/services/epoxy/src/workflows/replica/mod.rs index 85583b4619..36a29cfbde 100644 --- a/packages/services/epoxy/src/workflows/replica/mod.rs +++ b/packages/services/epoxy/src/workflows/replica/mod.rs @@ -14,14 +14,14 @@ pub use setup::*; pub struct Input {} #[workflow] -pub async fn replica(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> { +pub async fn epoxy_replica(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> { setup_replica(ctx, input).await?; // Main loop ctx.repeat(|ctx| { async move { // Noop for now - ctx.listen::().await?; + ctx.listen::
().await?; Ok(Loop::<()>::Continue) } .boxed() @@ -32,9 +32,9 @@ pub async fn replica(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> { } #[signal("epoxy_replica_begin_learning")] -pub struct BeginLearningSignal { +pub struct BeginLearning { pub config: types::ClusterConfig, } -#[signal("main")] -pub struct MainSignal {} +#[signal("epoxy_replica_main")] +pub struct Main {} diff --git a/packages/services/epoxy/src/workflows/replica/setup.rs b/packages/services/epoxy/src/workflows/replica/setup.rs index 278c77269d..31671b7341 100644 --- a/packages/services/epoxy/src/workflows/replica/setup.rs +++ b/packages/services/epoxy/src/workflows/replica/setup.rs @@ -10,14 +10,14 @@ use universaldb::{KeySelector, RangeOption, options::StreamingMode}; use crate::types; // IMPORTANT: Do not use `read_cluster_config`. Instead, use the config provided by -// `BeginLearningSignal`. This is because the value of `read_cluster_config` may change between +// `BeginLearning`. This is because the value of `read_cluster_config` may change between // activities which can cause the learning process to enter an invalid state. pub async fn setup_replica(ctx: &mut WorkflowCtx, _input: &super::Input) -> Result<()> { - // Wait for cooridinator to send begin learning signal - let begin_learning = ctx.listen::().await?; + // Wait for coordiinator to send begin learning signal + let begin_learning = ctx.listen::().await?; - // TODO: Paralellize replicas + // TODO: Parallelize replicas let total_replicas = begin_learning.config.replicas.len(); let mut replica_index = 0; @@ -134,7 +134,7 @@ pub async fn setup_replica(ctx: &mut WorkflowCtx, _input: &super::Input) -> Resu #[derive(Debug, Clone, Serialize, Deserialize, Hash)] pub struct DownloadInstancesChunkInput { - /// Config received from BeginLearningSignal + /// Config received from BeginLearning pub learning_config: types::ClusterConfig, pub from_replica_id: protocol::ReplicaId, pub after_instance: Option, @@ -312,7 +312,7 @@ async fn apply_log_entry( #[derive(Debug, Clone, Serialize, Deserialize, Hash)] pub struct RecoverKeysChunkInput { - /// Config received from BeginLearningSignal + /// Config received from BeginLearning pub learning_config: types::ClusterConfig, /// The last key value from the previous chunk, used for pagination pub after_key: Option>, @@ -728,7 +728,7 @@ struct CommittedEntry { #[derive(Debug, Clone, Serialize, Deserialize, Hash)] pub struct NotifyActiveInput { - /// Config received from BeginLearningSignal + /// Config received from BeginLearning pub learning_config: types::ClusterConfig, } diff --git a/packages/services/epoxy/tests/common/mod.rs b/packages/services/epoxy/tests/common/mod.rs index 2a108a3b63..5f60c2aa44 100644 --- a/packages/services/epoxy/tests/common/mod.rs +++ b/packages/services/epoxy/tests/common/mod.rs @@ -284,7 +284,7 @@ async fn setup_epoxy_coordinator_wf( ) .await?; leader_ctx - .signal(epoxy::workflows::coordinator::ReconfigureSignal {}) + .signal(epoxy::workflows::coordinator::Reconfigure {}) .to_workflow_id(workflow_id) .send() .await?; diff --git a/packages/services/epoxy/tests/reconfigure.rs b/packages/services/epoxy/tests/reconfigure.rs index 948e358b29..80369255b1 100644 --- a/packages/services/epoxy/tests/reconfigure.rs +++ b/packages/services/epoxy/tests/reconfigure.rs @@ -321,7 +321,7 @@ async fn test_inner(config: TestConfig) { .await .unwrap(); leader_ctx - .signal(epoxy::workflows::coordinator::ReconfigureSignal {}) + .signal(epoxy::workflows::coordinator::Reconfigure {}) .to_workflow_id(test_ctx.coordinator_workflow_id) .send() .await diff --git a/sdks/typescript/test-runner/src/main.ts b/sdks/typescript/test-runner/src/main.ts index fbe681326d..12e13fe7c7 100644 --- a/sdks/typescript/test-runner/src/main.ts +++ b/sdks/typescript/test-runner/src/main.ts @@ -61,8 +61,6 @@ app.get('/start', async (c) => { }); }); -app.get('*', (c) => c.text('ok')); - serve({ fetch: app.fetch, port: INTERNAL_SERVER_PORT, @@ -72,6 +70,8 @@ console.log(`Internal HTTP server listening on port ${INTERNAL_SERVER_PORT}`); if (AUTOSTART) runner = await startRunner(); async function startRunner(): Promise { + console.log("Starting runner"); + const config: RunnerConfig = { version: RIVET_RUNNER_VERSION, endpoint: RIVET_ENDPOINT, @@ -154,7 +154,7 @@ async function startRunner(): Promise { }, }; - runner = new Runner(config); + let runner = new Runner(config); // Start runner await runner.start(); @@ -163,5 +163,7 @@ async function startRunner(): Promise { console.log("Waiting runner start..."); await runnerStarted.promise; + console.log("Runner started"); + return runner; -} \ No newline at end of file +} From 8de25348c47a488aa9aaaab5c203779617c9cd2c Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Fri, 12 Sep 2025 13:01:43 -0700 Subject: [PATCH 14/17] fix(api): fix query serialization for remote req --- Cargo.toml | 1 + docker/template/src/docker-compose.ts | 1 - packages/common/api-builder/src/context.rs | 2 -- .../common/api-builder/src/error_response.rs | 1 + packages/common/api-client/Cargo.toml | 1 + packages/common/api-client/src/lib.rs | 36 ++++++++++++------- packages/common/types/src/datacenters.rs | 1 - packages/core/api-public/src/datacenters.rs | 1 - .../src/workflows/coordinator/reconfigure.rs | 22 +++++++----- 9 files changed, 41 insertions(+), 25 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index fd561bc4f5..798c872ede 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -59,6 +59,7 @@ rstest = "0.26.1" rustls-pemfile = "2.2.0" rustyline = "15.0.0" serde_bare = "0.5.0" +serde_html_form = "0.2.7" serde_yaml = "0.9.34" sha2 = "0.10" slog = "2.7" diff --git a/docker/template/src/docker-compose.ts b/docker/template/src/docker-compose.ts index d60cd53e27..19f044c5c8 100644 --- a/docker/template/src/docker-compose.ts +++ b/docker/template/src/docker-compose.ts @@ -320,7 +320,6 @@ export function generateDockerCompose(context: TemplateContext) { restart: "unless-stopped", environment: [ `RIVET_ENDPOINT=http://${context.getServiceHost("rivet-engine", datacenter.name, 0)}:6420`, - `RUNNER_HOST=${context.getServiceHost("runner", datacenter.name, i)}`, ], stop_grace_period: "4s", ports: isPrimary && i === 0 ? [`5050:5050`] : undefined, diff --git a/packages/common/api-builder/src/context.rs b/packages/common/api-builder/src/context.rs index acf2fc7cd1..5bfa83625d 100644 --- a/packages/common/api-builder/src/context.rs +++ b/packages/common/api-builder/src/context.rs @@ -32,8 +32,6 @@ impl fmt::Debug for ApiCtx { impl ApiCtx { pub fn new(global: GlobalApiCtx, ray_id: Id, req_id: Id) -> Result { - // Create StandaloneCtx synchronously by using a blocking call - // This is necessary because we need Clone support and async Clone is not possible let standalone_ctx = StandaloneCtx::new( global.db.clone(), global.config.clone(), diff --git a/packages/common/api-builder/src/error_response.rs b/packages/common/api-builder/src/error_response.rs index b946bb83ba..bc1faa1542 100644 --- a/packages/common/api-builder/src/error_response.rs +++ b/packages/common/api-builder/src/error_response.rs @@ -54,6 +54,7 @@ impl IntoResponse for ApiError { internal: if error_response.group == rivet_error::INTERNAL_ERROR.group && error_response.code == rivet_error::INTERNAL_ERROR.code { + tracing::debug!(err=?self.0, "internal debug error"); Some(format!("{}", self.0).into()) } else { None diff --git a/packages/common/api-client/Cargo.toml b/packages/common/api-client/Cargo.toml index 9c6e27c910..8075f82de2 100644 --- a/packages/common/api-client/Cargo.toml +++ b/packages/common/api-client/Cargo.toml @@ -16,6 +16,7 @@ rivet-config.workspace = true rivet-error.workspace = true rivet-pools.workspace = true rivet-util.workspace = true +serde_html_form.workspace = true serde.workspace = true tokio.workspace = true tracing.workspace = true diff --git a/packages/common/api-client/src/lib.rs b/packages/common/api-client/src/lib.rs index 634761ec02..deea077a71 100644 --- a/packages/common/api-client/src/lib.rs +++ b/packages/common/api-client/src/lib.rs @@ -25,20 +25,26 @@ pub async fn request_remote_datacenter_raw( .ok_or_else(|| errors::Datacenter::NotFound.build())?; let client = rivet_pools::reqwest::client().await?; - let url = dc.api_peer_url.join(endpoint)?; - - let mut request = client.request(method, url).headers(headers); + let mut url = dc.api_peer_url.join(endpoint)?; + // NOTE: We don't use reqwest's `.query` because it doesn't support list query parameters if let Some(q) = query { - request = request.query(q); + url.set_query(Some(&serde_html_form::to_string(q)?)); } + let mut request = client.request(method, url).headers(headers); + if let Some(b) = body { request = request.json(b); } - let res = request.send().await?; - rivet_api_util::reqwest_to_axum_response(res).await + let res = request + .send() + .await + .context("failed sending request to remote dc")?; + rivet_api_util::reqwest_to_axum_response(res) + .await + .context("failed parsing response from remote dc") } /// Generic function to make requests to a specific datacenter @@ -59,20 +65,26 @@ where .ok_or_else(|| errors::Datacenter::NotFound.build())?; let client = rivet_pools::reqwest::client().await?; - let url = dc.api_peer_url.join(endpoint)?; - - let mut request = client.request(method, url).headers(headers); + let mut url = dc.api_peer_url.join(endpoint)?; + // NOTE: We don't use reqwest's `.query` because it doesn't support list query parameters if let Some(q) = query { - request = request.query(q); + url.set_query(Some(&serde_html_form::to_string(q)?)); } + let mut request = client.request(method, url).headers(headers); + if let Some(b) = body { request = request.json(b); } - let res = request.send().await?; - rivet_api_util::parse_response::(res).await + let res = request + .send() + .await + .context("failed sending request to remote dc")?; + rivet_api_util::parse_response::(res) + .await + .context("failed parsing response from remote dc") } /// Generic function to fanout requests to all datacenters and aggregate results diff --git a/packages/common/types/src/datacenters.rs b/packages/common/types/src/datacenters.rs index b3eac04ded..4db7b95e65 100644 --- a/packages/common/types/src/datacenters.rs +++ b/packages/common/types/src/datacenters.rs @@ -6,5 +6,4 @@ use utoipa::ToSchema; pub struct Datacenter { pub datacenter_label: u16, pub name: String, - pub url: String, } diff --git a/packages/core/api-public/src/datacenters.rs b/packages/core/api-public/src/datacenters.rs index 22bfc86e5a..f3a507f826 100644 --- a/packages/core/api-public/src/datacenters.rs +++ b/packages/core/api-public/src/datacenters.rs @@ -21,7 +21,6 @@ pub async fn list(ctx: ApiCtx, _path: (), _query: ()) -> Result { .map(|dc| Datacenter { datacenter_label: dc.datacenter_label, name: dc.name.clone(), - url: dc.guard_url.to_string(), }) .collect(), pagination: Pagination { cursor: None }, diff --git a/packages/services/epoxy/src/workflows/coordinator/reconfigure.rs b/packages/services/epoxy/src/workflows/coordinator/reconfigure.rs index 1586e9c351..1c2a4748fd 100644 --- a/packages/services/epoxy/src/workflows/coordinator/reconfigure.rs +++ b/packages/services/epoxy/src/workflows/coordinator/reconfigure.rs @@ -153,19 +153,25 @@ pub async fn health_check_new_replicas( async move { tracing::info!(?replica_id, "sending health check to replica"); + let from_replica_id = ctx.config().epoxy_replica_id(); let request = protocol::Request { - from_replica_id: ctx.config().epoxy_replica_id(), + from_replica_id, to_replica_id: replica_id, kind: protocol::RequestKind::HealthCheckRequest, }; - crate::http_client::send_message_to_address( - replica.api_peer_url.clone(), - replica_id, - request, - ) - .await - .with_context(|| format!("health check failed for replica {}", replica_id))?; + // Call directly instead of sending request for same replica + if from_replica_id == replica_id { + tracing::info!(?replica_id, "skipping health check to self"); + } else { + crate::http_client::send_message_to_address( + replica.api_peer_url.clone(), + replica_id, + request, + ) + .await + .with_context(|| format!("health check failed for replica {}", replica_id))?; + } tracing::info!(?replica_id, "health check successful"); Ok(()) From 4f09c22a1da07290968450de7fea937ea76bb854 Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Fri, 12 Sep 2025 15:46:05 -0700 Subject: [PATCH 15/17] fix(epoxy): fix sending msgs to self --- Cargo.lock | 1 + out/openapi.json | 6 +-- packages/common/api-builder/src/context.rs | 20 ++++++++++ .../common/gasoline/core/src/ctx/activity.rs | 4 ++ .../common/gasoline/core/src/ctx/operation.rs | 4 ++ .../gasoline/core/src/ctx/standalone.rs | 40 ++++++++++++++----- packages/core/pegboard-gateway/src/lib.rs | 24 +++++++---- packages/services/epoxy/src/http_client.rs | 24 ++++++++--- .../epoxy/src/ops/explicit_prepare.rs | 15 +++++-- .../epoxy/src/ops/kv/get_optimistic.rs | 7 +++- packages/services/epoxy/src/ops/propose.rs | 25 +++++++++--- .../epoxy/src/replica/message_request.rs | 1 - .../src/workflows/coordinator/reconfigure.rs | 23 +++++------ .../coordinator/replica_status_change.rs | 3 +- .../epoxy/src/workflows/replica/setup.rs | 11 +++-- 15 files changed, 151 insertions(+), 57 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index cd762d64fa..21b23c5d82 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4122,6 +4122,7 @@ dependencies = [ "rivet-pools", "rivet-util", "serde", + "serde_html_form", "tokio", "tracing", ] diff --git a/out/openapi.json b/out/openapi.json index 637d75cf88..3eeca9e31d 100644 --- a/out/openapi.json +++ b/out/openapi.json @@ -1185,8 +1185,7 @@ "type": "object", "required": [ "datacenter_label", - "name", - "url" + "name" ], "properties": { "datacenter_label": { @@ -1196,9 +1195,6 @@ }, "name": { "type": "string" - }, - "url": { - "type": "string" } }, "additionalProperties": false diff --git a/packages/common/api-builder/src/context.rs b/packages/common/api-builder/src/context.rs index 5bfa83625d..7b22e3a5d4 100644 --- a/packages/common/api-builder/src/context.rs +++ b/packages/common/api-builder/src/context.rs @@ -49,6 +49,26 @@ impl ApiCtx { }) } + pub fn new_from_activity(ctx: &ActivityCtx) -> Result { + let req_id = Id::new_v1(ctx.config().dc_label()); + + Ok(Self { + ray_id: ctx.ray_id(), + req_id, + standalone_ctx: StandaloneCtx::new_from_activity(ctx, req_id)?, + }) + } + + pub fn new_from_operation(ctx: &OperationCtx) -> Result { + let req_id = Id::new_v1(ctx.config().dc_label()); + + Ok(Self { + ray_id: ctx.ray_id(), + req_id, + standalone_ctx: StandaloneCtx::new_from_operation(ctx, req_id)?, + }) + } + pub fn ray_id(&self) -> Id { self.ray_id } diff --git a/packages/common/gasoline/core/src/ctx/activity.rs b/packages/common/gasoline/core/src/ctx/activity.rs index 28650d73e8..ce1de28c55 100644 --- a/packages/common/gasoline/core/src/ctx/activity.rs +++ b/packages/common/gasoline/core/src/ctx/activity.rs @@ -144,6 +144,10 @@ impl ActivityCtx { } impl ActivityCtx { + pub(crate) fn db(&self) -> &DatabaseHandle { + &self.db + } + pub fn name(&self) -> &str { self.name } diff --git a/packages/common/gasoline/core/src/ctx/operation.rs b/packages/common/gasoline/core/src/ctx/operation.rs index a9695fec00..99fbd6071f 100644 --- a/packages/common/gasoline/core/src/ctx/operation.rs +++ b/packages/common/gasoline/core/src/ctx/operation.rs @@ -152,6 +152,10 @@ impl OperationCtx { } impl OperationCtx { + pub(crate) fn db(&self) -> &DatabaseHandle { + &self.db + } + pub fn name(&self) -> &str { self.name } diff --git a/packages/common/gasoline/core/src/ctx/standalone.rs b/packages/common/gasoline/core/src/ctx/standalone.rs index 6b5c476e32..089afc7428 100644 --- a/packages/common/gasoline/core/src/ctx/standalone.rs +++ b/packages/common/gasoline/core/src/ctx/standalone.rs @@ -7,7 +7,7 @@ use tracing::Instrument; use crate::{ builder::{WorkflowRepr, common as builder}, - ctx::{MessageCtx, common, message::SubscriptionHandle}, + ctx::{ActivityCtx, MessageCtx, OperationCtx, common, message::SubscriptionHandle}, db::{DatabaseHandle, WorkflowData}, error::WorkflowResult, message::Message, @@ -21,7 +21,7 @@ use crate::{ pub struct StandaloneCtx { ray_id: Id, req_id: Id, - name: &'static str, + name: String, ts: i64, db: DatabaseHandle, @@ -39,7 +39,7 @@ impl StandaloneCtx { config: rivet_config::Config, pools: rivet_pools::Pools, cache: rivet_cache::Cache, - name: &'static str, + name: &str, ray_id: Id, req_id: Id, ) -> WorkflowResult { @@ -54,7 +54,7 @@ impl StandaloneCtx { Ok(StandaloneCtx { ray_id, req_id, - name, + name: name.to_string(), ts, db, config, @@ -63,6 +63,32 @@ impl StandaloneCtx { msg_ctx, }) } + + #[tracing::instrument(skip_all)] + pub fn new_from_activity(ctx: &ActivityCtx, req_id: Id) -> WorkflowResult { + StandaloneCtx::new( + ctx.db().clone(), + ctx.config().clone(), + ctx.pools().clone(), + ctx.cache().clone(), + ctx.name(), + ctx.ray_id(), + req_id, + ) + } + + #[tracing::instrument(skip_all)] + pub fn new_from_operation(ctx: &OperationCtx, req_id: Id) -> WorkflowResult { + StandaloneCtx::new( + ctx.db().clone(), + ctx.config().clone(), + ctx.pools().clone(), + ctx.cache().clone(), + ctx.name(), + ctx.ray_id(), + req_id, + ) + } } impl StandaloneCtx { @@ -154,13 +180,9 @@ impl StandaloneCtx { impl StandaloneCtx { pub fn name(&self) -> &str { - self.name + &self.name } - // pub fn timeout(&self) -> Duration { - // self.timeout - // } - pub fn ray_id(&self) -> Id { self.ray_id } diff --git a/packages/core/pegboard-gateway/src/lib.rs b/packages/core/pegboard-gateway/src/lib.rs index 67a23eff2a..b1b89dd4ed 100644 --- a/packages/core/pegboard-gateway/src/lib.rs +++ b/packages/core/pegboard-gateway/src/lib.rs @@ -1,3 +1,13 @@ +use std::result::Result::Ok as ResultOk; +use std::{ + collections::HashMap, + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; + use anyhow::*; use async_trait::async_trait; use bytes::Bytes; @@ -108,9 +118,9 @@ impl PegboardGateway { let actor_id = req .headers() .get("x-rivet-actor") - .ok_or_else(|| anyhow!("missing x-rivet-actor header"))? + .context("missing x-rivet-actor header")? .to_str() - .map_err(|_| anyhow!("invalid x-rivet-actor header"))? + .context("invalid x-rivet-actor header")? .to_string(); // Extract request parts @@ -132,7 +142,7 @@ impl PegboardGateway { .into_body() .collect() .await - .map_err(|e| anyhow!("failed to read body: {}", e))? + .context("failed to read body")? .to_bytes(); // Build subject to publish to @@ -212,11 +222,9 @@ impl PegboardGateway { // Extract actor ID for the message let actor_id = match headers .get("x-rivet-actor") - .ok_or_else(|| anyhow!("missing x-rivet-actor header")) - .and_then(|v| { - v.to_str() - .map_err(|_| anyhow!("invalid x-rivet-actor header")) - }) { + .context("missing x-rivet-actor header") + .and_then(|v| v.to_str().context("invalid x-rivet-actor header")) + { Result::Ok(v) => v.to_string(), Err(err) => return Err((client_ws, err)), }; diff --git a/packages/services/epoxy/src/http_client.rs b/packages/services/epoxy/src/http_client.rs index 51864305c6..f3b440a79f 100644 --- a/packages/services/epoxy/src/http_client.rs +++ b/packages/services/epoxy/src/http_client.rs @@ -5,6 +5,7 @@ use epoxy_protocol::{ versioned, }; use futures_util::{StreamExt, stream::FuturesUnordered}; +use rivet_api_builder::ApiCtx; use std::future::Future; use versioned_data_util::OwnedVersionedData; @@ -56,7 +57,7 @@ where .await; tracing::info!(?quorum_size, len = ?responses.len(), ?quorum_type, "fanout quorum size"); - // Choow how many successful responses we need before considering a success + // Choose how many successful responses we need before considering a success let target_responses = match quorum_type { // Only require 1 response utils::QuorumType::Any => 1, @@ -93,19 +94,32 @@ where } pub async fn send_message( + ctx: &ApiCtx, config: &protocol::ClusterConfig, - to_replica_id: ReplicaId, request: protocol::Request, ) -> Result { - let replica_url = find_replica_address(config, to_replica_id)?; - send_message_to_address(replica_url, to_replica_id, request).await + let replica_url = find_replica_address(config, request.to_replica_id)?; + send_message_to_address(ctx, replica_url, request).await } pub async fn send_message_to_address( + ctx: &ApiCtx, replica_url: String, - to_replica_id: ReplicaId, request: protocol::Request, ) -> Result { + let from_replica_id = request.from_replica_id; + let to_replica_id = request.to_replica_id; + + if from_replica_id == to_replica_id { + tracing::info!( + to_replica = to_replica_id, + "sending message to replica directly" + ); + + return crate::replica::message_request::message_request(&ctx, from_replica_id, request) + .await; + } + let mut replica_url = url::Url::parse(&replica_url)?; replica_url.set_path(&format!("/v{PROTOCOL_VERSION}/epoxy/message")); diff --git a/packages/services/epoxy/src/ops/explicit_prepare.rs b/packages/services/epoxy/src/ops/explicit_prepare.rs index 14ca28852a..b75070d3f7 100644 --- a/packages/services/epoxy/src/ops/explicit_prepare.rs +++ b/packages/services/epoxy/src/ops/explicit_prepare.rs @@ -1,6 +1,7 @@ use anyhow::*; use epoxy_protocol::protocol::{self, ReplicaId}; use gas::prelude::*; +use rivet_api_builder::ApiCtx; use crate::{http_client, replica, types, utils}; @@ -48,8 +49,15 @@ pub async fn epoxy_explicit_prepare( let quorum_members = utils::get_quorum_members(&config); // EPaxos Step 26: Send Prepare to all replicas and wait for quorum - let prepare_responses = - send_prepares(&config, replica_id, &quorum_members, &new_ballot, instance).await?; + let prepare_responses = send_prepares( + ctx, + &config, + replica_id, + &quorum_members, + &new_ballot, + instance, + ) + .await?; // Check if we got enough responses for a quorum let required_quorum = utils::calculate_quorum(quorum_members.len(), utils::QuorumType::Slow); @@ -241,6 +249,7 @@ fn compare_ballots(a: &protocol::Ballot, b: &protocol::Ballot) -> std::cmp::Orde } async fn send_prepares( + ctx: &OperationCtx, config: &protocol::ClusterConfig, from_replica_id: ReplicaId, replica_ids: &[ReplicaId], @@ -257,8 +266,8 @@ async fn send_prepares( let instance = instance.clone(); async move { let response = http_client::send_message( + &ApiCtx::new_from_operation(&ctx)?, &config, - to_replica_id, protocol::Request { from_replica_id, to_replica_id, diff --git a/packages/services/epoxy/src/ops/kv/get_optimistic.rs b/packages/services/epoxy/src/ops/kv/get_optimistic.rs index 6335cbbb90..726b54656f 100644 --- a/packages/services/epoxy/src/ops/kv/get_optimistic.rs +++ b/packages/services/epoxy/src/ops/kv/get_optimistic.rs @@ -1,6 +1,7 @@ use anyhow::*; use epoxy_protocol::protocol::{self, ReplicaId}; use gas::prelude::*; +use rivet_api_builder::ApiCtx; use universaldb::utils::{FormalKey, IsolationLevel::*}; use crate::{http_client, keys, utils}; @@ -35,7 +36,7 @@ pub struct Output { /// /// We cannot use quorum reads for the fanout read because of the constraints of Epaxos. #[operation] -pub async fn epoxy_kv_get_optimistic(ctx: &OperationCtx, input: &Input) -> Result { +pub async fn epoxy_get_optimistic(ctx: &OperationCtx, input: &Input) -> Result { // Try to read locally let kv_key = keys::keys::KvValueKey::new(input.key.clone()); let cache_key = keys::keys::KvOptimisticCacheKey::new(input.key.clone()); @@ -113,7 +114,9 @@ pub async fn epoxy_kv_get_optimistic(ctx: &OperationCtx, input: &Input) -> Resul }; // Send the message and extract the KV response - let response = http_client::send_message(&config, replica_id, request).await?; + let response = + http_client::send_message(&ApiCtx::new_from_operation(&ctx)?, &config, request) + .await?; match response.kind { protocol::ResponseKind::KvGetResponse(kv_response) => Ok(kv_response.value), diff --git a/packages/services/epoxy/src/ops/propose.rs b/packages/services/epoxy/src/ops/propose.rs index 8c42747036..cc8a0fe7bc 100644 --- a/packages/services/epoxy/src/ops/propose.rs +++ b/packages/services/epoxy/src/ops/propose.rs @@ -49,7 +49,8 @@ pub async fn epoxy_propose(ctx: &OperationCtx, input: &Input) -> Result= utils::calculate_quorum(quorum_members.len(), utils::QuorumType::Slow) { @@ -137,12 +145,14 @@ pub async fn commit( // Send commits to all replicas (not just quorum members) let all_replicas = utils::get_all_replicas(config); tokio::spawn({ + let ctx = ctx.clone(); let config = config.clone(); let replica_id = replica_id; let all_replicas = all_replicas.to_vec(); let payload = payload.clone(); + async move { - let _ = send_commits(&config, replica_id, &all_replicas, &payload).await; + let _ = send_commits(&ctx, &config, replica_id, &all_replicas, &payload).await; } }); @@ -154,6 +164,7 @@ pub async fn commit( } async fn send_pre_accepts( + ctx: &OperationCtx, config: &protocol::ClusterConfig, from_replica_id: ReplicaId, replica_ids: &[ReplicaId], @@ -168,8 +179,8 @@ async fn send_pre_accepts( let payload = payload.clone(); async move { let response = http_client::send_message( + &ApiCtx::new_from_operation(&ctx)?, &config, - to_replica_id, protocol::Request { from_replica_id, to_replica_id, @@ -197,6 +208,7 @@ async fn send_pre_accepts( } async fn send_accepts( + ctx: &OperationCtx, config: &protocol::ClusterConfig, from_replica_id: ReplicaId, replica_ids: &[ReplicaId], @@ -211,8 +223,8 @@ async fn send_accepts( let payload = payload.clone(); async move { let response = http_client::send_message( + &ApiCtx::new_from_operation(&ctx)?, &config, - to_replica_id, protocol::Request { from_replica_id, to_replica_id, @@ -241,6 +253,7 @@ async fn send_accepts( } async fn send_commits( + ctx: &OperationCtx, config: &protocol::ClusterConfig, from_replica_id: ReplicaId, replica_ids: &[ReplicaId], @@ -255,8 +268,8 @@ async fn send_commits( let payload = payload.clone(); async move { let response = http_client::send_message( + &ApiCtx::new_from_operation(&ctx)?, &config, - to_replica_id, protocol::Request { from_replica_id, to_replica_id, diff --git a/packages/services/epoxy/src/replica/message_request.rs b/packages/services/epoxy/src/replica/message_request.rs index 44ec707be1..f3ca78426a 100644 --- a/packages/services/epoxy/src/replica/message_request.rs +++ b/packages/services/epoxy/src/replica/message_request.rs @@ -123,7 +123,6 @@ pub async fn message_request( }) .to_workflow::() .tag("replica", replica_id) - .to_workflow::() .send() .await?; diff --git a/packages/services/epoxy/src/workflows/coordinator/reconfigure.rs b/packages/services/epoxy/src/workflows/coordinator/reconfigure.rs index 1c2a4748fd..0adbe0fd6c 100644 --- a/packages/services/epoxy/src/workflows/coordinator/reconfigure.rs +++ b/packages/services/epoxy/src/workflows/coordinator/reconfigure.rs @@ -1,6 +1,7 @@ use anyhow::*; use epoxy_protocol::protocol::{self, ReplicaId}; use gas::prelude::*; +use rivet_api_builder::ApiCtx; use serde::{Deserialize, Serialize}; use crate::types; @@ -160,18 +161,13 @@ pub async fn health_check_new_replicas( kind: protocol::RequestKind::HealthCheckRequest, }; - // Call directly instead of sending request for same replica - if from_replica_id == replica_id { - tracing::info!(?replica_id, "skipping health check to self"); - } else { - crate::http_client::send_message_to_address( - replica.api_peer_url.clone(), - replica_id, - request, - ) - .await - .with_context(|| format!("health check failed for replica {}", replica_id))?; - } + crate::http_client::send_message_to_address( + &ApiCtx::new_from_activity(ctx)?, + replica.api_peer_url.clone(), + request, + ) + .await + .with_context(|| format!("health check failed for replica {}", replica_id))?; tracing::info!(?replica_id, "health check successful"); Ok(()) @@ -243,7 +239,8 @@ pub async fn send_begin_learning( }), }; - crate::http_client::send_message(&config, replica_id, request).await?; + crate::http_client::send_message(&ApiCtx::new_from_activity(&ctx)?, &config, request) + .await?; tracing::info!(?replica_id, "begin learning sent successfully"); Ok(()) diff --git a/packages/services/epoxy/src/workflows/coordinator/replica_status_change.rs b/packages/services/epoxy/src/workflows/coordinator/replica_status_change.rs index 41da43f595..d35f2c6706 100644 --- a/packages/services/epoxy/src/workflows/coordinator/replica_status_change.rs +++ b/packages/services/epoxy/src/workflows/coordinator/replica_status_change.rs @@ -1,6 +1,7 @@ use anyhow::*; use epoxy_protocol::protocol; use gas::prelude::*; +use rivet_api_builder::ApiCtx; use serde::{Deserialize, Serialize}; use super::State; @@ -126,7 +127,7 @@ pub async fn notify_all_replicas( }), }; - crate::http_client::send_message(&config, replica_id, request) + crate::http_client::send_message(&ApiCtx::new_from_activity(&ctx)?, &config, request) .await .with_context(|| format!("failed to update config for replica {}", replica_id))?; diff --git a/packages/services/epoxy/src/workflows/replica/setup.rs b/packages/services/epoxy/src/workflows/replica/setup.rs index 31671b7341..8d9b43b2f6 100644 --- a/packages/services/epoxy/src/workflows/replica/setup.rs +++ b/packages/services/epoxy/src/workflows/replica/setup.rs @@ -2,6 +2,7 @@ use anyhow::*; use epoxy_protocol::protocol; use futures_util::{FutureExt, TryStreamExt}; use gas::prelude::*; +use rivet_api_builder::ApiCtx; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, VecDeque}; use universaldb::prelude::*; @@ -14,10 +15,10 @@ use crate::types; // activities which can cause the learning process to enter an invalid state. pub async fn setup_replica(ctx: &mut WorkflowCtx, _input: &super::Input) -> Result<()> { - // Wait for coordiinator to send begin learning signal + // Wait for cooridinator to send begin learning signal let begin_learning = ctx.listen::().await?; - // TODO: Parallelize replicas + // TODO: Paralellize replicas let total_replicas = begin_learning.config.replicas.len(); let mut replica_index = 0; @@ -182,7 +183,8 @@ pub async fn download_instances_chunk( }; let response = - crate::http_client::send_message(&proto_config, input.from_replica_id, request).await?; + crate::http_client::send_message(&ApiCtx::new_from_activity(&ctx)?, &proto_config, request) + .await?; // Extract instances from response let protocol::ResponseKind::DownloadInstancesResponse(download_response) = response.kind else { @@ -751,7 +753,8 @@ pub async fn notify_active(ctx: &ActivityCtx, input: &NotifyActiveInput) -> Resu ), }; - crate::http_client::send_message(&proto_config, config.coordinator_replica_id, request).await?; + crate::http_client::send_message(&ApiCtx::new_from_activity(&ctx)?, &proto_config, request) + .await?; tracing::info!("notified coordinator of active status"); Ok(()) From a0f362cf5b77c678d2626c453f26e61a1c153efc Mon Sep 17 00:00:00 2001 From: MasterPtato Date: Fri, 12 Sep 2025 17:04:48 -0700 Subject: [PATCH 16/17] fix(api): fix creating actors in other dcs --- out/openapi.json | 3 +- .../common/api-types/src/actors/create.rs | 1 - packages/core/api-peer/src/actors/create.rs | 22 +++-- packages/core/api-public/src/actors/create.rs | 88 ++++++++----------- packages/core/api-public/src/router.rs | 2 +- packages/core/pegboard-gateway/src/lib.rs | 1 - .../services/pegboard/src/ops/actor/create.rs | 4 - .../pegboard/src/workflows/actor/runtime.rs | 4 +- 8 files changed, 55 insertions(+), 70 deletions(-) diff --git a/out/openapi.json b/out/openapi.json index 3eeca9e31d..fd8259a654 100644 --- a/out/openapi.json +++ b/out/openapi.json @@ -1011,7 +1011,8 @@ "actor": { "$ref": "#/components/schemas/Actor" } - } + }, + "additionalProperties": false }, "ActorsDeleteResponse": { "type": "object" diff --git a/packages/common/api-types/src/actors/create.rs b/packages/common/api-types/src/actors/create.rs index 082ed027eb..004058dc7f 100644 --- a/packages/common/api-types/src/actors/create.rs +++ b/packages/common/api-types/src/actors/create.rs @@ -13,7 +13,6 @@ pub struct CreateQuery { #[serde(deny_unknown_fields)] #[schema(as = ActorsCreateRequest)] pub struct CreateRequest { - pub actor_id: Id, pub name: String, pub key: Option, pub input: Option, diff --git a/packages/core/api-peer/src/actors/create.rs b/packages/core/api-peer/src/actors/create.rs index 1980262919..ebf91f98eb 100644 --- a/packages/core/api-peer/src/actors/create.rs +++ b/packages/core/api-peer/src/actors/create.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use gas::prelude::*; use rivet_api_builder::ApiCtx; use rivet_api_types::actors::create::{CreateQuery, CreateRequest, CreateResponse}; @@ -15,27 +16,24 @@ pub async fn create( .await? .ok_or_else(|| namespace::errors::Namespace::NotFound.build())?; + let actor_id = Id::new_v1(ctx.config().dc_label()); + let res = ctx .op(pegboard::ops::actor::create::Input { - actor_id: body.actor_id, + actor_id, namespace_id: namespace.namespace_id, name: body.name.clone(), - key: body.key.clone(), - runner_name_selector: body.runner_name_selector.clone(), + key: body.key, + runner_name_selector: body.runner_name_selector, input: body.input.clone(), crash_policy: body.crash_policy, - // Don't forward request since this request should be already forwarded if it is going - // to be forward. - // - // This should never throw a request needs to be forwarded error. If it does, something - // is broken. - forward_request: false, + // NOTE: This can forward if the user attempts to create an actor with a target dc and this dc + // ends up forwarding to another. + forward_request: true, // api-peer is always creating in its own datacenter datacenter_name: None, }) .await?; - let actor = res.actor; - - Ok(CreateResponse { actor }) + Ok(CreateResponse { actor: res.actor }) } diff --git a/packages/core/api-public/src/actors/create.rs b/packages/core/api-public/src/actors/create.rs index 2d037a05ee..f661b5b12c 100644 --- a/packages/core/api-public/src/actors/create.rs +++ b/packages/core/api-public/src/actors/create.rs @@ -1,11 +1,17 @@ use anyhow::Result; -use rivet_api_builder::ApiCtx; +use axum::{ + extract::{Extension, Query}, + http::HeaderMap, + response::{IntoResponse, Json, Response}, +}; +use rivet_api_builder::{ApiCtx, ApiError}; +use rivet_api_client::request_remote_datacenter; +use rivet_api_types::actors::create::{CreateRequest, CreateResponse}; use rivet_types::actors::CrashPolicy; -use rivet_util::Id; use serde::{Deserialize, Serialize}; use utoipa::{IntoParams, ToSchema}; -#[derive(Debug, Deserialize, IntoParams)] +#[derive(Debug, Serialize, Deserialize, IntoParams)] #[serde(deny_unknown_fields)] #[into_params(parameter_in = Query)] pub struct CreateQuery { @@ -13,23 +19,6 @@ pub struct CreateQuery { pub datacenter: Option, } -#[derive(Deserialize, ToSchema)] -#[serde(deny_unknown_fields)] -#[schema(as = ActorsCreateRequest)] -pub struct CreateRequest { - pub name: String, - pub key: Option, - pub input: Option, - pub runner_name_selector: String, - pub crash_policy: CrashPolicy, -} - -#[derive(Serialize, ToSchema)] -#[schema(as = ActorsCreateResponse)] -pub struct CreateResponse { - pub actor: rivet_types::actors::Actor, -} - /// ## Datacenter Round Trips /// /// **If actor is created in the current datacenter:** @@ -57,18 +46,23 @@ pub struct CreateResponse { ), )] pub async fn create( + Extension(ctx): Extension, + headers: HeaderMap, + Query(query): Query, + Json(body): Json, +) -> Response { + match create_inner(ctx, headers, query, body).await { + Ok(response) => Json(response).into_response(), + Err(err) => ApiError::from(err).into_response(), + } +} + +async fn create_inner( ctx: ApiCtx, - _path: (), + headers: HeaderMap, query: CreateQuery, body: CreateRequest, ) -> Result { - let namespace = ctx - .op(namespace::ops::resolve_for_name_global::Input { - name: query.namespace.clone(), - }) - .await? - .ok_or_else(|| namespace::errors::Namespace::NotFound.build())?; - // Determine which datacenter to create the actor in let target_dc_label = if let Some(dc_name) = &query.datacenter { ctx.config() @@ -79,26 +73,22 @@ pub async fn create( ctx.config().dc_label() }; - let actor_id = Id::new_v1(target_dc_label); - - let key: Option = body.key; - - let res = ctx - .op(pegboard::ops::actor::create::Input { - actor_id, - namespace_id: namespace.namespace_id, - name: body.name.clone(), - key, - runner_name_selector: body.runner_name_selector, - input: body.input.clone(), - crash_policy: body.crash_policy, - // Forward requests to the correct api-peer datacenter - forward_request: true, - datacenter_name: query.datacenter.clone(), - }) - .await?; - - let actor = res.actor; + let query = rivet_api_types::actors::create::CreateQuery { + namespace: query.namespace, + }; - Ok(CreateResponse { actor }) + if target_dc_label == ctx.config().dc_label() { + rivet_api_peer::actors::create::create(ctx, (), query, body).await + } else { + request_remote_datacenter::( + ctx.config(), + target_dc_label, + "/actors", + axum::http::Method::POST, + headers, + Some(&query), + Some(&body), + ) + .await + } } diff --git a/packages/core/api-public/src/router.rs b/packages/core/api-public/src/router.rs index 9c06fd15b8..a33b8b3243 100644 --- a/packages/core/api-public/src/router.rs +++ b/packages/core/api-public/src/router.rs @@ -68,7 +68,7 @@ pub async fn router( ) // MARK: Actors .route("/actors", axum::routing::get(actors::list::list)) - .route("/actors", post(actors::create::create)) + .route("/actors", axum::routing::post(actors::create::create)) .route( "/actors", axum::routing::put(actors::get_or_create::get_or_create), diff --git a/packages/core/pegboard-gateway/src/lib.rs b/packages/core/pegboard-gateway/src/lib.rs index b1b89dd4ed..8e45374927 100644 --- a/packages/core/pegboard-gateway/src/lib.rs +++ b/packages/core/pegboard-gateway/src/lib.rs @@ -26,7 +26,6 @@ use rivet_tunnel_protocol::{ ToServerWebSocketOpen, }; use rivet_util::serde::HashableMap; -use std::time::Duration; use tokio_tungstenite::tungstenite::Message; use crate::shared_state::{SharedState, TunnelMessageData}; diff --git a/packages/services/pegboard/src/ops/actor/create.rs b/packages/services/pegboard/src/ops/actor/create.rs index 66a7a1b42c..962cf1299a 100644 --- a/packages/services/pegboard/src/ops/actor/create.rs +++ b/packages/services/pegboard/src/ops/actor/create.rs @@ -133,9 +133,6 @@ async fn forward_to_datacenter( .next() .ok_or_else(|| namespace::errors::Namespace::NotFound.build())?; - // Generate a new actor ID with the correct datacenter label - let actor_id = Id::new_v1(datacenter_label); - // Make request to remote datacenter let response = request_remote_datacenter::( ctx.config(), @@ -147,7 +144,6 @@ async fn forward_to_datacenter( namespace: namespace.name.clone(), }), Some(&rivet_api_types::actors::create::CreateRequest { - actor_id, name, key, input, diff --git a/packages/services/pegboard/src/workflows/actor/runtime.rs b/packages/services/pegboard/src/workflows/actor/runtime.rs index d01755d520..ffd1c25fd4 100644 --- a/packages/services/pegboard/src/workflows/actor/runtime.rs +++ b/packages/services/pegboard/src/workflows/actor/runtime.rs @@ -106,10 +106,10 @@ async fn allocate_actor( .udb()? .run(|tx| async move { let ping_threshold_ts = util::timestamp::now() - RUNNER_ELIGIBLE_THRESHOLD_MS; - let tx = tx.with_subspace(keys::subspace()); // Check if runner is an serverless runner let for_serverless = tx + .with_subspace(namespace::keys::subspace()) .exists( &namespace::keys::RunnerConfigByVariantKey::new( namespace_id, @@ -120,6 +120,8 @@ async fn allocate_actor( ) .await?; + let tx = tx.with_subspace(keys::subspace()); + if for_serverless { tx.atomic_op( &rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new( From 3716f5a6dba8ba753339d074e46ad9279deb9d9a Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Fri, 12 Sep 2025 15:04:51 -0700 Subject: [PATCH 17/17] chore(util): increase max ident len --- packages/common/util/core/src/check.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/common/util/core/src/check.rs b/packages/common/util/core/src/check.rs index 86bbc4613b..9f56452638 100644 --- a/packages/common/util/core/src/check.rs +++ b/packages/common/util/core/src/check.rs @@ -1,7 +1,8 @@ use lazy_static::lazy_static; use regex::{Regex, RegexBuilder}; -pub const MAX_IDENT_LEN: usize = 16; +pub const MAX_IDENT_LEN: usize = 64; +// TODO: Do we still need long idents? pub const MAX_IDENT_LONG_LEN: usize = 64; pub const MAX_DISPLAY_NAME_LEN: usize = 24; pub const MAX_DISPLAY_NAME_LONG_LEN: usize = 128;