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 pub async fn started(&self) -> anyhow::Result<()> {
104 self.handle.started().await
105 }
106
107 pub async fn wait_actors_drained(&self) {
109 self.handle.wait_actors_drained().await
110 }
111
112 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 shutdown.cancelled().await;
642
643 let shutdown_envoy = async {
644 let stop_threshold = handle
647 .get_protocol_metadata()
648 .await
649 .map(|x| x.actor_stop_threshold)
650 .unwrap_or(30 * 60 * 1000);
651 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 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 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
1247fn map_envoy_stop_reason(reason: &protocol::StopActorReason) -> ShutdownKind {
1252 match reason {
1253 protocol::StopActorReason::SleepIntent => ShutdownKind::Sleep,
1255 protocol::StopActorReason::GoingAway => ShutdownKind::Sleep,
1258 protocol::StopActorReason::Lost => ShutdownKind::Sleep,
1261 protocol::StopActorReason::StopIntent => ShutdownKind::Destroy,
1263 protocol::StopActorReason::Destroy => ShutdownKind::Destroy,
1265 }
1266}
1267
1268#[cfg(test)]
1270#[path = "../../tests/registry.rs"]
1271pub(crate) mod tests;