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