From e98d4d0a9e25f0f9eaa4fdab48328c1f7ca21a6a Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Wed, 19 Nov 2025 12:17:06 -0800 Subject: [PATCH] chore(pegboard-gateway): add new message id format & add deprecated tunnel ack --- Cargo.lock | 2 + .../packages/guard-core/src/custom_serve.rs | 8 +- engine/packages/guard-core/src/lib.rs | 1 + .../packages/guard-core/src/proxy_service.rs | 220 +++-- engine/packages/pegboard-gateway/src/lib.rs | 37 +- .../pegboard-gateway/src/shared_state.rs | 167 ++-- engine/packages/pegboard-runner/src/lib.rs | 4 +- .../pegboard-runner/src/tunnel_to_ws_task.rs | 11 +- .../pegboard-runner/src/ws_to_tunnel_task.rs | 40 +- engine/packages/pegboard/Cargo.toml | 2 + engine/packages/pegboard/src/keys/actor.rs | 32 +- .../pegboard/src/keys/hibernating_request.rs | 40 +- engine/packages/pegboard/src/lib.rs | 1 + .../ops/actor/hibernating_request/delete.rs | 9 +- .../src/ops/actor/hibernating_request/list.rs | 15 +- .../ops/actor/hibernating_request/upsert.rs | 16 +- .../packages/pegboard/src/pubsub_subjects.rs | 12 +- engine/packages/pegboard/src/tunnel/id.rs | 86 ++ engine/packages/pegboard/src/tunnel/mod.rs | 1 + .../pegboard/src/workflows/actor/runtime.rs | 4 +- .../sdks/rust/runner-protocol/src/compat.rs | 7 + engine/sdks/rust/runner-protocol/src/lib.rs | 1 + .../rust/runner-protocol/src/versioned.rs | 78 +- engine/sdks/schemas/runner-protocol/v3.bare | 54 +- .../typescript/runner-protocol/src/index.ts | 218 +++-- engine/sdks/typescript/runner/src/actor.ts | 103 ++ engine/sdks/typescript/runner/src/mod.ts | 51 +- .../sdks/typescript/runner/src/stringify.ts | 12 +- .../sdks/typescript/runner/src/tunnel-id.ts | 104 ++ engine/sdks/typescript/runner/src/tunnel.ts | 532 +++++----- engine/sdks/typescript/runner/src/utils.ts | 10 + .../runner/src/websocket-tunnel-adapter.ts | 26 +- rivetkit-asyncapi/asyncapi.json | 923 +++++++++--------- .../rivetkit/schemas/actor-persist/v3.bare | 13 +- .../rivetkit/src/actor/conn/driver.ts | 14 +- .../rivetkit/src/actor/conn/drivers/http.ts | 3 - .../src/actor/conn/drivers/raw-request.ts | 3 - .../src/actor/conn/drivers/raw-websocket.ts | 6 +- .../src/actor/conn/drivers/websocket.ts | 7 +- .../packages/rivetkit/src/actor/conn/mod.ts | 6 +- .../rivetkit/src/actor/conn/persisted.ts | 21 +- .../rivetkit/src/actor/conn/state-manager.ts | 10 - .../src/actor/instance/connection-manager.ts | 31 +- .../rivetkit/src/actor/instance/mod.ts | 12 +- .../src/actor/instance/state-manager.ts | 12 +- .../src/actor/router-websocket-endpoints.ts | 16 +- .../packages/rivetkit/src/actor/router.ts | 4 +- .../src/drivers/engine/actor-driver.ts | 61 +- .../src/drivers/file-system/manager.ts | 5 +- .../packages/rivetkit/src/utils.ts | 4 - 50 files changed, 1865 insertions(+), 1190 deletions(-) create mode 100644 engine/packages/pegboard/src/tunnel/id.rs create mode 100644 engine/packages/pegboard/src/tunnel/mod.rs create mode 100644 engine/sdks/rust/runner-protocol/src/compat.rs create mode 100644 engine/sdks/typescript/runner/src/actor.ts create mode 100644 engine/sdks/typescript/runner/src/tunnel-id.ts diff --git a/Cargo.lock b/Cargo.lock index bbd0fe9db1..c8f22040e0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3428,6 +3428,7 @@ dependencies = [ "lazy_static", "namespace", "nix 0.30.1", + "rand 0.8.5", "rivet-api-types", "rivet-api-util", "rivet-data", @@ -3437,6 +3438,7 @@ dependencies = [ "rivet-types", "rivet-util", "serde", + "serde_bare", "serde_json", "strum", "tracing", diff --git a/engine/packages/guard-core/src/custom_serve.rs b/engine/packages/guard-core/src/custom_serve.rs index 1f57fbeac5..d7518d06ef 100644 --- a/engine/packages/guard-core/src/custom_serve.rs +++ b/engine/packages/guard-core/src/custom_serve.rs @@ -4,12 +4,13 @@ use bytes::Bytes; use http_body_util::Full; use hyper::{Request, Response}; use tokio_tungstenite::tungstenite::protocol::frame::CloseFrame; -use uuid::Uuid; use crate::WebSocketHandle; use crate::proxy_service::ResponseBody; use crate::request_context::RequestContext; +use pegboard::tunnel::id::RequestId; + pub enum HibernationResult { Continue, Close, @@ -23,6 +24,7 @@ pub trait CustomServeTrait: Send + Sync { &self, req: Request>, request_context: &mut RequestContext, + request_id: RequestId, ) -> Result>; /// Handle a WebSocket connection after upgrade. Supports connection retries. @@ -33,7 +35,7 @@ pub trait CustomServeTrait: Send + Sync { _path: &str, _request_context: &mut RequestContext, // Identifies the websocket across retries. - _unique_request_id: Uuid, + _unique_request_id: RequestId, // True if this websocket is reconnecting after hibernation. _after_hibernation: bool, ) -> Result> { @@ -44,7 +46,7 @@ pub trait CustomServeTrait: Send + Sync { async fn handle_websocket_hibernation( &self, _websocket: WebSocketHandle, - _unique_request_id: Uuid, + _unique_request_id: RequestId, ) -> Result { bail!("service does not support websocket hibernation"); } diff --git a/engine/packages/guard-core/src/lib.rs b/engine/packages/guard-core/src/lib.rs index 25f0d4a954..e2e03d8538 100644 --- a/engine/packages/guard-core/src/lib.rs +++ b/engine/packages/guard-core/src/lib.rs @@ -12,6 +12,7 @@ pub mod websocket_handle; pub use cert_resolver::CertResolverFn; pub use custom_serve::CustomServeTrait; +pub use pegboard::tunnel::id::{RequestId, generate_request_id}; pub use proxy_service::{ CacheKeyFn, MiddlewareFn, ProxyService, ProxyState, RouteTarget, RoutingFn, RoutingOutput, }; diff --git a/engine/packages/guard-core/src/proxy_service.rs b/engine/packages/guard-core/src/proxy_service.rs index 161b073ab7..8caa024a2e 100644 --- a/engine/packages/guard-core/src/proxy_service.rs +++ b/engine/packages/guard-core/src/proxy_service.rs @@ -13,9 +13,11 @@ use rivet_error::{INTERNAL_ERROR, RivetError}; use rivet_metrics::KeyValue; use rivet_util::Id; use serde_json; + +use pegboard::tunnel::id::{RequestId, generate_request_id}; use std::{ borrow::Cow, - collections::HashMap as StdHashMap, + collections::{HashMap as StdHashMap, HashSet}, net::SocketAddr, sync::Arc, time::{Duration, Instant}, @@ -28,7 +30,6 @@ use tokio_tungstenite::tungstenite::{ }; use tracing::Instrument; use url::Url; -use uuid::Uuid; use crate::{ WebSocketHandle, @@ -349,6 +350,7 @@ pub struct ProxyState { route_cache: RouteCache, rate_limiters: Cache<(Id, std::net::IpAddr), Arc>>, in_flight_counters: Cache<(Id, std::net::IpAddr), Arc>>, + inflight_requests: Arc>>, port_type: PortType, clickhouse_inserter: Option, tasks: Arc, @@ -377,6 +379,7 @@ impl ProxyState { .max_capacity(10_000) .time_to_live(PROXY_STATE_CACHE_TTL) .build(), + inflight_requests: Arc::new(Mutex::new(HashSet::new())), port_type, clickhouse_inserter, tasks: TaskGroup::new(), @@ -600,54 +603,94 @@ impl ProxyState { ip_addr: std::net::IpAddr, actor_id: &Option, headers: &hyper::HeaderMap, - ) -> Result { - let Some(actor_id) = *actor_id else { - // No in-flight limiting when actor_id is None - return Ok(true); - }; - - // Get actor-specific middleware config - let middleware_config = self.get_middleware_config(&actor_id, headers).await?; - - let cache_key = (actor_id, ip_addr); + ) -> Result> { + // Check in-flight limit if actor_id is present + if let Some(actor_id) = *actor_id { + // Get actor-specific middleware config + let middleware_config = self.get_middleware_config(&actor_id, headers).await?; + + let cache_key = (actor_id, ip_addr); + + // Get existing counter or create a new one + let counter_arc = + if let Some(existing_counter) = self.in_flight_counters.get(&cache_key).await { + existing_counter + } else { + let new_counter = Arc::new(Mutex::new(InFlightCounter::new( + middleware_config.max_in_flight.amount, + ))); + self.in_flight_counters + .insert(cache_key, new_counter.clone()) + .await; + metrics::IN_FLIGHT_COUNTER_COUNT.record(self.in_flight_counters.entry_count(), &[]); + new_counter + }; - // Get existing counter or create a new one - let counter_arc = - if let Some(existing_counter) = self.in_flight_counters.get(&cache_key).await { - existing_counter - } else { - let new_counter = Arc::new(Mutex::new(InFlightCounter::new( - middleware_config.max_in_flight.amount, - ))); - self.in_flight_counters - .insert(cache_key, new_counter.clone()) - .await; - metrics::IN_FLIGHT_COUNTER_COUNT.record(self.in_flight_counters.entry_count(), &[]); - new_counter + // Try to acquire from the counter + let acquired = { + let mut counter = counter_arc.lock().await; + counter.try_acquire() }; - // Try to acquire from the counter - let result = { - let mut counter = counter_arc.lock().await; - counter.try_acquire() - }; + if !acquired { + return Ok(None); // Rate limited + } + } - Ok(result) + // Generate unique request ID + let request_id = Some(self.generate_unique_request_id().await?); + Ok(request_id) } #[tracing::instrument(skip_all)] - async fn release_in_flight(&self, ip_addr: std::net::IpAddr, actor_id: &Option) { - // Skip if actor_id is None (no in-flight tracking) - let actor_id = match actor_id { - Some(id) => *id, - None => return, // No in-flight tracking when actor_id is None - }; + async fn release_in_flight( + &self, + ip_addr: std::net::IpAddr, + actor_id: &Option, + request_id: RequestId, + ) { + // Release in-flight counter if actor_id is present + if let Some(actor_id) = *actor_id { + let cache_key = (actor_id, ip_addr); + if let Some(counter_arc) = self.in_flight_counters.get(&cache_key).await { + let mut counter = counter_arc.lock().await; + counter.release(); + } + } - let cache_key = (actor_id, ip_addr); - if let Some(counter_arc) = self.in_flight_counters.get(&cache_key).await { - let mut counter = counter_arc.lock().await; - counter.release(); + // Release request ID + let mut requests = self.inflight_requests.lock().await; + requests.remove(&request_id); + } + + /// Generate a unique request ID that is not currently in flight + async fn generate_unique_request_id(&self) -> anyhow::Result { + const MAX_TRIES: u32 = 100; + let mut requests = self.inflight_requests.lock().await; + + for attempt in 0..MAX_TRIES { + let request_id = generate_request_id(); + + // Check if this ID is already in use + if !requests.contains(&request_id) { + // Insert the ID and return it + requests.insert(request_id); + return Ok(request_id); + } + + // Collision occurred (extremely rare with 4 bytes = 4 billion possibilities) + // Generate a new ID and try again + tracing::warn!( + ?request_id, + attempt, + "request id collision, generating new id" + ); } + + anyhow::bail!( + "failed to generate unique request id after {} attempts", + MAX_TRIES + ); } } @@ -766,20 +809,23 @@ impl ProxyService { .build()); } - // Check in-flight limit - if !self + // Acquire in-flight limit and generate request ID + let request_id = match self .state .acquire_in_flight(client_ip, &actor_id, req.headers()) .await? { - return Err(errors::RateLimit { - actor_id, - method: req.method().to_string(), - path: path.clone(), - ip: client_ip.to_string(), + Some(id) => id, + None => { + return Err(errors::RateLimit { + actor_id, + method: req.method().to_string(), + path: path.clone(), + ip: client_ip.to_string(), + } + .build()); } - .build()); - } + }; // Increment metrics metrics::PROXY_REQUEST_PENDING.add(1, &[]); @@ -791,10 +837,11 @@ impl ProxyService { } let res = if hyper_tungstenite::is_upgrade_request(&req) { - self.handle_websocket_upgrade(req, target, request_context) + self.handle_websocket_upgrade(req, target, request_context, client_ip, actor_id) .await } else { - self.handle_http_request(req, target, request_context).await + self.handle_http_request(req, target, request_context, client_ip, actor_id) + .await }; let status = match &res { @@ -809,11 +856,13 @@ impl ProxyService { metrics::PROXY_REQUEST_PENDING.add(-1, &[]); - // Release in-flight counter when done + // Release in-flight counter and request ID when done let state_clone = self.state.clone(); tokio::spawn( async move { - state_clone.release_in_flight(client_ip, &actor_id).await; + state_clone + .release_in_flight(client_ip, &actor_id, request_id) + .await; } .instrument(tracing::info_span!("release_in_flight_task")), ); @@ -827,6 +876,8 @@ impl ProxyService { req: Request, resolved_route: ResolveRouteOutput, request_context: &mut RequestContext, + client_ip: std::net::IpAddr, + actor_id: Option, ) -> Result> { // Get middleware config for this actor if it exists let middleware_config = if let ResolveRouteOutput::Target(target) = &resolved_route @@ -1043,6 +1094,25 @@ impl ProxyService { } ResolveRouteOutput::CustomServe(mut handler) => { let req_headers = req.headers().clone(); + let req_method = req.method().clone(); + + // Acquire in-flight limit and generate request ID + let request_id = match self + .state + .acquire_in_flight(client_ip, &actor_id, &req_headers) + .await? + { + Some(id) => id, + None => { + return Err(errors::RateLimit { + actor_id, + method: req_method.to_string(), + path: path.clone(), + ip: client_ip.to_string(), + } + .build()); + } + }; // Collect request body let (req_parts, body) = req.into_parts(); @@ -1064,7 +1134,7 @@ impl ProxyService { attempts += 1; let res = handler - .handle_request(req_collected.clone(), request_context) + .handle_request(req_collected.clone(), request_context, request_id) .await; if should_retry_request(&res) { // Request connect error, might retry @@ -1094,10 +1164,18 @@ impl ProxyService { continue; } + // Release in-flight counter and request ID before returning + self.state + .release_in_flight(client_ip, &actor_id, request_id) + .await; return res; } // If we get here, all attempts failed + // Release in-flight counter and request ID before returning error + self.state + .release_in_flight(client_ip, &actor_id, request_id) + .await; return Err(errors::RetryAttemptsExceeded { attempts: max_attempts, } @@ -1156,13 +1234,9 @@ impl ProxyService { req: Request, target: ResolveRouteOutput, request_context: &mut RequestContext, + client_ip: std::net::IpAddr, + actor_id: Option, ) -> Result> { - // Get actor and server IDs for metrics and middleware - let actor_id = match &target { - ResolveRouteOutput::Target(target) => target.actor_id, - _ => None, - }; - // Parsed for retries later let req_host = req .headers() @@ -1838,12 +1912,29 @@ impl ProxyService { tracing::debug!(%req_path, "Spawning task to handle WebSocket communication"); let mut request_context = request_context.clone(); let req_headers = req_headers.clone(); + let state = self.state.clone(); let req_path = req_path.clone(); let req_host = req_host.clone(); + let req_method = req_method.clone(); self.state.tasks.spawn( async move { - let request_id = Uuid::new_v4(); + let request_id = match state + .acquire_in_flight(client_ip, &actor_id, &req_headers) + .await? + { + Some(id) => id, + None => { + return Err(errors::RateLimit { + actor_id, + method: req_method.to_string(), + path: req_path.clone(), + ip: client_ip.to_string(), + } + .build() + .into()); + } + }; let mut ws_hibernation_close = false; let mut after_hibernation = false; let mut attempts = 0u32; @@ -2047,6 +2138,11 @@ impl ProxyService { } } + // Release in-flight counter and request ID when task completes + state + .release_in_flight(client_ip, &actor_id, request_id) + .await; + anyhow::Ok(()) } .instrument(tracing::info_span!("handle_ws_task_custom_serve")), diff --git a/engine/packages/pegboard-gateway/src/lib.rs b/engine/packages/pegboard-gateway/src/lib.rs index dbdc1c51b8..95cd6e8b35 100644 --- a/engine/packages/pegboard-gateway/src/lib.rs +++ b/engine/packages/pegboard-gateway/src/lib.rs @@ -5,29 +5,30 @@ use futures_util::TryStreamExt; use gas::prelude::*; use http_body_util::{BodyExt, Full}; use hyper::{Request, Response, StatusCode}; +use pegboard::tunnel::id::{self as tunnel_id, RequestId}; use rand::Rng; use rivet_error::*; use rivet_guard_core::{ - WebSocketHandle, custom_serve::{CustomServeTrait, HibernationResult}, errors::{ ServiceUnavailable, WebSocketServiceHibernate, WebSocketServiceTimeout, WebSocketServiceUnavailable, }, - proxy_service::{ResponseBody, is_ws_hibernate}, + proxy_service::{is_ws_hibernate, ResponseBody}, request_context::RequestContext, websocket_handle::WebSocketReceiver, + WebSocketHandle, }; use rivet_runner_protocol as protocol; use rivet_util::serde::HashableMap; use std::{sync::Arc, time::Duration}; use tokio::{ - sync::{Mutex, watch}, + sync::{watch, Mutex}, task::JoinHandle, }; use tokio_tungstenite::tungstenite::{ + protocol::frame::{coding::CloseCode, CloseFrame}, Message, - protocol::frame::{CloseFrame, coding::CloseCode}, }; use crate::shared_state::{InFlightRequestHandle, SharedState}; @@ -86,6 +87,7 @@ impl CustomServeTrait for PegboardGateway { &self, req: Request>, _request_context: &mut RequestContext, + request_id: RequestId, ) -> Result> { // Use the actor ID from the gateway instance let actor_id = self.actor_id.to_string(); @@ -163,7 +165,6 @@ impl CustomServeTrait for PegboardGateway { pegboard::pubsub_subjects::RunnerReceiverSubject::new(self.runner_id).to_string(); // Start listening for request responses - let request_id = Uuid::new_v4().into_bytes(); let InFlightRequestHandle { mut msg_rx, mut drop_rx, @@ -212,7 +213,7 @@ impl CustomServeTrait for PegboardGateway { } } else { tracing::warn!( - request_id=?Uuid::from_bytes(request_id), + request_id=?tunnel_id::request_id_to_string(&request_id), "received no message response during request init", ); break; @@ -274,7 +275,7 @@ impl CustomServeTrait for PegboardGateway { headers: &hyper::HeaderMap, _path: &str, _request_context: &mut RequestContext, - unique_request_id: Uuid, + unique_request_id: RequestId, after_hibernation: bool, ) -> Result> { // Use the actor ID from the gateway instance @@ -298,7 +299,7 @@ impl CustomServeTrait for PegboardGateway { pegboard::pubsub_subjects::RunnerReceiverSubject::new(self.runner_id).to_string(); // Start listening for WebSocket messages - let request_id = unique_request_id.into_bytes(); + let request_id = unique_request_id; let InFlightRequestHandle { mut msg_rx, mut drop_rx, @@ -348,7 +349,7 @@ impl CustomServeTrait for PegboardGateway { } } else { tracing::warn!( - request_id=?Uuid::from_bytes(request_id), + request_id=?tunnel_id::request_id_to_string(&request_id), "received no message response during ws init", ); break; @@ -414,7 +415,7 @@ impl CustomServeTrait for PegboardGateway { } protocol::ToServerTunnelMessageKind::ToServerWebSocketMessageAck(ack) => { tracing::debug!( - request_id=?Uuid::from_bytes(request_id), + request_id=?tunnel_id::request_id_to_string(&request_id), ack_index=?ack.index, "received WebSocketMessageAck from runner" ); @@ -477,8 +478,6 @@ impl CustomServeTrait for PegboardGateway { let ws_message = protocol::ToClientTunnelMessageKind::ToClientWebSocketMessage( protocol::ToClientWebSocketMessage { - // NOTE: This gets set in shared_state.ts - index: 0, data: data.into(), binary: true, }, @@ -491,8 +490,6 @@ impl CustomServeTrait for PegboardGateway { let ws_message = protocol::ToClientTunnelMessageKind::ToClientWebSocketMessage( protocol::ToClientWebSocketMessage { - // NOTE: This gets set in shared_state.ts - index: 0, data: text.as_bytes().to_vec(), binary: false, }, @@ -615,12 +612,14 @@ impl CustomServeTrait for PegboardGateway { async fn handle_websocket_hibernation( &self, client_ws: WebSocketHandle, - unique_request_id: Uuid, + unique_request_id: RequestId, ) -> Result { + let request_id = unique_request_id; + // Immediately rewake if we have pending messages if self .shared_state - .has_pending_websocket_messages(unique_request_id.into_bytes()) + .has_pending_websocket_messages(request_id) .await? { tracing::debug!( @@ -633,6 +632,8 @@ impl CustomServeTrait for PegboardGateway { // Start keepalive task let ctx = self.ctx.clone(); let actor_id = self.actor_id; + let gateway_id = self.shared_state.gateway_id(); + let request_id = unique_request_id; let keepalive_handle: JoinHandle> = tokio::spawn(async move { let mut ping_interval = tokio::time::interval(Duration::from_millis( (ctx.config() @@ -652,7 +653,8 @@ impl CustomServeTrait for PegboardGateway { ctx.op(pegboard::ops::actor::hibernating_request::upsert::Input { actor_id, - request_id: unique_request_id, + gateway_id, + request_id, }) .await?; } @@ -669,6 +671,7 @@ impl CustomServeTrait for PegboardGateway { self.ctx .op(pegboard::ops::actor::hibernating_request::delete::Input { actor_id: self.actor_id, + gateway_id: self.shared_state.gateway_id(), request_id: unique_request_id, }) .await?; diff --git a/engine/packages/pegboard-gateway/src/shared_state.rs b/engine/packages/pegboard-gateway/src/shared_state.rs index 49c610f0a7..fc47b50548 100644 --- a/engine/packages/pegboard-gateway/src/shared_state.rs +++ b/engine/packages/pegboard-gateway/src/shared_state.rs @@ -1,6 +1,7 @@ use anyhow::Result; use gas::prelude::*; -use rivet_runner_protocol::{self as protocol, MessageId, PROTOCOL_VERSION, RequestId, versioned}; +use pegboard::tunnel::id::{self as tunnel_id, GatewayId, RequestId}; +use rivet_runner_protocol::{self as protocol, PROTOCOL_VERSION, versioned}; use scc::{HashMap, hash_map::Entry}; use std::{ ops::Deref, @@ -14,7 +15,6 @@ use vbare::OwnedVersionedData; use crate::WebsocketPendingLimitReached; const GC_INTERVAL: Duration = Duration::from_secs(15); -const MESSAGE_ACK_TIMEOUT: Duration = Duration::from_secs(30); const HWS_MESSAGE_ACK_TIMEOUT: Duration = Duration::from_secs(30); const HWS_MAX_PENDING_MSGS_SIZE_PER_REQ: u64 = util::size::mebibytes(1); @@ -36,30 +36,26 @@ struct InFlightRequest { drop_tx: watch::Sender<()>, /// True once first message for this request has been sent (so runner learned reply_to). opened: bool, - pending_msgs: Vec, + /// Message index counter for this request. + message_index: tunnel_id::MessageIndex, hibernation_state: Option, stopping: bool, } -pub struct PendingMessage { - message_id: MessageId, - send_instant: Instant, -} - struct HibernationState { total_pending_ws_msgs_size: u64, - last_ws_msg_index: u16, pending_ws_msgs: Vec, } pub struct PendingWebsocketMessage { payload: Vec, send_instant: Instant, + message_index: tunnel_id::MessageIndex, } pub struct SharedStateInner { ups: PubSub, - gateway_id: Uuid, + gateway_id: GatewayId, receiver_subject: String, in_flight_requests: HashMap, } @@ -69,7 +65,7 @@ pub struct SharedState(Arc); impl SharedState { pub fn new(ups: PubSub) -> Self { - let gateway_id = Uuid::new_v4(); + let gateway_id = tunnel_id::generate_gateway_id(); let receiver_subject = pegboard::pubsub_subjects::GatewayReceiverSubject::new(gateway_id).to_string(); @@ -81,6 +77,10 @@ impl SharedState { })) } + pub fn gateway_id(&self) -> GatewayId { + self.gateway_id + } + pub async fn start(&self) -> Result<()> { let sub = self.ups.subscribe(&self.receiver_subject).await?; @@ -108,7 +108,7 @@ impl SharedState { msg_tx, drop_tx, opened: false, - pending_msgs: Vec::new(), + message_index: 0, hibernation_state: None, stopping: false, }); @@ -118,7 +118,7 @@ impl SharedState { entry.msg_tx = msg_tx; entry.drop_tx = drop_tx; entry.opened = false; - entry.pending_msgs.clear(); + entry.message_index = 0; if entry.stopping { entry.hibernation_state = None; @@ -133,60 +133,45 @@ impl SharedState { pub async fn send_message( &self, request_id: RequestId, - mut message_kind: protocol::ToClientTunnelMessageKind, + message_kind: protocol::ToClientTunnelMessageKind, ) -> Result<()> { - let message_id = Uuid::new_v4().as_bytes().clone(); - let mut req = self .in_flight_requests .get_async(&request_id) .await .context("request not in flight")?; + // Generate message ID + let message_id = + tunnel_id::build_message_id(self.gateway_id, request_id, req.message_index)?; + + // Increment message index for next message + let current_message_index = req.message_index; + req.message_index = req.message_index.wrapping_add(1); + let include_reply_to = !req.opened; if include_reply_to { // Mark as opened so subsequent messages skip reply_to req.opened = true; } - let ws_msg_index = - if let (Some(hs), protocol::ToClientTunnelMessageKind::ToClientWebSocketMessage(msg)) = - (&req.hibernation_state, &mut message_kind) - { - // TODO: This ends up skipping 0 as an index when initiated but whatever - msg.index = hs.last_ws_msg_index.wrapping_add(1); - - Some(msg.index) - } else { - None - }; + // Check if this is a WebSocket message for hibernation tracking + let is_ws_message = matches!( + message_kind, + protocol::ToClientTunnelMessageKind::ToClientWebSocketMessage(_) + ); let payload = protocol::ToClientTunnelMessage { - gateway_id: *self.gateway_id.as_bytes(), - request_id: request_id.clone(), message_id, - // Only send reply to subject on the first message for this request. This reduces - // overhead of subsequent messages. - gateway_reply_to: if include_reply_to { - Some(self.receiver_subject.clone()) - } else { - None - }, message_kind, }; - let now = Instant::now(); - req.pending_msgs.push(PendingMessage { - message_id, - send_instant: now, - }); - // Send message let message = protocol::ToRunner::ToClientTunnelMessage(payload); let message_serialized = versioned::ToRunner::wrap_latest(message) .serialize_with_embedded_version(PROTOCOL_VERSION)?; - if let (Some(hs), Some(ws_msg_index)) = (&mut req.hibernation_state, ws_msg_index) { + if let (Some(hs), true) = (&mut req.hibernation_state, is_ws_message) { hs.total_pending_ws_msgs_size += message_serialized.len() as u64; if hs.total_pending_ws_msgs_size > HWS_MAX_PENDING_MSGS_SIZE_PER_REQ @@ -195,11 +180,10 @@ impl SharedState { return Err(WebsocketPendingLimitReached {}.build()); } - hs.last_ws_msg_index = ws_msg_index; - let pending_ws_msg = PendingWebsocketMessage { payload: message_serialized.clone(), - send_instant: now, + send_instant: Instant::now(), + message_index: current_message_index, }; hs.pending_ws_msgs.push(pending_ws_msg); @@ -225,35 +209,25 @@ impl SharedState { match versioned::ToGateway::deserialize_with_embedded_version(&msg.payload) { Ok(protocol::ToGateway::ToGatewayKeepAlive) => { - // TODO: - // let prev_len = in_flight.pending_msgs.len(); - // - // tracing::debug!(message_id=?Uuid::from_bytes(msg.message_id), "received tunnel ack"); - // - // in_flight - // .pending_msgs - // .retain(|m| m.message_id != msg.message_id); - // - // if prev_len == in_flight.pending_msgs.len() { - // tracing::warn!( - // request_id=?Uuid::from_bytes(msg.request_id), - // message_id=?Uuid::from_bytes(msg.message_id), - // "pending message does not exist or ack received after message body" - // ) - // } else { - // tracing::debug!( - // request_id=?Uuid::from_bytes(msg.request_id), - // message_id=?Uuid::from_bytes(msg.message_id), - // "received TunnelAck, removed from pending" - // ); - // } + // No-op } Ok(protocol::ToGateway::ToServerTunnelMessage(msg)) => { - let Some(in_flight) = self.in_flight_requests.get_async(&msg.request_id).await + // Parse message ID to extract components + let parts = match tunnel_id::parse_message_id(msg.message_id) { + Ok(p) => p, + Err(err) => { + tracing::error!(?err, message_id=?msg.message_id, "failed to parse message id"); + continue; + } + }; + + let Some(in_flight) = + self.in_flight_requests.get_async(&parts.request_id).await else { tracing::warn!( - request_id=?Uuid::from_bytes(msg.request_id), - message_id=?Uuid::from_bytes(msg.message_id), + gateway_id=?tunnel_id::gateway_id_to_string(&parts.gateway_id), + request_id=?tunnel_id::request_id_to_string(&parts.request_id), + message_index=parts.message_index, "in flight has already been disconnected, dropping message" ); continue; @@ -261,8 +235,9 @@ impl SharedState { // Send message to the request handler to emulate the real network action tracing::debug!( - request_id=?Uuid::from_bytes(msg.request_id), - message_id=?Uuid::from_bytes(msg.message_id), + gateway_id=?tunnel_id::gateway_id_to_string(&parts.gateway_id), + request_id=?tunnel_id::request_id_to_string(&parts.request_id), + message_index=parts.message_index, "forwarding message to request handler" ); let _ = in_flight.msg_tx.send(msg.message_kind.clone()).await; @@ -287,7 +262,6 @@ impl SharedState { (false, true) => { req.hibernation_state = Some(HibernationState { total_pending_ws_msgs_size: 0, - last_ws_msg_index: 0, pending_ws_msgs: Vec::new(), }); } @@ -306,7 +280,7 @@ impl SharedState { if let Some(hs) = &mut req.hibernation_state { if !hs.pending_ws_msgs.is_empty() { - tracing::debug!(request_id=?Uuid::from_bytes(request_id.clone()), len=?hs.pending_ws_msgs.len(), "resending pending messages"); + tracing::debug!(request_id=?tunnel_id::request_id_to_string(&request_id), len=?hs.pending_ws_msgs.len(), "resending pending messages"); for pending_msg in &hs.pending_ws_msgs { self.ups @@ -342,31 +316,21 @@ impl SharedState { let Some(hs) = &mut req.hibernation_state else { tracing::warn!( - request_id=?Uuid::from_bytes(request_id), + request_id=?tunnel_id::request_id_to_string(&request_id), "cannot ack ws messages, hibernation is not enabled" ); return Ok(()); }; + // Retain messages with index > ack_index (messages that haven't been acknowledged yet) let len_before = hs.pending_ws_msgs.len(); - let len = len_before.try_into()?; - let mut iter_index = 0u16; - hs.pending_ws_msgs.retain(|_| { - let msg_index = hs - .last_ws_msg_index - .wrapping_sub(len) - .wrapping_add(1) - .wrapping_add(iter_index); - let keep = wrapping_gt(msg_index, ack_index); - - iter_index += 1; - - keep + hs.pending_ws_msgs.retain(|msg| { + wrapping_gt(msg.message_index, ack_index) }); let len_after = hs.pending_ws_msgs.len(); tracing::debug!( - request_id=?Uuid::from_bytes(request_id), + request_id=?tunnel_id::request_id_to_string(&request_id), ack_index, removed_count=len_before - len_after, remaining_count=len_after, @@ -412,10 +376,8 @@ impl SharedState { enum MsgGcReason { /// Gateway channel is closed and there are no pending messages GatewayClosed, - /// Any tunnel message not acked (TunnelAck) - MessageNotAcked { message_id: Uuid }, /// WebSocket pending messages (ToServerWebSocketMessageAck) - WebSocketMessageNotAcked { last_ws_msg_index: u16 }, + WebSocketMessageNotAcked { message_index: u16 }, } let now = Instant::now(); @@ -434,7 +396,6 @@ impl SharedState { // If we have no pending messages of any kind and the channel is closed, remove the // in flight req if req.msg_tx.is_closed() - && req.pending_msgs.is_empty() && req .hibernation_state .as_ref() @@ -444,21 +405,13 @@ impl SharedState { break 'reason Some(MsgGcReason::GatewayClosed); } - if let Some(earliest_pending_msg) = req.pending_msgs.first() { - if now.duration_since(earliest_pending_msg.send_instant) - <= MESSAGE_ACK_TIMEOUT - { - break 'reason Some(MsgGcReason::MessageNotAcked{message_id:Uuid::from_bytes(earliest_pending_msg.message_id)}); - } - } - if let Some(hs) = &req.hibernation_state && let Some(earliest_pending_ws_msg) = hs.pending_ws_msgs.first() { if now.duration_since(earliest_pending_ws_msg.send_instant) - <= HWS_MESSAGE_ACK_TIMEOUT + > HWS_MESSAGE_ACK_TIMEOUT { - break 'reason Some(MsgGcReason::WebSocketMessageNotAcked{last_ws_msg_index: hs.last_ws_msg_index}); + break 'reason Some(MsgGcReason::WebSocketMessageNotAcked{message_index: earliest_pending_ws_msg.message_index}); } } @@ -467,13 +420,13 @@ impl SharedState { if let Some(reason) = &reason { tracing::debug!( - request_id=?Uuid::from_bytes(*request_id), + request_id=?tunnel_id::request_id_to_string(request_id), ?reason, "gc stopping in flight request" ); if req.drop_tx.send(()).is_err() { - tracing::debug!(request_id=?Uuid::from_bytes(*request_id), "failed to send timeout msg to tunnel"); + tracing::debug!(request_id=?tunnel_id::request_id_to_string(request_id), "failed to send timeout msg to tunnel"); } // Mark req as stopping to skip this loop next time the gc is run @@ -491,7 +444,7 @@ impl SharedState { // When the websocket reconnects a new channel will be created if req.stopping && req.drop_tx.is_closed() { tracing::debug!( - request_id=?Uuid::from_bytes(*request_id), + request_id=?tunnel_id::request_id_to_string(request_id), "gc removing in flight request" ); diff --git a/engine/packages/pegboard-runner/src/lib.rs b/engine/packages/pegboard-runner/src/lib.rs index 99fed0928e..895e9aee19 100644 --- a/engine/packages/pegboard-runner/src/lib.rs +++ b/engine/packages/pegboard-runner/src/lib.rs @@ -5,6 +5,7 @@ use gas::prelude::*; use http_body_util::Full; use hyper::{Response, StatusCode}; use pegboard::ops::runner::update_alloc_idx::Action; +use pegboard::tunnel::id::RequestId; use rivet_guard_core::{ WebSocketHandle, custom_serve::CustomServeTrait, proxy_service::ResponseBody, request_context::RequestContext, @@ -47,6 +48,7 @@ impl CustomServeTrait for PegboardRunnerWsCustomServe { &self, _req: hyper::Request>, _request_context: &mut RequestContext, + _request_id: RequestId, ) -> Result> { // Pegboard runner ws doesn't handle regular HTTP requests // Return a simple status response @@ -66,7 +68,7 @@ impl CustomServeTrait for PegboardRunnerWsCustomServe { _headers: &hyper::HeaderMap, path: &str, _request_context: &mut RequestContext, - _unique_request_id: Uuid, + _unique_request_id: pegboard::tunnel::id::RequestId, _after_hibernation: bool, ) -> Result> { // Get UPS diff --git a/engine/packages/pegboard-runner/src/tunnel_to_ws_task.rs b/engine/packages/pegboard-runner/src/tunnel_to_ws_task.rs index 842bf99e46..2a029f54bf 100644 --- a/engine/packages/pegboard-runner/src/tunnel_to_ws_task.rs +++ b/engine/packages/pegboard-runner/src/tunnel_to_ws_task.rs @@ -64,7 +64,7 @@ pub async fn task( for command_wrapper in &mut command_wrappers { if let protocol::Command::CommandStartActor(protocol::CommandStartActor { actor_id, - hibernating_request_ids, + hibernating_requests, .. }) = &mut command_wrapper.inner { @@ -74,8 +74,13 @@ pub async fn task( }) .await?; - *hibernating_request_ids = - ids.into_iter().map(|x| x.into_bytes().to_vec()).collect(); + *hibernating_requests = ids + .into_iter() + .map(|x| protocol::HibernatingRequest { + gateway_id: x.gateway_id, + request_id: x.request_id, + }) + .collect(); } } diff --git a/engine/packages/pegboard-runner/src/ws_to_tunnel_task.rs b/engine/packages/pegboard-runner/src/ws_to_tunnel_task.rs index 3c5511823f..1758695d6c 100644 --- a/engine/packages/pegboard-runner/src/ws_to_tunnel_task.rs +++ b/engine/packages/pegboard-runner/src/ws_to_tunnel_task.rs @@ -5,6 +5,7 @@ use gas::prelude::*; use hyper_tungstenite::tungstenite::Message as WsMessage; use hyper_tungstenite::tungstenite::Message; use pegboard::pubsub_subjects::GatewayReceiverSubject; +use pegboard::tunnel::id as tunnel_id; use pegboard_actor_kv as kv; use rivet_guard_core::websocket_handle::WebSocketReceiver; use rivet_runner_protocol::{self as protocol, PROTOCOL_VERSION, versioned}; @@ -334,7 +335,7 @@ async fn handle_message( } } protocol::ToServer::ToServerTunnelMessage(tunnel_msg) => { - handle_tunnel_message(&ctx, tunnel_msg) + handle_tunnel_message(&ctx, &conn, tunnel_msg) .await .context("failed to handle tunnel message")?; } @@ -362,11 +363,44 @@ async fn handle_message( #[tracing::instrument(skip_all)] async fn handle_tunnel_message( ctx: &StandaloneCtx, + conn: &Arc, msg: protocol::ToServerTunnelMessage, ) -> Result<()> { + // Ignore DeprecatedTunnelAck messages (used only for backwards compatibility) + if matches!( + msg.message_kind, + protocol::ToServerTunnelMessageKind::DeprecatedTunnelAck + ) { + return Ok(()); + } + + // Send DeprecatedTunnelAck back to runner for older protocol versions + if protocol::compat::version_needs_tunnel_ack(conn.protocol_version) { + let ack_msg = versioned::ToClient::wrap_latest(protocol::ToClient::ToClientTunnelMessage( + protocol::ToClientTunnelMessage { + message_id: msg.message_id, + message_kind: protocol::ToClientTunnelMessageKind::DeprecatedTunnelAck, + }, + )); + + let ack_serialized = ack_msg + .serialize(conn.protocol_version) + .context("failed to serialize DeprecatedTunnelAck response")?; + + conn.ws_handle + .send(hyper_tungstenite::tungstenite::Message::Binary( + ack_serialized.into(), + )) + .await + .context("failed to send DeprecatedTunnelAck to runner")?; + } + + // Parse message ID to extract gateway_id + let parts = + tunnel_id::parse_message_id(msg.message_id).context("failed to parse message id")?; + // Publish message to UPS - let gateway_reply_to = - GatewayReceiverSubject::new(Uuid::from_bytes(msg.gateway_id)).to_string(); + let gateway_reply_to = GatewayReceiverSubject::new(parts.gateway_id).to_string(); let msg_serialized = versioned::ToGateway::wrap_latest(protocol::ToGateway::ToServerTunnelMessage(msg)) .serialize_with_embedded_version(PROTOCOL_VERSION) diff --git a/engine/packages/pegboard/Cargo.toml b/engine/packages/pegboard/Cargo.toml index 9204944e14..ed443b4ebd 100644 --- a/engine/packages/pegboard/Cargo.toml +++ b/engine/packages/pegboard/Cargo.toml @@ -13,6 +13,7 @@ gas.workspace = true lazy_static.workspace = true namespace.workspace = true nix.workspace = true +rand.workspace = true rivet-api-types.workspace = true rivet-api-util.workspace = true rivet-data.workspace = true @@ -21,6 +22,7 @@ rivet-metrics.workspace = true rivet-runner-protocol.workspace = true rivet-types.workspace = true rivet-util.workspace = true +serde_bare.workspace = true serde_json.workspace = true serde.workspace = true strum.workspace = true diff --git a/engine/packages/pegboard/src/keys/actor.rs b/engine/packages/pegboard/src/keys/actor.rs index 00a889506e..5348b9b191 100644 --- a/engine/packages/pegboard/src/keys/actor.rs +++ b/engine/packages/pegboard/src/keys/actor.rs @@ -1,7 +1,8 @@ use anyhow::Result; use gas::prelude::*; use universaldb::prelude::*; -use uuid::Uuid; + +use crate::tunnel::id::{GatewayId, RequestId}; #[derive(Debug)] pub struct CreateTsKey { @@ -317,14 +318,21 @@ impl<'de> TupleUnpack<'de> for NamespaceIdKey { pub struct HibernatingRequestKey { actor_id: Id, last_ping_ts: i64, - pub request_id: Uuid, + pub gateway_id: GatewayId, + pub request_id: RequestId, } impl HibernatingRequestKey { - pub fn new(actor_id: Id, last_ping_ts: i64, request_id: Uuid) -> Self { + pub fn new( + actor_id: Id, + last_ping_ts: i64, + gateway_id: GatewayId, + request_id: RequestId, + ) -> Self { HibernatingRequestKey { actor_id, last_ping_ts, + gateway_id, request_id, } } @@ -361,7 +369,8 @@ impl TuplePack for HibernatingRequestKey { HIBERNATING_REQUEST, self.actor_id, self.last_ping_ts, - self.request_id, + &self.gateway_id[..], + &self.request_id[..], ); t.pack(w, tuple_depth) } @@ -369,12 +378,23 @@ impl TuplePack for HibernatingRequestKey { impl<'de> TupleUnpack<'de> for HibernatingRequestKey { fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let (input, (_, _, actor_id, last_ping_ts, request_id)) = - <(usize, usize, Id, i64, Uuid)>::unpack(input, tuple_depth)?; + let (input, (_, _, actor_id, last_ping_ts, gateway_id_bytes, request_id_bytes)) = + <(usize, usize, Id, i64, Vec, Vec)>::unpack(input, tuple_depth)?; + + let gateway_id: GatewayId = gateway_id_bytes + .as_slice() + .try_into() + .expect("invalid gateway_id length"); + + let request_id: RequestId = request_id_bytes + .as_slice() + .try_into() + .expect("invalid request_id length"); let v = HibernatingRequestKey { actor_id, last_ping_ts, + gateway_id, request_id, }; diff --git a/engine/packages/pegboard/src/keys/hibernating_request.rs b/engine/packages/pegboard/src/keys/hibernating_request.rs index 49e47b3069..7f91883f60 100644 --- a/engine/packages/pegboard/src/keys/hibernating_request.rs +++ b/engine/packages/pegboard/src/keys/hibernating_request.rs @@ -1,15 +1,20 @@ use anyhow::Result; use universaldb::prelude::*; -use uuid::Uuid; + +use crate::tunnel::id::{GatewayId, RequestId}; #[derive(Debug)] pub struct LastPingTsKey { - request_id: Uuid, + gateway_id: GatewayId, + request_id: RequestId, } impl LastPingTsKey { - pub fn new(request_id: Uuid) -> Self { - LastPingTsKey { request_id } + pub fn new(gateway_id: GatewayId, request_id: RequestId) -> Self { + LastPingTsKey { + gateway_id, + request_id, + } } } @@ -32,17 +37,36 @@ impl TuplePack for LastPingTsKey { w: &mut W, tuple_depth: TupleDepth, ) -> std::io::Result { - let t = (HIBERNATING_REQUEST, DATA, self.request_id, LAST_PING_TS); + let t = ( + HIBERNATING_REQUEST, + DATA, + &self.gateway_id[..], + &self.request_id[..], + LAST_PING_TS, + ); t.pack(w, tuple_depth) } } impl<'de> TupleUnpack<'de> for LastPingTsKey { fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { - let (input, (_, _, request_id, _)) = - <(usize, usize, Uuid, usize)>::unpack(input, tuple_depth)?; + let (input, (_, _, gateway_id_bytes, request_id_bytes, _)) = + <(usize, usize, Vec, Vec, usize)>::unpack(input, tuple_depth)?; + + let gateway_id: GatewayId = gateway_id_bytes + .as_slice() + .try_into() + .expect("invalid gateway_id length"); + + let request_id: RequestId = request_id_bytes + .as_slice() + .try_into() + .expect("invalid request_id length"); - let v = LastPingTsKey { request_id }; + let v = LastPingTsKey { + gateway_id, + request_id, + }; Ok((input, v)) } diff --git a/engine/packages/pegboard/src/lib.rs b/engine/packages/pegboard/src/lib.rs index a776a3d227..1574ff4503 100644 --- a/engine/packages/pegboard/src/lib.rs +++ b/engine/packages/pegboard/src/lib.rs @@ -5,6 +5,7 @@ pub mod keys; mod metrics; pub mod ops; pub mod pubsub_subjects; +pub mod tunnel; mod utils; pub mod workflows; diff --git a/engine/packages/pegboard/src/ops/actor/hibernating_request/delete.rs b/engine/packages/pegboard/src/ops/actor/hibernating_request/delete.rs index 026d7ab207..3af721741f 100644 --- a/engine/packages/pegboard/src/ops/actor/hibernating_request/delete.rs +++ b/engine/packages/pegboard/src/ops/actor/hibernating_request/delete.rs @@ -1,13 +1,14 @@ use gas::prelude::*; use universaldb::utils::IsolationLevel::*; -use uuid::Uuid; use crate::keys; +use crate::tunnel::id::{GatewayId, RequestId}; #[derive(Debug, Default)] pub struct Input { pub actor_id: Id, - pub request_id: Uuid, + pub gateway_id: GatewayId, + pub request_id: RequestId, } #[operation] @@ -19,12 +20,14 @@ pub async fn pegboard_actor_hibernating_request_delete( .run(|tx| async move { let tx = tx.with_subspace(keys::subspace()); - let last_ping_ts_key = keys::hibernating_request::LastPingTsKey::new(input.request_id); + let last_ping_ts_key = + keys::hibernating_request::LastPingTsKey::new(input.gateway_id, input.request_id); if let Some(last_ping_ts) = tx.read_opt(&last_ping_ts_key, Serializable).await? { tx.delete(&keys::actor::HibernatingRequestKey::new( input.actor_id, last_ping_ts, + input.gateway_id, input.request_id, )); } diff --git a/engine/packages/pegboard/src/ops/actor/hibernating_request/list.rs b/engine/packages/pegboard/src/ops/actor/hibernating_request/list.rs index 31baa2baf7..875ae86ca3 100644 --- a/engine/packages/pegboard/src/ops/actor/hibernating_request/list.rs +++ b/engine/packages/pegboard/src/ops/actor/hibernating_request/list.rs @@ -2,20 +2,26 @@ use futures_util::{StreamExt, TryStreamExt}; use gas::prelude::*; use universaldb::options::StreamingMode; use universaldb::utils::IsolationLevel::*; -use uuid::Uuid; use crate::keys; +use crate::tunnel::id::{GatewayId, RequestId}; #[derive(Debug, Default)] pub struct Input { pub actor_id: Id, } +#[derive(Debug)] +pub struct HibernatingRequestItem { + pub gateway_id: GatewayId, + pub request_id: RequestId, +} + #[operation] pub async fn pegboard_actor_hibernating_request_list( ctx: &OperationCtx, input: &Input, -) -> Result> { +) -> Result> { let hibernating_request_eligible_threshold = ctx .config() .pegboard() @@ -46,7 +52,10 @@ pub async fn pegboard_actor_hibernating_request_list( ) .map(|res| { let key = tx.unpack::(res?.key())?; - Ok(key.request_id) + Ok(HibernatingRequestItem { + gateway_id: key.gateway_id, + request_id: key.request_id, + }) }) .try_collect::>() .await diff --git a/engine/packages/pegboard/src/ops/actor/hibernating_request/upsert.rs b/engine/packages/pegboard/src/ops/actor/hibernating_request/upsert.rs index 9fd18b6dbf..26f8b8cb7d 100644 --- a/engine/packages/pegboard/src/ops/actor/hibernating_request/upsert.rs +++ b/engine/packages/pegboard/src/ops/actor/hibernating_request/upsert.rs @@ -1,13 +1,14 @@ use gas::prelude::*; use universaldb::utils::IsolationLevel::*; -use uuid::Uuid; use crate::keys; +use crate::tunnel::id::{GatewayId, RequestId}; #[derive(Debug, Default)] pub struct Input { pub actor_id: Id, - pub request_id: Uuid, + pub gateway_id: GatewayId, + pub request_id: RequestId, } #[operation] @@ -19,12 +20,14 @@ pub async fn pegboard_actor_hibernating_request_upsert( .run(|tx| async move { let tx = tx.with_subspace(keys::subspace()); - let last_ping_ts_key = keys::hibernating_request::LastPingTsKey::new(input.request_id); + let last_ping_ts_key = + keys::hibernating_request::LastPingTsKey::new(input.gateway_id, input.request_id); if let Some(last_ping_ts) = tx.read_opt(&last_ping_ts_key, Serializable).await? { tx.delete(&keys::actor::HibernatingRequestKey::new( input.actor_id, last_ping_ts, + input.gateway_id, input.request_id, )); } @@ -32,7 +35,12 @@ pub async fn pegboard_actor_hibernating_request_upsert( let now = util::timestamp::now(); tx.write(&last_ping_ts_key, now)?; tx.write( - &keys::actor::HibernatingRequestKey::new(input.actor_id, now, input.request_id), + &keys::actor::HibernatingRequestKey::new( + input.actor_id, + now, + input.gateway_id, + input.request_id, + ), (), )?; diff --git a/engine/packages/pegboard/src/pubsub_subjects.rs b/engine/packages/pegboard/src/pubsub_subjects.rs index 30b2a5becd..27f40dba5b 100644 --- a/engine/packages/pegboard/src/pubsub_subjects.rs +++ b/engine/packages/pegboard/src/pubsub_subjects.rs @@ -1,5 +1,7 @@ use gas::prelude::*; +use crate::tunnel::id as tunnel_id; + pub struct RunnerReceiverSubject { runner_id: Id, } @@ -59,17 +61,21 @@ impl std::fmt::Display for RunnerEvictionByNameSubject { } pub struct GatewayReceiverSubject { - gateway_id: Uuid, + gateway_id: tunnel_id::GatewayId, } impl GatewayReceiverSubject { - pub fn new(gateway_id: Uuid) -> Self { + pub fn new(gateway_id: tunnel_id::GatewayId) -> Self { Self { gateway_id } } } impl std::fmt::Display for GatewayReceiverSubject { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "pegboard.gateway.{}", self.gateway_id) + write!( + f, + "pegboard.gateway.{}", + tunnel_id::gateway_id_to_string(&self.gateway_id) + ) } } diff --git a/engine/packages/pegboard/src/tunnel/id.rs b/engine/packages/pegboard/src/tunnel/id.rs new file mode 100644 index 0000000000..6bad4de4ed --- /dev/null +++ b/engine/packages/pegboard/src/tunnel/id.rs @@ -0,0 +1,86 @@ +use anyhow::{Context, Result, ensure}; +use base64::{Engine, engine::general_purpose::STANDARD as BASE64}; +use rivet_runner_protocol as protocol; + +// Type aliases for the message ID components +pub type GatewayId = [u8; 4]; +pub type RequestId = [u8; 4]; +pub type MessageIndex = u16; +pub type MessageId = [u8; 12]; + +/// Generate a new 4-byte gateway ID from a random u32 +pub fn generate_gateway_id() -> GatewayId { + rand::random::().to_le_bytes() +} + +/// Build a MessageId from its components +pub fn build_message_id( + gateway_id: GatewayId, + request_id: RequestId, + message_index: MessageIndex, +) -> Result { + let parts = protocol::MessageIdParts { + gateway_id, + request_id, + message_index, + }; + + // Serialize directly to a fixed-size buffer on the stack + let mut message_id = [0u8; 12]; + let mut cursor = std::io::Cursor::new(&mut message_id[..]); + serde_bare::to_writer(&mut cursor, &parts).context("failed to serialize message id parts")?; + + // Verify we wrote exactly 12 bytes + let written = cursor.position() as usize; + ensure!( + written == 12, + "message id serialization produced wrong size: expected 12 bytes, got {}", + written + ); + + Ok(message_id) +} + +/// Parse a MessageId into its components +pub fn parse_message_id(message_id: MessageId) -> Result { + serde_bare::from_slice(&message_id).context("failed to deserialize message id") +} + +/// Convert a GatewayId to a base64 string +pub fn gateway_id_to_string(gateway_id: &GatewayId) -> String { + BASE64.encode(gateway_id) +} + +/// Parse a GatewayId from a base64 string +pub fn gateway_id_from_string(s: &str) -> Result { + let bytes = BASE64.decode(s).context("failed to decode base64")?; + let gateway_id: GatewayId = bytes.try_into().map_err(|v: Vec| { + anyhow::anyhow!( + "invalid gateway id length: expected 4 bytes, got {}", + v.len() + ) + })?; + Ok(gateway_id) +} + +/// Generate a new 4-byte request ID from a random u32 +pub fn generate_request_id() -> RequestId { + rand::random::().to_le_bytes() +} + +/// Convert a RequestId to a base64 string +pub fn request_id_to_string(request_id: &RequestId) -> String { + BASE64.encode(request_id) +} + +/// Parse a RequestId from a base64 string +pub fn request_id_from_string(s: &str) -> Result { + let bytes = BASE64.decode(s).context("failed to decode base64")?; + let request_id: RequestId = bytes.try_into().map_err(|v: Vec| { + anyhow::anyhow!( + "invalid request id length: expected 4 bytes, got {}", + v.len() + ) + })?; + Ok(request_id) +} diff --git a/engine/packages/pegboard/src/tunnel/mod.rs b/engine/packages/pegboard/src/tunnel/mod.rs new file mode 100644 index 0000000000..fd6bb6c432 --- /dev/null +++ b/engine/packages/pegboard/src/tunnel/mod.rs @@ -0,0 +1 @@ +pub mod id; diff --git a/engine/packages/pegboard/src/workflows/actor/runtime.rs b/engine/packages/pegboard/src/workflows/actor/runtime.rs index ce9bc1acec..7f062d03e6 100644 --- a/engine/packages/pegboard/src/workflows/actor/runtime.rs +++ b/engine/packages/pegboard/src/workflows/actor/runtime.rs @@ -532,7 +532,7 @@ pub async fn spawn_actor( }, // Empty because request ids are ephemeral. This is intercepted by guard and // populated before it reaches the runner - hibernating_request_ids: Vec::new(), + hibernating_requests: Vec::new(), }), }) .to_workflow_id(runner_workflow_id) @@ -579,7 +579,7 @@ pub async fn spawn_actor( }, // Empty because request ids are ephemeral. This is intercepted by guard and // populated before it reaches the runner - hibernating_request_ids: Vec::new(), + hibernating_requests: Vec::new(), }), }) .to_workflow_id(sig.runner_workflow_id) diff --git a/engine/sdks/rust/runner-protocol/src/compat.rs b/engine/sdks/rust/runner-protocol/src/compat.rs new file mode 100644 index 0000000000..b8e9f2ef37 --- /dev/null +++ b/engine/sdks/rust/runner-protocol/src/compat.rs @@ -0,0 +1,7 @@ +/// Returns whether the given protocol version needs tunnel ack messages. +/// +/// Older protocols have GC cycles to check for tunnel ack, so we need to send DeprecatedTunnelAck +/// for backwards compatibility. +pub fn version_needs_tunnel_ack(version: u16) -> bool { + version <= 2 +} diff --git a/engine/sdks/rust/runner-protocol/src/lib.rs b/engine/sdks/rust/runner-protocol/src/lib.rs index cd8d059452..2c0151c639 100644 --- a/engine/sdks/rust/runner-protocol/src/lib.rs +++ b/engine/sdks/rust/runner-protocol/src/lib.rs @@ -1,3 +1,4 @@ +pub mod compat; pub mod generated; pub mod versioned; diff --git a/engine/sdks/rust/runner-protocol/src/versioned.rs b/engine/sdks/rust/runner-protocol/src/versioned.rs index db2a24e300..aa4dfae195 100644 --- a/engine/sdks/rust/runner-protocol/src/versioned.rs +++ b/engine/sdks/rust/runner-protocol/src/versioned.rs @@ -149,7 +149,7 @@ impl ToClient { create_ts: start.config.create_ts, input: start.config.input, }, - hibernating_request_ids: Vec::new(), + hibernating_requests: Vec::new(), }) } v2::Command::CommandStopActor(stop) => { @@ -174,14 +174,17 @@ impl ToClient { }) } v2::ToClient::ToClientTunnelMessage(msg) => { + // Extract v3 message_id from v2's message_id + // v3: gateway_id (4) + request_id (4) + message_index (4) = 12 bytes + // v2.message_id contains: entire v3 message_id (12 bytes) + padding (4 bytes) + let mut message_id = [0u8; 12]; + message_id.copy_from_slice(&msg.message_id[..12]); + v3::ToClient::ToClientTunnelMessage(v3::ToClientTunnelMessage { - gateway_id: [0; 16], - request_id: msg.request_id, - message_id: msg.message_id, + message_id, message_kind: convert_to_client_tunnel_message_kind_v2_to_v3( msg.message_kind, ), - gateway_reply_to: msg.gateway_reply_to, }) } }; @@ -243,13 +246,23 @@ impl ToClient { }) } v3::ToClient::ToClientTunnelMessage(msg) => { + // Split v3 message_id into v2's request_id and message_id + // v3: gateway_id (4) + request_id (4) + message_index (4) = 12 bytes + // v2.request_id = gateway_id (4) + request_id (4) + padding (8 zeros) + // v2.message_id = entire v3 message_id (12 bytes) + padding (4 zeros) + let mut request_id = [0u8; 16]; + let mut message_id = [0u8; 16]; + request_id[..8].copy_from_slice(&msg.message_id[..8]); // gateway_id + request_id + message_id[..12].copy_from_slice(&msg.message_id); // entire v3 message_id + v2::ToClient::ToClientTunnelMessage(v2::ToClientTunnelMessage { - request_id: msg.request_id, - message_id: msg.message_id, + request_id, + message_id, message_kind: convert_to_client_tunnel_message_kind_v3_to_v2( msg.message_kind, + &msg.message_id, )?, - gateway_reply_to: msg.gateway_reply_to, + gateway_reply_to: None, }) } }; @@ -489,10 +502,14 @@ impl ToServer { }) } v2::ToServer::ToServerTunnelMessage(msg) => { + // Extract v3 message_id from v2's message_id + // v3: gateway_id (4) + request_id (4) + message_index (4) = 12 bytes + // v2.message_id contains: entire v3 message_id (12 bytes) + padding (4 bytes) + let mut message_id = [0u8; 12]; + message_id.copy_from_slice(&msg.message_id[..12]); + v3::ToServer::ToServerTunnelMessage(v3::ToServerTunnelMessage { - gateway_id: [0; 16], - request_id: msg.request_id, - message_id: msg.message_id, + message_id, message_kind: convert_to_server_tunnel_message_kind_v2_to_v3( msg.message_kind, ), @@ -554,9 +571,18 @@ impl ToServer { }) } v3::ToServer::ToServerTunnelMessage(msg) => { + // Split v3 message_id into v2's request_id and message_id + // v3: gateway_id (4) + request_id (4) + message_index (4) = 12 bytes + // v2.request_id = gateway_id (4) + request_id (4) + padding (8 zeros) + // v2.message_id = entire v3 message_id (12 bytes) + padding (4 zeros) + let mut request_id = [0u8; 16]; + let mut message_id = [0u8; 16]; + request_id[..8].copy_from_slice(&msg.message_id[..8]); // gateway_id + request_id + message_id[..12].copy_from_slice(&msg.message_id); // entire v3 message_id + v2::ToServer::ToServerTunnelMessage(v2::ToServerTunnelMessage { - request_id: msg.request_id, - message_id: msg.message_id, + request_id, + message_id, message_kind: convert_to_server_tunnel_message_kind_v3_to_v2( msg.message_kind, )?, @@ -1217,7 +1243,6 @@ fn convert_to_client_tunnel_message_kind_v2_to_v3( } v2::ToClientTunnelMessageKind::ToClientWebSocketMessage(msg) => { v3::ToClientTunnelMessageKind::ToClientWebSocketMessage(v3::ToClientWebSocketMessage { - index: msg.index, data: msg.data, binary: msg.binary, }) @@ -1228,18 +1253,16 @@ fn convert_to_client_tunnel_message_kind_v2_to_v3( reason: close.reason, }) } - // TunnelAck was removed in v3 + // DeprecatedTunnelAck is kept for backwards compatibility v2::ToClientTunnelMessageKind::TunnelAck => { - // TunnelAck is deprecated and should not be used - // For backwards compatibility, we skip it - // This shouldn't happen in practice as TunnelAck was removed - v3::ToClientTunnelMessageKind::ToClientRequestAbort + v3::ToClientTunnelMessageKind::DeprecatedTunnelAck } } } fn convert_to_client_tunnel_message_kind_v3_to_v2( kind: v3::ToClientTunnelMessageKind, + message_id: &[u8; 12], ) -> Result { Ok(match kind { v3::ToClientTunnelMessageKind::ToClientRequestStart(req) => { @@ -1269,10 +1292,12 @@ fn convert_to_client_tunnel_message_kind_v3_to_v2( }) } v3::ToClientTunnelMessageKind::ToClientWebSocketMessage(msg) => { + // Extract message index from message_id (bytes 8-9, u16 little-endian per BARE spec) + let index = u16::from_le_bytes([message_id[8], message_id[9]]); v2::ToClientTunnelMessageKind::ToClientWebSocketMessage(v2::ToClientWebSocketMessage { data: msg.data, binary: msg.binary, - index: msg.index, + index, }) } v3::ToClientTunnelMessageKind::ToClientWebSocketClose(close) => { @@ -1281,6 +1306,9 @@ fn convert_to_client_tunnel_message_kind_v3_to_v2( reason: close.reason, }) } + v3::ToClientTunnelMessageKind::DeprecatedTunnelAck => { + v2::ToClientTunnelMessageKind::TunnelAck + } }) } @@ -1328,12 +1356,9 @@ fn convert_to_server_tunnel_message_kind_v2_to_v3( hibernate: close.retry, }) } - // TunnelAck was removed in v3 + // DeprecatedTunnelAck is kept for backwards compatibility v2::ToServerTunnelMessageKind::TunnelAck => { - // TunnelAck is deprecated and should not be used - // For backwards compatibility, we skip it - // This shouldn't happen in practice as TunnelAck was removed - v3::ToServerTunnelMessageKind::ToServerResponseAbort + v3::ToServerTunnelMessageKind::DeprecatedTunnelAck } } } @@ -1383,6 +1408,9 @@ fn convert_to_server_tunnel_message_kind_v3_to_v2( retry: close.hibernate, }) } + v3::ToServerTunnelMessageKind::DeprecatedTunnelAck => { + v2::ToServerTunnelMessageKind::TunnelAck + } }) } diff --git a/engine/sdks/schemas/runner-protocol/v3.bare b/engine/sdks/schemas/runner-protocol/v3.bare index d2085e59eb..53c258cece 100644 --- a/engine/sdks/schemas/runner-protocol/v3.bare +++ b/engine/sdks/schemas/runner-protocol/v3.bare @@ -5,6 +5,10 @@ type Id str type Json str +type GatewayId data[4] +type RequestId data[4] +type MessageIndex u16 + # MARK: KV # Basic types @@ -168,12 +172,17 @@ type EventWrapper struct { } # MARK: Commands -# + +type HibernatingRequest struct { + gatewayId: GatewayId + requestId: RequestId +} + type CommandStartActor struct { actorId: Id generation: u32 config: ActorConfig - hibernatingRequestIds: list + hibernatingRequests: list } type CommandStopActor struct { @@ -193,9 +202,22 @@ type CommandWrapper struct { # MARK: Tunnel -type GatewayId data[16] # UUIDv4 -type RequestId data[16] # UUIDv4 -type MessageId data[16] # UUIDv4 +# Message ID + +type MessageIdParts struct { + # Globally unique ID + gatewayId: GatewayId + # Unique ID to the gateway + requestId: RequestId + # Unique ID to the request + messageIndex: MessageIndex +} + +type MessageId data[12] + +# Ack (deprecated, older protocols that have gc cycles to check for tunnel ack) + +type DeprecatedTunnelAck void # HTTP type ToClientRequestStart struct { @@ -236,7 +258,6 @@ type ToClientWebSocketOpen struct { } type ToClientWebSocketMessage struct { - index: u16 data: data binary: bool } @@ -256,7 +277,7 @@ type ToServerWebSocketMessage struct { } type ToServerWebSocketMessageAck struct { - index: u16 + index: MessageIndex } type ToServerWebSocketClose struct { @@ -267,11 +288,13 @@ type ToServerWebSocketClose struct { # To Server type ToServerTunnelMessageKind union { + DeprecatedTunnelAck | + # HTTP ToServerResponseStart | ToServerResponseChunk | ToServerResponseAbort | - + # WebSocket ToServerWebSocketOpen | ToServerWebSocketMessage | @@ -280,19 +303,19 @@ type ToServerTunnelMessageKind union { } type ToServerTunnelMessage struct { - gatewayId: GatewayId - requestId: RequestId messageId: MessageId messageKind: ToServerTunnelMessageKind } # To Client type ToClientTunnelMessageKind union { + DeprecatedTunnelAck | + # HTTP ToClientRequestStart | ToClientRequestChunk | ToClientRequestAbort | - + # WebSocket ToClientWebSocketOpen | ToClientWebSocketMessage | @@ -300,17 +323,8 @@ type ToClientTunnelMessageKind union { } type ToClientTunnelMessage struct { - gatewayId: GatewayId - requestId: RequestId messageId: MessageId messageKind: ToClientTunnelMessageKind - - # Subject to send replies to. - # - # Only sent when opening a new request from gateway -> pegboard-runner-ws. - # - # Should be stripped before sending to the runner. - gatewayReplyTo: optional } # MARK: To Server diff --git a/engine/sdks/typescript/runner-protocol/src/index.ts b/engine/sdks/typescript/runner-protocol/src/index.ts index 0531f2c169..3e1d938483 100644 --- a/engine/sdks/typescript/runner-protocol/src/index.ts +++ b/engine/sdks/typescript/runner-protocol/src/index.ts @@ -28,6 +28,38 @@ export function writeJson(bc: bare.ByteCursor, x: Json): void { bare.writeString(bc, x) } +export type GatewayId = ArrayBuffer + +export function readGatewayId(bc: bare.ByteCursor): GatewayId { + return bare.readFixedData(bc, 4) +} + +export function writeGatewayId(bc: bare.ByteCursor, x: GatewayId): void { + assert(x.byteLength === 4) + bare.writeFixedData(bc, x) +} + +export type RequestId = ArrayBuffer + +export function readRequestId(bc: bare.ByteCursor): RequestId { + return bare.readFixedData(bc, 4) +} + +export function writeRequestId(bc: bare.ByteCursor, x: RequestId): void { + assert(x.byteLength === 4) + bare.writeFixedData(bc, x) +} + +export type MessageIndex = u16 + +export function readMessageIndex(bc: bare.ByteCursor): MessageIndex { + return bare.readU16(bc) +} + +export function writeMessageIndex(bc: bare.ByteCursor, x: MessageIndex): void { + bare.writeU16(bc, x) +} + /** * Basic types */ @@ -808,22 +840,39 @@ export function writeEventWrapper(bc: bare.ByteCursor, x: EventWrapper): void { writeEvent(bc, x.inner) } -function read8(bc: bare.ByteCursor): readonly ArrayBuffer[] { +export type HibernatingRequest = { + readonly gatewayId: GatewayId + readonly requestId: RequestId +} + +export function readHibernatingRequest(bc: bare.ByteCursor): HibernatingRequest { + return { + gatewayId: readGatewayId(bc), + requestId: readRequestId(bc), + } +} + +export function writeHibernatingRequest(bc: bare.ByteCursor, x: HibernatingRequest): void { + writeGatewayId(bc, x.gatewayId) + writeRequestId(bc, x.requestId) +} + +function read8(bc: bare.ByteCursor): readonly HibernatingRequest[] { const len = bare.readUintSafe(bc) if (len === 0) { return [] } - const result = [bare.readData(bc)] + const result = [readHibernatingRequest(bc)] for (let i = 1; i < len; i++) { - result[i] = bare.readData(bc) + result[i] = readHibernatingRequest(bc) } return result } -function write8(bc: bare.ByteCursor, x: readonly ArrayBuffer[]): void { +function write8(bc: bare.ByteCursor, x: readonly HibernatingRequest[]): void { bare.writeUintSafe(bc, x.length) for (let i = 0; i < x.length; i++) { - bare.writeData(bc, x[i]) + writeHibernatingRequest(bc, x[i]) } } @@ -831,7 +880,7 @@ export type CommandStartActor = { readonly actorId: Id readonly generation: u32 readonly config: ActorConfig - readonly hibernatingRequestIds: readonly ArrayBuffer[] + readonly hibernatingRequests: readonly HibernatingRequest[] } export function readCommandStartActor(bc: bare.ByteCursor): CommandStartActor { @@ -839,7 +888,7 @@ export function readCommandStartActor(bc: bare.ByteCursor): CommandStartActor { actorId: readId(bc), generation: bare.readU32(bc), config: readActorConfig(bc), - hibernatingRequestIds: read8(bc), + hibernatingRequests: read8(bc), } } @@ -847,7 +896,7 @@ export function writeCommandStartActor(bc: bare.ByteCursor, x: CommandStartActor writeId(bc, x.actorId) bare.writeU32(bc, x.generation) writeActorConfig(bc, x.config) - write8(bc, x.hibernatingRequestIds) + write8(bc, x.hibernatingRequests) } export type CommandStopActor = { @@ -918,45 +967,67 @@ export function writeCommandWrapper(bc: bare.ByteCursor, x: CommandWrapper): voi writeCommand(bc, x.inner) } -export type GatewayId = ArrayBuffer - -export function readGatewayId(bc: bare.ByteCursor): GatewayId { - return bare.readFixedData(bc, 16) +export type MessageIdParts = { + /** + * Globally unique ID + */ + readonly gatewayId: GatewayId + /** + * Unique ID to the gateway + */ + readonly requestId: RequestId + /** + * Unique ID to the request + */ + readonly messageIndex: MessageIndex } -export function writeGatewayId(bc: bare.ByteCursor, x: GatewayId): void { - assert(x.byteLength === 16) - bare.writeFixedData(bc, x) +export function readMessageIdParts(bc: bare.ByteCursor): MessageIdParts { + return { + gatewayId: readGatewayId(bc), + requestId: readRequestId(bc), + messageIndex: readMessageIndex(bc), + } } -/** - * UUIDv4 - */ -export type RequestId = ArrayBuffer +export function writeMessageIdParts(bc: bare.ByteCursor, x: MessageIdParts): void { + writeGatewayId(bc, x.gatewayId) + writeRequestId(bc, x.requestId) + writeMessageIndex(bc, x.messageIndex) +} -export function readRequestId(bc: bare.ByteCursor): RequestId { - return bare.readFixedData(bc, 16) +export function encodeMessageIdParts(x: MessageIdParts, config?: Partial): Uint8Array { + const fullConfig = config != null ? bare.Config(config) : DEFAULT_CONFIG + const bc = new bare.ByteCursor( + new Uint8Array(fullConfig.initialBufferLength), + fullConfig, + ) + writeMessageIdParts(bc, x) + return new Uint8Array(bc.view.buffer, bc.view.byteOffset, bc.offset) } -export function writeRequestId(bc: bare.ByteCursor, x: RequestId): void { - assert(x.byteLength === 16) - bare.writeFixedData(bc, x) +export function decodeMessageIdParts(bytes: Uint8Array): MessageIdParts { + const bc = new bare.ByteCursor(bytes, DEFAULT_CONFIG) + const result = readMessageIdParts(bc) + if (bc.offset < bc.view.byteLength) { + throw new bare.BareError(bc.offset, "remaining bytes") + } + return result } -/** - * UUIDv4 - */ export type MessageId = ArrayBuffer export function readMessageId(bc: bare.ByteCursor): MessageId { - return bare.readFixedData(bc, 16) + return bare.readFixedData(bc, 12) } export function writeMessageId(bc: bare.ByteCursor, x: MessageId): void { - assert(x.byteLength === 16) + assert(x.byteLength === 12) bare.writeFixedData(bc, x) } +export type DeprecatedTunnelAck = null + function read9(bc: bare.ByteCursor): ReadonlyMap { const len = bare.readUintSafe(bc) const result = new Map() @@ -1097,21 +1168,18 @@ export function writeToClientWebSocketOpen(bc: bare.ByteCursor, x: ToClientWebSo } export type ToClientWebSocketMessage = { - readonly index: u16 readonly data: ArrayBuffer readonly binary: boolean } export function readToClientWebSocketMessage(bc: bare.ByteCursor): ToClientWebSocketMessage { return { - index: bare.readU16(bc), data: bare.readData(bc), binary: bare.readBool(bc), } } export function writeToClientWebSocketMessage(bc: bare.ByteCursor, x: ToClientWebSocketMessage): void { - bare.writeU16(bc, x.index) bare.writeData(bc, x.data) bare.writeBool(bc, x.binary) } @@ -1176,17 +1244,17 @@ export function writeToServerWebSocketMessage(bc: bare.ByteCursor, x: ToServerWe } export type ToServerWebSocketMessageAck = { - readonly index: u16 + readonly index: MessageIndex } export function readToServerWebSocketMessageAck(bc: bare.ByteCursor): ToServerWebSocketMessageAck { return { - index: bare.readU16(bc), + index: readMessageIndex(bc), } } export function writeToServerWebSocketMessageAck(bc: bare.ByteCursor, x: ToServerWebSocketMessageAck): void { - bare.writeU16(bc, x.index) + writeMessageIndex(bc, x.index) } export type ToServerWebSocketClose = { @@ -1213,6 +1281,7 @@ export function writeToServerWebSocketClose(bc: bare.ByteCursor, x: ToServerWebS * To Server */ export type ToServerTunnelMessageKind = + | { readonly tag: "DeprecatedTunnelAck"; readonly val: DeprecatedTunnelAck } /** * HTTP */ @@ -1232,18 +1301,20 @@ export function readToServerTunnelMessageKind(bc: bare.ByteCursor): ToServerTunn const tag = bare.readU8(bc) switch (tag) { case 0: - return { tag: "ToServerResponseStart", val: readToServerResponseStart(bc) } + return { tag: "DeprecatedTunnelAck", val: null } case 1: - return { tag: "ToServerResponseChunk", val: readToServerResponseChunk(bc) } + return { tag: "ToServerResponseStart", val: readToServerResponseStart(bc) } case 2: - return { tag: "ToServerResponseAbort", val: null } + return { tag: "ToServerResponseChunk", val: readToServerResponseChunk(bc) } case 3: - return { tag: "ToServerWebSocketOpen", val: readToServerWebSocketOpen(bc) } + return { tag: "ToServerResponseAbort", val: null } case 4: - return { tag: "ToServerWebSocketMessage", val: readToServerWebSocketMessage(bc) } + return { tag: "ToServerWebSocketOpen", val: readToServerWebSocketOpen(bc) } case 5: - return { tag: "ToServerWebSocketMessageAck", val: readToServerWebSocketMessageAck(bc) } + return { tag: "ToServerWebSocketMessage", val: readToServerWebSocketMessage(bc) } case 6: + return { tag: "ToServerWebSocketMessageAck", val: readToServerWebSocketMessageAck(bc) } + case 7: return { tag: "ToServerWebSocketClose", val: readToServerWebSocketClose(bc) } default: { bc.offset = offset @@ -1254,37 +1325,41 @@ export function readToServerTunnelMessageKind(bc: bare.ByteCursor): ToServerTunn export function writeToServerTunnelMessageKind(bc: bare.ByteCursor, x: ToServerTunnelMessageKind): void { switch (x.tag) { - case "ToServerResponseStart": { + case "DeprecatedTunnelAck": { bare.writeU8(bc, 0) + break + } + case "ToServerResponseStart": { + bare.writeU8(bc, 1) writeToServerResponseStart(bc, x.val) break } case "ToServerResponseChunk": { - bare.writeU8(bc, 1) + bare.writeU8(bc, 2) writeToServerResponseChunk(bc, x.val) break } case "ToServerResponseAbort": { - bare.writeU8(bc, 2) + bare.writeU8(bc, 3) break } case "ToServerWebSocketOpen": { - bare.writeU8(bc, 3) + bare.writeU8(bc, 4) writeToServerWebSocketOpen(bc, x.val) break } case "ToServerWebSocketMessage": { - bare.writeU8(bc, 4) + bare.writeU8(bc, 5) writeToServerWebSocketMessage(bc, x.val) break } case "ToServerWebSocketMessageAck": { - bare.writeU8(bc, 5) + bare.writeU8(bc, 6) writeToServerWebSocketMessageAck(bc, x.val) break } case "ToServerWebSocketClose": { - bare.writeU8(bc, 6) + bare.writeU8(bc, 7) writeToServerWebSocketClose(bc, x.val) break } @@ -1292,24 +1367,18 @@ export function writeToServerTunnelMessageKind(bc: bare.ByteCursor, x: ToServerT } export type ToServerTunnelMessage = { - readonly gatewayId: GatewayId - readonly requestId: RequestId readonly messageId: MessageId readonly messageKind: ToServerTunnelMessageKind } export function readToServerTunnelMessage(bc: bare.ByteCursor): ToServerTunnelMessage { return { - gatewayId: readGatewayId(bc), - requestId: readRequestId(bc), messageId: readMessageId(bc), messageKind: readToServerTunnelMessageKind(bc), } } export function writeToServerTunnelMessage(bc: bare.ByteCursor, x: ToServerTunnelMessage): void { - writeGatewayId(bc, x.gatewayId) - writeRequestId(bc, x.requestId) writeMessageId(bc, x.messageId) writeToServerTunnelMessageKind(bc, x.messageKind) } @@ -1318,6 +1387,7 @@ export function writeToServerTunnelMessage(bc: bare.ByteCursor, x: ToServerTunne * To Client */ export type ToClientTunnelMessageKind = + | { readonly tag: "DeprecatedTunnelAck"; readonly val: DeprecatedTunnelAck } /** * HTTP */ @@ -1336,16 +1406,18 @@ export function readToClientTunnelMessageKind(bc: bare.ByteCursor): ToClientTunn const tag = bare.readU8(bc) switch (tag) { case 0: - return { tag: "ToClientRequestStart", val: readToClientRequestStart(bc) } + return { tag: "DeprecatedTunnelAck", val: null } case 1: - return { tag: "ToClientRequestChunk", val: readToClientRequestChunk(bc) } + return { tag: "ToClientRequestStart", val: readToClientRequestStart(bc) } case 2: - return { tag: "ToClientRequestAbort", val: null } + return { tag: "ToClientRequestChunk", val: readToClientRequestChunk(bc) } case 3: - return { tag: "ToClientWebSocketOpen", val: readToClientWebSocketOpen(bc) } + return { tag: "ToClientRequestAbort", val: null } case 4: - return { tag: "ToClientWebSocketMessage", val: readToClientWebSocketMessage(bc) } + return { tag: "ToClientWebSocketOpen", val: readToClientWebSocketOpen(bc) } case 5: + return { tag: "ToClientWebSocketMessage", val: readToClientWebSocketMessage(bc) } + case 6: return { tag: "ToClientWebSocketClose", val: readToClientWebSocketClose(bc) } default: { bc.offset = offset @@ -1356,32 +1428,36 @@ export function readToClientTunnelMessageKind(bc: bare.ByteCursor): ToClientTunn export function writeToClientTunnelMessageKind(bc: bare.ByteCursor, x: ToClientTunnelMessageKind): void { switch (x.tag) { - case "ToClientRequestStart": { + case "DeprecatedTunnelAck": { bare.writeU8(bc, 0) + break + } + case "ToClientRequestStart": { + bare.writeU8(bc, 1) writeToClientRequestStart(bc, x.val) break } case "ToClientRequestChunk": { - bare.writeU8(bc, 1) + bare.writeU8(bc, 2) writeToClientRequestChunk(bc, x.val) break } case "ToClientRequestAbort": { - bare.writeU8(bc, 2) + bare.writeU8(bc, 3) break } case "ToClientWebSocketOpen": { - bare.writeU8(bc, 3) + bare.writeU8(bc, 4) writeToClientWebSocketOpen(bc, x.val) break } case "ToClientWebSocketMessage": { - bare.writeU8(bc, 4) + bare.writeU8(bc, 5) writeToClientWebSocketMessage(bc, x.val) break } case "ToClientWebSocketClose": { - bare.writeU8(bc, 5) + bare.writeU8(bc, 6) writeToClientWebSocketClose(bc, x.val) break } @@ -1389,32 +1465,20 @@ export function writeToClientTunnelMessageKind(bc: bare.ByteCursor, x: ToClientT } export type ToClientTunnelMessage = { - readonly gatewayId: GatewayId - readonly requestId: RequestId readonly messageId: MessageId readonly messageKind: ToClientTunnelMessageKind - /** - * Should be stripped before sending to the runner. - */ - readonly gatewayReplyTo: string | null } export function readToClientTunnelMessage(bc: bare.ByteCursor): ToClientTunnelMessage { return { - gatewayId: readGatewayId(bc), - requestId: readRequestId(bc), messageId: readMessageId(bc), messageKind: readToClientTunnelMessageKind(bc), - gatewayReplyTo: read5(bc), } } export function writeToClientTunnelMessage(bc: bare.ByteCursor, x: ToClientTunnelMessage): void { - writeGatewayId(bc, x.gatewayId) - writeRequestId(bc, x.requestId) writeMessageId(bc, x.messageId) writeToClientTunnelMessageKind(bc, x.messageKind) - write5(bc, x.gatewayReplyTo) } function read11(bc: bare.ByteCursor): ReadonlyMap { diff --git a/engine/sdks/typescript/runner/src/actor.ts b/engine/sdks/typescript/runner/src/actor.ts new file mode 100644 index 0000000000..8db2c29aa7 --- /dev/null +++ b/engine/sdks/typescript/runner/src/actor.ts @@ -0,0 +1,103 @@ +import type * as protocol from "@rivetkit/engine-runner-protocol"; +import type { PendingRequest } from "./tunnel"; +import type { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter"; +import { arraysEqual } from "./utils"; + +export interface ActorConfig { + name: string; + key: string | null; + createTs: bigint; + input: Uint8Array | null; +} + +export class RunnerActor { + actorId: string; + generation: number; + config: ActorConfig; + pendingRequests: Array<{ + gatewayId: protocol.GatewayId; + requestId: protocol.RequestId; + request: PendingRequest; + }> = []; + webSockets: Array<{ + gatewayId: protocol.GatewayId; + requestId: protocol.RequestId; + ws: WebSocketTunnelAdapter; + }> = []; + + constructor(actorId: string, generation: number, config: ActorConfig) { + this.actorId = actorId; + this.generation = generation; + this.config = config; + } + + // Pending request methods + getPendingRequest( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + ): PendingRequest | undefined { + return this.pendingRequests.find( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + )?.request; + } + + setPendingRequest( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + request: PendingRequest, + ) { + this.deletePendingRequest(gatewayId, requestId); + this.pendingRequests.push({ gatewayId, requestId, request }); + } + + deletePendingRequest( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + ) { + const index = this.pendingRequests.findIndex( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + ); + if (index !== -1) { + this.pendingRequests.splice(index, 1); + } + } + + // WebSocket methods + getWebSocket( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + ): WebSocketTunnelAdapter | undefined { + return this.webSockets.find( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + )?.ws; + } + + setWebSocket( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + ws: WebSocketTunnelAdapter, + ) { + this.deleteWebSocket(gatewayId, requestId); + this.webSockets.push({ gatewayId, requestId, ws }); + } + + deleteWebSocket( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + ) { + const index = this.webSockets.findIndex( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + ); + if (index !== -1) { + this.webSockets.splice(index, 1); + } + } +} diff --git a/engine/sdks/typescript/runner/src/mod.ts b/engine/sdks/typescript/runner/src/mod.ts index bf76be7e77..86a223e76b 100644 --- a/engine/sdks/typescript/runner/src/mod.ts +++ b/engine/sdks/typescript/runner/src/mod.ts @@ -3,7 +3,7 @@ import type { Logger } from "pino"; import type WebSocket from "ws"; import { logger, setLogger } from "./log.js"; import { stringifyCommandWrapper, stringifyEvent } from "./stringify"; -import type { PendingRequest, PendingTunnelMessage } from "./tunnel"; +import type { PendingRequest } from "./tunnel"; import { type HibernatingWebSocketMetadata, Tunnel } from "./tunnel"; import { calculateBackoff, @@ -12,8 +12,11 @@ import { } from "./utils"; import { importWebSocket } from "./websocket.js"; import type { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter"; +import { RunnerActor, type ActorConfig } from "./actor"; export type { HibernatingWebSocketMetadata }; +export * as tunnelId from "./tunnel-id"; +export { RunnerActor, type ActorConfig }; const KV_EXPIRE: number = 30_000; const PROTOCOL_VERSION: number = 3; @@ -23,22 +26,6 @@ const RUNNER_PING_INTERVAL = 3_000; const EVENT_BACKLOG_WARN_THRESHOLD = 10_000; const SIGNAL_HANDLERS: (() => void)[] = []; -export interface RunnerActor { - actorId: string; - generation: number; - config: ActorConfig; - pendingRequests: Map; - webSockets: Map; - pendingTunnelMessages: Map; -} - -export interface ActorConfig { - name: string; - key: string | null; - createTs: bigint; - input: Uint8Array | null; -} - export interface RunnerConfig { logger?: Logger; version: number; @@ -60,6 +47,7 @@ export interface RunnerConfig { fetch: ( runner: Runner, actorId: string, + gatewayId: protocol.GatewayId, requestId: protocol.RequestId, request: Request, ) => Promise; @@ -139,6 +127,7 @@ export interface RunnerConfig { runner: Runner, actorId: string, ws: any, + gatewayId: protocol.GatewayId, requestId: protocol.RequestId, request: Request, path: string, @@ -154,6 +143,7 @@ export interface RunnerConfig { */ canHibernate: ( actorId: string, + gatewayId: ArrayBuffer, requestId: ArrayBuffer, request: Request, ) => boolean; @@ -232,7 +222,7 @@ export class Runner { #ackInterval?: NodeJS.Timeout; // KV operations - #nextRequestId: number = 0; + #nextKvRequestId: number = 0; #kvRequests: Map = new Map(); #kvCleanupInterval?: NodeJS.Timeout; @@ -932,14 +922,7 @@ export class Runner { input: config.input ? new Uint8Array(config.input) : null, }; - const instance: RunnerActor = { - actorId, - generation, - config: actorConfig, - pendingRequests: new Map(), - webSockets: new Map(), - pendingTunnelMessages: new Map(), - }; + const instance = new RunnerActor(actorId, generation, actorConfig); this.#actors.set(actorId, instance); @@ -958,7 +941,7 @@ export class Runner { // Restore hibernating requests await this.#tunnel.restoreHibernatingRequests( actorId, - startCommand.hibernatingRequestIds, + startCommand.hibernatingRequests, ); } catch (err) { this.log?.error({ @@ -1451,7 +1434,7 @@ export class Runner { return; } - const requestId = this.#nextRequestId++; + const requestId = this.#nextKvRequestId++; const isConnected = this.#pegboardWebSocket && this.#pegboardWebSocket.readyState === 1; @@ -1548,10 +1531,18 @@ export class Runner { } } - sendHibernatableWebSocketMessageAck(requestId: ArrayBuffer, index: number) { + sendHibernatableWebSocketMessageAck( + gatewayId: ArrayBuffer, + requestId: ArrayBuffer, + index: number, + ) { if (!this.#tunnel) throw new Error("missing tunnel to send message ack"); - this.#tunnel.sendHibernatableWebSocketMessageAck(requestId, index); + this.#tunnel.sendHibernatableWebSocketMessageAck( + gatewayId, + requestId, + index, + ); } getServerlessInitPacket(): string | undefined { diff --git a/engine/sdks/typescript/runner/src/stringify.ts b/engine/sdks/typescript/runner/src/stringify.ts index f29923d0b2..25f345a789 100644 --- a/engine/sdks/typescript/runner/src/stringify.ts +++ b/engine/sdks/typescript/runner/src/stringify.ts @@ -32,8 +32,8 @@ export function stringifyToServerTunnelMessageKind( kind: protocol.ToServerTunnelMessageKind, ): string { switch (kind.tag) { - case "TunnelAck": - return "TunnelAck"; + case "DeprecatedTunnelAck": + return "DeprecatedTunnelAck"; case "ToServerResponseStart": { const { status, headers, body, stream } = kind.val; const bodyStr = body === null ? "null" : stringifyArrayBuffer(body); @@ -74,8 +74,8 @@ export function stringifyToClientTunnelMessageKind( kind: protocol.ToClientTunnelMessageKind, ): string { switch (kind.tag) { - case "TunnelAck": - return "TunnelAck"; + case "DeprecatedTunnelAck": + return "DeprecatedTunnelAck"; case "ToClientRequestStart": { const { actorId, method, path, headers, body, stream } = kind.val; const bodyStr = body === null ? "null" : stringifyArrayBuffer(body); @@ -92,8 +92,8 @@ export function stringifyToClientTunnelMessageKind( return `ToClientWebSocketOpen{actorId: "${actorId}", path: "${path}", headers: ${stringifyMap(headers)}}`; } case "ToClientWebSocketMessage": { - const { index, data, binary } = kind.val; - return `ToClientWebSocketMessage{index: ${index}, data: ${stringifyArrayBuffer(data)}, binary: ${binary}}`; + const { data, binary } = kind.val; + return `ToClientWebSocketMessage{data: ${stringifyArrayBuffer(data)}, binary: ${binary}}`; } case "ToClientWebSocketClose": { const { code, reason } = kind.val; diff --git a/engine/sdks/typescript/runner/src/tunnel-id.ts b/engine/sdks/typescript/runner/src/tunnel-id.ts new file mode 100644 index 0000000000..bea04a2453 --- /dev/null +++ b/engine/sdks/typescript/runner/src/tunnel-id.ts @@ -0,0 +1,104 @@ +import * as protocol from "@rivetkit/engine-runner-protocol"; + +// Type aliases for the message ID components +export type GatewayId = ArrayBuffer; +export type RequestId = ArrayBuffer; +export type MessageIndex = number; +export type MessageId = ArrayBuffer; + +/** + * Build a MessageId from its components + */ +export function buildMessageId( + gatewayId: GatewayId, + requestId: RequestId, + messageIndex: MessageIndex, +): MessageId { + if (gatewayId.byteLength !== 4) { + throw new Error( + `invalid gateway id length: expected 4 bytes, got ${gatewayId.byteLength}`, + ); + } + if (requestId.byteLength !== 4) { + throw new Error( + `invalid request id length: expected 4 bytes, got ${requestId.byteLength}`, + ); + } + if (messageIndex < 0 || messageIndex > 0xffff) { + throw new Error( + `invalid message index: must be u16, got ${messageIndex}`, + ); + } + + const parts: protocol.MessageIdParts = { + gatewayId, + requestId, + messageIndex, + }; + + const encoded = protocol.encodeMessageIdParts(parts); + + if (encoded.byteLength !== 12) { + throw new Error( + `message id serialization produced wrong size: expected 12 bytes, got ${encoded.byteLength}`, + ); + } + + // Create a new ArrayBuffer from the Uint8Array + const messageId = new ArrayBuffer(12); + new Uint8Array(messageId).set(encoded); + return messageId; +} + +/** + * Parse a MessageId into its components + */ +export function parseMessageId(messageId: MessageId): protocol.MessageIdParts { + if (messageId.byteLength !== 12) { + throw new Error( + `invalid message id length: expected 12 bytes, got ${messageId.byteLength}`, + ); + } + const uint8Array = new Uint8Array(messageId); + return protocol.decodeMessageIdParts(uint8Array); +} + +/** + * Convert a GatewayId to a base64 string + */ +export function gatewayIdToString(gatewayId: GatewayId): string { + const uint8Array = new Uint8Array(gatewayId); + return bufferToBase64(uint8Array); +} + +/** + * Convert a RequestId to a base64 string + */ +export function requestIdToString(requestId: RequestId): string { + const uint8Array = new Uint8Array(requestId); + return bufferToBase64(uint8Array); +} + +/** + * Convert a MessageId to a base64 string + */ +export function messageIdToString(messageId: MessageId): string { + const uint8Array = new Uint8Array(messageId); + return bufferToBase64(uint8Array); +} + +// Helper functions for base64 encoding/decoding + +function bufferToBase64(buffer: Uint8Array): string { + // Use Node.js Buffer if available, otherwise use browser btoa + if (typeof Buffer !== "undefined") { + return Buffer.from(buffer).toString("base64"); + } else { + // Browser environment + let binary = ""; + for (let i = 0; i < buffer.byteLength; i++) { + binary += String.fromCharCode(buffer[i]); + } + return btoa(binary); + } +} diff --git a/engine/sdks/typescript/runner/src/tunnel.ts b/engine/sdks/typescript/runner/src/tunnel.ts index f71c9cb89e..16d836003d 100644 --- a/engine/sdks/typescript/runner/src/tunnel.ts +++ b/engine/sdks/typescript/runner/src/tunnel.ts @@ -1,5 +1,9 @@ import type * as protocol from "@rivetkit/engine-runner-protocol"; -import type { MessageId, RequestId } from "@rivetkit/engine-runner-protocol"; +import type { + MessageId, + RequestId, + GatewayId, +} from "@rivetkit/engine-runner-protocol"; import type { Logger } from "pino"; import { parse as uuidparse, @@ -11,32 +15,31 @@ import { stringifyToClientTunnelMessageKind, stringifyToServerTunnelMessageKind, } from "./stringify"; -import { unreachable } from "./utils"; +import * as tunnelId from "./tunnel-id"; +import { arraysEqual, unreachable } from "./utils"; import { HIBERNATABLE_SYMBOL, WebSocketTunnelAdapter, } from "./websocket-tunnel-adapter"; -const GC_INTERVAL = 60000; // 60 seconds -const MESSAGE_ACK_TIMEOUT = 5000; // 5 seconds - export interface PendingRequest { resolve: (response: Response) => void; reject: (error: Error) => void; streamController?: ReadableStreamDefaultController; actorId?: string; + gatewayId?: GatewayId; + requestId?: RequestId; + clientMessageIndex: number; } export interface HibernatingWebSocketMetadata { + gatewayId: GatewayId; requestId: RequestId; + clientMessageIndex: number; + serverMessageIndex: number; + path: string; headers: Record; - messageIndex: number; -} - -export interface PendingTunnelMessage { - sentAt: number; - requestIdStr: string; } class RunnerShutdownError extends Error { @@ -49,9 +52,11 @@ export class Tunnel { #runner: Runner; /** Maps request IDs to actor IDs for lookup */ - #requestToActor: Map = new Map(); - - #gcInterval?: NodeJS.Timeout; + #requestToActor: Array<{ + gatewayId: GatewayId; + requestId: RequestId; + actorId: string; + }> = []; get log(): Logger | undefined { return this.#runner.log; @@ -62,79 +67,68 @@ export class Tunnel { } start(): void { - this.#startGarbageCollector(); + // No-op - kept for compatibility } shutdown() { // NOTE: Pegboard WS already closed at this point, cannot send // anything. All teardown logic is handled by pegboard-runner. - if (this.#gcInterval) { - clearInterval(this.#gcInterval); - this.#gcInterval = undefined; - } - // Reject all pending requests and close all WebSockets for all actors // RunnerShutdownError will be explicitly ignored for (const [_actorId, actor] of this.#runner.actors) { // Reject all pending requests for this actor - for (const [_, request] of actor.pendingRequests) { - request.reject(new RunnerShutdownError()); + for (const entry of actor.pendingRequests) { + entry.request.reject(new RunnerShutdownError()); } - actor.pendingRequests.clear(); + actor.pendingRequests = []; // Close all WebSockets for this actor // The WebSocket close event with retry is automatically sent when the // runner WS closes, so we only need to notify the client that the WS // closed: // https://github.com/rivet-dev/rivet/blob/00d4f6a22da178a6f8115e5db50d96c6f8387c2e/engine/packages/pegboard-runner/src/lib.rs#L157 - for (const [_, ws] of actor.webSockets) { + for (const entry of actor.webSockets) { // Only close non-hibernatable websockets to prevent sending // unnecessary close messages for websockets that will be hibernated - if (!ws[HIBERNATABLE_SYMBOL]) { - ws._closeWithoutCallback(1000, "ws.tunnel_shutdown"); + if (!entry.ws[HIBERNATABLE_SYMBOL]) { + entry.ws._closeWithoutCallback(1000, "ws.tunnel_shutdown"); } } - actor.webSockets.clear(); + actor.webSockets = []; } // Clear the request-to-actor mapping - this.#requestToActor.clear(); + this.#requestToActor = []; } async restoreHibernatingRequests( actorId: string, - requestIds: readonly RequestId[], + hibernatingRequests: readonly protocol.HibernatingRequest[], ) { this.log?.debug({ msg: "restoring hibernating requests", actorId, - requests: requestIds.length, + requests: hibernatingRequests.length, }); // Load all persisted metadata const metaEntries = await this.#runner.config.hibernatableWebSocket.loadAll(actorId); - // Create maps for efficient lookup - const requestIdMap = new Map(); - for (const requestId of requestIds) { - requestIdMap.set(idToStr(requestId), requestId); - } - - const metaMap = new Map(); - for (const meta of metaEntries) { - metaMap.set(idToStr(meta.requestId), meta); - } - // Track all background operations const backgroundOperations: Promise[] = []; // Process connected WebSockets let connectedButNotLoadedCount = 0; let restoredCount = 0; - for (const [requestIdStr, requestId] of requestIdMap) { - const meta = metaMap.get(requestIdStr); + for (const { gatewayId, requestId } of hibernatingRequests) { + const requestIdStr = tunnelId.requestIdToString(requestId); + const meta = metaEntries.find( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + ); if (!meta) { // Connected but not loaded (not persisted) - close it @@ -145,7 +139,7 @@ export class Tunnel { requestId: requestIdStr, }); - this.#sendMessage(requestId, { + this.#sendMessage(gatewayId, requestId, { tag: "ToServerWebSocketClose", val: { code: 1000, @@ -167,17 +161,35 @@ export class Tunnel { // Track this operation to ensure it completes const restoreOperation = this.#createWebSocket( actorId, + gatewayId, requestId, requestIdStr, + meta.serverMessageIndex, true, true, - meta.messageIndex, request, meta.path, meta.headers, false, ) .then(() => { + // Create a PendingRequest entry to track the message index + const actor = this.#runner.getActor(actorId); + if (actor) { + actor.pendingRequests.push({ + gatewayId, + requestId, + request: { + resolve: () => {}, + reject: () => {}, + actorId: actorId, + gatewayId: gatewayId, + requestId: requestId, + clientMessageIndex: meta.clientMessageIndex, + }, + }); + } + this.log?.info({ msg: "connection successfully restored", actorId, @@ -192,7 +204,7 @@ export class Tunnel { }); // Close the WebSocket on error - this.#sendMessage(requestId, { + this.#sendMessage(gatewayId, requestId, { tag: "ToServerWebSocketClose", val: { code: 1011, @@ -209,8 +221,14 @@ export class Tunnel { // Process loaded but not connected (stale) - remove them let loadedButNotConnectedCount = 0; - for (const [requestIdStr, meta] of metaMap) { - if (!requestIdMap.has(requestIdStr)) { + for (const meta of metaEntries) { + const requestIdStr = tunnelId.requestIdToString(meta.requestId); + const isConnected = hibernatingRequests.some( + (req) => + arraysEqual(req.gatewayId, meta.gatewayId) && + arraysEqual(req.requestId, meta.requestId), + ); + if (!isConnected) { this.log?.warn({ msg: "removing stale persisted websocket", requestId: requestIdStr, @@ -226,11 +244,12 @@ export class Tunnel { // Track this operation to ensure it completes const cleanupOperation = this.#createWebSocket( actorId, + meta.gatewayId, meta.requestId, requestIdStr, + meta.serverMessageIndex, true, true, - meta.messageIndex, request, meta.path, meta.headers, @@ -276,11 +295,12 @@ export class Tunnel { */ async #createWebSocket( actorId: string, + gatewayId: GatewayId, requestId: RequestId, requestIdStr: string, + serverMessageIndex: number, isHibernatable: boolean, isRestoringHibernatable: boolean, - messageIndex: number, request: Request, path: string, headers: Record, @@ -298,8 +318,8 @@ export class Tunnel { this, actorId, requestIdStr, + serverMessageIndex, isHibernatable, - messageIndex, isRestoringHibernatable, request, (data: ArrayBuffer | string, isBinary: boolean) => { @@ -309,7 +329,7 @@ export class Tunnel { ? (new TextEncoder().encode(data).buffer as ArrayBuffer) : data; - this.#sendMessage(requestId, { + this.#sendMessage(gatewayId, requestId, { tag: "ToServerWebSocketMessage", val: { data: dataBuffer, @@ -320,7 +340,7 @@ export class Tunnel { (code?: number, reason?: string) => { // Send close through tunnel if engine doesn't already know it's closed if (!engineAlreadyClosed) { - this.#sendMessage(requestId, { + this.#sendMessage(gatewayId, requestId, { tag: "ToServerWebSocketClose", val: { code: code || null, @@ -333,11 +353,12 @@ export class Tunnel { // Clean up actor tracking const actor = this.#runner.getActor(actorId); if (actor) { - actor.webSockets.delete(requestIdStr); + actor.deleteWebSocket(gatewayId, requestId); + actor.deletePendingRequest(gatewayId, requestId); } // Clean up request-to-actor mapping - this.#requestToActor.delete(requestIdStr); + this.#removeRequestToActor(gatewayId, requestId); }, ); @@ -347,8 +368,8 @@ export class Tunnel { throw new Error(`Actor ${actorId} not found`); } - actor.webSockets.set(requestIdStr, adapter); - this.#requestToActor.set(requestIdStr, actorId); + actor.setWebSocket(gatewayId, requestId, adapter); + this.#addRequestToActor(gatewayId, requestId, actorId); // Call WebSocket handler. This handler will add event listeners // for `open`, etc. @@ -356,6 +377,7 @@ export class Tunnel { this.#runner, actorId, adapter, + gatewayId, requestId, request, path, @@ -367,22 +389,49 @@ export class Tunnel { return adapter; } - getRequestActor(requestIdStr: string): RunnerActor | undefined { - const actorId = this.#requestToActor.get(requestIdStr); - if (!actorId) { + #addRequestToActor( + gatewayId: GatewayId, + requestId: RequestId, + actorId: string, + ) { + this.#requestToActor.push({ gatewayId, requestId, actorId }); + } + + #removeRequestToActor(gatewayId: GatewayId, requestId: RequestId) { + const index = this.#requestToActor.findIndex( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + ); + if (index !== -1) { + this.#requestToActor.splice(index, 1); + } + } + + getRequestActor( + gatewayId: GatewayId, + requestId: RequestId, + ): RunnerActor | undefined { + const entry = this.#requestToActor.find( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + ); + + if (!entry) { this.log?.warn({ msg: "missing requestToActor entry", - requestId: requestIdStr, + requestId: tunnelId.requestIdToString(requestId), }); return undefined; } - const actor = this.#runner.getActor(actorId); + const actor = this.#runner.getActor(entry.actorId); if (!actor) { this.log?.warn({ msg: "missing actor for requestToActor lookup", - requestId: requestIdStr, - actorId, + requestId: tunnelId.requestIdToString(requestId), + actorId: entry.actorId, }); return undefined; } @@ -391,6 +440,7 @@ export class Tunnel { } #sendMessage( + gatewayId: GatewayId, requestId: RequestId, messageKind: protocol.ToServerTunnelMessageKind, ) { @@ -398,27 +448,45 @@ export class Tunnel { if (!this.#runner.__webSocketReady()) { this.log?.warn({ msg: "cannot send tunnel message, socket not connected to engine. tunnel data dropped.", - requestId: idToStr(requestId), + requestId: tunnelId.requestIdToString(requestId), message: stringifyToServerTunnelMessageKind(messageKind), }); return; } - // Build message - const messageId = generateUuidBuffer(); - - const requestIdStr = idToStr(requestId); - const messageIdStr = idToStr(messageId); + // Get or initialize message index for this request + const requestIdStr = tunnelId.requestIdToString(requestId); + const actor = this.getRequestActor(gatewayId, requestId); + if (!actor) { + this.log?.warn({ + msg: "cannot send tunnel message, actor not found", + requestId: requestIdStr, + }); + return; + } - // Store the pending message in the actor's map - const actor = this.getRequestActor(requestIdStr); - if (actor) { - actor.pendingTunnelMessages.set(messageIdStr, { - sentAt: Date.now(), - requestIdStr, + // Get message index from pending request + let clientMessageIndex: number; + const pending = actor.getPendingRequest(gatewayId, requestId); + if (pending) { + clientMessageIndex = pending.clientMessageIndex; + pending.clientMessageIndex++; + } else { + // No pending request + this.log?.warn({ + msg: "missing pending request for send message, defaulting to message index 0", }); + clientMessageIndex = 0; } + // Build message ID from gatewayId + requestId + messageIndex + const messageId = tunnelId.buildMessageId( + gatewayId, + requestId, + clientMessageIndex, + ); + const messageIdStr = tunnelId.messageIdToString(messageId); + this.log?.debug({ msg: "send tunnel msg", requestId: requestIdStr, @@ -430,7 +498,6 @@ export class Tunnel { const message: protocol.ToServer = { tag: "ToServerTunnelMessage", val: { - requestId, messageId, messageKind, }, @@ -438,110 +505,36 @@ export class Tunnel { this.#runner.__sendToServer(message); } - #startGarbageCollector() { - if (this.#gcInterval) { - clearInterval(this.#gcInterval); - } - - this.#gcInterval = setInterval(() => { - this.#gc(); - }, GC_INTERVAL); - } - - #gc() { - const now = Date.now(); - let totalMessagesToDelete = 0; - - // Iterate through all actors - for (const [_actorId, actor] of this.#runner.actors) { - const messagesToDelete: string[] = []; - - for (const [ - messageId, - pendingMessage, - ] of actor.pendingTunnelMessages) { - // 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 = - actor.pendingRequests.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 pendingRequests map - actor.pendingRequests.delete(requestIdStr); - } - - // Check if this is a WebSocket - const webSocket = actor.webSockets.get(requestIdStr); - if (webSocket) { - // Close the WebSocket connection - webSocket.close(1000, "ws.ack_timeout"); - - // Clean up from webSockets map - actor.webSockets.delete(requestIdStr); - } - - // Clean up request-to-actor mapping - this.#requestToActor.delete(requestIdStr); - } - } - - // Remove timed out messages for this actor - for (const messageId of messagesToDelete) { - actor.pendingTunnelMessages.delete(messageId); - } - - totalMessagesToDelete += messagesToDelete.length; - } - - // Log if we purged any messages - if (totalMessagesToDelete > 0) { - this.log?.warn({ - msg: "purging unacked tunnel messages, this indicates that the Rivet Engine is disconnected or not responding", - count: totalMessagesToDelete, - }); - } - } - closeActiveRequests(actor: RunnerActor) { const actorId = actor.actorId; // Terminate all requests for this actor. This will no send a // ToServerResponse* message since the actor will no longer be loaded. // The gateway is responsible for closing the request. - for (const [requestIdStr, pending] of actor.pendingRequests) { - pending.reject(new Error(`Actor ${actorId} stopped`)); - this.#requestToActor.delete(requestIdStr); + for (const entry of actor.pendingRequests) { + entry.request.reject(new Error(`Actor ${actorId} stopped`)); + if (entry.gatewayId && entry.requestId) { + this.#removeRequestToActor( + entry.gatewayId, + entry.requestId, + ); + } } // Close all WebSockets. Only send close event to non-HWS. The gateway is // responsible for hibernating HWS and closing regular WS. - for (const [requestIdStr, ws] of actor.webSockets) { - const isHibernatable = ws[HIBERNATABLE_SYMBOL]; + for (const entry of actor.webSockets) { + const isHibernatable = entry.ws[HIBERNATABLE_SYMBOL]; if (!isHibernatable) { - ws._closeWithoutCallback(1000, "actor.stopped"); + entry.ws._closeWithoutCallback(1000, "actor.stopped"); } - this.#requestToActor.delete(requestIdStr); + // Note: request-to-actor mapping is cleaned up in the close callback } } async #fetch( actorId: string, + gatewayId: protocol.GatewayId, requestId: protocol.RequestId, request: Request, ): Promise { @@ -565,6 +558,7 @@ export class Tunnel { const fetchHandler = this.#runner.config.fetch( this.#runner, actorId, + gatewayId, requestId, request, ); @@ -577,8 +571,13 @@ export class Tunnel { } async handleTunnelMessage(message: protocol.ToClientTunnelMessage) { - const requestIdStr = idToStr(message.requestId); - const messageIdStr = idToStr(message.messageId); + // Parse the gateway ID, request ID, and message index from the messageId + const messageIdParts = tunnelId.parseMessageId(message.messageId); + const gatewayId = messageIdParts.gatewayId; + const requestId = messageIdParts.requestId; + + const requestIdStr = tunnelId.requestIdToString(requestId); + const messageIdStr = tunnelId.messageIdToString(message.messageId); this.log?.debug({ msg: "receive tunnel msg", requestId: requestIdStr, @@ -589,49 +588,59 @@ export class Tunnel { switch (message.messageKind.tag) { case "ToClientRequestStart": await this.#handleRequestStart( - message.requestId, + gatewayId, + requestId, message.messageKind.val, ); break; case "ToClientRequestChunk": await this.#handleRequestChunk( - message.requestId, + gatewayId, + requestId, message.messageKind.val, ); break; case "ToClientRequestAbort": - await this.#handleRequestAbort(message.requestId); + await this.#handleRequestAbort(gatewayId, requestId); break; case "ToClientWebSocketOpen": await this.#handleWebSocketOpen( - message.requestId, + gatewayId, + requestId, message.messageKind.val, ); break; case "ToClientWebSocketMessage": { this.#handleWebSocketMessage( - message.requestId, + gatewayId, + requestId, + messageIdParts.messageIndex, message.messageKind.val, ); break; } case "ToClientWebSocketClose": await this.#handleWebSocketClose( - message.requestId, + gatewayId, + requestId, message.messageKind.val, ); break; + case "DeprecatedTunnelAck": + // Ignore deprecated tunnel ack messages + break; default: unreachable(message.messageKind); } } async #handleRequestStart( - requestId: ArrayBuffer, + gatewayId: GatewayId, + requestId: RequestId, req: protocol.ToClientRequestStart, ) { // Track this request for the actor - const requestIdStr = idToStr(requestId); + const requestIdStr = tunnelId.requestIdToString(requestId); const actor = this.#runner.getActor(req.actorId); if (!actor) { this.log?.warn({ @@ -643,7 +652,7 @@ export class Tunnel { } // Add to request-to-actor mapping - this.#requestToActor.set(requestIdStr, req.actorId); + this.#addRequestToActor(gatewayId, requestId, req.actorId); try { // Convert headers map to Headers object @@ -666,16 +675,21 @@ export class Tunnel { start: (controller) => { // Store controller for chunks const existing = - actor.pendingRequests.get(requestIdStr); + actor.getPendingRequest(gatewayId, requestId); if (existing) { existing.streamController = controller; existing.actorId = req.actorId; + existing.gatewayId = gatewayId; + existing.requestId = requestId; } else { - actor.pendingRequests.set(requestIdStr, { + actor.setPendingRequest(gatewayId, requestId, { resolve: () => {}, reject: () => {}, streamController: controller, actorId: req.actorId, + gatewayId: gatewayId, + requestId: requestId, + clientMessageIndex: 0, }); } }, @@ -690,25 +704,39 @@ export class Tunnel { // Call fetch handler with validation const response = await this.#fetch( req.actorId, + gatewayId, requestId, streamingRequest, ); await this.#sendResponse( actor.actorId, actor.generation, + gatewayId, requestId, response, ); } else { // Non-streaming request + // Create a pending request entry to track messageIndex for the response + actor.setPendingRequest(gatewayId, requestId, { + resolve: () => {}, + reject: () => {}, + actorId: req.actorId, + gatewayId: gatewayId, + requestId: requestId, + clientMessageIndex: 0, + }); + const response = await this.#fetch( req.actorId, + gatewayId, requestId, request, ); await this.#sendResponse( actor.actorId, actor.generation, + gatewayId, requestId, response, ); @@ -721,6 +749,7 @@ export class Tunnel { this.#sendResponseError( actor.actorId, actor.generation, + gatewayId, requestId, 500, "Internal Server Error", @@ -729,47 +758,49 @@ export class Tunnel { } finally { // Clean up request tracking if (this.#runner.hasActor(req.actorId, actor.generation)) { - actor.pendingRequests.delete(requestIdStr); - this.#requestToActor.delete(requestIdStr); + actor.deletePendingRequest(gatewayId, requestId); + this.#removeRequestToActor(gatewayId, requestId); } } } async #handleRequestChunk( - requestId: ArrayBuffer, + gatewayId: GatewayId, + requestId: RequestId, chunk: protocol.ToClientRequestChunk, ) { - const requestIdStr = idToStr(requestId); - const actor = this.getRequestActor(requestIdStr); + const requestIdStr = tunnelId.requestIdToString(requestId); + const actor = this.getRequestActor(gatewayId, requestId); if (actor) { - const pending = actor.pendingRequests.get(requestIdStr); + const pending = actor.getPendingRequest(gatewayId, requestId); if (pending?.streamController) { pending.streamController.enqueue(new Uint8Array(chunk.body)); if (chunk.finish) { pending.streamController.close(); - actor.pendingRequests.delete(requestIdStr); - this.#requestToActor.delete(requestIdStr); + actor.deletePendingRequest(gatewayId, requestId); + this.#removeRequestToActor(gatewayId, requestId); } } } } - async #handleRequestAbort(requestId: ArrayBuffer) { - const requestIdStr = idToStr(requestId); - const actor = this.getRequestActor(requestIdStr); + async #handleRequestAbort(gatewayId: GatewayId, requestId: RequestId) { + const requestIdStr = tunnelId.requestIdToString(requestId); + const actor = this.getRequestActor(gatewayId, requestId); if (actor) { - const pending = actor.pendingRequests.get(requestIdStr); + const pending = actor.getPendingRequest(gatewayId, requestId); if (pending?.streamController) { pending.streamController.error(new Error("Request aborted")); } - actor.pendingRequests.delete(requestIdStr); - this.#requestToActor.delete(requestIdStr); + actor.deletePendingRequest(gatewayId, requestId); + this.#removeRequestToActor(gatewayId, requestId); } } async #sendResponse( actorId: string, generation: number, + gatewayId: GatewayId, requestId: ArrayBuffer, response: Response, ) { @@ -804,7 +835,7 @@ export class Tunnel { } // Send as non-streaming response if actor has not stopped - this.#sendMessage(requestId, { + this.#sendMessage(gatewayId, requestId, { tag: "ToServerResponseStart", val: { status: response.status as protocol.u16, @@ -818,6 +849,7 @@ export class Tunnel { #sendResponseError( actorId: string, generation: number, + gatewayId: GatewayId, requestId: ArrayBuffer, status: number, message: string, @@ -835,7 +867,7 @@ export class Tunnel { const headers = new Map(); headers.set("content-type", "text/plain"); - this.#sendMessage(requestId, { + this.#sendMessage(gatewayId, requestId, { tag: "ToServerResponseStart", val: { status: status as protocol.u16, @@ -847,16 +879,17 @@ export class Tunnel { } async #handleWebSocketOpen( - requestId: protocol.RequestId, + gatewayId: GatewayId, + requestId: RequestId, open: protocol.ToClientWebSocketOpen, ) { // NOTE: This method is safe to be async since we will not receive any // further WebSocket events until we send a ToServerWebSocketOpen // tunnel message. We can do any async logic we need to between thoes two events. // - // Sedning a ToServerWebSocketClose will terminate the WebSocket early. + // Sending a ToServerWebSocketClose will terminate the WebSocket early. - const requestIdStr = idToStr(requestId); + const requestIdStr = tunnelId.requestIdToString(requestId); // Validate actor exists const actor = this.#runner.getActor(open.actorId); @@ -871,7 +904,7 @@ export class Tunnel { // // See // https://github.com/rivet-dev/rivet/blob/222dae87e3efccaffa2b503de40ecf8afd4e31eb/engine/packages/pegboard-gateway/src/lib.rs#L238 - this.#sendMessage(requestId, { + this.#sendMessage(gatewayId, requestId, { tag: "ToServerWebSocketClose", val: { code: 1011, @@ -885,7 +918,7 @@ export class Tunnel { // Close existing WebSocket if one already exists for this request ID. // This should never happen, but prevents any potential duplicate // WebSockets from retransmits. - const existingAdapter = actor.webSockets.get(requestIdStr); + const existingAdapter = actor.getWebSocket(gatewayId, requestId); if (existingAdapter) { this.log?.warn({ msg: "closing existing websocket for duplicate open event for the same request id", @@ -906,6 +939,7 @@ export class Tunnel { const canHibernate = this.#runner.config.hibernatableWebSocket.canHibernate( actor.actorId, + gatewayId, requestId, request, ); @@ -916,21 +950,32 @@ export class Tunnel { // open event. const adapter = await this.#createWebSocket( actor.actorId, + gatewayId, requestId, requestIdStr, + 0, canHibernate, false, - 0, request, open.path, Object.fromEntries(open.headers), false, ); + // Create a PendingRequest entry to track the message index + actor.setPendingRequest(gatewayId, requestId, { + resolve: () => {}, + reject: () => {}, + actorId: actor.actorId, + gatewayId: gatewayId, + requestId: requestId, + clientMessageIndex: 0, + }); + // Open the WebSocket after `config.socket` so (a) the event // handlers can be added and (b) any errors in `config.websocket` // will cause the WebSocket to terminate before the open event. - this.#sendMessage(requestId, { + this.#sendMessage(gatewayId, requestId, { tag: "ToServerWebSocketOpen", val: { canHibernate, @@ -945,7 +990,7 @@ export class Tunnel { // TODO: Call close event on adapter if needed // Send close on error - this.#sendMessage(requestId, { + this.#sendMessage(gatewayId, requestId, { tag: "ToServerWebSocketClose", val: { code: 1011, @@ -955,28 +1000,36 @@ export class Tunnel { }); // Clean up actor tracking - actor.webSockets.delete(requestIdStr); - this.#requestToActor.delete(requestIdStr); + actor.deleteWebSocket(gatewayId, requestId); + actor.deletePendingRequest(gatewayId, requestId); + this.#removeRequestToActor(gatewayId, requestId); } } #handleWebSocketMessage( - requestId: ArrayBuffer, + gatewayId: GatewayId, + requestId: RequestId, + serverMessageIndex: number, msg: protocol.ToClientWebSocketMessage, ) { // NOTE: This method cannot be async in order to ensure in-order // message processing. - const requestIdStr = idToStr(requestId); - const actor = this.getRequestActor(requestIdStr); + const requestIdStr = tunnelId.requestIdToString(requestId); + const actor = this.getRequestActor(gatewayId, requestId); if (actor) { - const adapter = actor.webSockets.get(requestIdStr); + const adapter = actor.getWebSocket(gatewayId, requestId); if (adapter) { const data = msg.binary ? new Uint8Array(msg.data) : new TextDecoder().decode(new Uint8Array(msg.data)); - adapter._handleMessage(requestId, data, msg.index, msg.binary); + adapter._handleMessage( + requestId, + data, + serverMessageIndex, + msg.binary, + ); return; } } @@ -988,33 +1041,60 @@ export class Tunnel { }); } - sendHibernatableWebSocketMessageAck(requestId: ArrayBuffer, index: number) { + sendHibernatableWebSocketMessageAck( + gatewayId: ArrayBuffer, + requestId: ArrayBuffer, + clientMessageIndex: number, + ) { + const requestIdStr = tunnelId.requestIdToString(requestId); + this.log?.debug({ msg: "ack ws msg", - requestId: idToStr(requestId), - index, + requestId: requestIdStr, + index: clientMessageIndex, }); - if (index < 0 || index > 65535) + if (clientMessageIndex < 0 || clientMessageIndex > 65535) throw new Error("invalid websocket ack index"); + // Get the actor to find the gatewayId + const actor = this.getRequestActor(gatewayId, requestId); + if (!actor) { + this.log?.warn({ + msg: "cannot send websocket ack, actor not found", + requestId: requestIdStr, + }); + return; + } + + // Get gatewayId from the pending request + const pending = actor.getPendingRequest(gatewayId, requestId); + if (!pending?.gatewayId) { + this.log?.warn({ + msg: "cannot send websocket ack, gatewayId not found in pending request", + requestId: requestIdStr, + }); + return; + } + // Send the ack message - this.#sendMessage(requestId, { + this.#sendMessage(pending.gatewayId, requestId, { tag: "ToServerWebSocketMessageAck", val: { - index, + index: clientMessageIndex, }, }); } async #handleWebSocketClose( - requestId: ArrayBuffer, + gatewayId: GatewayId, + requestId: RequestId, close: protocol.ToClientWebSocketClose, ) { - const requestIdStr = idToStr(requestId); - const actor = this.getRequestActor(requestIdStr); + const requestIdStr = tunnelId.requestIdToString(requestId); + const actor = this.getRequestActor(gatewayId, requestId); if (actor) { - const adapter = actor.webSockets.get(requestIdStr); + const adapter = actor.getWebSocket(gatewayId, requestId); if (adapter) { // We don't need to send a close response adapter._handleClose( @@ -1022,24 +1102,14 @@ export class Tunnel { close.code || undefined, close.reason || undefined, ); - actor.webSockets.delete(requestIdStr); - this.#requestToActor.delete(requestIdStr); + actor.deleteWebSocket(gatewayId, requestId); + actor.deletePendingRequest(gatewayId, requestId); + this.#removeRequestToActor(gatewayId, requestId); } } } } -/** Generates a UUID as bytes. */ -function generateUuidBuffer(): ArrayBuffer { - const buffer = new Uint8Array(16); - uuidv4(undefined, buffer); - return buffer.buffer; -} - -function idToStr(id: ArrayBuffer): string { - return uuidstringify(new Uint8Array(id)); -} - /** * Builds a request that represents the incoming request for a given WebSocket. * diff --git a/engine/sdks/typescript/runner/src/utils.ts b/engine/sdks/typescript/runner/src/utils.ts index c6a9c5e7b3..f9c1278f3e 100644 --- a/engine/sdks/typescript/runner/src/utils.ts +++ b/engine/sdks/typescript/runner/src/utils.ts @@ -121,3 +121,13 @@ function wrappingSub(a: number, b: number, max: number): number { } return result; } + +export function arraysEqual(a: ArrayBuffer, b: ArrayBuffer): boolean { + const ua = new Uint8Array(a); + const ub = new Uint8Array(b); + if (ua.length !== ub.length) return false; + for (let i = 0; i < ua.length; i++) { + if (ua[i] !== ub[i]) return false; + } + return true; +} diff --git a/engine/sdks/typescript/runner/src/websocket-tunnel-adapter.ts b/engine/sdks/typescript/runner/src/websocket-tunnel-adapter.ts index 3eb7b6c068..a2eec4dd23 100644 --- a/engine/sdks/typescript/runner/src/websocket-tunnel-adapter.ts +++ b/engine/sdks/typescript/runner/src/websocket-tunnel-adapter.ts @@ -27,7 +27,7 @@ export class WebSocketTunnelAdapter { #actorId: string; #requestId: string; #hibernatable: boolean; - #messageIndex: number; + #serverMessageIndex: number; get [HIBERNATABLE_SYMBOL](): boolean { return this.#hibernatable; @@ -59,8 +59,8 @@ export class WebSocketTunnelAdapter { tunnel: Tunnel, actorId: string, requestId: string, + serverMessageIndex: number, hibernatable: boolean, - messageIndex: number, isRestoringHibernatable: boolean, /** @experimental */ public readonly request: Request, @@ -71,7 +71,7 @@ export class WebSocketTunnelAdapter { this.#actorId = actorId; this.#requestId = requestId; this.#hibernatable = hibernatable; - this.#messageIndex = messageIndex; + this.#serverMessageIndex = serverMessageIndex; this.#sendCallback = sendCallback; this.#closeCallback = closeCallback; @@ -112,7 +112,7 @@ export class WebSocketTunnelAdapter { _handleMessage( requestId: ArrayBuffer, data: string | Uint8Array, - messageIndex: number, + serverMessageIndex: number, isBinary: boolean, ): boolean { if (this.#readyState !== 1) { @@ -122,7 +122,7 @@ export class WebSocketTunnelAdapter { actorId: this.#actorId, currentReadyState: this.#readyState, expectedReadyState: 1, - messageIndex, + serverMessageIndex: serverMessageIndex, hibernatable: this.#hibernatable, }); return true; @@ -130,7 +130,7 @@ export class WebSocketTunnelAdapter { // Validate message index if (this.#hibernatable) { - const previousIndex = this.#messageIndex; + const previousIndex = this.#serverMessageIndex; // Ignore duplicate old messages // @@ -140,14 +140,14 @@ export class WebSocketTunnelAdapter { // received by the gateway (due to a crash or network // issue), the gateway will resend all messages from // the last ack on reconnect. - if (wrappingLteU16(messageIndex, previousIndex)) { + if (wrappingLteU16(serverMessageIndex, previousIndex)) { this.#log?.info({ msg: "received duplicate hibernating websocket message, this indicates the actor failed to ack the message index before restarting", requestId, actorId: this.#actorId, previousIndex, expectedIndex: wrappingAddU16(previousIndex, 1), - receivedIndex: messageIndex, + receivedIndex: serverMessageIndex, }); return true; @@ -157,7 +157,7 @@ export class WebSocketTunnelAdapter { // // There is no scenario where this should ever happen const expectedIndex = wrappingAddU16(previousIndex, 1); - if (messageIndex !== expectedIndex) { + if (serverMessageIndex !== expectedIndex) { const closeReason = "ws.message_index_skip"; this.#log?.warn({ @@ -166,10 +166,10 @@ export class WebSocketTunnelAdapter { actorId: this.#actorId, previousIndex, expectedIndex, - receivedIndex: messageIndex, + receivedIndex: serverMessageIndex, closeReason, gap: wrappingSubU16( - wrappingSubU16(messageIndex, previousIndex), + wrappingSubU16(serverMessageIndex, previousIndex), 1, ), }); @@ -181,7 +181,7 @@ export class WebSocketTunnelAdapter { } // Update to the next index - this.#messageIndex = messageIndex; + this.#serverMessageIndex = serverMessageIndex; } // Dispatch event @@ -215,7 +215,7 @@ export class WebSocketTunnelAdapter { type: "message", data: messageData, rivetRequestId: requestId, - rivetMessageIndex: messageIndex, + rivetMessageIndex: serverMessageIndex, target: this, }; diff --git a/rivetkit-asyncapi/asyncapi.json b/rivetkit-asyncapi/asyncapi.json index 0d9c7c7f98..e6074c7754 100644 --- a/rivetkit-asyncapi/asyncapi.json +++ b/rivetkit-asyncapi/asyncapi.json @@ -1,436 +1,489 @@ { - "asyncapi": "3.0.0", - "info": { - "title": "RivetKit WebSocket Protocol", - "version": "2.0.24-rc.1", - "description": "WebSocket protocol for bidirectional communication between RivetKit clients and actors" - }, - "channels": { - "/gateway/{actorId}/connect": { - "address": "/gateway/{actorId}/connect", - "parameters": { - "actorId": { - "description": "The unique identifier for the actor instance" - } - }, - "messages": { - "toClient": { - "$ref": "#/components/messages/ToClient" - }, - "toServer": { - "$ref": "#/components/messages/ToServer" - } - } - } - }, - "operations": { - "sendToClient": { - "action": "send", - "channel": { - "$ref": "#/channels/~1gateway~1{actorId}~1connect" - }, - "messages": [ - { - "$ref": "#/channels/~1gateway~1{actorId}~1connect/messages/toClient" - } - ], - "summary": "Send messages from server to client", - "description": "Messages sent from the RivetKit actor to connected clients" - }, - "receiveFromClient": { - "action": "receive", - "channel": { - "$ref": "#/channels/~1gateway~1{actorId}~1connect" - }, - "messages": [ - { - "$ref": "#/channels/~1gateway~1{actorId}~1connect/messages/toServer" - } - ], - "summary": "Receive messages from client", - "description": "Messages received by the RivetKit actor from connected clients" - } - }, - "components": { - "messages": { - "ToClient": { - "name": "ToClient", - "title": "Message To Client", - "summary": "A message sent from the server to the client", - "contentType": "application/json", - "payload": { - "type": "object", - "properties": { - "body": { - "anyOf": [ - { - "type": "object", - "properties": { - "tag": { - "type": "string", - "const": "Init" - }, - "val": { - "type": "object", - "properties": { - "actorId": { - "type": "string" - }, - "connectionId": { - "type": "string" - } - }, - "required": [ - "actorId", - "connectionId" - ], - "additionalProperties": false - } - }, - "required": ["tag", "val"], - "additionalProperties": false - }, - { - "type": "object", - "properties": { - "tag": { - "type": "string", - "const": "Error" - }, - "val": { - "type": "object", - "properties": { - "group": { - "type": "string" - }, - "code": { - "type": "string" - }, - "message": { - "type": "string" - }, - "metadata": {}, - "actionId": { - "type": ["integer", "null"] - } - }, - "required": [ - "group", - "code", - "message", - "actionId" - ], - "additionalProperties": false - } - }, - "required": ["tag", "val"], - "additionalProperties": false - }, - { - "type": "object", - "properties": { - "tag": { - "type": "string", - "const": "ActionResponse" - }, - "val": { - "type": "object", - "properties": { - "id": { - "type": "integer", - "format": "int64" - }, - "output": {} - }, - "required": ["id"], - "additionalProperties": false - } - }, - "required": ["tag", "val"], - "additionalProperties": false - }, - { - "type": "object", - "properties": { - "tag": { - "type": "string", - "const": "Event" - }, - "val": { - "type": "object", - "properties": { - "name": { - "type": "string" - }, - "args": {} - }, - "required": ["name"], - "additionalProperties": false - } - }, - "required": ["tag", "val"], - "additionalProperties": false - } - ] - } - }, - "required": ["body"], - "additionalProperties": false - }, - "examples": [ - { - "name": "Init message", - "summary": "Initial connection message", - "payload": { - "body": { - "tag": "Init", - "val": { - "actorId": "actor_123", - "connectionId": "conn_456" - } - } - } - }, - { - "name": "Error message", - "summary": "Error response", - "payload": { - "body": { - "tag": "Error", - "val": { - "group": "auth", - "code": "unauthorized", - "message": "Authentication failed", - "actionId": null - } - } - } - }, - { - "name": "Action response", - "summary": "Response to an action request", - "payload": { - "body": { - "tag": "ActionResponse", - "val": { - "id": "123", - "output": { - "result": "success" - } - } - } - } - }, - { - "name": "Event", - "summary": "Event broadcast to subscribed clients", - "payload": { - "body": { - "tag": "Event", - "val": { - "name": "stateChanged", - "args": { - "newState": "active" - } - } - } - } - } - ] - }, - "ToServer": { - "name": "ToServer", - "title": "Message To Server", - "summary": "A message sent from the client to the server", - "contentType": "application/json", - "payload": { - "type": "object", - "properties": { - "body": { - "anyOf": [ - { - "type": "object", - "properties": { - "tag": { - "type": "string", - "const": "ActionRequest" - }, - "val": { - "type": "object", - "properties": { - "id": { - "type": "integer", - "format": "int64" - }, - "name": { - "type": "string" - }, - "args": {} - }, - "required": ["id", "name"], - "additionalProperties": false - } - }, - "required": ["tag", "val"], - "additionalProperties": false - }, - { - "type": "object", - "properties": { - "tag": { - "type": "string", - "const": "SubscriptionRequest" - }, - "val": { - "type": "object", - "properties": { - "eventName": { - "type": "string" - }, - "subscribe": { - "type": "boolean" - } - }, - "required": [ - "eventName", - "subscribe" - ], - "additionalProperties": false - } - }, - "required": ["tag", "val"], - "additionalProperties": false - } - ] - } - }, - "required": ["body"], - "additionalProperties": false - }, - "examples": [ - { - "name": "Action request", - "summary": "Request to execute an action", - "payload": { - "body": { - "tag": "ActionRequest", - "val": { - "id": "123", - "name": "updateState", - "args": { - "key": "value" - } - } - } - } - }, - { - "name": "Subscription request", - "summary": "Request to subscribe/unsubscribe from an event", - "payload": { - "body": { - "tag": "SubscriptionRequest", - "val": { - "eventName": "stateChanged", - "subscribe": true - } - } - } - } - ] - } - }, - "schemas": { - "Init": { - "type": "object", - "properties": { - "actorId": { - "type": "string" - }, - "connectionId": { - "type": "string" - } - }, - "required": ["actorId", "connectionId"], - "additionalProperties": false, - "description": "Initial connection message sent from server to client" - }, - "Error": { - "type": "object", - "properties": { - "group": { - "type": "string" - }, - "code": { - "type": "string" - }, - "message": { - "type": "string" - }, - "metadata": {}, - "actionId": { - "type": ["integer", "null"] - } - }, - "required": ["group", "code", "message", "actionId"], - "additionalProperties": false, - "description": "Error message sent from server to client" - }, - "ActionResponse": { - "type": "object", - "properties": { - "id": { - "type": "integer", - "format": "int64" - }, - "output": {} - }, - "required": ["id"], - "additionalProperties": false, - "description": "Response to an action request" - }, - "Event": { - "type": "object", - "properties": { - "name": { - "type": "string" - }, - "args": {} - }, - "required": ["name"], - "additionalProperties": false, - "description": "Event broadcast to subscribed clients" - }, - "ActionRequest": { - "type": "object", - "properties": { - "id": { - "type": "integer", - "format": "int64" - }, - "name": { - "type": "string" - }, - "args": {} - }, - "required": ["id", "name"], - "additionalProperties": false, - "description": "Request to execute an action on the actor" - }, - "SubscriptionRequest": { - "type": "object", - "properties": { - "eventName": { - "type": "string" - }, - "subscribe": { - "type": "boolean" - } - }, - "required": ["eventName", "subscribe"], - "additionalProperties": false, - "description": "Request to subscribe or unsubscribe from an event" - } - } - } -} + "asyncapi": "3.0.0", + "info": { + "title": "RivetKit WebSocket Protocol", + "version": "2.0.24-rc.1", + "description": "WebSocket protocol for bidirectional communication between RivetKit clients and actors" + }, + "channels": { + "/gateway/{actorId}/connect": { + "address": "/gateway/{actorId}/connect", + "parameters": { + "actorId": { + "description": "The unique identifier for the actor instance" + } + }, + "messages": { + "toClient": { + "$ref": "#/components/messages/ToClient" + }, + "toServer": { + "$ref": "#/components/messages/ToServer" + } + } + } + }, + "operations": { + "sendToClient": { + "action": "send", + "channel": { + "$ref": "#/channels/~1gateway~1{actorId}~1connect" + }, + "messages": [ + { + "$ref": "#/channels/~1gateway~1{actorId}~1connect/messages/toClient" + } + ], + "summary": "Send messages from server to client", + "description": "Messages sent from the RivetKit actor to connected clients" + }, + "receiveFromClient": { + "action": "receive", + "channel": { + "$ref": "#/channels/~1gateway~1{actorId}~1connect" + }, + "messages": [ + { + "$ref": "#/channels/~1gateway~1{actorId}~1connect/messages/toServer" + } + ], + "summary": "Receive messages from client", + "description": "Messages received by the RivetKit actor from connected clients" + } + }, + "components": { + "messages": { + "ToClient": { + "name": "ToClient", + "title": "Message To Client", + "summary": "A message sent from the server to the client", + "contentType": "application/json", + "payload": { + "type": "object", + "properties": { + "body": { + "anyOf": [ + { + "type": "object", + "properties": { + "tag": { + "type": "string", + "const": "Init" + }, + "val": { + "type": "object", + "properties": { + "actorId": { + "type": "string" + }, + "connectionId": { + "type": "string" + } + }, + "required": [ + "actorId", + "connectionId" + ], + "additionalProperties": false + } + }, + "required": [ + "tag", + "val" + ], + "additionalProperties": false + }, + { + "type": "object", + "properties": { + "tag": { + "type": "string", + "const": "Error" + }, + "val": { + "type": "object", + "properties": { + "group": { + "type": "string" + }, + "code": { + "type": "string" + }, + "message": { + "type": "string" + }, + "metadata": {}, + "actionId": { + "type": [ + "integer", + "null" + ] + } + }, + "required": [ + "group", + "code", + "message", + "actionId" + ], + "additionalProperties": false + } + }, + "required": [ + "tag", + "val" + ], + "additionalProperties": false + }, + { + "type": "object", + "properties": { + "tag": { + "type": "string", + "const": "ActionResponse" + }, + "val": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "format": "int64" + }, + "output": {} + }, + "required": [ + "id" + ], + "additionalProperties": false + } + }, + "required": [ + "tag", + "val" + ], + "additionalProperties": false + }, + { + "type": "object", + "properties": { + "tag": { + "type": "string", + "const": "Event" + }, + "val": { + "type": "object", + "properties": { + "name": { + "type": "string" + }, + "args": {} + }, + "required": [ + "name" + ], + "additionalProperties": false + } + }, + "required": [ + "tag", + "val" + ], + "additionalProperties": false + } + ] + } + }, + "required": [ + "body" + ], + "additionalProperties": false + }, + "examples": [ + { + "name": "Init message", + "summary": "Initial connection message", + "payload": { + "body": { + "tag": "Init", + "val": { + "actorId": "actor_123", + "connectionId": "conn_456" + } + } + } + }, + { + "name": "Error message", + "summary": "Error response", + "payload": { + "body": { + "tag": "Error", + "val": { + "group": "auth", + "code": "unauthorized", + "message": "Authentication failed", + "actionId": null + } + } + } + }, + { + "name": "Action response", + "summary": "Response to an action request", + "payload": { + "body": { + "tag": "ActionResponse", + "val": { + "id": "123", + "output": { + "result": "success" + } + } + } + } + }, + { + "name": "Event", + "summary": "Event broadcast to subscribed clients", + "payload": { + "body": { + "tag": "Event", + "val": { + "name": "stateChanged", + "args": { + "newState": "active" + } + } + } + } + } + ] + }, + "ToServer": { + "name": "ToServer", + "title": "Message To Server", + "summary": "A message sent from the client to the server", + "contentType": "application/json", + "payload": { + "type": "object", + "properties": { + "body": { + "anyOf": [ + { + "type": "object", + "properties": { + "tag": { + "type": "string", + "const": "ActionRequest" + }, + "val": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "format": "int64" + }, + "name": { + "type": "string" + }, + "args": {} + }, + "required": [ + "id", + "name" + ], + "additionalProperties": false + } + }, + "required": [ + "tag", + "val" + ], + "additionalProperties": false + }, + { + "type": "object", + "properties": { + "tag": { + "type": "string", + "const": "SubscriptionRequest" + }, + "val": { + "type": "object", + "properties": { + "eventName": { + "type": "string" + }, + "subscribe": { + "type": "boolean" + } + }, + "required": [ + "eventName", + "subscribe" + ], + "additionalProperties": false + } + }, + "required": [ + "tag", + "val" + ], + "additionalProperties": false + } + ] + } + }, + "required": [ + "body" + ], + "additionalProperties": false + }, + "examples": [ + { + "name": "Action request", + "summary": "Request to execute an action", + "payload": { + "body": { + "tag": "ActionRequest", + "val": { + "id": "123", + "name": "updateState", + "args": { + "key": "value" + } + } + } + } + }, + { + "name": "Subscription request", + "summary": "Request to subscribe/unsubscribe from an event", + "payload": { + "body": { + "tag": "SubscriptionRequest", + "val": { + "eventName": "stateChanged", + "subscribe": true + } + } + } + } + ] + } + }, + "schemas": { + "Init": { + "type": "object", + "properties": { + "actorId": { + "type": "string" + }, + "connectionId": { + "type": "string" + } + }, + "required": [ + "actorId", + "connectionId" + ], + "additionalProperties": false, + "description": "Initial connection message sent from server to client" + }, + "Error": { + "type": "object", + "properties": { + "group": { + "type": "string" + }, + "code": { + "type": "string" + }, + "message": { + "type": "string" + }, + "metadata": {}, + "actionId": { + "type": [ + "integer", + "null" + ] + } + }, + "required": [ + "group", + "code", + "message", + "actionId" + ], + "additionalProperties": false, + "description": "Error message sent from server to client" + }, + "ActionResponse": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "format": "int64" + }, + "output": {} + }, + "required": [ + "id" + ], + "additionalProperties": false, + "description": "Response to an action request" + }, + "Event": { + "type": "object", + "properties": { + "name": { + "type": "string" + }, + "args": {} + }, + "required": [ + "name" + ], + "additionalProperties": false, + "description": "Event broadcast to subscribed clients" + }, + "ActionRequest": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "format": "int64" + }, + "name": { + "type": "string" + }, + "args": {} + }, + "required": [ + "id", + "name" + ], + "additionalProperties": false, + "description": "Request to execute an action on the actor" + }, + "SubscriptionRequest": { + "type": "object", + "properties": { + "eventName": { + "type": "string" + }, + "subscribe": { + "type": "boolean" + } + }, + "required": [ + "eventName", + "subscribe" + ], + "additionalProperties": false, + "description": "Request to subscribe or unsubscribe from an event" + } + } + } +} \ No newline at end of file diff --git a/rivetkit-typescript/packages/rivetkit/schemas/actor-persist/v3.bare b/rivetkit-typescript/packages/rivetkit/schemas/actor-persist/v3.bare index 698121ea90..a2054e7146 100644 --- a/rivetkit-typescript/packages/rivetkit/schemas/actor-persist/v3.bare +++ b/rivetkit-typescript/packages/rivetkit/schemas/actor-persist/v3.bare @@ -1,4 +1,7 @@ -type RequestId data +type GatewayId data[4] +type RequestId data[4] +type MessageIndex u16 + type Cbor data # MARK: Connection @@ -14,10 +17,10 @@ type Conn struct { state: Cbor subscriptions: list - # Request ID of the hibernatable WebSocket - hibernatableRequestId: RequestId - # Last seem message index for this WebSocket - msgIndex: i64 + gatewayId: GatewayId + requestId: RequestId + serverMessageIndex: u16 + clientMessageIndex: u16 requestPath: str requestHeaders: map diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/conn/driver.ts b/rivetkit-typescript/packages/rivetkit/src/actor/conn/driver.ts index 8530841a3e..3543b20b40 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/conn/driver.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/conn/driver.ts @@ -30,19 +30,13 @@ export interface ConnDriver { ): void; }; - /** - * Unique request ID provided by the underlying provider. If none is - * available for this conn driver, a random UUID is generated. - **/ - requestId: string; - - /** ArrayBuffer version of requestId if relevant. */ - requestIdBuf: ArrayBuffer | undefined; - /** * If the connection can be hibernated. If true, this will allow the actor to go to sleep while the connection is still active. **/ - hibernatable: boolean; + hibernatable?: { + gatewayId: ArrayBuffer; + requestId: ArrayBuffer; + }; /** * This returns a promise since we commonly disconnect at the end of a program, and not waiting will cause the socket to not close cleanly. diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/http.ts b/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/http.ts index c8610bd024..3ba0338405 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/http.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/http.ts @@ -5,9 +5,6 @@ export type ConnHttpState = Record; export function createHttpDriver(): ConnDriver { return { type: "http", - requestId: crypto.randomUUID(), - requestIdBuf: undefined, - hibernatable: false, getConnectionReadyState(_actor, _conn) { // TODO: This might not be the correct logic return DriverReadyState.OPEN; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/raw-request.ts b/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/raw-request.ts index 00d525533f..f1165be2ee 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/raw-request.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/raw-request.ts @@ -11,9 +11,6 @@ import { DriverReadyState } from "../driver"; export function createRawRequestDriver(): ConnDriver { return { type: "raw-request", - requestId: crypto.randomUUID(), - requestIdBuf: undefined, - hibernatable: false, disconnect: async () => { // Noop diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/raw-websocket.ts b/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/raw-websocket.ts index cfe907ee02..4a505516b7 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/raw-websocket.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/raw-websocket.ts @@ -13,17 +13,13 @@ import { type ConnDriver, DriverReadyState } from "../driver"; * actor's onWebSocket handler. */ export function createRawWebSocketDriver( - requestId: string, - requestIdBuf: ArrayBuffer | undefined, - hibernatable: boolean, + hibernatable: ConnDriver['hibernatable'], closePromise: Promise, ): { driver: ConnDriver; setWebSocket(ws: UniversalWebSocket): void } { let websocket: UniversalWebSocket | undefined; const driver: ConnDriver = { type: "raw-websocket", - requestId, - requestIdBuf, hibernatable, // No sendMessage implementation since this is a raw WebSocket that doesn't diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/websocket.ts b/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/websocket.ts index 1b6fadc610..d156c86a23 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/websocket.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/websocket.ts @@ -8,15 +8,12 @@ import { type ConnDriver, DriverReadyState } from "../driver"; export type ConnDriverWebSocketState = Record; export function createWebSocketDriver( - requestId: string, - requestIdBuf: ArrayBuffer | undefined, - hibernatable: boolean, + hibernatable: ConnDriver['hibernatable'], encoding: Encoding, closePromise: Promise, ): { driver: ConnDriver; setWebSocket(ws: WSContext): void } { loggerWithoutContext().debug({ msg: "createWebSocketDriver creating driver", - requestId, hibernatable, }); // Wait for WS to open @@ -24,8 +21,6 @@ export function createWebSocketDriver( const driver: ConnDriver = { type: "websocket", - requestId, - requestIdBuf, hibernatable, rivetKitProtocol: { sendMessage: ( diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/conn/mod.ts b/rivetkit-typescript/packages/rivetkit/src/actor/conn/mod.ts index f662fabe71..985ddfa670 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/conn/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/conn/mod.ts @@ -13,10 +13,6 @@ import { CachedSerializer } from "../protocol/serde"; import type { ConnDriver } from "./driver"; import { type ConnDataInput, StateManager } from "./state-manager"; -export function generateConnRequestId(): string { - return crypto.randomUUID(); -} - export type ConnId = string; export type AnyConn = Conn; @@ -112,7 +108,7 @@ export class Conn { * If the underlying connection can hibernate. */ get isHibernatable(): boolean { - return this[CONN_DRIVER_SYMBOL]?.hibernatable ?? false; + return this[CONN_DRIVER_SYMBOL]?.hibernatable !== undefined; } /** diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/conn/persisted.ts b/rivetkit-typescript/packages/rivetkit/src/actor/conn/persisted.ts index 0fed69fd8f..1484495576 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/conn/persisted.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/conn/persisted.ts @@ -8,6 +8,7 @@ import * as cbor from "cbor-x"; import type * as persistSchema from "@/schemas/actor-persist/mod"; import { bufferToArrayBuffer } from "@/utils"; +export type GatewayId = ArrayBuffer; export type RequestId = ArrayBuffer; export type Cbor = ArrayBuffer; @@ -25,10 +26,10 @@ export interface PersistedConn { parameters: CP; state: CS; subscriptions: PersistedSubscription[]; - /** Request ID of the hibernatable WebSocket */ - hibernatableRequestId: RequestId; - /** Last seen message index for this WebSocket */ - msgIndex: number; + gatewayId: GatewayId; + requestId: RequestId; + serverMessageIndex: number; + clientMessageIndex: number; requestPath: string; requestHeaders: Record; } @@ -47,8 +48,10 @@ export function convertConnToBarePersistedConn( subscriptions: persist.subscriptions.map((sub) => ({ eventName: sub.eventName, })), - hibernatableRequestId: persist.hibernatableRequestId, - msgIndex: BigInt(persist.msgIndex), + gatewayId: persist.gatewayId, + requestId: persist.requestId, + serverMessageIndex: persist.serverMessageIndex, + clientMessageIndex: persist.clientMessageIndex, requestPath: persist.requestPath, requestHeaders: new Map(Object.entries(persist.requestHeaders)), }; @@ -68,8 +71,10 @@ export function convertConnFromBarePersistedConn( subscriptions: bareData.subscriptions.map((sub) => ({ eventName: sub.eventName, })), - hibernatableRequestId: bareData.hibernatableRequestId, - msgIndex: Number(bareData.msgIndex), + gatewayId: bareData.gatewayId, + requestId: bareData.requestId, + serverMessageIndex: bareData.serverMessageIndex, + clientMessageIndex: bareData.clientMessageIndex, requestPath: bareData.requestPath, requestHeaders: Object.fromEntries(bareData.requestHeaders), }; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/conn/state-manager.ts b/rivetkit-typescript/packages/rivetkit/src/actor/conn/state-manager.ts index d41a1124f9..8c79502d60 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/conn/state-manager.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/conn/state-manager.ts @@ -193,14 +193,4 @@ export class StateManager { } return subIdx !== -1; } - - buildHwsMeta(): HibernatingWebSocketMetadata { - const hibernatable = this.hibernatableDataOrError(); - return { - requestId: hibernatable.hibernatableRequestId, - path: hibernatable.requestPath, - headers: hibernatable.requestHeaders, - messageIndex: hibernatable.msgIndex, - }; - } } diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/instance/connection-manager.ts b/rivetkit-typescript/packages/rivetkit/src/actor/instance/connection-manager.ts index f561378efd..91644254a5 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/instance/connection-manager.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/instance/connection-manager.ts @@ -125,10 +125,8 @@ export class ConnectionManager< // Create connection persist data let connData: ConnDataInput; if (isHibernatable) { - invariant( - driver.requestIdBuf, - "must have requestIdBuf if hibernatable", - ); + const hibernatable = driver.hibernatable; + invariant(hibernatable, "must have hibernatable"); invariant(requestPath, "missing requestPath for hibernatable ws"); invariant( requestHeaders, @@ -140,12 +138,11 @@ export class ConnectionManager< parameters: params, state: connState as CS, subscriptions: [], - // Fallback to empty buf if not provided since we don't use this value - hibernatableRequestId: driver.hibernatable - ? driver.requestIdBuf - : new ArrayBuffer(), + gatewayId: hibernatable.gatewayId, + requestId: hibernatable.requestId, + clientMessageIndex: 0, // First message index will be 1, so we start at 0 - msgIndex: 0, + serverMessageIndex: 0, requestPath, requestHeaders, }, @@ -228,8 +225,11 @@ export class ConnectionManager< } #reconnectHibernatableConn(driver: ConnDriver): Conn { - invariant(driver.requestIdBuf, "missing requestIdBuf"); - const existingConn = this.findHibernatableConn(driver.requestIdBuf); + invariant(driver.hibernatable, "missing requestIdBuf"); + const existingConn = this.findHibernatableConn( + driver.hibernatable.gatewayId, + driver.hibernatable.requestId, + ); invariant( existingConn, "cannot find connection for restoring connection", @@ -238,7 +238,6 @@ export class ConnectionManager< this.#actor.rLog.debug({ msg: "reconnecting hibernatable websocket connection", connectionId: existingConn.id, - requestId: driver.requestId, }); // Clean up existing driver state if present @@ -393,14 +392,16 @@ export class ConnectionManager< // MARK: - Private Helpers findHibernatableConn( + gatewayIdBuf: ArrayBuffer, requestIdBuf: ArrayBuffer, ): Conn | undefined { return Array.from(this.#connections.values()).find((conn) => { const connStateManager = conn[CONN_STATE_MANAGER_SYMBOL]; - const connRequestId = - connStateManager.hibernatableDataRaw?.hibernatableRequestId; + const h = connStateManager.hibernatableDataRaw; return ( - connRequestId && arrayBuffersEqual(connRequestId, requestIdBuf) + h && + arrayBuffersEqual(h.gatewayId, gatewayIdBuf) && + arrayBuffersEqual(h.requestId, requestIdBuf) ); }); } diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts b/rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts index bc16251a3a..7ad78d8c1a 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts @@ -14,7 +14,7 @@ import { import type * as protocol from "@/schemas/client-protocol/mod"; import { TO_CLIENT_VERSIONED } from "@/schemas/client-protocol/versioned"; import { ToClientSchema } from "@/schemas/client-protocol-zod/mod"; -import { EXTRA_ERROR_LOG, idToStr } from "@/utils"; +import { EXTRA_ERROR_LOG } from "@/utils"; import type { ActorConfig, InitContext } from "../config"; import type { ConnDriver } from "../conn/driver"; import { createHttpDriver } from "../conn/drivers/http"; @@ -168,14 +168,8 @@ export class ActorInstance { : undefined, subscriptions: conn.subscriptions.size, isHibernatable: conn.isHibernatable, - hibernatableRequestId: connStateManager - .hibernatableDataRaw?.hibernatableRequestId - ? idToStr( - connStateManager.hibernatableDataRaw - .hibernatableRequestId, - ) - : undefined, - // TODO: Include the underlying request for path & headers? + // TODO: Include underlying hibernatable metadata + + // path + headers }; }); }, diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/instance/state-manager.ts b/rivetkit-typescript/packages/rivetkit/src/actor/instance/state-manager.ts index 2f4bdf2b08..ce6086595f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/instance/state-manager.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/instance/state-manager.ts @@ -19,6 +19,7 @@ import { isConnStatePath, isStatePath } from "../utils"; import { KEYS, makeConnKey } from "./kv"; import type { ActorInstance } from "./mod"; import { convertActorToBarePersisted, type PersistedActor } from "./persisted"; +import { tunnelId } from "@rivetkit/engine-runner"; export interface SaveStateOptions { /** @@ -430,9 +431,14 @@ export class StateManager { this.#actor.rLog.info({ msg: "persisting connection", connId, - hibernatableRequestId: - hibernatableDataRaw.hibernatableRequestId, - msgIndex: hibernatableDataRaw.msgIndex, + gatewayId: tunnelId.gatewayIdToString( + hibernatableDataRaw.requestId, + ), + requestId: tunnelId.requestIdToString( + hibernatableDataRaw.requestId, + ), + serverMessageIndex: hibernatableDataRaw.serverMessageIndex, + clientMessageIndex: hibernatableDataRaw.clientMessageIndex, hasState: hibernatableDataRaw.state !== undefined, }); diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/router-websocket-endpoints.ts b/rivetkit-typescript/packages/rivetkit/src/actor/router-websocket-endpoints.ts index c96862e3e4..0cac7f8f2f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/router-websocket-endpoints.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/router-websocket-endpoints.ts @@ -62,8 +62,8 @@ export async function routeWebSocket( actorId: string, encoding: Encoding, parameters: unknown, - requestId: string, - requestIdBuf: ArrayBuffer | undefined, + gatewayId: ArrayBuffer | undefined, + requestId: ArrayBuffer | undefined, isHibernatable: boolean, isRestoringHibernatable: boolean, ): Promise { @@ -90,9 +90,9 @@ export async function routeWebSocket( let connDriver: ConnDriver; if (requestPath === PATH_CONNECT) { const { driver, setWebSocket } = createWebSocketDriver( - requestId, - requestIdBuf, - isHibernatable, + isHibernatable + ? { gatewayId: gatewayId!, requestId: requestId! } + : undefined, encoding, closePromiseResolvers.promise, ); @@ -103,9 +103,9 @@ export async function routeWebSocket( requestPath.startsWith(PATH_WEBSOCKET_PREFIX) ) { const { driver, setWebSocket } = createRawWebSocketDriver( - requestId, - requestIdBuf, - isHibernatable, + isHibernatable + ? { gatewayId: gatewayId!, requestId: requestId! } + : undefined, closePromiseResolvers.promise, ); handler = handleRawWebSocket.bind(undefined, setWebSocket); diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/router.ts b/rivetkit-typescript/packages/rivetkit/src/actor/router.ts index d50f88996a..bace613cbc 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/router.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/router.ts @@ -24,7 +24,7 @@ import { } from "@/inspector/actor"; import { isInspectorEnabled, secureInspector } from "@/inspector/utils"; import type { RunnerConfig } from "@/registry/run-config"; -import { CONN_DRIVER_SYMBOL, generateConnRequestId } from "./conn/mod"; +import { CONN_DRIVER_SYMBOL } from "./conn/mod"; import type { ActorDriver } from "./driver"; import { loggerWithoutContext } from "./log"; import { @@ -125,7 +125,7 @@ export function createActorRouter( c.env.actorId, encoding, connParams, - generateConnRequestId(), + undefined, undefined, false, false, diff --git a/rivetkit-typescript/packages/rivetkit/src/drivers/engine/actor-driver.ts b/rivetkit-typescript/packages/rivetkit/src/drivers/engine/actor-driver.ts index cd90975380..bed6a47f9f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/drivers/engine/actor-driver.ts +++ b/rivetkit-typescript/packages/rivetkit/src/drivers/engine/actor-driver.ts @@ -3,7 +3,7 @@ import type { RunnerConfig as EngineRunnerConfig, HibernatingWebSocketMetadata, } from "@rivetkit/engine-runner"; -import { Runner } from "@rivetkit/engine-runner"; +import { Runner, tunnelId } from "@rivetkit/engine-runner"; import * as cbor from "cbor-x"; import type { Context as HonoContext } from "hono"; import { streamSSE } from "hono/streaming"; @@ -11,7 +11,6 @@ import { WSContext, type WSContextInit } from "hono/ws"; import invariant from "invariant"; import { type AnyConn, - CONN_ACTOR_SYMBOL, CONN_STATE_MANAGER_SYMBOL, } from "@/actor/conn/mod"; import { lookupInRegistry } from "@/actor/definition"; @@ -50,7 +49,6 @@ import type { RequestId } from "@/schemas/actor-persist/mod"; import { arrayBuffersEqual, assertUnreachable, - idToStr, type LongTimeoutHandle, promiseWithResolvers, setLongTimeout, @@ -104,10 +102,10 @@ export class EngineActorDriver implements ActorDriver { // Map of conn IDs to message index waiting to be persisted before sending // an ack // - // messageIndex is updated and pendingAck is flagged in needed in + // serverMessageIndex is updated and pendingAck is flagged in needed in // onBeforePersistConnect, then the HWS ack message is sent in // onAfterPersistConn. This allows us to track what's about to be written - // to storage to prevent race conditions with the messageIndex being + // to storage to prevent race conditions with the serverMessageIndex being // updated while writing the existing state. // // bufferedMessageSize tracks the total bytes received since last persist @@ -116,7 +114,7 @@ export class EngineActorDriver implements ActorDriver { #hwsMessageIndex = new Map< string, { - messageIndex: number; + serverMessageIndex: number; bufferedMessageSize: number; pendingAckFromMessageIndex: boolean; pendingAckFromBufferSize: boolean; @@ -518,6 +516,7 @@ export class EngineActorDriver implements ActorDriver { async #runnerFetch( _runner: Runner, actorId: string, + _gatewayIdBuf: ArrayBuffer, _requestIdBuf: ArrayBuffer, request: Request, ): Promise { @@ -534,6 +533,7 @@ export class EngineActorDriver implements ActorDriver { _runner: Runner, actorId: string, websocketRaw: any, + gatewayIdBuf: ArrayBuffer, requestIdBuf: ArrayBuffer, request: Request, requestPath: string, @@ -542,7 +542,6 @@ export class EngineActorDriver implements ActorDriver { isRestoringHibernatable: boolean, ): Promise { const websocket = websocketRaw as UniversalWebSocket; - const requestId = idToStr(requestIdBuf); // Add a unique ID to track this WebSocket object const wsUniqueId = `ws_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`; @@ -581,7 +580,7 @@ export class EngineActorDriver implements ActorDriver { actorId, encoding, connParams, - requestId, + gatewayIdBuf, requestIdBuf, isHibernatable, isRestoringHibernatable, @@ -673,8 +672,8 @@ export class EngineActorDriver implements ActorDriver { ); // Persist message index - const previousMsgIndex = hibernate.msgIndex; - hibernate.msgIndex = event.rivetMessageIndex; + const previousMsgIndex = hibernate.serverMessageIndex; + hibernate.serverMessageIndex = event.rivetMessageIndex; logger().info({ msg: "persisting message index", connId: conn.id, @@ -742,7 +741,8 @@ export class EngineActorDriver implements ActorDriver { msg: "event listeners attached to restored websocket", actorId, connId: conn?.id, - requestId, + gatewayId: tunnelId.gatewayIdToString(gatewayIdBuf), + requestId: tunnelId.requestIdToString(requestIdBuf), websocketType: websocket?.constructor?.name, hasMessageListener: !!websocket.addEventListener, }); @@ -752,6 +752,7 @@ export class EngineActorDriver implements ActorDriver { // MARK: - Hibernating WebSockets #hwsCanHibernate( actorId: string, + gatewayId: ArrayBuffer, requestId: ArrayBuffer, request: Request, ): boolean { @@ -788,7 +789,8 @@ export class EngineActorDriver implements ActorDriver { // Determine configuration for new WS logger().debug({ msg: "no existing hibernatable websocket found", - requestId: idToStr(requestId), + gatewayId: tunnelId.gatewayIdToString(gatewayId), + requestId: tunnelId.requestIdToString(requestId), }); if (path === PATH_CONNECT) { return true; @@ -853,10 +855,12 @@ export class EngineActorDriver implements ActorDriver { const hibernatable = connStateManager.hibernatableData; if (!hibernatable) return undefined; return { - requestId: hibernatable.hibernatableRequestId, + gatewayId: hibernatable.gatewayId, + requestId: hibernatable.requestId, + serverMessageIndex: hibernatable.serverMessageIndex, + clientMessageIndex: hibernatable.clientMessageIndex, path: hibernatable.requestPath, headers: hibernatable.requestHeaders, - messageIndex: hibernatable.msgIndex, } satisfies HibernatingWebSocketMetadata; }) .filter((x) => x !== undefined) @@ -865,31 +869,29 @@ export class EngineActorDriver implements ActorDriver { onCreateConn(conn: AnyConn) { const hibernatable = conn[CONN_STATE_MANAGER_SYMBOL].hibernatableData; - logger().info({ - msg: "EngineActorDriver.onCreateConn called", - connId: conn.id, - hasHibernatable: !!hibernatable, - msgIndex: hibernatable?.msgIndex, - }); - if (!hibernatable) return; this.#hwsMessageIndex.set(conn.id, { - messageIndex: hibernatable.msgIndex, + serverMessageIndex: hibernatable.serverMessageIndex, bufferedMessageSize: 0, pendingAckFromMessageIndex: false, pendingAckFromBufferSize: false, }); - logger().info({ - msg: "EngineActorDriver: created #hwsMessageIndex entry", + logger().debug({ + msg: "created #hwsMessageIndex entry", connId: conn.id, - msgIndex: hibernatable.msgIndex, + serverMessageIndex: hibernatable.serverMessageIndex, }); } onDestroyConn(conn: AnyConn) { this.#hwsMessageIndex.delete(conn.id); + + logger().debug({ + msg: "removed #hwsMessageIndex entry", + connId: conn.id, + }); } onBeforePersistConn(conn: AnyConn) { @@ -907,8 +909,8 @@ export class EngineActorDriver implements ActorDriver { // There is a newer message index entry.pendingAckFromMessageIndex = - hibernatable.msgIndex > entry.messageIndex; - entry.messageIndex = hibernatable.msgIndex; + hibernatable.serverMessageIndex > entry.serverMessageIndex; + entry.serverMessageIndex = hibernatable.serverMessageIndex; } onAfterPersistConn(conn: AnyConn) { @@ -930,8 +932,9 @@ export class EngineActorDriver implements ActorDriver { entry.pendingAckFromBufferSize ) { this.#runner.sendHibernatableWebSocketMessageAck( - hibernatable.hibernatableRequestId, - entry.messageIndex, + hibernatable.gatewayId, + hibernatable.requestId, + entry.serverMessageIndex, ); entry.pendingAckFromMessageIndex = false; entry.pendingAckFromBufferSize = false; diff --git a/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/manager.ts b/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/manager.ts index 35fe31c166..6db4d36312 100644 --- a/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/manager.ts +++ b/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/manager.ts @@ -1,6 +1,5 @@ import type { Context as HonoContext } from "hono"; import invariant from "invariant"; -import { generateConnRequestId } from "@/actor/conn/mod"; import { ActorStopping } from "@/actor/errors"; import { type ActorRouter, createActorRouter } from "@/actor/router"; import { routeWebSocket } from "@/actor/router-websocket-endpoints"; @@ -170,7 +169,7 @@ export class FileSystemManagerDriver implements ManagerDriver { actorId, encoding, params, - generateConnRequestId(), + undefined, undefined, false, false, @@ -213,7 +212,7 @@ export class FileSystemManagerDriver implements ManagerDriver { actorId, encoding, params, - generateConnRequestId(), + undefined, undefined, false, false, diff --git a/rivetkit-typescript/packages/rivetkit/src/utils.ts b/rivetkit-typescript/packages/rivetkit/src/utils.ts index 7f6e63487a..c9021ffa27 100644 --- a/rivetkit-typescript/packages/rivetkit/src/utils.ts +++ b/rivetkit-typescript/packages/rivetkit/src/utils.ts @@ -279,7 +279,3 @@ export const EXTRA_ERROR_LOG = { support: "https://rivet.dev/discord", version: VERSION, }; - -export function idToStr(id: ArrayBuffer): string { - return uuidstringify(new Uint8Array(id)); -}