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	pub fn complete(self) -> bool {
37		self.finish(Ok(()))
38	}
39
40	pub fn fail(self, error: anyhow::Error) -> bool {
41		self.finish(Err(error))
42	}
43
44	pub fn finish(self, result: anyhow::Result<()>) -> bool {
45		let mut guard = match self.tx.lock() {
46			Ok(guard) => guard,
47			Err(poisoned) => poisoned.into_inner(),
48		};
49
50		let Some(tx) = guard.take() else {
51			return false;
52		};
53
54		tx.send(result).is_ok()
55	}
56}
57
58/// Callbacks that the consumer of the envoy client must implement.
59pub trait EnvoyCallbacks: Send + Sync + 'static {
60	fn on_connect(&self, _handle: EnvoyHandle) {}
61
62	fn on_disconnect(&self, _handle: EnvoyHandle) {}
63
64	fn on_actor_start(
65		&self,
66		handle: EnvoyHandle,
67		actor_id: String,
68		generation: u32,
69		config: protocol::ActorConfig,
70		preloaded_kv: Option<protocol::PreloadedKv>,
71	) -> BoxFuture<anyhow::Result<()>>;
72
73	fn on_actor_stop(
74		&self,
75		_handle: EnvoyHandle,
76		_actor_id: String,
77		_generation: u32,
78		_reason: protocol::StopActorReason,
79	) -> BoxFuture<anyhow::Result<()>> {
80		Box::pin(async { Ok(()) })
81	}
82
83	fn on_actor_stop_with_completion(
84		&self,
85		handle: EnvoyHandle,
86		actor_id: String,
87		generation: u32,
88		reason: protocol::StopActorReason,
89		stop_handle: ActorStopHandle,
90	) -> BoxFuture<anyhow::Result<()>> {
91		let stop_future = self.on_actor_stop(handle, actor_id, generation, reason);
92
93		Box::pin(async move {
94			stop_future.await?;
95			stop_handle.complete();
96			Ok(())
97		})
98	}
99
100	fn on_shutdown(&self);
101
102	fn fetch(
103		&self,
104		handle: EnvoyHandle,
105		actor_id: String,
106		gateway_id: protocol::GatewayId,
107		request_id: protocol::RequestId,
108		request: HttpRequest,
109	) -> BoxFuture<anyhow::Result<HttpResponse>>;
110
111	fn websocket(
112		&self,
113		handle: EnvoyHandle,
114		actor_id: String,
115		gateway_id: protocol::GatewayId,
116		request_id: protocol::RequestId,
117		request: HttpRequest,
118		path: String,
119		headers: HashMap<String, String>,
120		is_hibernatable: bool,
121		is_restoring_hibernatable: bool,
122		sender: WebSocketSender,
123	) -> BoxFuture<anyhow::Result<WebSocketHandler>>;
124
125	fn can_hibernate(
126		&self,
127		actor_id: &str,
128		gateway_id: &protocol::GatewayId,
129		request_id: &protocol::RequestId,
130		request: &HttpRequest,
131	) -> BoxFuture<anyhow::Result<bool>>;
132}
rivet_envoy_client/callbacks.rs

rivet_envoy_client/
callbacks.rs
