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