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