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