rivet_envoy_client/
callbacks.rs

1use std::{
2	collections::HashMap,
3	future::Future,
4	pin::Pin,
5	sync::{Arc, Mutex},
6};
7
8use rivet_envoy_protocol as protocol;
9use tokio::sync::oneshot;
10
11use crate::{
12	handle::EnvoyHandle,
13	http::{HttpRequest, HttpResponse},
14	websocket::{WebSocketHandler, WebSocketSender},
15};
16
17#[cfg(not(target_arch = "wasm32"))]
18pub type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
19
20#[cfg(target_arch = "wasm32")]
21pub type BoxFuture<T> = Pin<Box<dyn Future<Output = T>>>;
22
23/// One-shot completion handle used to defer the final stopped event until teardown is done.
24#[derive(Clone)]
25pub struct ActorStopHandle {
26	tx: Arc<Mutex<Option<oneshot::Sender<anyhow::Result<()>>>>>,
27}
28
29impl ActorStopHandle {
30	pub(crate) fn new(tx: oneshot::Sender<anyhow::Result<()>>) -> Self {
31		Self {
32			tx: Arc::new(Mutex::new(Some(tx))),
33		}
34	}
35
36	/// Test-only constructor that discards the completion signal. Not for production use.
37	#[doc(hidden)]
38	pub fn detached() -> Self {
39		let (tx, _rx) = oneshot::channel();
40		Self::new(tx)
41	}
42
43	pub fn complete(self) -> bool {
44		self.finish(Ok(()))
45	}
46
47	pub fn fail(self, error: anyhow::Error) -> bool {
48		self.finish(Err(error))
49	}
50
51	pub fn finish(self, result: anyhow::Result<()>) -> bool {
52		let mut guard = match self.tx.lock() {
53			Ok(guard) => guard,
54			Err(poisoned) => poisoned.into_inner(),
55		};
56
57		let Some(tx) = guard.take() else {
58			return false;
59		};
60
61		tx.send(result).is_ok()
62	}
63}
64
65/// Callbacks that the consumer of the envoy client must implement.
66pub trait EnvoyCallbacks: Send + Sync + 'static {
67	fn on_connect(&self, _handle: EnvoyHandle) {}
68
69	fn on_disconnect(&self, _handle: EnvoyHandle) {}
70
71	fn on_actor_start(
72		&self,
73		handle: EnvoyHandle,
74		actor_id: String,
75		generation: u32,
76		config: protocol::ActorConfig,
77		preloaded_kv: Option<protocol::PreloadedKv>,
78	) -> BoxFuture<anyhow::Result<()>>;
79
80	fn on_actor_stop(
81		&self,
82		_handle: EnvoyHandle,
83		_actor_id: String,
84		_generation: u32,
85		_reason: protocol::StopActorReason,
86	) -> BoxFuture<anyhow::Result<()>> {
87		Box::pin(async { Ok(()) })
88	}
89
90	fn on_actor_stop_with_completion(
91		&self,
92		handle: EnvoyHandle,
93		actor_id: String,
94		generation: u32,
95		reason: protocol::StopActorReason,
96		stop_handle: ActorStopHandle,
97	) -> BoxFuture<anyhow::Result<()>> {
98		let stop_future = self.on_actor_stop(handle, actor_id, generation, reason);
99
100		Box::pin(async move {
101			stop_future.await?;
102			stop_handle.complete();
103			Ok(())
104		})
105	}
106
107	fn on_shutdown(&self);
108
109	fn fetch(
110		&self,
111		handle: EnvoyHandle,
112		actor_id: String,
113		gateway_id: protocol::GatewayId,
114		request_id: protocol::RequestId,
115		request: HttpRequest,
116	) -> BoxFuture<anyhow::Result<HttpResponse>>;
117
118	fn websocket(
119		&self,
120		handle: EnvoyHandle,
121		actor_id: String,
122		gateway_id: protocol::GatewayId,
123		request_id: protocol::RequestId,
124		request: HttpRequest,
125		path: String,
126		headers: HashMap<String, String>,
127		is_hibernatable: bool,
128		is_restoring_hibernatable: bool,
129		sender: WebSocketSender,
130	) -> BoxFuture<anyhow::Result<WebSocketHandler>>;
131
132	fn can_hibernate(
133		&self,
134		actor_id: &str,
135		gateway_id: &protocol::GatewayId,
136		request_id: &protocol::RequestId,
137		request: &HttpRequest,
138	) -> BoxFuture<anyhow::Result<bool>>;
139}
rivet_envoy_client/callbacks.rs

rivet_envoy_client/
callbacks.rs
