Skip to main content

rivetkit_core/registry/
mod.rs

1use std::collections::HashMap;
2use std::env;
3use std::io::Cursor;
4use std::path::PathBuf;
5use std::sync::Arc;
6use std::sync::atomic::{AtomicBool, Ordering};
7use std::time::Duration;
8
9use crate::time::{Instant, timeout};
10
11use ::http::StatusCode;
12use anyhow::{Context, Result};
13use parking_lot::Mutex;
14use rivet_envoy_client::config::{
15	ActorStopHandle, BoxFuture as EnvoyBoxFuture, EnvoyCallbacks, HttpRequest, HttpResponse,
16	WebSocketHandler, WebSocketMessage, WebSocketSender,
17};
18use rivet_envoy_client::envoy::start_envoy;
19use rivet_envoy_client::handle::EnvoyHandle;
20use rivet_envoy_client::protocol;
21use rivet_error::{ActorSpecifier, RivetError};
22use rivetkit_client_protocol as client_protocol;
23use rivetkit_shared_types::serverless_metadata::{
24	ActorName, ServerlessMetadataEnvoy, ServerlessMetadataEnvoyKind, ServerlessMetadataPayload,
25};
26use scc::{HashMap as SccHashMap, hash_map::Entry as SccEntry};
27use serde::{Deserialize, Serialize};
28use serde_bytes::ByteBuf;
29use serde_json::{Value as JsonValue, json};
30use tokio::sync::{Mutex as TokioMutex, Notify, broadcast, mpsc, oneshot};
31use tokio::task::JoinHandle;
32use tokio_util::sync::CancellationToken;
33use url::Url;
34use vbare::OwnedVersionedData;
35
36use crate::actor::action::ActionDispatchError;
37use crate::actor::config::CanHibernateWebSocket;
38use crate::actor::connection::{ConnHandle, HibernatableConnectionMetadata};
39use crate::actor::context::{ActorContext, InspectorAttachGuard};
40use crate::actor::factory::ActorFactory;
41use crate::actor::keys::PERSIST_DATA_KEY;
42use crate::actor::lifecycle_hooks::Reply;
43use crate::actor::messages::{ActorEvent, QueueSendResult, Request, Response, StateDelta};
44use crate::actor::preload::{PreloadedKv, PreloadedPersistedActor};
45use crate::actor::state::decode_persisted_actor;
46use crate::actor::task::{
47	ActorTask, DispatchCommand, LifecycleCommand, try_send_dispatch_command,
48	try_send_lifecycle_command,
49};
50use crate::actor::task_types::ShutdownKind;
51#[cfg(feature = "native-runtime")]
52use crate::engine_process::{EngineProcessManager, EngineResolverConfig};
53use crate::error::{ActorLifecycle as ActorLifecycleError, ActorRuntime};
54use crate::inspector::protocol::{
55	self as inspector_protocol, ServerMessage as InspectorServerMessage,
56};
57use crate::inspector::{Inspector, InspectorAuth, InspectorSignal, InspectorSubscription};
58use crate::kv::Kv;
59use crate::runtime::RuntimeSpawner;
60use crate::sqlite::SqliteDb;
61use crate::types::{ActorKey, ActorKeySegment, WsMessage, format_actor_key};
62use crate::websocket::WebSocket;
63
64mod actor_connect;
65mod dispatch;
66mod envoy_callbacks;
67mod http;
68mod inspector;
69mod inspector_ws;
70#[cfg(feature = "native-runtime")]
71mod runner_config;
72mod websocket;
73
74use inspector::build_actor_inspector;
75use websocket::is_actor_connect_path;
76
77#[derive(Default)]
78pub struct CoreRegistry {
79	factories: HashMap<String, Arc<ActorFactory>>,
80}
81
82#[derive(Clone)]
83pub struct CoreEnvoyHandle {
84	handle: EnvoyHandle,
85}
86
87#[derive(Clone, Debug)]
88pub struct CoreEnvoyStatus {
89	pub active_actor_count: usize,
90	pub ping_healthy: bool,
91}
92
93impl CoreEnvoyHandle {
94	pub(crate) fn new(handle: EnvoyHandle) -> Self {
95		Self { handle }
96	}
97
98	pub fn status(&self) -> CoreEnvoyStatus {
99		CoreEnvoyStatus {
100			active_actor_count: self.handle.active_actor_count(),
101			ping_healthy: self.handle.is_ping_healthy(),
102		}
103	}
104
105	/// Engine-reported drain threshold in milliseconds. `None` until the
106	/// envoy has completed its first protocol-metadata exchange with the
107	/// engine.
108	pub async fn actor_stop_threshold_ms(&self) -> Option<i64> {
109		self.handle
110			.get_protocol_metadata()
111			.await
112			.map(|metadata| metadata.actor_stop_threshold)
113	}
114}
115
116#[derive(Clone)]
117struct ActorTaskHandle {
118	actor_id: String,
119	actor_name: String,
120	generation: u32,
121	ctx: ActorContext,
122	factory: Arc<ActorFactory>,
123	inspector: Inspector,
124	lifecycle: mpsc::UnboundedSender<LifecycleCommand>,
125	dispatch: mpsc::UnboundedSender<DispatchCommand>,
126	join: Arc<TokioMutex<Option<JoinHandle<Result<()>>>>>,
127}
128
129type ActiveActorInstance = Arc<ActorTaskHandle>;
130
131enum ActorInstanceState {
132	Active(ActiveActorInstance),
133	Stopping {
134		instance: ActiveActorInstance,
135		reason: ShutdownKind,
136	},
137}
138
139impl ActorInstanceState {
140	fn instance(&self) -> ActiveActorInstance {
141		match self {
142			Self::Active(instance) | Self::Stopping { instance, .. } => instance.clone(),
143		}
144	}
145
146	fn active_instance(&self) -> Option<ActiveActorInstance> {
147		match self {
148			Self::Active(instance) => Some(instance.clone()),
149			Self::Stopping { .. } => None,
150		}
151	}
152}
153
154#[derive(Clone)]
155struct PendingStop {
156	reason: protocol::StopActorReason,
157	stop_handle: ActorStopHandle,
158}
159
160pub(crate) struct RegistryDispatcher {
161	pub(crate) factories: HashMap<String, Arc<ActorFactory>>,
162	actor_instances: SccHashMap<String, ActorInstanceState>,
163	starting_instances: SccHashMap<String, Arc<Notify>>,
164	pending_stops: SccHashMap<String, PendingStop>,
165	region: String,
166	handle_inspector_http_in_runtime: bool,
167}
168
169pub(crate) struct RegistryCallbacks {
170	pub(crate) dispatcher: Arc<RegistryDispatcher>,
171}
172
173#[derive(Clone, Debug)]
174struct StartActorRequest {
175	actor_id: String,
176	generation: u32,
177	actor_name: String,
178	input: Option<Vec<u8>>,
179	preload_persisted_actor: PreloadedPersistedActor,
180	preloaded_kv: Option<PreloadedKv>,
181	ctx: ActorContext,
182}
183
184#[derive(Clone, Debug)]
185struct ServeSettings {
186	version: u32,
187	endpoint: String,
188	token: Option<String>,
189	namespace: String,
190	pool_name: String,
191	engine_binary_path: Option<PathBuf>,
192	engine_host: Option<String>,
193	engine_port: Option<u16>,
194	engine_spawn: EngineSpawnMode,
195	engine_auto_download: bool,
196	handle_inspector_http_in_runtime: bool,
197	serverless_base_path: Option<String>,
198	serverless_package_version: String,
199	serverless_client_endpoint: Option<String>,
200	serverless_client_namespace: Option<String>,
201	serverless_client_token: Option<String>,
202	serverless_validate_endpoint: bool,
203	serverless_max_start_payload_bytes: usize,
204}
205
206#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
207pub enum EngineSpawnMode {
208	#[default]
209	Auto,
210	Always,
211	Never,
212}
213
214impl EngineSpawnMode {
215	pub(crate) fn from_env() -> Self {
216		match env::var("RIVETKIT_ENGINE_SPAWN") {
217			Ok(value) if value.eq_ignore_ascii_case("always") => Self::Always,
218			Ok(value) if value.eq_ignore_ascii_case("never") => Self::Never,
219			_ => Self::Auto,
220		}
221	}
222}
223
224#[derive(Clone, Debug, Default)]
225pub struct ServeConfig {
226	pub version: u32,
227	pub endpoint: String,
228	pub token: Option<String>,
229	pub namespace: String,
230	pub pool_name: String,
231	pub engine_binary_path: Option<PathBuf>,
232	pub engine_host: Option<String>,
233	pub engine_port: Option<u16>,
234	pub engine_spawn: EngineSpawnMode,
235	pub engine_auto_download: bool,
236	pub handle_inspector_http_in_runtime: bool,
237	pub serverless_base_path: Option<String>,
238	pub serverless_package_version: String,
239	pub serverless_client_endpoint: Option<String>,
240	pub serverless_client_namespace: Option<String>,
241	pub serverless_client_token: Option<String>,
242	pub serverless_validate_endpoint: bool,
243	pub serverless_max_start_payload_bytes: usize,
244	pub serverless_cache_envoy: bool,
245}
246
247#[derive(Debug, Default, Deserialize)]
248#[serde(default)]
249struct InspectorPatchStateBody {
250	state: JsonValue,
251}
252
253#[derive(Debug, Default, Deserialize)]
254#[serde(default)]
255struct InspectorActionBody {
256	args: Vec<JsonValue>,
257	properties: Option<JsonValue>,
258}
259
260#[derive(Debug, Default, Deserialize)]
261#[serde(default)]
262struct InspectorDatabaseExecuteBody {
263	sql: String,
264	args: Vec<JsonValue>,
265	properties: Option<JsonValue>,
266}
267
268#[derive(Debug, Default, Deserialize)]
269#[serde(default, rename_all = "camelCase")]
270struct InspectorWorkflowReplayBody {
271	entry_id: Option<String>,
272}
273
274#[derive(Debug, Serialize)]
275#[serde(rename_all = "camelCase")]
276struct InspectorQueueMessageJson {
277	id: u64,
278	name: String,
279	created_at_ms: i64,
280}
281
282#[derive(Debug, Serialize)]
283#[serde(rename_all = "camelCase")]
284struct InspectorQueueResponseJson {
285	size: u32,
286	max_size: u32,
287	truncated: bool,
288	messages: Vec<InspectorQueueMessageJson>,
289}
290
291#[derive(Debug, Deserialize)]
292#[serde(default)]
293struct HttpActionRequestJson {
294	args: JsonValue,
295}
296
297impl Default for HttpActionRequestJson {
298	fn default() -> Self {
299		Self {
300			args: JsonValue::Array(Vec::new()),
301		}
302	}
303}
304
305pub(crate) fn should_manage_engine(endpoint: &str, spawn_mode: EngineSpawnMode) -> Result<bool> {
306	match spawn_mode {
307		EngineSpawnMode::Always => Ok(true),
308		EngineSpawnMode::Never => Ok(false),
309		EngineSpawnMode::Auto => is_loopback_endpoint(endpoint),
310	}
311}
312
313fn is_loopback_endpoint(endpoint: &str) -> Result<bool> {
314	let url =
315		Url::parse(endpoint).with_context(|| format!("parse engine endpoint `{endpoint}`"))?;
316	let Some(host) = url.host_str() else {
317		anyhow::bail!("engine endpoint `{endpoint}` is invalid: missing host");
318	};
319
320	if host == "localhost" || host.ends_with(".localhost") {
321		return Ok(true);
322	}
323
324	let ip_host = host
325		.strip_prefix('[')
326		.and_then(|value| value.strip_suffix(']'))
327		.unwrap_or(host);
328
329	Ok(ip_host
330		.parse::<std::net::IpAddr>()
331		.map(|ip| ip.is_loopback() || ip.is_unspecified())
332		.unwrap_or(false))
333}
334
335#[cfg(test)]
336mod engine_spawn_tests {
337	use super::{EngineSpawnMode, should_manage_engine};
338
339	#[test]
340	fn auto_manages_loopback_endpoints() {
341		assert!(should_manage_engine("http://127.0.0.1:6420", EngineSpawnMode::Auto).unwrap());
342		assert!(should_manage_engine("http://localhost:6420", EngineSpawnMode::Auto).unwrap());
343		assert!(should_manage_engine("http://dev.localhost:6420", EngineSpawnMode::Auto).unwrap());
344		assert!(should_manage_engine("http://[::1]:6420", EngineSpawnMode::Auto).unwrap());
345	}
346
347	#[test]
348	fn auto_leaves_remote_endpoints_connect_only() {
349		assert!(!should_manage_engine("https://api.rivet.dev", EngineSpawnMode::Auto).unwrap());
350		assert!(!should_manage_engine("http://192.0.2.10:6420", EngineSpawnMode::Auto).unwrap());
351	}
352
353	#[test]
354	fn explicit_spawn_mode_overrides_endpoint_shape() {
355		assert!(should_manage_engine("https://api.rivet.dev", EngineSpawnMode::Always).unwrap());
356		assert!(!should_manage_engine("http://127.0.0.1:6420", EngineSpawnMode::Never).unwrap());
357	}
358}
359
360#[derive(Debug, Deserialize)]
361#[serde(default, rename_all = "camelCase")]
362struct HttpQueueSendRequestJson {
363	body: JsonValue,
364	wait: Option<bool>,
365	timeout: Option<u64>,
366}
367
368impl Default for HttpQueueSendRequestJson {
369	fn default() -> Self {
370		Self {
371			body: JsonValue::Null,
372			wait: None,
373			timeout: None,
374		}
375	}
376}
377
378#[derive(RivetError)]
379#[error("message", "incoming_too_long", "Incoming message too long")]
380struct IncomingMessageTooLong;
381
382#[derive(RivetError)]
383#[error("message", "outgoing_too_long", "Outgoing message too long")]
384struct OutgoingMessageTooLong;
385
386#[derive(RivetError)]
387#[error("actor", "action_timed_out", "Action timed out")]
388struct ActionTimedOut;
389
390#[derive(RivetError, Serialize)]
391#[error("actor", "method_not_allowed", "Method not allowed")]
392struct MethodNotAllowed {
393	method: String,
394	path: String,
395}
396
397#[derive(Debug, Serialize)]
398#[serde(rename_all = "camelCase")]
399struct InspectorConnectionJson {
400	#[serde(rename = "type")]
401	connection_type: Option<String>,
402	id: String,
403	details: InspectorConnectionDetailsJson,
404}
405
406#[derive(Debug, Serialize)]
407#[serde(rename_all = "camelCase")]
408struct InspectorConnectionDetailsJson {
409	#[serde(rename = "type")]
410	connection_type: Option<String>,
411	params: JsonValue,
412	state_enabled: bool,
413	state: JsonValue,
414	subscriptions: usize,
415	is_hibernatable: bool,
416}
417
418#[derive(Debug, Serialize)]
419#[serde(rename_all = "camelCase")]
420struct InspectorSummaryJson {
421	state: JsonValue,
422	is_state_enabled: bool,
423	connections: Vec<InspectorConnectionJson>,
424	rpcs: Vec<String>,
425	queue_size: u32,
426	is_database_enabled: bool,
427	#[serde(rename = "isWorkflowEnabled")]
428	workflow_supported: bool,
429	workflow_history: Option<JsonValue>,
430}
431
432const WS_PROTOCOL_ENCODING: &str = "rivet_encoding.";
433const WS_PROTOCOL_CONN_PARAMS: &str = "rivet_conn_params.";
434
435#[derive(Debug)]
436struct ActorConnectInit {
437	actor_id: String,
438	connection_id: String,
439}
440
441#[derive(Debug)]
442struct ActorConnectError {
443	group: String,
444	code: String,
445	message: String,
446	metadata: Option<ByteBuf>,
447	action_id: Option<u64>,
448	actor: Option<ActorSpecifier>,
449}
450
451#[derive(Debug)]
452struct ActorConnectActionResponse {
453	id: u64,
454	output: ByteBuf,
455}
456
457#[derive(Debug)]
458struct ActorConnectEvent {
459	name: String,
460	args: ByteBuf,
461}
462
463#[derive(Clone, Copy, Debug, PartialEq, Eq)]
464enum ActorConnectEncoding {
465	Json,
466	Cbor,
467	Bare,
468}
469
470#[derive(Debug)]
471enum ActorConnectToClient {
472	Init(ActorConnectInit),
473	Error(ActorConnectError),
474	ActionResponse(ActorConnectActionResponse),
475	Event(ActorConnectEvent),
476}
477
478#[derive(Debug)]
479struct ActorConnectActionRequest {
480	id: u64,
481	name: String,
482	args: ByteBuf,
483}
484
485#[derive(Debug)]
486enum ActorConnectSendError {
487	OutgoingTooLong,
488	Encode(anyhow::Error),
489}
490
491#[derive(Debug, Deserialize)]
492struct ActorConnectSubscriptionRequest {
493	#[serde(rename = "eventName")]
494	event_name: String,
495	subscribe: bool,
496}
497
498#[derive(Debug)]
499enum ActorConnectToServer {
500	ActionRequest(ActorConnectActionRequest),
501	SubscriptionRequest(ActorConnectSubscriptionRequest),
502}
503
504#[derive(Debug, Deserialize)]
505struct ActorConnectActionRequestJson {
506	id: u64,
507	name: String,
508	args: JsonValue,
509}
510
511#[derive(Debug, Deserialize)]
512#[serde(tag = "tag", content = "val")]
513enum ActorConnectToServerJsonBody {
514	ActionRequest(ActorConnectActionRequestJson),
515	SubscriptionRequest(ActorConnectSubscriptionRequest),
516}
517
518#[derive(Debug, Deserialize)]
519struct ActorConnectToServerJsonEnvelope {
520	body: ActorConnectToServerJsonBody,
521}
522
523impl CoreRegistry {
524	pub fn new() -> Self {
525		Self::default()
526	}
527
528	pub fn register(&mut self, name: &str, factory: ActorFactory) {
529		self.factories.insert(name.to_owned(), Arc::new(factory));
530	}
531
532	pub fn register_shared(&mut self, name: &str, factory: Arc<ActorFactory>) {
533		self.factories.insert(name.to_owned(), factory);
534	}
535
536	pub fn normal_metadata_payload(&self, config: &ServeConfig) -> ServerlessMetadataPayload {
537		serverless_metadata_payload(
538			build_actor_metadata_map_from_factories(&self.factories),
539			config,
540			ServerlessMetadataEnvoyKind::Normal {},
541		)
542	}
543
544	pub fn serverless_metadata_payload(&self, config: &ServeConfig) -> ServerlessMetadataPayload {
545		serverless_metadata_payload(
546			build_actor_metadata_map_from_factories(&self.factories),
547			config,
548			ServerlessMetadataEnvoyKind::Serverless {},
549		)
550	}
551
552	pub async fn serve(self, shutdown: CancellationToken) -> Result<()> {
553		self.serve_with_config(ServeConfig::from_env(), shutdown)
554			.await
555	}
556
557	pub async fn serve_with_config(
558		self,
559		config: ServeConfig,
560		shutdown: CancellationToken,
561	) -> Result<()> {
562		self.serve_with_config_and_handle_observer(config, shutdown, |_| {})
563			.await
564	}
565
566	pub async fn serve_with_config_and_handle_observer(
567		self,
568		config: ServeConfig,
569		shutdown: CancellationToken,
570		on_handle: impl FnOnce(CoreEnvoyHandle) + Send + 'static,
571	) -> Result<()> {
572		crate::metrics_endpoint::record_rivetkit_info(
573			config.serverless_package_version.clone(),
574			config.version,
575			"serverful",
576			config.pool_name.clone(),
577		);
578
579		let dispatcher = self.into_dispatcher(&config);
580		#[cfg(feature = "native-runtime")]
581		let _engine_process = if should_manage_engine(&config.endpoint, config.engine_spawn)? {
582			Some(
583				EngineProcessManager::start_or_reuse(EngineResolverConfig::from_parts(
584					&config.endpoint,
585					config.engine_binary_path.clone(),
586					config.engine_host.clone(),
587					config.engine_port,
588					config.engine_auto_download,
589				))
590				.await?,
591			)
592		} else {
593			None
594		};
595		#[cfg(not(feature = "native-runtime"))]
596		if should_manage_engine(&config.endpoint, config.engine_spawn)? {
597			anyhow::bail!("engine process spawning requires the `native-runtime` feature");
598		}
599
600		#[cfg(feature = "native-runtime")]
601		runner_config::ensure_local_normal_runner_config(&config).await?;
602		let callbacks = Arc::new(RegistryCallbacks {
603			dispatcher: dispatcher.clone(),
604		});
605
606		let prepopulate_actor_names = dispatcher
607			.build_actor_metadata_map()
608			.into_iter()
609			.map(|(name, metadata)| (name, rivet_envoy_client::config::ActorName { metadata }))
610			.collect();
611		let handle = start_envoy(rivet_envoy_client::config::EnvoyConfig {
612			version: config.version,
613			endpoint: config.endpoint,
614			token: config.token,
615			namespace: config.namespace,
616			pool_name: config.pool_name,
617			prepopulate_actor_names,
618			metadata: Some(json!({
619				"rivetkit": { "version": config.serverless_package_version },
620			})),
621			not_global: false,
622			debug_latency_ms: None,
623			callbacks,
624		})
625		.await;
626		on_handle(CoreEnvoyHandle::new(handle.clone()));
627
628		// Do not install `tokio::signal::ctrl_c()` here. It calls
629		// `sigaction(SIGINT, ...)` at the POSIX level, which overrides the
630		// host's default SIGINT handling when rivetkit-core is embedded in
631		// Node via NAPI and leaves the host process unable to exit. Callers
632		// trip the `shutdown` token instead.
633		shutdown.cancelled().await;
634
635		// TODO: Move into envoy-client since timing out has to do with protocol compliance
636		// Read threshold from protocol metadata, fall back to 30 min
637		let stop_threshold = handle
638			.get_protocol_metadata()
639			.await
640			.map(|x| x.actor_stop_threshold)
641			.unwrap_or(30 * 60 * 1000);
642		// Bounded drain. If envoy cannot reach the engine (reconnect loop stuck),
643		// we fall back to immediate `Stop` rather than hanging indefinitely.
644		// The outer host (TS signal handler / Rust binary) is the backstop.
645		match timeout(
646			Duration::from_millis(stop_threshold as u64),
647			handle.shutdown_and_wait(false),
648		)
649		.await
650		{
651			Ok(()) => {}
652			Err(_) => {
653				tracing::warn!("envoy shutdown drain exceeded timeout; forcing immediate stop");
654				handle.shutdown(true);
655				handle.wait_stopped().await;
656			}
657		}
658
659		Ok(())
660	}
661
662	fn into_dispatcher(self, config: &ServeConfig) -> Arc<RegistryDispatcher> {
663		Arc::new(RegistryDispatcher::new(
664			self.factories,
665			config.handle_inspector_http_in_runtime,
666		))
667	}
668
669	pub async fn into_serverless_runtime(
670		self,
671		config: ServeConfig,
672	) -> Result<crate::serverless::CoreServerlessRuntime> {
673		crate::serverless::CoreServerlessRuntime::new(self.factories, config).await
674	}
675}
676
677impl RegistryDispatcher {
678	pub(crate) fn new(
679		factories: HashMap<String, Arc<ActorFactory>>,
680		handle_inspector_http_in_runtime: bool,
681	) -> Self {
682		Self {
683			factories,
684			actor_instances: SccHashMap::new(),
685			starting_instances: SccHashMap::new(),
686			pending_stops: SccHashMap::new(),
687			region: env::var("RIVET_REGION").unwrap_or_default(),
688			handle_inspector_http_in_runtime,
689		}
690	}
691
692	pub(crate) fn build_actor_metadata_map(&self) -> HashMap<String, JsonValue> {
693		build_actor_metadata_map_from_factories(&self.factories)
694	}
695}
696
697pub(crate) fn serverless_metadata_payload(
698	actor_metadata: HashMap<String, JsonValue>,
699	config: &ServeConfig,
700	envoy_kind: ServerlessMetadataEnvoyKind,
701) -> ServerlessMetadataPayload {
702	let actor_names = actor_metadata
703		.into_iter()
704		.map(|(name, metadata)| {
705			(
706				name,
707				ActorName {
708					metadata: Some(metadata),
709				},
710			)
711		})
712		.collect::<HashMap<_, _>>();
713
714	ServerlessMetadataPayload {
715		runtime: "rivetkit".to_owned(),
716		version: config.serverless_package_version.clone(),
717		envoy_protocol_version: Some(protocol::PROTOCOL_VERSION),
718		actor_names,
719		envoy: Some(ServerlessMetadataEnvoy {
720			kind: Some(envoy_kind),
721			version: Some(config.version),
722		}),
723		runner: None,
724		client_endpoint: config.serverless_client_endpoint.clone(),
725		client_namespace: config.serverless_client_namespace.clone(),
726		client_token: config.serverless_client_token.clone(),
727	}
728}
729
730fn build_actor_metadata_map_from_factories(
731	factories: &HashMap<String, Arc<ActorFactory>>,
732) -> HashMap<String, JsonValue> {
733	factories
734		.iter()
735		.map(|(actor_name, factory)| {
736			let config = factory.config();
737			let mut metadata = serde_json::Map::new();
738			if let Some(icon) = &config.icon {
739				metadata.insert("icon".to_owned(), json!(icon));
740			}
741			if let Some(name) = &config.name {
742				metadata.insert("name".to_owned(), json!(name));
743			}
744			metadata.insert(
745				"preload".to_owned(),
746				json!({
747					"keys": [
748						[1],
749						[3],
750						[5, 1, 1],
751						[6],
752					],
753					"prefixes": [
754						{
755							"prefix": [6, 1],
756							"maxBytes": config.preload_max_workflow_bytes.unwrap_or(131_072),
757							"partial": false,
758						},
759						{
760							"prefix": [2],
761							"maxBytes": config.preload_max_connections_bytes.unwrap_or(65_536),
762							"partial": false,
763						},
764						{
765							"prefix": [5, 1, 2],
766							"maxBytes": 65_536,
767							"partial": false,
768						},
769					],
770				}),
771			);
772			(actor_name.clone(), JsonValue::Object(metadata))
773		})
774		.collect()
775}
776
777impl RegistryDispatcher {
778	async fn start_actor(self: &Arc<Self>, request: StartActorRequest) -> Result<()> {
779		let startup_notify = Arc::new(Notify::new());
780		let _ = self
781			.starting_instances
782			.insert_async(request.actor_id.clone(), startup_notify.clone())
783			.await;
784		let factory = self
785			.factories
786			.get(&request.actor_name)
787			.cloned()
788			.ok_or_else(|| {
789				ActorRuntime::NotRegistered {
790					actor_name: request.actor_name.clone(),
791				}
792				.build()
793			})?;
794		let (lifecycle_tx, lifecycle_rx) = mpsc::unbounded_channel();
795		let (dispatch_tx, dispatch_rx) = mpsc::unbounded_channel();
796		let (lifecycle_events_tx, lifecycle_events_rx) = mpsc::unbounded_channel();
797		request
798			.ctx
799			.configure_lifecycle_events(Some(lifecycle_events_tx));
800		request.ctx.cancel_sleep_timer();
801		request.ctx.set_local_alarm_callback(Some(Arc::new({
802			let lifecycle_tx = lifecycle_tx.clone();
803			move || {
804				let lifecycle_tx = lifecycle_tx.clone();
805				Box::pin(async move {
806					let (reply_tx, reply_rx) = oneshot::channel();
807					if let Err(error) = try_send_lifecycle_command(
808						&lifecycle_tx,
809						LifecycleCommand::FireAlarm { reply: reply_tx },
810					) {
811						tracing::warn!(?error, "failed to enqueue actor alarm");
812						return;
813					}
814					let _ = reply_rx.await;
815				})
816			}
817		})));
818		let task = ActorTask::new(
819			request.actor_id.clone(),
820			request.generation,
821			lifecycle_rx,
822			dispatch_rx,
823			lifecycle_events_rx,
824			factory.clone(),
825			request.ctx.clone(),
826			request.input,
827			None,
828		)
829		.with_preloaded_persisted_actor(request.preload_persisted_actor)
830		.with_preloaded_kv(request.preloaded_kv);
831		let join = RuntimeSpawner::spawn(task.run());
832
833		let (start_tx, start_rx) = oneshot::channel();
834		let result: Result<Arc<ActorTaskHandle>> = async {
835			try_send_lifecycle_command(&lifecycle_tx, LifecycleCommand::Start { reply: start_tx })
836				.context("send actor task start command")?;
837			start_rx
838				.await
839				.context("receive actor task start reply")?
840				.context("actor task start")?;
841			let inspector = build_actor_inspector();
842			request.ctx.configure_inspector(Some(inspector.clone()));
843			Ok::<Arc<ActorTaskHandle>, anyhow::Error>(Arc::new(ActorTaskHandle {
844				actor_id: request.actor_id.clone(),
845				actor_name: request.actor_name.clone(),
846				generation: request.generation,
847				ctx: request.ctx.clone(),
848				factory,
849				inspector,
850				lifecycle: lifecycle_tx,
851				dispatch: dispatch_tx,
852				join: Arc::new(TokioMutex::new(Some(join))),
853			}))
854		}
855		.await
856		.with_context(|| format!("start actor `{}`", request.actor_id));
857
858		match result {
859			Ok(instance) => {
860				let pending_stop = self
861					.pending_stops
862					.remove_async(&request.actor_id.clone())
863					.await
864					.map(|(_, pending_stop)| pending_stop);
865				if let Some(pending_stop) = pending_stop {
866					let actor_id = request.actor_id.clone();
867					let stop_reason = map_envoy_stop_reason(&pending_stop.reason);
868					if matches!(stop_reason, ShutdownKind::Destroy) {
869						instance.ctx.mark_destroy_requested();
870					}
871					self.set_actor_instance_state(
872						actor_id.clone(),
873						ActorInstanceState::Stopping {
874							instance: instance.clone(),
875							reason: stop_reason,
876						},
877					)
878					.await;
879					let _ = self
880						.starting_instances
881						.remove_async(&request.actor_id.clone())
882						.await;
883
884					let dispatcher = self.clone();
885					RuntimeSpawner::spawn(async move {
886						if let Err(error) = dispatcher
887							.shutdown_started_instance(
888								&actor_id,
889								instance.clone(),
890								pending_stop.reason,
891								pending_stop.stop_handle,
892							)
893							.await
894						{
895							tracing::error!(
896								actor_id,
897								?error,
898								"failed to stop actor queued during startup"
899							);
900						}
901						dispatcher
902							.remove_stopping_actor_instance(&actor_id, &instance)
903							.await;
904					});
905					startup_notify.notify_waiters();
906
907					Ok(())
908				} else {
909					self.set_actor_instance_state(
910						request.actor_id.clone(),
911						ActorInstanceState::Active(instance),
912					)
913					.await;
914					let _ = self
915						.starting_instances
916						.remove_async(&request.actor_id.clone())
917						.await;
918					startup_notify.notify_waiters();
919					Ok(())
920				}
921			}
922			Err(error) => {
923				let _ = self
924					.starting_instances
925					.remove_async(&request.actor_id.clone())
926					.await;
927				startup_notify.notify_waiters();
928				Err(error)
929			}
930		}
931	}
932
933	async fn set_actor_instance_state(&self, actor_id: String, state: ActorInstanceState) {
934		match self.actor_instances.entry_async(actor_id).await {
935			SccEntry::Occupied(mut entry) => {
936				entry.insert(state);
937			}
938			SccEntry::Vacant(entry) => {
939				entry.insert_entry(state);
940			}
941		}
942	}
943
944	async fn transition_actor_to_stopping(
945		&self,
946		actor_id: &str,
947		reason: ShutdownKind,
948	) -> Option<ActiveActorInstance> {
949		match self.actor_instances.entry_async(actor_id.to_owned()).await {
950			SccEntry::Occupied(mut entry) => {
951				let instance = entry.get().instance();
952				if matches!(entry.get(), ActorInstanceState::Active(_)) {
953					entry.insert(ActorInstanceState::Stopping {
954						instance: instance.clone(),
955						reason,
956					});
957				} else {
958					instance
959						.ctx
960						.warn_work_sent_to_stopping_instance("stop_actor");
961				}
962				Some(instance)
963			}
964			SccEntry::Vacant(entry) => {
965				drop(entry);
966				None
967			}
968		}
969	}
970
971	async fn remove_stopping_actor_instance(&self, actor_id: &str, expected: &ActiveActorInstance) {
972		match self.actor_instances.entry_async(actor_id.to_owned()).await {
973			SccEntry::Occupied(entry) => {
974				let should_remove = match entry.get() {
975					ActorInstanceState::Stopping { instance, .. } => {
976						Arc::ptr_eq(instance, expected)
977					}
978					ActorInstanceState::Active(_) => false,
979				};
980				if should_remove {
981					let _ = entry.remove_entry();
982				}
983			}
984			SccEntry::Vacant(entry) => {
985				drop(entry);
986			}
987		}
988	}
989
990	async fn active_actor(&self, actor_id: &str) -> Result<Arc<ActorTaskHandle>> {
991		if let Some(instance) = self.actor_instances.get_async(&actor_id.to_owned()).await {
992			match instance.get() {
993				ActorInstanceState::Active(instance) => {
994					let instance = instance.clone();
995					// TODO: Share admission policy with ActorTask::dispatch_lifecycle_error.
996					if instance.ctx.started() {
997						if instance.ctx.destroy_requested() {
998							instance
999								.ctx
1000								.warn_work_sent_to_stopping_instance("active_actor");
1001							return Err(ActorLifecycleError::Destroying.build());
1002						}
1003						return Ok(instance);
1004					}
1005
1006					instance
1007						.ctx
1008						.warn_work_sent_to_stopping_instance("active_actor");
1009					return Err(if instance.ctx.destroy_requested() {
1010						ActorLifecycleError::Destroying.build()
1011					} else if instance.ctx.sleep_requested() {
1012						ActorLifecycleError::Stopping.build()
1013					} else {
1014						ActorLifecycleError::Starting.build()
1015					});
1016				}
1017				ActorInstanceState::Stopping { instance, reason } => {
1018					let instance = instance.clone();
1019					match reason {
1020						ShutdownKind::Sleep if instance.ctx.started() => return Ok(instance),
1021						ShutdownKind::Sleep => {
1022							instance
1023								.ctx
1024								.warn_work_sent_to_stopping_instance("active_actor");
1025							return Err(ActorLifecycleError::Stopping.build());
1026						}
1027						ShutdownKind::Destroy => {
1028							instance
1029								.ctx
1030								.warn_work_sent_to_stopping_instance("active_actor");
1031							return Err(ActorLifecycleError::Destroying.build());
1032						}
1033					}
1034				}
1035			}
1036		}
1037
1038		tracing::warn!(actor_id, "actor instance not found");
1039		Err(ActorRuntime::NotFound {
1040			resource: "instance".to_owned(),
1041			id: actor_id.to_owned(),
1042		}
1043		.build())
1044	}
1045
1046	async fn stop_actor(
1047		&self,
1048		actor_id: &str,
1049		reason: protocol::StopActorReason,
1050		stop_handle: ActorStopHandle,
1051	) -> Result<()> {
1052		if self
1053			.starting_instances
1054			.get_async(&actor_id.to_owned())
1055			.await
1056			.is_some()
1057		{
1058			let _ = self
1059				.pending_stops
1060				.insert_async(
1061					actor_id.to_owned(),
1062					PendingStop {
1063						reason,
1064						stop_handle,
1065					},
1066				)
1067				.await;
1068			return Ok(());
1069		}
1070
1071		let task_stop_reason = map_envoy_stop_reason(&reason);
1072		let instance = match self
1073			.transition_actor_to_stopping(actor_id, task_stop_reason)
1074			.await
1075		{
1076			Some(instance) => instance,
1077			None => {
1078				let _ = self
1079					.pending_stops
1080					.insert_async(
1081						actor_id.to_owned(),
1082						PendingStop {
1083							reason,
1084							stop_handle,
1085						},
1086					)
1087					.await;
1088				return Ok(());
1089			}
1090		};
1091		let result = self
1092			.shutdown_started_instance(actor_id, instance.clone(), reason, stop_handle)
1093			.await;
1094		self.remove_stopping_actor_instance(actor_id, &instance)
1095			.await;
1096		result
1097	}
1098
1099	async fn shutdown_started_instance(
1100		&self,
1101		actor_id: &str,
1102		instance: Arc<ActorTaskHandle>,
1103		reason: protocol::StopActorReason,
1104		stop_handle: ActorStopHandle,
1105	) -> Result<()> {
1106		let task_stop_reason = map_envoy_stop_reason(&reason);
1107
1108		if matches!(task_stop_reason, ShutdownKind::Destroy) {
1109			instance.ctx.mark_destroy_requested();
1110		}
1111
1112		tracing::debug!(
1113			actor_id,
1114			handle_actor_id = %instance.actor_id,
1115			actor_name = %instance.actor_name,
1116			generation = instance.generation,
1117			?reason,
1118			?task_stop_reason,
1119			"stopping actor instance"
1120		);
1121
1122		let (reply_tx, reply_rx) = oneshot::channel();
1123		let shutdown_result = match try_send_lifecycle_command(
1124			&instance.lifecycle,
1125			LifecycleCommand::Stop {
1126				reason: task_stop_reason,
1127				reply: reply_tx,
1128			},
1129		) {
1130			Ok(()) => reply_rx
1131				.await
1132				.context("receive actor task stop reply")
1133				.and_then(|result| result),
1134			Err(error) => Err(error),
1135		};
1136
1137		if matches!(task_stop_reason, ShutdownKind::Destroy) {
1138			let shutdown_deadline =
1139				Instant::now() + instance.factory.config().effective_sleep_grace_period();
1140			if !instance
1141				.ctx
1142				.wait_for_internal_keep_awake_idle(shutdown_deadline)
1143				.await
1144			{
1145				instance.ctx.record_direct_subsystem_shutdown_warning(
1146					"internal_keep_awake",
1147					"destroy_drain",
1148				);
1149				tracing::warn!(
1150					actor_id,
1151					"destroy shutdown timed out waiting for in-flight actions"
1152				);
1153			}
1154			if !instance
1155				.ctx
1156				.wait_for_http_requests_drained(shutdown_deadline)
1157				.await
1158			{
1159				instance
1160					.ctx
1161					.record_direct_subsystem_shutdown_warning("http_requests", "destroy_drain");
1162				tracing::warn!(
1163					actor_id,
1164					"destroy shutdown timed out waiting for in-flight http requests"
1165				);
1166			}
1167		}
1168
1169		let mut join_guard = instance.join.lock().await;
1170		if let Some(join) = join_guard.take() {
1171			join.await
1172				.context("join actor task")?
1173				.context("actor task failed")?;
1174		}
1175		instance.ctx.configure_lifecycle_events(None);
1176
1177		match shutdown_result {
1178			Ok(_) => {
1179				let _ = stop_handle.complete();
1180				Ok(())
1181			}
1182			Err(error) => {
1183				let _ = stop_handle.fail(anyhow::Error::new(RivetError::extract(&error)));
1184				Err(error).with_context(|| format!("stop actor `{actor_id}`"))
1185			}
1186		}
1187	}
1188}
1189
1190impl RegistryDispatcher {
1191	fn can_hibernate(&self, actor_id: &str, request: &HttpRequest) -> bool {
1192		if matches!(is_actor_connect_path(&request.path), Ok(true)) {
1193			return true;
1194		}
1195
1196		let Some(instance) = self
1197			.actor_instances
1198			.read_sync(actor_id, |_, state| state.active_instance())
1199			.flatten()
1200		else {
1201			return false;
1202		};
1203
1204		match &instance.factory.config().can_hibernate_websocket {
1205			CanHibernateWebSocket::Bool(value) => *value,
1206			CanHibernateWebSocket::Callback(callback) => callback(request),
1207		}
1208	}
1209
1210	#[allow(clippy::too_many_arguments)]
1211	fn build_actor_context(
1212		&self,
1213		handle: EnvoyHandle,
1214		actor_id: &str,
1215		generation: u32,
1216		actor_name: &str,
1217		key: ActorKey,
1218		factory: &ActorFactory,
1219	) -> ActorContext {
1220		let formatted_key = format_actor_key(&key);
1221		let ctx = ActorContext::build(
1222			actor_id.to_owned(),
1223			actor_name.to_owned(),
1224			key,
1225			self.region.clone(),
1226			Some(generation),
1227			handle.get_envoy_key().to_owned(),
1228			factory.config().clone(),
1229			Kv::new(handle.clone(), actor_id.to_owned()),
1230			SqliteDb::new_with_remote_sqlite(
1231				handle.clone(),
1232				actor_id.to_owned(),
1233				Some(formatted_key),
1234				Some(generation as u64),
1235				factory.config().has_database,
1236				factory.config().remote_sqlite,
1237			),
1238		);
1239		ctx.configure_envoy(handle, Some(generation));
1240		ctx
1241	}
1242}
1243
1244/// Maps an envoy-protocol stop reason to the lifecycle `ShutdownKind` used by
1245/// `ActorTask`. Reallocation paths (the actor will resurrect on a new envoy)
1246/// are routed through `Sleep` so user `onSleep` runs and durable state is
1247/// preserved without firing a permanent destroy.
1248fn map_envoy_stop_reason(reason: &protocol::StopActorReason) -> ShutdownKind {
1249	match reason {
1250		// Idle sleep requested by the actor itself.
1251		protocol::StopActorReason::SleepIntent => ShutdownKind::Sleep,
1252		// Runner is being drained; engine will reallocate the actor on a new
1253		// envoy. Treat as sleep so persistent state and onSleep semantics hold.
1254		protocol::StopActorReason::GoingAway => ShutdownKind::Sleep,
1255		// Runner connection lost; once reconnected (or another runner is
1256		// allocated) the actor resurrects with the same id.
1257		protocol::StopActorReason::Lost => ShutdownKind::Sleep,
1258		// User-initiated stop intent (`ctx.destroy()` and equivalents).
1259		protocol::StopActorReason::StopIntent => ShutdownKind::Destroy,
1260		// Engine-initiated permanent destroy.
1261		protocol::StopActorReason::Destroy => ShutdownKind::Destroy,
1262	}
1263}