Skip to main content

rivet_envoy_client/
context.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3use std::sync::Mutex as StdMutex;
4use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64};
5
6use crate::async_counter::AsyncCounter;
7use rivet_envoy_protocol as protocol;
8use tokio::sync::{Mutex, Notify, Semaphore, mpsc, oneshot, watch};
9
10use crate::actor::ToActor;
11use crate::config::EnvoyConfig;
12use crate::envoy::ToEnvoyMessage;
13use crate::tunnel::HibernatingWebSocketMetadata;
14
15pub struct SharedActorEntry {
16	pub handle: mpsc::UnboundedSender<ToActor>,
17	pub active_http_request_count: Arc<AsyncCounter>,
18}
19
20pub struct SharedContext {
21	pub config: EnvoyConfig,
22	pub envoy_key: String,
23	pub envoy_tx: mpsc::UnboundedSender<ToEnvoyMessage>,
24	pub actors: Arc<StdMutex<HashMap<String, HashMap<u32, SharedActorEntry>>>>,
25	pub actors_notify: Arc<Notify>,
26	pub live_tunnel_requests: Arc<StdMutex<HashMap<[u8; 8], String>>>,
27	pub pending_hibernation_restores:
28		Arc<StdMutex<HashMap<String, Vec<HibernatingWebSocketMetadata>>>>,
29	pub ws_tx: Arc<Mutex<Option<mpsc::UnboundedSender<WsTxMessage>>>>,
30	pub http_ws_tx: Arc<Mutex<Option<HttpConnectionTx>>>,
31	/// The currently connected WebSocket session, or zero while disconnected.
32	///
33	/// Session IDs are actor-client-local and never cross the wire. Remote SQLite
34	/// transactions use them to prevent a statement from being admitted on a
35	/// replacement WebSocket after the server-side connection (and its SQLite
36	/// handle) has already been torn down.
37	pub connection_session: AtomicU64,
38	pub next_connection_session: AtomicU64,
39	pub connection_session_tx: watch::Sender<u64>,
40	pub protocol_metadata: Arc<Mutex<Option<protocol::ProtocolMetadata>>>,
41	pub shutting_down: AtomicBool,
42	/// Epoch ms timestamp of the most recent ping packet received from the engine. Used by
43	/// `EnvoyHandle::is_ping_healthy` to surface a dead engine link to upstream health checks.
44	/// Zero means no ping has been received yet.
45	pub last_ping_ts: AtomicI64,
46	// Latched signal fired by `envoy_loop` after its cleanup block completes.
47	// Waiters observing `true` are guaranteed that the loop has exited and
48	// every pending KV/SQLite request has been resolved (with `EnvoyShutdownError`
49	// if it didn't complete naturally).
50	pub stopped_tx: watch::Sender<bool>,
51}
52
53#[derive(Debug)]
54pub enum WsTxMessage {
55	Send(Vec<u8>),
56	Close,
57}
58
59#[derive(Clone)]
60pub struct HttpConnectionTx {
61	pub session: u64,
62	pub tx: mpsc::Sender<HttpWsTxMessage>,
63	pub byte_budget: Arc<Semaphore>,
64}
65
66pub struct HttpWsTxMessage {
67	pub data: Vec<u8>,
68	pub _byte_permit: tokio::sync::OwnedSemaphorePermit,
69	pub written: oneshot::Sender<anyhow::Result<()>>,
70}