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