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	on_reply: Option<Box<dyn FnOnce(&Result<T>) + Send>>,
11}
12
13impl<T> Reply<T> {
14	/// Runs `on_reply` with the result just before it is sent. A reply dropped
15	/// unsent runs it with the dropped-reply error, so every outcome is seen.
16	pub(crate) fn on_reply(mut self, on_reply: impl FnOnce(&Result<T>) + Send + 'static) -> Self {
17		self.on_reply = Some(Box::new(on_reply));
18		self
19	}
20
21	pub fn send(mut self, result: Result<T>) {
22		self.deliver(result);
23	}
24
25	fn deliver(&mut self, result: Result<T>) {
26		if let Some(on_reply) = self.on_reply.take() {
27			on_reply(&result);
28		}
29		if let Some(tx) = self.tx.take() {
30			let _ = tx.send(result);
31		}
32	}
33}
34
35impl<T> Drop for Reply<T> {
36	fn drop(&mut self) {
37		if self.tx.is_some() {
38			self.deliver(Err(crate::error::ActorLifecycle::DroppedReply.build()));
39		}
40	}
41}
42
43impl<T> std::fmt::Debug for Reply<T> {
44	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45		f.debug_struct("Reply")
46			.field("pending", &self.tx.is_some())
47			.finish()
48	}
49}
50
51impl<T> From<oneshot::Sender<Result<T>>> for Reply<T> {
52	fn from(tx: oneshot::Sender<Result<T>>) -> Self {
53		Self {
54			tx: Some(tx),
55			on_reply: None,
56		}
57	}
58}
59
60pub struct ActorEvents {
61	actor_id: String,
62	inner: mpsc::UnboundedReceiver<ActorEvent>,
63}
64
65impl ActorEvents {
66	pub(crate) fn new(actor_id: String, inner: mpsc::UnboundedReceiver<ActorEvent>) -> Self {
67		Self { actor_id, inner }
68	}
69
70	pub async fn recv(&mut self) -> Option<ActorEvent> {
71		let event = self.inner.recv().await;
72		if let Some(event) = &event {
73			tracing::debug!(
74				actor_id = %self.actor_id,
75				event = event.kind(),
76				"actor event drained"
77			);
78		}
79		event
80	}
81
82	pub fn try_recv(&mut self) -> Option<ActorEvent> {
83		let event = self.inner.try_recv().ok();
84		if let Some(event) = &event {
85			tracing::debug!(
86				actor_id = %self.actor_id,
87				event = event.kind(),
88				"actor event drained"
89			);
90		}
91		event
92	}
93}
94
95impl std::fmt::Debug for ActorEvents {
96	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97		f.write_str("ActorEvents(..)")
98	}
99}
100
101impl From<mpsc::UnboundedReceiver<ActorEvent>> for ActorEvents {
102	fn from(value: mpsc::UnboundedReceiver<ActorEvent>) -> Self {
103		Self::new("unknown".to_owned(), value)
104	}
105}
106
107#[derive(Debug)]
108pub struct ActorStart {
109	pub ctx: ActorContext,
110	pub input: Option<Vec<u8>>,
111	pub is_new: bool,
112	pub snapshot: Option<Vec<u8>>,
113	pub hibernated: Vec<(ConnHandle, Vec<u8>)>,
114	pub events: ActorEvents,
115	pub startup_ready: Option<oneshot::Sender<Result<()>>>,
116}