Skip to main content

rivetkit_core/actor/
context.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::future::Future;
3use std::sync::Weak;
4#[cfg(any(test, feature = "test-support"))]
5use std::sync::atomic::AtomicI64;
6use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering};
7use std::sync::{Arc, OnceLock};
8use std::time::Duration;
9
10use crate::time::Instant;
11
12use anyhow::{Context as AnyhowContext, Result};
13use futures::future::BoxFuture;
14use parking_lot::{Mutex, RwLock};
15use rivet_envoy_client::handle::EnvoyHandle;
16use rivet_envoy_client::tunnel::HibernatingWebSocketMetadata;
17use rivet_error::ActorSpecifier;
18use scc::HashMap as SccHashMap;
19use scc::HashSet as SccHashSet;
20use tokio::runtime::Handle;
21use tokio::sync::{Mutex as AsyncMutex, Notify, OnceCell, broadcast, mpsc, oneshot};
22use tokio::task::JoinHandle;
23use tokio_util::sync::CancellationToken;
24
25use crate::ActorConfig;
26#[cfg(feature = "sqlite-local")]
27use crate::actor::actor_runtime_socket::{
28	ActorRuntimeSocketEndpoint, ActorRuntimeSocketEndpointInfo,
29};
30use crate::actor::connection::{
31	ConnHandle, ConnHandles, HibernatableConnectionMetadata, PendingHibernationChanges,
32	hibernatable_id_from_slice,
33};
34use crate::actor::diagnostics::ActorDiagnostics;
35use crate::actor::internal_storage;
36use crate::actor::kv::LegacyActorKv;
37use crate::actor::lifecycle_hooks::Reply;
38use crate::actor::messages::{ActorEvent, Request, StateDelta, WorkflowKvWrite};
39use crate::actor::metrics::ActorMetrics;
40use crate::actor::queue::{QueueInspectorUpdateCallback, QueueMetadata, QueueWaitActivityCallback};
41use crate::actor::schedule::{InternalKeepAwakeCallback, LocalAlarmCallback};
42use crate::actor::sleep::{CanSleep, SleepState};
43use crate::actor::state::{PendingSave, PersistedActor, RequestSaveOpts};
44use crate::actor::task::LifecycleEvent;
45use crate::actor::task_types::UserTaskKind;
46use crate::actor::work_registry::{ActorWorkKind, CountGuard, RegionGuard};
47use crate::error::{ActorLifecycle as ActorLifecycleError, ActorRuntime};
48use crate::inspector::{Inspector, InspectorSnapshot};
49use crate::sqlite::SqliteDb;
50use crate::types::{ActorKey, ConnId, ListOpts, format_actor_key};
51
52/// Shared actor runtime context.
53///
54/// This public surface is the foreign-runtime contract for `rivetkit-core`.
55/// Native Rust, NAPI-backed TypeScript, and future V8 runtimes should be able
56/// to drive actor behavior through `ActorFactory` plus the methods exposed here
57/// and on the returned runtime objects like `SqliteDb`, schedule APIs,
58/// queue APIs, `ConnHandle`, and `WebSocket`.
59#[derive(Clone)]
60pub struct ActorContext(pub(crate) Arc<ActorContextInner>);
61
62#[derive(Clone)]
63pub struct ActorKv {
64	sql: SqliteDb,
65}
66
67pub(crate) struct ActorContextInner {
68	pub(super) legacy_kv: LegacyActorKv,
69	user_kv: ActorKv,
70	sql: SqliteDb,
71	#[cfg(feature = "sqlite-local")]
72	actor_runtime_socket: ActorRuntimeSocketEndpoint,
73	// Forced-sync: actor state snapshots are exposed through synchronous
74	// accessors and are never held across `.await`.
75	pub(super) current_state: RwLock<Vec<u8>>,
76	pub(super) persisted: RwLock<PersistedActor>,
77	pub(super) last_pushed_alarm: RwLock<Option<i64>>,
78	pub(super) state_save_interval: Duration,
79	pub(super) state_dirty: AtomicBool,
80	pub(super) state_revision: AtomicU64,
81	pub(super) save_request_revision: AtomicU64,
82	pub(super) save_completed_revision: AtomicU64,
83	pub(super) save_completion: Notify,
84	pub(super) save_requested: AtomicBool,
85	pub(super) save_requested_immediate: AtomicBool,
86	// Forced-sync: debounce bookkeeping is updated from sync save-request paths.
87	pub(super) save_requested_within_deadline: Mutex<Option<crate::time::Instant>>,
88	pub(super) last_save_at: Mutex<Option<crate::time::Instant>>,
89	pub(super) pending_save: Mutex<Option<PendingSave>>,
90	pub(super) tracked_persist: Mutex<Option<JoinHandle<()>>>,
91	pub(super) save_guard: AsyncMutex<()>,
92	pub(super) in_flight_state_writes: AtomicUsize,
93	pub(super) state_write_completion: Notify,
94	pub(super) on_state_change_in_flight: AtomicUsize,
95	pub(super) on_state_change_idle: Notify,
96	// Forced-sync: hooks are registered and cloned from synchronous runtime
97	// wiring slots before use.
98	pub(super) request_save_hooks: RwLock<Vec<Arc<dyn Fn(RequestSaveOpts) + Send + Sync>>>,
99	// Forced-sync: schedule runtime handles and callbacks are synchronous
100	// wiring slots cloned before actor/envoy I/O.
101	pub(super) schedule_generation: Mutex<Option<u32>>,
102	pub(super) schedule_envoy_handle: Mutex<Option<EnvoyHandle>>,
103	pub(super) client_endpoint: OnceLock<String>,
104	pub(super) client_token: OnceLock<String>,
105	pub(super) client_namespace: OnceLock<String>,
106	pub(super) client_pool_name: OnceLock<String>,
107	pub(super) schedule_internal_keep_awake: Mutex<Option<InternalKeepAwakeCallback>>,
108	pub(super) schedule_local_alarm_callback: Mutex<Option<LocalAlarmCallback>>,
109	// Forced-sync: the local alarm timer is aborted from sync paths.
110	pub(super) schedule_local_alarm_task: Mutex<Option<JoinHandle<()>>>,
111	// Forced-sync: receivers are pushed/taken from sync paths and awaited after
112	// being moved out of the lock.
113	pub(super) schedule_pending_alarm_writes: Mutex<Vec<oneshot::Receiver<()>>>,
114	pub(super) schedule_local_alarm_epoch: AtomicU64,
115	pub(super) schedule_alarm_dispatch_enabled: AtomicBool,
116	pub(super) schedule_dirty_since_push: AtomicBool,
117	pub(super) schedule_mutation_lock: AsyncMutex<()>,
118	pub(super) schedule_running: SccHashSet<String>,
119	pub(super) schedule_history_insert_count: AtomicUsize,
120	pub(super) max_schedules: u32,
121	#[cfg(any(test, feature = "test-support"))]
122	pub(super) schedule_now_override: AtomicI64,
123	#[cfg(test)]
124	pub(super) schedule_driver_alarm_cancel_count: AtomicUsize,
125	#[cfg(test)]
126	pub(super) schedule_sync_alarm_failures: AtomicUsize,
127	// Forced-sync: queue config is read from sync public methods before blocking
128	// on async queue work.
129	pub(super) queue_config: Mutex<ActorConfig>,
130	pub(super) queue_abort_signal: Mutex<CancellationToken>,
131	pub(super) queue_initialize: OnceCell<()>,
132	pub(super) queue_metadata: AsyncMutex<QueueMetadata>,
133	pub(super) queue_receive_lock: AsyncMutex<()>,
134	pub(super) queue_completion_waiters: SccHashMap<u64, oneshot::Sender<Option<Vec<u8>>>>,
135	pub(super) queue_notify: Notify,
136	pub(super) active_queue_wait_count: AtomicU32,
137	// Forced-sync: callbacks are registered and cloned from synchronous hooks.
138	pub(super) queue_wait_activity_callback: Mutex<Option<QueueWaitActivityCallback>>,
139	pub(super) queue_inspector_update_callback: Mutex<Option<QueueInspectorUpdateCallback>>,
140	// Forced-sync: connection operations expose sync accessors or clone handles
141	// before awaiting; connection_disconnect_state serializes disconnect
142	// bookkeeping with pending hibernation snapshots.
143	pub(super) connection_config: RwLock<ActorConfig>,
144	pub(super) connections: RwLock<BTreeMap<ConnId, ConnHandle>>,
145	pub(super) pending_hibernation_updates: RwLock<BTreeSet<ConnId>>,
146	pub(super) pending_hibernation_removals: RwLock<BTreeSet<ConnId>>,
147	pub(super) connection_disconnect_state: Mutex<()>,
148	pub(super) sleep: SleepState,
149	activity: ActivityState,
150	sleep_requested: AtomicBool,
151	destroy_requested: AtomicBool,
152	destroy_completed: AtomicBool,
153	destroy_completion_notify: Notify,
154	abort_signal: Mutex<CancellationToken>,
155	shutdown_deadline: CancellationToken,
156	// Forced-sync: runtime wiring slots are configured through synchronous
157	// lifecycle setup and cloned before sending events.
158	inspector: RwLock<Option<Inspector>>,
159	inspector_attach_count: RwLock<Option<Arc<AtomicU32>>>,
160	inspector_overlay_tx: RwLock<Option<broadcast::Sender<Arc<Vec<u8>>>>>,
161	actor_events: RwLock<Option<mpsc::UnboundedSender<ActorEvent>>>,
162	pub(super) lifecycle_events: RwLock<Option<mpsc::UnboundedSender<LifecycleEvent>>>,
163	hibernated_connection_liveness_override: RwLock<Option<BTreeSet<(Vec<u8>, Vec<u8>)>>>,
164	pub(super) metrics: ActorMetrics,
165	diagnostics: ActorDiagnostics,
166	actor_id: String,
167	name: String,
168	key: ActorKey,
169	region: String,
170}
171
172#[derive(Debug, Default)]
173pub(crate) struct ActivityState {
174	dirty: AtomicBool,
175}
176
177impl ActivityState {
178	fn mark_dirty(&self) -> bool {
179		!self.dirty.swap(true, Ordering::AcqRel)
180	}
181
182	fn take_dirty(&self) -> bool {
183		self.dirty.swap(false, Ordering::AcqRel)
184	}
185}
186
187impl ActorKv {
188	pub async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
189		let mut values = self.batch_get(&[key]).await?;
190		Ok(values.pop().flatten())
191	}
192
193	pub async fn put(&self, key: &[u8], value: &[u8]) -> Result<()> {
194		self.batch_put(&[(key, value)]).await
195	}
196
197	pub async fn delete(&self, key: &[u8]) -> Result<()> {
198		self.batch_delete(&[key]).await
199	}
200
201	pub async fn batch_get(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>> {
202		internal_storage::user_kv_batch_get(&self.sql, keys).await
203	}
204
205	pub async fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<()> {
206		internal_storage::user_kv_batch_put(&self.sql, entries).await
207	}
208
209	pub async fn batch_delete(&self, keys: &[&[u8]]) -> Result<()> {
210		internal_storage::user_kv_batch_delete(&self.sql, keys).await
211	}
212
213	pub async fn delete_range(&self, start: &[u8], end: &[u8]) -> Result<()> {
214		internal_storage::user_kv_delete_range(&self.sql, start, end).await
215	}
216
217	pub async fn list_prefix(
218		&self,
219		prefix: &[u8],
220		opts: ListOpts,
221	) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
222		internal_storage::user_kv_list_prefix(&self.sql, prefix, opts).await
223	}
224
225	pub async fn list_range(
226		&self,
227		start: &[u8],
228		end: &[u8],
229		opts: ListOpts,
230	) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
231		internal_storage::user_kv_list_range(&self.sql, start, end, opts).await
232	}
233}
234
235impl ActorContext {
236	#[cfg(test)]
237	pub(crate) fn new(
238		actor_id: impl Into<String>,
239		name: impl Into<String>,
240		key: ActorKey,
241		region: impl Into<String>,
242	) -> Self {
243		tests::new_with_kv(actor_id, name, key, region, LegacyActorKv::new_in_memory())
244	}
245
246	#[cfg(test)]
247	pub(crate) fn new_with_kv(
248		actor_id: impl Into<String>,
249		name: impl Into<String>,
250		key: ActorKey,
251		region: impl Into<String>,
252		legacy_kv: LegacyActorKv,
253	) -> Self {
254		tests::new_with_kv(actor_id, name, key, region, legacy_kv)
255	}
256
257	#[cfg(test)]
258	pub(crate) fn new_for_state_tests(legacy_kv: LegacyActorKv, config: ActorConfig) -> Self {
259		let ctx = tests::new_with_kv("state-test", "state-test", Vec::new(), "local", legacy_kv);
260		ctx.configure_connection_runtime(config.clone());
261		ctx.configure_queue(config);
262		ctx
263	}
264
265	pub(crate) fn build(
266		actor_id: String,
267		name: String,
268		key: ActorKey,
269		region: String,
270		_generation: Option<u32>,
271		_envoy_key: String,
272		config: ActorConfig,
273		legacy_kv: LegacyActorKv,
274		sql: SqliteDb,
275	) -> Self {
276		let metrics = ActorMetrics::new(name.clone());
277		#[cfg(feature = "sqlite-local")]
278		let mut sql = sql;
279		#[cfg(feature = "sqlite-local")]
280		sql.set_vfs_metrics(Arc::new(metrics.clone()));
281		let diagnostics = ActorDiagnostics::new(actor_id.clone());
282		#[cfg(feature = "sqlite-local")]
283		let actor_runtime_socket =
284			ActorRuntimeSocketEndpoint::new(config.enable_actor_runtime_socket, sql.clone());
285		let state_save_interval = config.state_save_interval;
286		let max_schedules = config.max_schedules;
287		let abort_signal = CancellationToken::new();
288		let shutdown_deadline = CancellationToken::new();
289		let sleep = SleepState::new(config.clone());
290		let user_kv = ActorKv { sql: sql.clone() };
291		let ctx = Self(Arc::new(ActorContextInner {
292			legacy_kv,
293			user_kv,
294			sql,
295			#[cfg(feature = "sqlite-local")]
296			actor_runtime_socket,
297			current_state: RwLock::new(Vec::new()),
298			persisted: RwLock::new(PersistedActor::default()),
299			last_pushed_alarm: RwLock::new(None),
300			state_save_interval,
301			state_dirty: AtomicBool::new(false),
302			state_revision: AtomicU64::new(0),
303			save_request_revision: AtomicU64::new(0),
304			save_completed_revision: AtomicU64::new(0),
305			save_completion: Notify::new(),
306			save_requested: AtomicBool::new(false),
307			save_requested_immediate: AtomicBool::new(false),
308			save_requested_within_deadline: Mutex::new(None),
309			last_save_at: Mutex::new(None),
310			pending_save: Mutex::new(None),
311			tracked_persist: Mutex::new(None),
312			save_guard: AsyncMutex::new(()),
313			in_flight_state_writes: AtomicUsize::new(0),
314			state_write_completion: Notify::new(),
315			on_state_change_in_flight: AtomicUsize::new(0),
316			on_state_change_idle: Notify::new(),
317			request_save_hooks: RwLock::new(Vec::new()),
318			schedule_generation: Mutex::new(None),
319			schedule_envoy_handle: Mutex::new(None),
320			client_endpoint: OnceLock::new(),
321			client_token: OnceLock::new(),
322			client_namespace: OnceLock::new(),
323			client_pool_name: OnceLock::new(),
324			schedule_internal_keep_awake: Mutex::new(None),
325			schedule_local_alarm_callback: Mutex::new(None),
326			schedule_local_alarm_task: Mutex::new(None),
327			schedule_pending_alarm_writes: Mutex::new(Vec::new()),
328			schedule_local_alarm_epoch: AtomicU64::new(0),
329			schedule_alarm_dispatch_enabled: AtomicBool::new(true),
330			// A fresh actor context has no in-process record of a successful
331			// envoy alarm push yet, so the first sync must always push.
332			schedule_dirty_since_push: AtomicBool::new(true),
333			schedule_mutation_lock: AsyncMutex::new(()),
334			schedule_running: SccHashSet::new(),
335			schedule_history_insert_count: AtomicUsize::new(0),
336			max_schedules,
337			#[cfg(any(test, feature = "test-support"))]
338			schedule_now_override: AtomicI64::new(i64::MIN),
339			#[cfg(test)]
340			schedule_driver_alarm_cancel_count: AtomicUsize::new(0),
341			#[cfg(test)]
342			schedule_sync_alarm_failures: AtomicUsize::new(0),
343			queue_config: Mutex::new(config.clone()),
344			queue_abort_signal: Mutex::new(abort_signal.clone()),
345			queue_initialize: OnceCell::new(),
346			queue_metadata: AsyncMutex::new(QueueMetadata::default()),
347			queue_receive_lock: AsyncMutex::new(()),
348			queue_completion_waiters: SccHashMap::new(),
349			queue_notify: Notify::new(),
350			active_queue_wait_count: AtomicU32::new(0),
351			queue_wait_activity_callback: Mutex::new(None),
352			queue_inspector_update_callback: Mutex::new(None),
353			connection_config: RwLock::new(config),
354			connections: RwLock::new(BTreeMap::new()),
355			pending_hibernation_updates: RwLock::new(BTreeSet::new()),
356			pending_hibernation_removals: RwLock::new(BTreeSet::new()),
357			connection_disconnect_state: Mutex::new(()),
358			sleep,
359			activity: ActivityState::default(),
360			sleep_requested: AtomicBool::new(false),
361			destroy_requested: AtomicBool::new(false),
362			destroy_completed: AtomicBool::new(false),
363			destroy_completion_notify: Notify::new(),
364			abort_signal: Mutex::new(abort_signal),
365			shutdown_deadline,
366			inspector: RwLock::new(None),
367			inspector_attach_count: RwLock::new(None),
368			inspector_overlay_tx: RwLock::new(None),
369			actor_events: RwLock::new(None),
370			lifecycle_events: RwLock::new(None),
371			hibernated_connection_liveness_override: RwLock::new(None),
372			metrics,
373			diagnostics,
374			actor_id,
375			name,
376			key,
377			region,
378		}));
379		ctx.configure_sleep_hooks();
380		ctx
381	}
382
383	#[deprecated(
384		note = "Actor KV is deprecated. Use embedded SQLite (`sql()`) or actor state instead."
385	)]
386	pub async fn kv_batch_get(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>> {
387		self.0.user_kv.batch_get(keys).await
388	}
389
390	#[deprecated(
391		note = "Actor KV is deprecated. Use embedded SQLite (`sql()`) or actor state instead."
392	)]
393	pub async fn kv_batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<()> {
394		self.0.user_kv.batch_put(entries).await
395	}
396
397	#[deprecated(
398		note = "Actor KV is deprecated. Use embedded SQLite (`sql()`) or actor state instead."
399	)]
400	pub async fn kv_batch_delete(&self, keys: &[&[u8]]) -> Result<()> {
401		self.0.user_kv.batch_delete(keys).await
402	}
403
404	#[deprecated(
405		note = "Actor KV is deprecated. Use embedded SQLite (`sql()`) or actor state instead."
406	)]
407	pub async fn kv_delete_range(&self, start: &[u8], end: &[u8]) -> Result<()> {
408		self.0.user_kv.delete_range(start, end).await
409	}
410
411	#[deprecated(
412		note = "Actor KV is deprecated. Use embedded SQLite (`sql()`) or actor state instead."
413	)]
414	pub async fn kv_list_prefix(
415		&self,
416		prefix: &[u8],
417		opts: ListOpts,
418	) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
419		self.0.user_kv.list_prefix(prefix, opts).await
420	}
421
422	#[deprecated(
423		note = "Actor KV is deprecated. Use embedded SQLite (`sql()`) or actor state instead."
424	)]
425	pub async fn kv_list_range(
426		&self,
427		start: &[u8],
428		end: &[u8],
429		opts: ListOpts,
430	) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
431		self.0.user_kv.list_range(start, end, opts).await
432	}
433
434	#[deprecated(
435		note = "Actor KV is deprecated. Use embedded SQLite (`sql()`) or actor state instead."
436	)]
437	pub fn kv(&self) -> &ActorKv {
438		&self.0.user_kv
439	}
440
441	/// Legacy actor KV access used only by the one-time KV-to-SQLite importer and
442	/// the temporary inspector-token mirror consumed by the dashboard.
443	pub(crate) fn legacy_kv(&self) -> &LegacyActorKv {
444		&self.0.legacy_kv
445	}
446
447	pub fn sql(&self) -> &SqliteDb {
448		&self.0.sql
449	}
450
451	#[cfg(feature = "sqlite-local")]
452	pub async fn provision_actor_runtime_socket(&self) -> Result<ActorRuntimeSocketEndpointInfo> {
453		self.0.actor_runtime_socket.provision().await
454	}
455
456	#[cfg(feature = "sqlite-local")]
457	pub(crate) async fn shutdown_actor_runtime_socket(&self) {
458		self.0.actor_runtime_socket.shutdown().await;
459	}
460
461	pub async fn db_exec(&self, sql: &str) -> Result<Vec<u8>> {
462		self.0.sql.exec_rows_cbor(sql).await
463	}
464
465	pub async fn db_query(&self, sql: &str, params: Option<&[u8]>) -> Result<Vec<u8>> {
466		self.0.sql.query_rows_cbor(sql, params).await
467	}
468
469	pub async fn db_execute(&self, sql: &str, params: Option<&[u8]>) -> Result<Vec<u8>> {
470		self.0.sql.execute_rows_cbor(sql, params).await
471	}
472
473	pub async fn db_run(&self, sql: &str, params: Option<&[u8]>) -> Result<()> {
474		self.0.sql.run_cbor(sql, params).await?;
475		Ok(())
476	}
477
478	pub fn set_alarm(&self, timestamp_ms: Option<i64>) -> Result<()> {
479		self.set_schedule_alarm(timestamp_ms)
480	}
481
482	/// Resync persisted alarms with the runtime's alarm transport.
483	///
484	/// Foreign-runtime adapters should call this during startup after loading
485	/// any persisted schedule state and before accepting user callbacks that rely
486	/// on future alarms being armed.
487	pub async fn init_alarms(&self) {
488		if let Err(error) = self.recover_interrupted_schedule_history().await {
489			tracing::error!(?error, "failed to recover interrupted schedule history");
490		}
491		self.sync_future_alarm_logged().await;
492	}
493
494	pub fn queue(&self) -> &Self {
495		self
496	}
497
498	pub fn sleep(&self) -> Result<()> {
499		// `started` is cleared when the lifecycle state machine transitions
500		// into SleepGrace / DestroyGrace, so `started=false` covers both
501		// "never started" and "already shutting down". Distinguish with the
502		// request flags for an accurate diagnostic.
503		if !self.0.sleep.lifecycle_started.load(Ordering::SeqCst) {
504			let already_stopping = self.0.sleep_requested.load(Ordering::SeqCst)
505				|| self.0.destroy_requested.load(Ordering::SeqCst);
506			return if already_stopping {
507				Err(ActorLifecycleError::Stopping.build()).context("actor is already shutting down")
508			} else {
509				Err(ActorLifecycleError::Starting.build())
510					.context("cannot request sleep before actor startup completes")
511			};
512		}
513		if self.0.sleep_requested.swap(true, Ordering::SeqCst) {
514			return Ok(());
515		}
516		self.cancel_sleep_timer();
517		if Handle::try_current().is_ok() {
518			let ctx = self.clone();
519			let tracked = self.track_shutdown_task(async move {
520				ctx.record_user_task_started(UserTaskKind::SleepFinalize);
521				let started_at = Instant::now();
522				ctx.request_sleep_from_envoy();
523				ctx.record_user_task_finished(UserTaskKind::SleepFinalize, started_at.elapsed());
524			});
525			if tracked {
526				return Ok(());
527			}
528		}
529
530		self.request_sleep_from_envoy();
531		Ok(())
532	}
533
534	pub fn destroy(&self) -> Result<()> {
535		self.request_stop(None)
536	}
537
538	/// Request a stop with an error attached. Behaves like [`Self::destroy`]
539	/// locally (destroy grace hooks run for this generation), but the envoy
540	/// reports the stop to the engine with `StopCode::Error` and the message,
541	/// so the engine records the crash and applies its crash handling instead
542	/// of unconditionally destroying the actor.
543	pub fn stop_with_error(&self, message: impl Into<String>) -> Result<()> {
544		self.request_stop(Some(truncate_stop_error_message(message.into())))
545	}
546
547	fn request_stop(&self, error: Option<String>) -> Result<()> {
548		// See `sleep` for why the request flags disambiguate `started=false`.
549		// destroy() is allowed after sleep() has been requested because
550		// destroy is a stronger signal that escalates an in-flight sleep.
551		if !self.0.sleep.lifecycle_started.load(Ordering::SeqCst)
552			&& !self.0.sleep_requested.load(Ordering::SeqCst)
553			&& !self.0.destroy_requested.load(Ordering::SeqCst)
554		{
555			return Err(ActorLifecycleError::Starting.build())
556				.context("cannot request destroy before actor startup completes");
557		}
558		if self.0.destroy_requested.swap(true, Ordering::SeqCst) {
559			return Err(ActorLifecycleError::Stopping.build())
560				.context("destroy already requested for this generation");
561		}
562		// Winning the swap above makes this the only writer of the error slot
563		// for this generation. The slot is consumed by
564		// `request_destroy_from_envoy` when the stop intent is sent.
565		if error.is_some() {
566			*self.0.sleep.destroy_error.lock() = error;
567		}
568		// Reuse the shared teardown sequence used by the registry shutdown path
569		// so future changes to `mark_destroy_requested` cannot drift.
570		// `destroy_requested` is already true from the swap above. The redundant
571		// `store(true)` inside is harmless.
572		#[cfg(not(feature = "wasm-runtime"))]
573		self.mark_destroy_requested();
574		#[cfg(feature = "wasm-runtime")]
575		self.mark_destroy_requested_without_spawn();
576
577		let ctx = self.clone();
578		if Handle::try_current().is_ok() {
579			let tracked = self.track_shutdown_task(async move {
580				ctx.record_user_task_started(UserTaskKind::DestroyRequest);
581				let started_at = Instant::now();
582				ctx.request_destroy_from_envoy();
583				ctx.record_user_task_finished(UserTaskKind::DestroyRequest, started_at.elapsed());
584			});
585			if tracked {
586				return Ok(());
587			}
588		}
589
590		self.request_destroy_from_envoy();
591		Ok(())
592	}
593
594	pub fn mark_destroy_requested(&self) {
595		self.cancel_sleep_timer();
596		self.flush_on_shutdown();
597		self.0.destroy_requested.store(true, Ordering::SeqCst);
598		self.0.destroy_completed.store(false, Ordering::SeqCst);
599	}
600
601	#[cfg(feature = "wasm-runtime")]
602	fn mark_destroy_requested_without_spawn(&self) {
603		self.cancel_sleep_timer();
604		self.0.destroy_requested.store(true, Ordering::SeqCst);
605		self.0.destroy_completed.store(false, Ordering::SeqCst);
606	}
607
608	#[doc(hidden)]
609	pub fn cancel_actor_abort_signal(&self) {
610		self.0.abort_signal.lock().cancel();
611	}
612
613	pub(crate) fn reset_abort_signal_for_start(&self) {
614		let mut abort_signal = self.0.abort_signal.lock();
615		if !abort_signal.is_cancelled() {
616			return;
617		}
618
619		// Sleep or destroy cancels the generation abort signal to break actor
620		// scoped waits. A restarted actor needs a fresh signal so the next
621		// generation can wait normally.
622		let next_signal = CancellationToken::new();
623		*abort_signal = next_signal.clone();
624		*self.0.queue_abort_signal.lock() = next_signal;
625	}
626
627	#[doc(hidden)]
628	pub fn actor_abort_signal(&self) -> CancellationToken {
629		self.0.abort_signal.lock().clone()
630	}
631
632	#[doc(hidden)]
633	pub fn actor_aborted(&self) -> bool {
634		self.0.abort_signal.lock().is_cancelled()
635	}
636
637	/// Fires when the shutdown grace deadline has elapsed and core is forcing
638	/// cleanup. Foreign-runtime adapters should abort any in-flight shutdown
639	/// work (for example `onSleep` / `onDestroy`) when this token is cancelled
640	/// so resources like SQLite are not torn down mid-operation.
641	#[doc(hidden)]
642	pub fn shutdown_deadline_token(&self) -> CancellationToken {
643		self.0.shutdown_deadline.clone()
644	}
645
646	#[doc(hidden)]
647	pub fn cancel_shutdown_deadline(&self) {
648		self.0.shutdown_deadline.cancel();
649	}
650
651	/// Deprecated no-op. Use `keep_awake` to hold the actor awake for the
652	/// duration of a future, or `wait_until` to keep work alive across the
653	/// sleep grace period. Retained only for NAPI bridge compatibility.
654	#[deprecated(note = "no-op: use `keep_awake` or `wait_until` instead")]
655	pub fn set_prevent_sleep(&self, _enabled: bool) {}
656
657	#[deprecated(note = "no-op: always returns false")]
658	pub fn prevent_sleep(&self) -> bool {
659		false
660	}
661
662	#[cfg(not(feature = "wasm-runtime"))]
663	pub fn wait_until(&self, future: impl Future<Output = ()> + Send + 'static) {
664		self.spawn_work(ActorWorkKind::WaitUntil, future);
665	}
666
667	#[cfg(not(feature = "wasm-runtime"))]
668	pub fn register_task(&self, future: impl Future<Output = ()> + Send + 'static) {
669		self.spawn_work(ActorWorkKind::RegisteredTask, future);
670	}
671
672	#[cfg(feature = "wasm-runtime")]
673	pub fn wait_until(&self, future: impl Future<Output = ()> + 'static) {
674		self.spawn_work(ActorWorkKind::WaitUntil, future);
675	}
676
677	#[cfg(feature = "wasm-runtime")]
678	pub fn register_task(&self, future: impl Future<Output = ()> + 'static) {
679		self.spawn_work(ActorWorkKind::RegisteredTask, future);
680	}
681
682	pub async fn keep_awake<F>(&self, future: F) -> F::Output
683	where
684		F: Future,
685	{
686		self.track_work(ActorWorkKind::KeepAwake, future).await
687	}
688
689	pub fn keep_awake_region(&self) -> KeepAwakeRegion {
690		KeepAwakeRegion {
691			region: Some(self.begin_work_region(ActorWorkKind::KeepAwake)),
692		}
693	}
694
695	pub async fn internal_keep_awake<F>(&self, future: F) -> F::Output
696	where
697		F: Future,
698	{
699		self.track_work(ActorWorkKind::InternalKeepAwake, future)
700			.await
701	}
702
703	pub fn keep_awake_count(&self) -> usize {
704		self.sleep_keep_awake_count()
705	}
706
707	pub fn internal_keep_awake_count(&self) -> usize {
708		self.sleep_internal_keep_awake_count()
709	}
710
711	pub async fn track_work<F>(&self, kind: ActorWorkKind, future: F) -> F::Output
712	where
713		F: Future,
714	{
715		let _region = self.begin_work_region(kind);
716		future.await
717	}
718
719	#[cfg(not(feature = "wasm-runtime"))]
720	pub fn spawn_work<F>(&self, kind: ActorWorkKind, future: F)
721	where
722		F: Future<Output = ()> + Send + 'static,
723	{
724		self.spawn_work_inner(kind, future);
725	}
726
727	#[cfg(feature = "wasm-runtime")]
728	pub fn spawn_work<F>(&self, kind: ActorWorkKind, future: F)
729	where
730		F: Future<Output = ()> + 'static,
731	{
732		self.spawn_work_inner(kind, future);
733	}
734
735	pub fn begin_work_region(&self, kind: ActorWorkKind) -> ActorWorkRegion {
736		ActorWorkRegion {
737			guard: Some(ActorWorkGuard::new(self.clone(), kind)),
738		}
739	}
740
741	pub fn actor_id(&self) -> &str {
742		&self.0.actor_id
743	}
744
745	pub fn name(&self) -> &str {
746		&self.0.name
747	}
748
749	pub fn key(&self) -> &ActorKey {
750		&self.0.key
751	}
752
753	pub(crate) fn actor_specifier(&self) -> Option<ActorSpecifier> {
754		Some(
755			ActorSpecifier::new(self.actor_id().to_owned(), self.sleep_generation()? as u64)
756				.with_key(format_actor_key(self.key())),
757		)
758	}
759
760	pub(crate) fn attach_actor_to_error(&self, error: anyhow::Error) -> anyhow::Error {
761		match self.actor_specifier() {
762			Some(actor) => error.context(actor),
763			None => error,
764		}
765	}
766
767	pub fn region(&self) -> &str {
768		&self.0.region
769	}
770
771	pub fn has_state(&self) -> bool {
772		self.0.connection_config.read().has_state
773	}
774
775	#[doc(hidden)]
776	pub fn record_startup_create_state(&self, duration: Duration) {
777		self.0.metrics.observe_create_state(duration);
778	}
779
780	#[doc(hidden)]
781	pub fn record_startup_create_vars(&self, duration: Duration) {
782		self.0.metrics.observe_create_vars(duration);
783	}
784
785	pub fn broadcast(&self, name: &str, args: &[u8]) {
786		for connection in self.conns() {
787			if connection.is_subscribed(name) {
788				connection.send(name, args);
789			}
790		}
791	}
792
793	/// Returns a lock-backed iterator over live connections.
794	///
795	/// Do not hold the returned iterator across `.await`. It keeps a read lock
796	/// on the connection map until dropped, which blocks connection writers.
797	#[must_use]
798	pub fn conns(&self) -> ConnHandles<'_> {
799		self.iter_connections()
800	}
801
802	pub fn client_endpoint(&self) -> Option<&str> {
803		self.0.client_endpoint.get().map(String::as_str)
804	}
805
806	pub fn client_token(&self) -> Option<&str> {
807		self.0.client_token.get().map(String::as_str)
808	}
809
810	pub fn client_namespace(&self) -> Option<&str> {
811		self.0.client_namespace.get().map(String::as_str)
812	}
813
814	pub fn client_pool_name(&self) -> Option<&str> {
815		self.0.client_pool_name.get().map(String::as_str)
816	}
817
818	pub fn ack_hibernatable_websocket_message(
819		&self,
820		gateway_id: &[u8],
821		request_id: &[u8],
822		server_message_index: u16,
823	) -> Result<()> {
824		let gateway_id = hibernatable_id_from_slice("gateway_id", gateway_id)?;
825		let request_id = hibernatable_id_from_slice("request_id", request_id)?;
826		let envoy_handle = self.sleep_envoy_handle().ok_or_else(|| {
827			ActorRuntime::NotConfigured {
828				component: "hibernatable websocket ack".to_owned(),
829			}
830			.build()
831		})?;
832		envoy_handle.send_hibernatable_ws_message_ack(gateway_id, request_id, server_message_index);
833		Ok(())
834	}
835
836	pub(crate) fn load_persisted_actor(&self, persisted: PersistedActor) {
837		self.load_persisted(persisted);
838	}
839
840	pub(crate) fn persisted_actor(&self) -> PersistedActor {
841		self.persisted()
842	}
843
844	/// Dispatches any scheduled actions whose deadline has already passed.
845	///
846	/// Foreign-runtime adapters should call this after startup callbacks complete
847	/// so overdue scheduled work enters the normal actor event loop.
848	pub async fn drain_overdue_scheduled_events(&self) -> Result<()> {
849		for dispatch in self.take_due_schedule_dispatches().await? {
850			self.dispatch_scheduled_action(dispatch).await;
851		}
852		Ok(())
853	}
854
855	pub(crate) fn metrics(&self) -> &ActorMetrics {
856		&self.0.metrics
857	}
858
859	pub(crate) fn record_user_task_started(&self, kind: UserTaskKind) {
860		self.0.metrics.begin_user_task(kind);
861	}
862
863	pub(crate) fn record_user_task_finished(&self, kind: UserTaskKind, duration: Duration) {
864		self.0.metrics.end_user_task(kind, duration);
865	}
866
867	pub(crate) fn record_shutdown_wait(
868		&self,
869		reason: crate::actor::task_types::ShutdownKind,
870		duration: Duration,
871	) {
872		self.0.metrics.observe_shutdown_wait(reason, duration);
873	}
874
875	pub(crate) fn record_shutdown_timeout(&self, reason: crate::actor::task_types::ShutdownKind) {
876		self.0.metrics.inc_shutdown_timeout(reason);
877	}
878
879	pub(crate) fn record_direct_subsystem_shutdown_warning(
880		&self,
881		subsystem: &str,
882		operation: &str,
883	) {
884		self.0
885			.metrics
886			.inc_direct_subsystem_shutdown_warning(subsystem, operation);
887	}
888
889	pub(crate) fn warn_work_sent_to_stopping_instance(&self, operation: &'static str) {
890		if let Some(suppression) = self.0.diagnostics.record("work_sent_to_stopping_instance") {
891			tracing::warn!(
892				actor_id = %suppression.actor_id,
893				operation,
894				per_actor_suppressed = suppression.per_actor_suppressed,
895				global_suppressed = suppression.global_suppressed,
896				"work sent to stopping actor instance"
897			);
898		}
899	}
900
901	pub(crate) fn warn_self_call_risk(&self, operation: &'static str) {
902		if let Some(suppression) = self.0.diagnostics.record("self_call_risk") {
903			tracing::warn!(
904				actor_id = %suppression.actor_id,
905				operation,
906				per_actor_suppressed = suppression.per_actor_suppressed,
907				global_suppressed = suppression.global_suppressed,
908				"actor dispatch may be parked behind the current instance"
909			);
910		}
911	}
912
913	#[cfg(test)]
914	pub(crate) fn add_conn(&self, conn: ConnHandle) {
915		self.insert_existing(conn);
916		self.record_connections_updated();
917		self.reset_sleep_timer();
918	}
919
920	pub(crate) fn remove_conn(&self, conn_id: &str) -> Option<ConnHandle> {
921		let removed = self.remove_existing(conn_id);
922		if removed.is_some() {
923			self.record_connections_updated();
924			self.reset_sleep_timer();
925		}
926		removed
927	}
928
929	pub(crate) fn configure_connection_runtime(&self, config: ActorConfig) {
930		self.configure_sleep_state(config.clone());
931		self.configure_connection_storage(config);
932	}
933
934	pub(crate) fn configure_actor_events(&self, sender: Option<mpsc::UnboundedSender<ActorEvent>>) {
935		*self.0.actor_events.write() = sender;
936	}
937
938	pub(crate) fn try_send_actor_event(
939		&self,
940		event: ActorEvent,
941		operation: &'static str,
942	) -> Result<()> {
943		let sender = self.0.actor_events.read().clone().ok_or_else(|| {
944			ActorRuntime::NotConfigured {
945				component: "actor event inbox".to_owned(),
946			}
947			.build()
948		})?;
949		tracing::debug!(
950			actor_id = %self.actor_id(),
951			operation,
952			event = event.kind(),
953			"actor event enqueued"
954		);
955		sender.send(event).map_err(|_| {
956			ActorRuntime::NotConfigured {
957				component: "actor event inbox".to_owned(),
958			}
959			.build()
960		})
961	}
962
963	#[doc(hidden)]
964	pub fn configure_envoy(&self, envoy_handle: EnvoyHandle, generation: Option<u32>) {
965		let _ = self
966			.0
967			.client_endpoint
968			.set(envoy_handle.endpoint().to_owned());
969		if let Some(token) = envoy_handle.token() {
970			let _ = self.0.client_token.set(token.to_owned());
971		}
972		let _ = self
973			.0
974			.client_namespace
975			.set(envoy_handle.namespace().to_owned());
976		let _ = self
977			.0
978			.client_pool_name
979			.set(envoy_handle.pool_name().to_owned());
980		self.configure_sleep_envoy(envoy_handle.clone(), generation);
981		self.configure_schedule_envoy(envoy_handle, generation);
982	}
983
984	pub(crate) async fn connect_conn<F>(
985		&self,
986		params: Vec<u8>,
987		is_hibernatable: bool,
988		hibernation: Option<HibernatableConnectionMetadata>,
989		request: Option<Request>,
990		create_state: F,
991	) -> Result<ConnHandle>
992	where
993		F: Future<Output = Result<Vec<u8>>> + Send,
994	{
995		let conn = self
996			.connect_with_state(params, is_hibernatable, hibernation, request, create_state)
997			.await?;
998		self.record_connections_updated();
999		self.reset_sleep_timer();
1000		Ok(conn)
1001	}
1002
1003	pub(crate) async fn connect_conn_with_prepare<F, P>(
1004		&self,
1005		params: Vec<u8>,
1006		is_hibernatable: bool,
1007		hibernation: Option<HibernatableConnectionMetadata>,
1008		request: Option<Request>,
1009		create_state: F,
1010		prepare_connection: P,
1011	) -> Result<ConnHandle>
1012	where
1013		F: Future<Output = Result<Vec<u8>>> + Send,
1014		P: FnOnce(&ConnHandle) -> Result<()>,
1015	{
1016		let conn = self
1017			.connect_with_state_and_prepare(
1018				params,
1019				is_hibernatable,
1020				hibernation,
1021				request,
1022				create_state,
1023				prepare_connection,
1024			)
1025			.await?;
1026		self.record_connections_updated();
1027		self.reset_sleep_timer();
1028		Ok(conn)
1029	}
1030
1031	pub async fn connect_conn_with_request<F>(
1032		&self,
1033		params: Vec<u8>,
1034		request: Option<Request>,
1035		create_state: F,
1036	) -> Result<ConnHandle>
1037	where
1038		F: Future<Output = Result<Vec<u8>>> + Send,
1039	{
1040		self.connect_conn(params, false, None, request, create_state)
1041			.await
1042	}
1043
1044	pub(crate) fn reconnect_hibernatable_conn(
1045		&self,
1046		gateway_id: &[u8],
1047		request_id: &[u8],
1048	) -> Result<ConnHandle> {
1049		self.reconnect_hibernatable(gateway_id, request_id)
1050	}
1051
1052	pub async fn disconnect_conn(&self, id: ConnId) -> Result<()> {
1053		self.disconnect_transport_only(|conn| conn.id() == id).await
1054	}
1055
1056	pub async fn disconnect_conns<F>(&self, predicate: F) -> Result<()>
1057	where
1058		F: FnMut(&ConnHandle) -> bool,
1059	{
1060		self.disconnect_transport_only(predicate).await
1061	}
1062
1063	pub(crate) fn request_hibernation_transport_save(&self, conn_id: &str) {
1064		self.queue_hibernation_update(conn_id.to_owned());
1065		self.request_save(RequestSaveOpts::default());
1066	}
1067
1068	pub(crate) fn request_hibernation_transport_removal(&self, conn_id: impl Into<String>) {
1069		self.queue_hibernation_removal_inner(conn_id.into());
1070		self.request_save(RequestSaveOpts::default());
1071	}
1072
1073	pub fn queue_hibernation_removal(&self, conn_id: impl Into<String>) {
1074		self.request_hibernation_transport_removal(conn_id);
1075	}
1076
1077	pub fn has_pending_hibernation_changes(&self) -> bool {
1078		self.has_pending_hibernation_changes_inner()
1079	}
1080
1081	pub fn take_pending_hibernation_changes(&self) -> Vec<ConnId> {
1082		self.pending_hibernation_removals()
1083	}
1084
1085	pub fn dirty_hibernatable_conns(&self) -> Vec<ConnHandle> {
1086		self.dirty_hibernatable_conns_inner()
1087	}
1088
1089	pub(crate) fn hibernated_connection_is_live(
1090		&self,
1091		gateway_id: &[u8],
1092		request_id: &[u8],
1093	) -> Result<bool> {
1094		let gateway_id = hibernatable_id_from_slice("gateway_id", gateway_id)?;
1095		let request_id = hibernatable_id_from_slice("request_id", request_id)?;
1096
1097		if let Some(override_pairs) = self
1098			.0
1099			.hibernated_connection_liveness_override
1100			.read()
1101			.as_ref()
1102		{
1103			return Ok(override_pairs.contains(&(gateway_id.to_vec(), request_id.to_vec())));
1104		}
1105
1106		let Some(envoy_handle) = self.sleep_envoy_handle() else {
1107			return Ok(false);
1108		};
1109		let is_live = envoy_handle.hibernatable_connection_is_live(
1110			self.actor_id(),
1111			self.sleep_generation(),
1112			gateway_id,
1113			request_id,
1114		);
1115		Ok(is_live)
1116	}
1117
1118	#[cfg(test)]
1119	pub(crate) fn set_hibernated_connection_liveness_override<I>(&self, pairs: I)
1120	where
1121		I: IntoIterator<Item = (Vec<u8>, Vec<u8>)>,
1122	{
1123		*self.0.hibernated_connection_liveness_override.write() = Some(pairs.into_iter().collect());
1124	}
1125
1126	fn prepare_state_deltas(
1127		&self,
1128		deltas: Vec<StateDelta>,
1129	) -> Result<(Vec<StateDelta>, PendingHibernationChanges)> {
1130		fn finish_with_error(
1131			ctx: &ActorContext,
1132			pending: PendingHibernationChanges,
1133			error: anyhow::Error,
1134		) -> Result<(Vec<StateDelta>, PendingHibernationChanges)> {
1135			ctx.restore_pending_hibernation_changes(pending);
1136			Err(error)
1137		}
1138
1139		let mut next_deltas = Vec::new();
1140		let mut explicit_updates = std::collections::BTreeMap::new();
1141		let mut explicit_removals = std::collections::BTreeSet::new();
1142
1143		for delta in deltas {
1144			match delta {
1145				StateDelta::ConnHibernation { conn, bytes } => {
1146					explicit_updates.insert(conn, bytes);
1147				}
1148				StateDelta::ConnHibernationRemoved(conn) => {
1149					explicit_removals.insert(conn);
1150				}
1151				other => next_deltas.push(other),
1152			}
1153		}
1154
1155		let mut pending = self.take_pending_hibernation_changes_inner();
1156		let mut removal_ids = pending.removed.clone();
1157		removal_ids.extend(explicit_removals.iter().cloned());
1158
1159		let explicit_update_ids: std::collections::BTreeSet<_> =
1160			explicit_updates.keys().cloned().collect();
1161		pending.updated.extend(explicit_update_ids.iter().cloned());
1162		pending.removed.extend(explicit_removals.iter().cloned());
1163
1164		for (conn, bytes) in explicit_updates {
1165			if removal_ids.contains(&conn) {
1166				continue;
1167			}
1168			let encoded = match self.encode_hibernation_delta(&conn, bytes) {
1169				Ok(encoded) => encoded,
1170				Err(error) => {
1171					return finish_with_error(self, pending, error);
1172				}
1173			};
1174			next_deltas.push(StateDelta::ConnHibernation {
1175				conn,
1176				bytes: encoded,
1177			});
1178		}
1179
1180		for conn in &pending.updated {
1181			if removal_ids.contains(conn)
1182				|| explicit_removals.contains(conn)
1183				|| explicit_update_ids.contains(conn)
1184			{
1185				continue;
1186			}
1187			let Some(handle) = self.connection(conn) else {
1188				continue;
1189			};
1190			if !handle.is_hibernatable() || handle.hibernation().is_none() {
1191				continue;
1192			}
1193			let encoded = match self.encode_hibernation_delta(conn, handle.state()) {
1194				Ok(encoded) => encoded,
1195				Err(error) => {
1196					return finish_with_error(self, pending, error);
1197				}
1198			};
1199			next_deltas.push(StateDelta::ConnHibernation {
1200				conn: conn.clone(),
1201				bytes: encoded,
1202			});
1203		}
1204
1205		for conn in removal_ids {
1206			next_deltas.push(StateDelta::ConnHibernationRemoved(conn));
1207		}
1208
1209		Ok((next_deltas, pending))
1210	}
1211
1212	pub(crate) async fn restore_hibernatable_connections(&self) -> Result<Vec<ConnHandle>> {
1213		let restored = self.restore_persisted().await?;
1214		if !restored.is_empty() {
1215			if let Some(envoy_handle) = self.sleep_envoy_handle() {
1216				let meta_entries: Vec<_> = restored
1217					.iter()
1218					.filter_map(|conn| {
1219						let hibernation = conn.hibernation()?;
1220						Some(HibernatingWebSocketMetadata {
1221							gateway_id: hibernation.gateway_id,
1222							request_id: hibernation.request_id,
1223							envoy_message_index: hibernation.client_message_index,
1224							rivet_message_index: hibernation.server_message_index,
1225							path: hibernation.request_path,
1226							headers: hibernation.request_headers.into_iter().collect(),
1227						})
1228					})
1229					.collect();
1230				envoy_handle.restore_hibernating_requests(self.actor_id().to_owned(), meta_entries);
1231			}
1232			self.record_connections_updated();
1233			self.reset_sleep_timer();
1234		}
1235		Ok(restored)
1236	}
1237
1238	pub(crate) fn configure_inspector(&self, inspector: Option<Inspector>) {
1239		*self.0.inspector.write() = inspector;
1240	}
1241
1242	pub(crate) fn inspector(&self) -> Option<Inspector> {
1243		self.0.inspector.read().clone()
1244	}
1245
1246	pub fn inspector_snapshot(&self) -> InspectorSnapshot {
1247		self.inspector()
1248			.map(|inspector| inspector.snapshot())
1249			.unwrap_or_default()
1250	}
1251
1252	pub(crate) fn configure_inspector_runtime(
1253		&self,
1254		attach_count: Arc<AtomicU32>,
1255		overlay_tx: broadcast::Sender<Arc<Vec<u8>>>,
1256	) {
1257		*self.0.inspector_attach_count.write() = Some(attach_count);
1258		*self.0.inspector_overlay_tx.write() = Some(overlay_tx);
1259	}
1260
1261	pub(crate) fn inspector_attach(&self) -> Option<InspectorAttachGuard> {
1262		InspectorAttachGuard::new(self.clone())
1263	}
1264
1265	#[cfg(test)]
1266	pub(crate) fn inspector_attach_count(&self) -> u32 {
1267		self.inspector_attach_count_arc()
1268			.map(|attach_count| attach_count.load(Ordering::SeqCst))
1269			.unwrap_or(0)
1270	}
1271
1272	pub(crate) fn subscribe_inspector(&self) -> Option<broadcast::Receiver<Arc<Vec<u8>>>> {
1273		self.0
1274			.inspector_overlay_tx
1275			.read()
1276			.clone()
1277			.map(|overlay_tx| overlay_tx.subscribe())
1278	}
1279
1280	pub(crate) fn downgrade(&self) -> Weak<ActorContextInner> {
1281		Arc::downgrade(&self.0)
1282	}
1283
1284	pub(crate) fn from_weak(weak: &Weak<ActorContextInner>) -> Option<Self> {
1285		weak.upgrade().map(Self)
1286	}
1287
1288	#[doc(hidden)]
1289	pub fn set_started(&self, started: bool) {
1290		self.set_lifecycle_started(started);
1291		self.reset_sleep_timer();
1292	}
1293
1294	#[doc(hidden)]
1295	pub fn started(&self) -> bool {
1296		self.lifecycle_started()
1297	}
1298
1299	pub(crate) fn destroy_requested(&self) -> bool {
1300		self.0.destroy_requested.load(Ordering::SeqCst)
1301	}
1302
1303	pub fn is_destroy_requested(&self) -> bool {
1304		self.destroy_requested()
1305	}
1306
1307	pub(crate) async fn wait_for_destroy_completion(&self) {
1308		if self.0.destroy_completed.load(Ordering::SeqCst) {
1309			return;
1310		}
1311
1312		loop {
1313			let notified = self.0.destroy_completion_notify.notified();
1314			if self.0.destroy_completed.load(Ordering::SeqCst) {
1315				return;
1316			}
1317			notified.await;
1318			if self.0.destroy_completed.load(Ordering::SeqCst) {
1319				return;
1320			}
1321		}
1322	}
1323
1324	pub async fn wait_for_destroy_completion_public(&self) {
1325		self.wait_for_destroy_completion().await;
1326	}
1327
1328	pub(crate) fn mark_destroy_completed(&self) {
1329		self.0.destroy_completed.store(true, Ordering::SeqCst);
1330		self.0.destroy_completion_notify.notify_waiters();
1331	}
1332
1333	pub(crate) async fn can_sleep(&self) -> CanSleep {
1334		self.can_arm_sleep_timer().await
1335	}
1336
1337	pub(crate) fn pending_disconnect_count(&self) -> usize {
1338		self.0.sleep.work.disconnect_callback.load()
1339	}
1340
1341	pub async fn with_disconnect_callback<F, Fut, T>(&self, run: F) -> T
1342	where
1343		F: FnOnce() -> Fut,
1344		Fut: Future<Output = T>,
1345	{
1346		self.track_work(ActorWorkKind::DisconnectCallback, run())
1347			.await
1348	}
1349
1350	pub(crate) fn configure_lifecycle_events(
1351		&self,
1352		sender: Option<mpsc::UnboundedSender<LifecycleEvent>>,
1353	) {
1354		*self.0.lifecycle_events.write() = sender;
1355	}
1356
1357	pub(crate) fn notify_inspector_serialize_requested(&self) {
1358		self.try_send_lifecycle_event(
1359			LifecycleEvent::InspectorSerializeRequested,
1360			"inspector_serialize_requested",
1361		);
1362	}
1363
1364	pub(crate) fn notify_activity_dirty(&self) -> bool {
1365		if self.0.lifecycle_events.read().is_none() {
1366			return false;
1367		}
1368		if self.0.activity.mark_dirty() {
1369			self.sleep_activity_notify().notify_one();
1370		}
1371		true
1372	}
1373
1374	pub(crate) fn acknowledge_activity_dirty(&self) -> bool {
1375		self.0.activity.take_dirty()
1376	}
1377
1378	/// Notify the ActorTask that a `can_sleep` input has changed so the sleep
1379	/// deadline gets re-evaluated. Falls back to the detached compat timer
1380	/// when the actor has no wired `ActorTask` (test-only contexts).
1381	pub(crate) fn reset_sleep_timer(&self) {
1382		if self.notify_activity_dirty() {
1383			return;
1384		}
1385
1386		#[cfg(feature = "wasm-runtime")]
1387		return;
1388
1389		#[cfg(not(feature = "wasm-runtime"))]
1390		self.reset_sleep_timer_state();
1391	}
1392
1393	fn notify_inspector_attachments_changed(&self) {
1394		self.try_send_lifecycle_event(
1395			LifecycleEvent::InspectorAttachmentsChanged,
1396			"inspector_attachments_changed",
1397		);
1398	}
1399
1400	pub(crate) fn configure_sleep(&self, config: ActorConfig) {
1401		self.configure_sleep_state(config.clone());
1402		self.configure_queue(config);
1403		self.reset_sleep_timer();
1404	}
1405
1406	pub(crate) fn sleep_config(&self) -> ActorConfig {
1407		self.sleep_state_config()
1408	}
1409
1410	pub(crate) fn sleep_requested(&self) -> bool {
1411		self.0.sleep_requested.load(Ordering::SeqCst)
1412	}
1413
1414	pub(crate) fn clear_sleep_requested(&self) {
1415		self.0.sleep_requested.store(false, Ordering::SeqCst);
1416	}
1417
1418	pub(crate) async fn internal_keep_awake_task(
1419		&self,
1420		future: BoxFuture<'static, Result<()>>,
1421	) -> Result<()> {
1422		self.internal_keep_awake(future).await
1423	}
1424
1425	pub fn websocket_callback_region(&self) -> WebSocketCallbackRegion {
1426		WebSocketCallbackRegion {
1427			region: Some(self.begin_work_region(ActorWorkKind::WebSocketCallback)),
1428		}
1429	}
1430
1431	pub(crate) async fn with_websocket_callback<F, Fut, T>(&self, run: F) -> T
1432	where
1433		F: FnOnce() -> Fut,
1434		Fut: Future<Output = T>,
1435	{
1436		let _guard = self.websocket_callback_region();
1437		run().await
1438	}
1439
1440	fn idle_work_region(&self, kind: ActorWorkKind) -> Option<RegionGuard> {
1441		if !kind.policy().blocks_idle_sleep {
1442			return None;
1443		}
1444		let region = match kind {
1445			ActorWorkKind::Action => self.internal_keep_awake_region(),
1446			ActorWorkKind::KeepAwake => self.keep_awake_region_state(),
1447			ActorWorkKind::InternalKeepAwake => self.internal_keep_awake_region(),
1448			ActorWorkKind::WaitUntil => return None,
1449			ActorWorkKind::RegisteredTask => return None,
1450			ActorWorkKind::WebSocketCallback => self.websocket_callback_region_state(),
1451			ActorWorkKind::DisconnectCallback => self.disconnect_callback_region_state(),
1452		};
1453		Some(region.with_log_fields(kind.label(), Some(self.actor_id().to_owned())))
1454	}
1455
1456	fn shutdown_work_region(&self) -> CountGuard {
1457		let counter = self.0.sleep.work.shutdown_counter.clone();
1458		counter.increment();
1459		CountGuard::from_incremented(counter)
1460	}
1461
1462	fn configure_sleep_hooks(&self) {
1463		let keep_awake_ctx = self.downgrade();
1464		self.0
1465			.sleep
1466			.work
1467			.keep_awake
1468			.register_change_callback(Arc::new(move || {
1469				if let Some(ctx) = ActorContext::from_weak(&keep_awake_ctx) {
1470					ctx.0
1471						.metrics
1472						.set_keep_awake_active(ctx.sleep_keep_awake_count());
1473				}
1474			}));
1475
1476		let internal_keep_awake_metric_ctx = self.downgrade();
1477		self.0
1478			.sleep
1479			.work
1480			.internal_keep_awake
1481			.register_change_callback(Arc::new(move || {
1482				if let Some(ctx) = ActorContext::from_weak(&internal_keep_awake_metric_ctx) {
1483					ctx.0
1484						.metrics
1485						.set_internal_keep_awake_active(ctx.sleep_internal_keep_awake_count());
1486				}
1487			}));
1488
1489		let shutdown_tasks_ctx = self.downgrade();
1490		self.0
1491			.sleep
1492			.work
1493			.shutdown_counter
1494			.register_change_callback(Arc::new(move || {
1495				if let Some(ctx) = ActorContext::from_weak(&shutdown_tasks_ctx) {
1496					ctx.0
1497						.metrics
1498						.set_shutdown_tasks_active(ctx.shutdown_task_count());
1499				}
1500			}));
1501
1502		let internal_keep_awake_ctx = self.downgrade();
1503		self.set_internal_keep_awake(Some(Arc::new(move |future| {
1504			let ctx = ActorContext::from_weak(&internal_keep_awake_ctx);
1505			Box::pin(async move {
1506				let Some(ctx) = ctx else {
1507					return Err(ActorRuntime::NotConfigured {
1508						component: "actor context".to_owned(),
1509					}
1510					.build());
1511				};
1512				ctx.internal_keep_awake_task(future).await
1513			})
1514		})));
1515
1516		let queue_ctx = self.downgrade();
1517		self.set_wait_activity_callback(Some(Arc::new(move || {
1518			if let Some(ctx) = ActorContext::from_weak(&queue_ctx) {
1519				ctx.reset_sleep_timer();
1520			}
1521		})));
1522
1523		let queue_ctx = self.downgrade();
1524		self.set_inspector_update_callback(Some(Arc::new(move |queue_size| {
1525			if let Some(ctx) = ActorContext::from_weak(&queue_ctx) {
1526				ctx.record_queue_updated(queue_size);
1527			}
1528		})));
1529	}
1530
1531	pub(crate) fn record_state_updated(&self) {
1532		if let Some(inspector) = self.inspector() {
1533			inspector.record_state_updated();
1534		}
1535	}
1536
1537	pub(crate) fn record_connections_updated(&self) {
1538		let Some(inspector) = self.inspector() else {
1539			return;
1540		};
1541		let active_connections = self.active_connection_count();
1542		inspector.record_connections_updated(active_connections);
1543	}
1544
1545	fn record_queue_updated(&self, queue_size: u32) {
1546		if let Some(inspector) = self.inspector() {
1547			inspector.record_queue_updated(queue_size);
1548		}
1549	}
1550
1551	pub(crate) fn record_schedules_updated(&self) {
1552		if let Some(inspector) = self.inspector() {
1553			inspector.record_schedules_updated();
1554		}
1555	}
1556
1557	pub(crate) async fn save_state_with_revision(
1558		&self,
1559		deltas: Vec<StateDelta>,
1560		save_request_revision: u64,
1561	) -> Result<()> {
1562		let (deltas, pending_hibernation_changes) = match self.prepare_state_deltas(deltas) {
1563			Ok(prepared) => prepared,
1564			Err(error) => return Err(error),
1565		};
1566		if let Err(error) = self.apply_state_deltas(deltas, save_request_revision).await {
1567			self.restore_pending_hibernation_changes(pending_hibernation_changes);
1568			return Err(error);
1569		}
1570		self.record_state_updated();
1571		Ok(())
1572	}
1573
1574	pub(crate) async fn save_state_and_workflow_batch_with_revision(
1575		&self,
1576		deltas: Vec<StateDelta>,
1577		workflow_writes: Vec<WorkflowKvWrite>,
1578		save_request_revision: u64,
1579	) -> Result<()> {
1580		let (deltas, pending_hibernation_changes) = match self.prepare_state_deltas(deltas) {
1581			Ok(prepared) => prepared,
1582			Err(error) => return Err(error),
1583		};
1584		if let Err(error) = self
1585			.apply_state_deltas_and_workflow(deltas, workflow_writes, save_request_revision)
1586			.await
1587		{
1588			self.restore_pending_hibernation_changes(pending_hibernation_changes);
1589			return Err(error);
1590		}
1591		self.record_state_updated();
1592		Ok(())
1593	}
1594
1595	async fn dispatch_scheduled_action(
1596		&self,
1597		dispatch: crate::actor::schedule::DueScheduleDispatch,
1598	) {
1599		let ctx = self.clone();
1600		let event_id = dispatch.event_id;
1601		let action = dispatch.action;
1602		let args = dispatch.args;
1603		let scheduled_fire = dispatch.fire;
1604		let recurring_name = scheduled_fire.name.clone();
1605		let history_id = dispatch.history_id;
1606		let internal_keep_awake_region = self.begin_work_region(ActorWorkKind::InternalKeepAwake);
1607
1608		self.track_shutdown_task(async move {
1609			let _internal_keep_awake_region = internal_keep_awake_region;
1610			ctx.record_user_task_started(UserTaskKind::ScheduledAction);
1611			let started_at = Instant::now();
1612			let action_name = action.clone();
1613			let (reply_tx, reply_rx) = oneshot::channel();
1614
1615			let mut dispatch_error = None;
1616			match ctx.try_send_actor_event(
1617				ActorEvent::Action {
1618					name: action.clone(),
1619					args,
1620					conn: None,
1621					scheduled_fire: Some(scheduled_fire),
1622					reply: Reply::from(reply_tx),
1623				},
1624				"scheduled_action",
1625			) {
1626				Ok(()) => match reply_rx.await {
1627					Ok(Ok(_)) => {}
1628					Ok(Err(error)) => {
1629						dispatch_error = Some(error);
1630						tracing::error!(
1631							error = ?dispatch_error.as_ref().expect("just assigned"),
1632							event_id,
1633							action_name,
1634							"scheduled event execution failed"
1635						);
1636					}
1637					Err(error) => {
1638						dispatch_error = Some(error.into());
1639						tracing::error!(
1640							error = ?dispatch_error.as_ref().expect("just assigned"),
1641							event_id,
1642							action_name,
1643							"scheduled event reply dropped"
1644						);
1645					}
1646				},
1647				Err(error) => {
1648					dispatch_error = Some(error);
1649					tracing::error!(
1650						error = ?dispatch_error.as_ref().expect("just assigned"),
1651						event_id,
1652						action_name,
1653						"failed to enqueue scheduled event"
1654					);
1655				}
1656			}
1657
1658			ctx.finish_schedule_dispatch(&event_id, history_id, dispatch_error.as_ref())
1659				.await;
1660			if let (Some(name), Some(error)) = (recurring_name, dispatch_error.as_ref()) {
1661				let structured = rivet_error::RivetError::extract(error);
1662				if structured.group() == "actor" && structured.code() == "action_not_found" {
1663					if let Err(delete_error) = ctx.cron_delete_if_action(&name, &action_name).await
1664					{
1665						tracing::error!(
1666							?delete_error,
1667							%name,
1668							"failed to delete recurring schedule for missing action"
1669						);
1670					}
1671				}
1672			}
1673
1674			ctx.record_user_task_finished(UserTaskKind::ScheduledAction, started_at.elapsed());
1675		});
1676	}
1677
1678	fn inspector_attach_count_arc(&self) -> Option<Arc<AtomicU32>> {
1679		self.0.inspector_attach_count.read().clone()
1680	}
1681
1682	fn try_send_lifecycle_event(&self, event: LifecycleEvent, operation: &'static str) {
1683		let Some(sender) = self.0.lifecycle_events.read().clone() else {
1684			return;
1685		};
1686
1687		if sender.send(event).is_err() {
1688			tracing::warn!(operation, "failed to enqueue actor lifecycle event");
1689		}
1690	}
1691}
1692
1693/// Cap on the stop error message forwarded to the engine. The message is
1694/// persisted in engine workflow state and rendered in the dashboard, so
1695/// unbounded anyhow chains must be truncated at this boundary.
1696const MAX_STOP_ERROR_MESSAGE_LEN: usize = 2048;
1697
1698fn truncate_stop_error_message(mut message: String) -> String {
1699	if message.len() <= MAX_STOP_ERROR_MESSAGE_LEN {
1700		return message;
1701	}
1702	let mut end = MAX_STOP_ERROR_MESSAGE_LEN;
1703	while !message.is_char_boundary(end) {
1704		end -= 1;
1705	}
1706	message.truncate(end);
1707	message.push_str("... (truncated)");
1708	message
1709}
1710
1711struct ActorWorkGuard {
1712	ctx: ActorContext,
1713	kind: ActorWorkKind,
1714	started_at: Option<Instant>,
1715	idle_region: Option<RegionGuard>,
1716	shutdown_region: Option<CountGuard>,
1717}
1718
1719#[must_use]
1720pub struct ActorWorkRegion {
1721	guard: Option<ActorWorkGuard>,
1722}
1723
1724impl ActorWorkGuard {
1725	fn new(ctx: ActorContext, kind: ActorWorkKind) -> Self {
1726		let policy = kind.policy();
1727		let idle_region = ctx.idle_work_region(kind);
1728		let shutdown_region = if policy.drains_shutdown_grace {
1729			Some(ctx.shutdown_work_region())
1730		} else {
1731			None
1732		};
1733		let started_at = if let Some(user_task_kind) = policy.user_task_kind {
1734			ctx.record_user_task_started(user_task_kind);
1735			Some(Instant::now())
1736		} else {
1737			None
1738		};
1739		ctx.reset_sleep_timer();
1740		Self {
1741			ctx,
1742			kind,
1743			started_at,
1744			idle_region,
1745			shutdown_region,
1746		}
1747	}
1748}
1749
1750impl Drop for ActorWorkGuard {
1751	fn drop(&mut self) {
1752		if let Some(started_at) = self.started_at.take()
1753			&& let Some(user_task_kind) = self.kind.policy().user_task_kind
1754		{
1755			self.ctx
1756				.record_user_task_finished(user_task_kind, started_at.elapsed());
1757		}
1758		self.idle_region.take();
1759		self.shutdown_region.take();
1760		self.ctx.reset_sleep_timer();
1761	}
1762}
1763
1764impl Drop for ActorWorkRegion {
1765	fn drop(&mut self) {
1766		self.guard.take();
1767	}
1768}
1769
1770#[must_use]
1771#[derive(Debug)]
1772pub(crate) struct InspectorAttachGuard {
1773	ctx: ActorContext,
1774}
1775
1776impl InspectorAttachGuard {
1777	fn new(ctx: ActorContext) -> Option<Self> {
1778		let attach_count = ctx.inspector_attach_count_arc()?;
1779		let previous = attach_count.fetch_add(1, Ordering::SeqCst);
1780		let current = previous.saturating_add(1);
1781		tracing::debug!(
1782			actor_id = %ctx.actor_id(),
1783			previous_count = previous,
1784			current_count = current,
1785			"inspector attached"
1786		);
1787		if previous == 0 {
1788			ctx.notify_inspector_attachments_changed();
1789		}
1790		Some(Self { ctx })
1791	}
1792}
1793
1794impl Drop for InspectorAttachGuard {
1795	fn drop(&mut self) {
1796		let Some(attach_count) = self.ctx.inspector_attach_count_arc() else {
1797			return;
1798		};
1799		let Ok(previous) =
1800			attach_count.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |current| {
1801				current.checked_sub(1)
1802			})
1803		else {
1804			return;
1805		};
1806		let current = previous.saturating_sub(1);
1807		tracing::debug!(
1808			actor_id = %self.ctx.actor_id(),
1809			previous_count = previous,
1810			current_count = current,
1811			"inspector detached"
1812		);
1813		if previous == 1 {
1814			self.ctx.notify_inspector_attachments_changed();
1815		}
1816	}
1817}
1818
1819pub struct WebSocketCallbackRegion {
1820	region: Option<ActorWorkRegion>,
1821}
1822
1823pub struct KeepAwakeRegion {
1824	region: Option<ActorWorkRegion>,
1825}
1826
1827impl Drop for WebSocketCallbackRegion {
1828	fn drop(&mut self) {
1829		self.region.take();
1830	}
1831}
1832
1833impl Drop for KeepAwakeRegion {
1834	fn drop(&mut self) {
1835		// Take the region explicitly to mirror WebSocketCallbackRegion.
1836		self.region.take();
1837	}
1838}
1839
1840impl std::fmt::Debug for ActorContext {
1841	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1842		f.debug_struct("ActorContext")
1843			.field("actor_id", &self.0.actor_id)
1844			.field("name", &self.0.name)
1845			.field("key", &self.0.key)
1846			.field("region", &self.0.region)
1847			.finish()
1848	}
1849}
1850
1851// Test shim keeps moved tests in crate-root tests/ with private-module access.
1852#[cfg(test)]
1853#[path = "../../tests/context.rs"]
1854pub(crate) mod tests;