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