rivetkit_core/actor/
lifecycle_hooks.rs1use 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 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}