Skip to main content

rivetkit_core/actor/
lifecycle_hooks.rs

1use anyhow::Result;
2use tokio::sync::{mpsc, oneshot};
3
4use crate::actor::connection::ConnHandle;
5use crate::actor::context::ActorContext;
6use crate::actor::messages::ActorEvent;
7
8pub struct Reply<T> {
9	tx: Option<oneshot::Sender<Result<T>>>,
10}
11
12impl<T> Reply<T> {
13	pub fn send(mut self, result: Result<T>) {
14		if let Some(tx) = self.tx.take() {
15			let _ = tx.send(result);
16		}
17	}
18}
19
20impl<T> Drop for Reply<T> {
21	fn drop(&mut self) {
22		if let Some(tx) = self.tx.take() {
23			let _ = tx.send(Err(crate::error::ActorLifecycle::DroppedReply.build()));
24		}
25	}
26}
27
28impl<T> std::fmt::Debug for Reply<T> {
29	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30		f.debug_struct("Reply")
31			.field("pending", &self.tx.is_some())
32			.finish()
33	}
34}
35
36impl<T> From<oneshot::Sender<Result<T>>> for Reply<T> {
37	fn from(tx: oneshot::Sender<Result<T>>) -> Self {
38		Self { tx: Some(tx) }
39	}
40}
41
42pub struct ActorEvents {
43	actor_id: String,
44	inner: mpsc::UnboundedReceiver<ActorEvent>,
45}
46
47impl ActorEvents {
48	pub(crate) fn new(actor_id: String, inner: mpsc::UnboundedReceiver<ActorEvent>) -> Self {
49		Self { actor_id, inner }
50	}
51
52	pub async fn recv(&mut self) -> Option<ActorEvent> {
53		let event = self.inner.recv().await;
54		if let Some(event) = &event {
55			tracing::debug!(
56				actor_id = %self.actor_id,
57				event = event.kind(),
58				"actor event drained"
59			);
60		}
61		event
62	}
63
64	pub fn try_recv(&mut self) -> Option<ActorEvent> {
65		let event = self.inner.try_recv().ok();
66		if let Some(event) = &event {
67			tracing::debug!(
68				actor_id = %self.actor_id,
69				event = event.kind(),
70				"actor event drained"
71			);
72		}
73		event
74	}
75}
76
77impl std::fmt::Debug for ActorEvents {
78	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79		f.write_str("ActorEvents(..)")
80	}
81}
82
83impl From<mpsc::UnboundedReceiver<ActorEvent>> for ActorEvents {
84	fn from(value: mpsc::UnboundedReceiver<ActorEvent>) -> Self {
85		Self::new("unknown".to_owned(), value)
86	}
87}
88
89#[derive(Debug)]
90pub struct ActorStart {
91	pub ctx: ActorContext,
92	pub input: Option<Vec<u8>>,
93	pub is_new: bool,
94	pub snapshot: Option<Vec<u8>>,
95	pub hibernated: Vec<(ConnHandle, Vec<u8>)>,
96	pub events: ActorEvents,
97	pub startup_ready: Option<oneshot::Sender<Result<()>>>,
98}