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