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