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