rivet_envoy_client/
kv.rs

1use rivet_envoy_protocol as protocol;
2use tokio::sync::oneshot;
3
4use crate::connection::ws_send;
5use crate::envoy::EnvoyContext;
6use crate::metrics::METRICS;
7
8pub struct KvRequestEntry {
9	pub actor_id: String,
10	pub data: protocol::KvRequestData,
11	pub response_tx: oneshot::Sender<anyhow::Result<protocol::KvResponseData>>,
12	pub sent: bool,
13	pub timestamp: crate::time::Instant,
14}
15
16pub const KV_EXPIRE_MS: u64 = 30_000;
17pub const KV_CLEANUP_INTERVAL_MS: u64 = 15_000;
18
19pub async fn handle_kv_request(
20	ctx: &mut EnvoyContext,
21	actor_id: String,
22	data: protocol::KvRequestData,
23	response_tx: oneshot::Sender<anyhow::Result<protocol::KvResponseData>>,
24) {
25	let request_id = ctx.next_kv_request_id;
26	ctx.next_kv_request_id += 1;
27
28	let entry = KvRequestEntry {
29		actor_id,
30		data,
31		response_tx,
32		sent: false,
33		timestamp: crate::time::Instant::now(),
34	};
35
36	ctx.kv_requests.insert(request_id, entry);
37	METRICS.kv_requests_inflight.inc();
38
39	let ws_available = {
40		let guard = ctx.shared.ws_tx.lock().await;
41		guard.is_some()
42	};
43
44	if ws_available {
45		send_single_kv_request(ctx, request_id).await;
46	}
47}
48
49pub async fn handle_kv_response(ctx: &mut EnvoyContext, response: protocol::ToEnvoyKvResponse) {
50	let request = ctx.kv_requests.remove(&response.request_id);
51
52	if let Some(request) = request {
53		METRICS.kv_requests_inflight.dec();
54		match response.data {
55			protocol::KvResponseData::KvErrorResponse(ref e) => {
56				let _ = request
57					.response_tx
58					.send(Err(anyhow::anyhow!("{}", e.message)));
59			}
60			_ => {
61				let _ = request.response_tx.send(Ok(response.data));
62			}
63		}
64	} else {
65		tracing::error!(
66			request_id = response.request_id,
67			"received kv response for unknown request id"
68		);
69	}
70}
71
72pub async fn send_single_kv_request(ctx: &mut EnvoyContext, request_id: u32) {
73	let request = ctx.kv_requests.get_mut(&request_id);
74	let Some(request) = request else { return };
75	if request.sent {
76		return;
77	}
78
79	ws_send(
80		&ctx.shared,
81		protocol::ToRivet::ToRivetKvRequest(protocol::ToRivetKvRequest {
82			actor_id: request.actor_id.clone(),
83			request_id,
84			data: request.data.clone(),
85		}),
86	)
87	.await;
88
89	// Re-get after async call
90	if let Some(request) = ctx.kv_requests.get_mut(&request_id) {
91		request.sent = true;
92		request.timestamp = crate::time::Instant::now();
93	}
94}
95
96pub async fn process_unsent_kv_requests(ctx: &mut EnvoyContext) {
97	let ws_available = {
98		let guard = ctx.shared.ws_tx.lock().await;
99		guard.is_some()
100	};
101
102	if !ws_available {
103		return;
104	}
105
106	let unsent: Vec<u32> = ctx
107		.kv_requests
108		.iter()
109		.filter(|(_, req)| !req.sent)
110		.map(|(id, _)| *id)
111		.collect();
112
113	for request_id in unsent {
114		send_single_kv_request(ctx, request_id).await;
115	}
116}
117
118pub fn cleanup_old_kv_requests(ctx: &mut EnvoyContext) {
119	let now = crate::time::Instant::now();
120	let mut to_delete = Vec::new();
121
122	for (request_id, request) in &ctx.kv_requests {
123		if now.duration_since(request.timestamp).as_millis() > KV_EXPIRE_MS as u128 {
124			to_delete.push(*request_id);
125		}
126	}
127
128	for request_id in to_delete {
129		if let Some(request) = ctx.kv_requests.remove(&request_id) {
130			METRICS.kv_requests_inflight.dec();
131			tracing::warn!(
132				request_id,
133				was_sent = request.sent,
134				age_ms = now.duration_since(request.timestamp).as_millis() as u64,
135				"kv request expired by cleanup"
136			);
137			let _ = request
138				.response_tx
139				.send(Err(anyhow::anyhow!("KV request timed out")));
140		}
141	}
142}
rivet_envoy_client/kv.rs

rivet_envoy_client/
kv.rs
