Skip to main content

rivet_envoy_client/
handle.rs

1use std::sync::Arc;
2use std::sync::atomic::Ordering;
3
4use crate::async_counter::AsyncCounter;
5use rivet_envoy_protocol as protocol;
6use tokio::sync::oneshot;
7
8use crate::context::SharedContext;
9use crate::envoy::{ActorInfo, ToEnvoyMessage};
10use crate::metrics::METRICS;
11use crate::sqlite::{RemoteSqliteRequest, RemoteSqliteResponse, SqliteRequest, SqliteResponse};
12use crate::tunnel::HibernatingWebSocketMetadata;
13
14/// Handle for interacting with the envoy from callbacks.
15#[derive(Clone)]
16pub struct EnvoyHandle {
17	pub(crate) shared: Arc<SharedContext>,
18	pub(crate) started_rx: tokio::sync::watch::Receiver<()>,
19}
20
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct ServerlessActorStart {
23	pub actor_id: String,
24	pub generation: u32,
25}
26
27impl EnvoyHandle {
28	#[doc(hidden)]
29	pub fn from_shared(shared: Arc<SharedContext>) -> Self {
30		Self {
31			shared,
32			started_rx: tokio::sync::watch::channel(()).1,
33		}
34	}
35
36	pub fn shutdown(&self, immediate: bool) {
37		self.shared.shutting_down.store(true, Ordering::Release);
38
39		if immediate {
40			let _ = crate::envoy::send_to_envoy_tx(&self.shared, ToEnvoyMessage::Stop);
41		} else {
42			let _ = crate::envoy::send_to_envoy_tx(&self.shared, ToEnvoyMessage::Shutdown);
43		}
44	}
45
46	/// True once the envoy loop has finished its cleanup block. Latched: stays
47	/// true forever after the loop exits.
48	pub fn is_stopped(&self) -> bool {
49		*self.shared.stopped_tx.borrow()
50	}
51
52	/// Resolves when the envoy loop has finished its cleanup block.
53	///
54	/// Returning does NOT imply successful delivery of pending KV/SQLite/tunnel
55	/// requests. The cleanup block errors out every outstanding request with
56	/// `EnvoyShutdownError`. Callers needing durability must wait on individual
57	/// request acks before invoking shutdown.
58	///
59	/// Latched: safe to call before, during, or after the envoy loop exits.
60	/// A waiter arriving after the loop already exited resolves immediately.
61	pub async fn wait_stopped(&self) {
62		let mut rx = self.shared.stopped_tx.subscribe();
63		if *rx.borrow_and_update() {
64			return;
65		}
66		let _ = rx.changed().await;
67	}
68
69	/// Convenience: signal shutdown then await `wait_stopped`.
70	pub async fn shutdown_and_wait(&self, immediate: bool) {
71		self.shutdown(immediate);
72		self.wait_stopped().await;
73	}
74
75	pub async fn get_protocol_metadata(&self) -> Option<protocol::ProtocolMetadata> {
76		self.shared.protocol_metadata.lock().await.clone()
77	}
78
79	/// Threshold for `is_ping_healthy`.
80	pub const PING_HEALTHY_THRESHOLD_MS: i64 = 20_000;
81
82	/// True after the engine has sent at least one ping and the most recent ping is within
83	/// `PING_HEALTHY_THRESHOLD_MS`. Returns false when the engine link has never completed
84	/// the ping handshake or has gone silently dead long enough that an upstream health check
85	/// should treat this envoy as unhealthy and recycle it.
86	pub fn is_ping_healthy(&self) -> bool {
87		let last = self.shared.last_ping_ts.load(Ordering::Acquire);
88		if last == 0 {
89			return false;
90		}
91		crate::time::now_millis() - last < Self::PING_HEALTHY_THRESHOLD_MS
92	}
93
94	pub fn get_envoy_key(&self) -> &str {
95		&self.shared.envoy_key
96	}
97
98	pub fn endpoint(&self) -> &str {
99		&self.shared.config.endpoint
100	}
101
102	pub fn token(&self) -> Option<&str> {
103		self.shared.config.token.as_deref()
104	}
105
106	/// Returns the current WebSocket session ID, or `None` while disconnected.
107	/// This is an internal client-side affinity token; it is never sent over the
108	/// Envoy protocol.
109	#[doc(hidden)]
110	pub fn connection_session(&self) -> Option<u64> {
111		let session = self.shared.connection_session.load(Ordering::Acquire);
112		(session != 0).then_some(session)
113	}
114
115	#[doc(hidden)]
116	pub fn subscribe_connection_session(&self) -> tokio::sync::watch::Receiver<u64> {
117		self.shared.connection_session_tx.subscribe()
118	}
119
120	pub fn namespace(&self) -> &str {
121		&self.shared.config.namespace
122	}
123
124	pub fn active_actor_count(&self) -> usize {
125		let guard = self
126			.shared
127			.actors
128			.lock()
129			.expect("shared actor registry poisoned");
130		guard
131			.values()
132			.map(|generations| {
133				generations
134					.values()
135					.filter(|actor| !actor.handle.is_closed())
136					.count()
137			})
138			.sum()
139	}
140
141	pub fn pool_name(&self) -> &str {
142		&self.shared.config.pool_name
143	}
144
145	pub async fn started(&self) -> anyhow::Result<()> {
146		self.started_rx
147			.clone()
148			.changed()
149			.await
150			.map_err(|_| anyhow::anyhow!("envoy stopped before startup completed"))?;
151		Ok(())
152	}
153
154	pub fn sleep_actor(&self, actor_id: String, generation: Option<u32>) {
155		let _ = crate::envoy::send_to_envoy_tx(
156			&self.shared,
157			ToEnvoyMessage::ActorIntent {
158				actor_id,
159				generation,
160				intent: protocol::ActorIntent::ActorIntentSleep,
161				error: None,
162			},
163		);
164	}
165
166	pub fn stop_actor(&self, actor_id: String, generation: Option<u32>, error: Option<String>) {
167		let _ = crate::envoy::send_to_envoy_tx(
168			&self.shared,
169			ToEnvoyMessage::ActorIntent {
170				actor_id,
171				generation,
172				intent: protocol::ActorIntent::ActorIntentStop,
173				error,
174			},
175		);
176	}
177
178	pub async fn get_actor(&self, actor_id: &str, generation: Option<u32>) -> Option<ActorInfo> {
179		let (tx, rx) = tokio::sync::oneshot::channel();
180		crate::envoy::send_to_envoy_tx(
181			&self.shared,
182			ToEnvoyMessage::GetActor {
183				actor_id: actor_id.to_string(),
184				generation,
185				response_tx: tx,
186			},
187		)
188		.ok()?;
189		rx.await.ok().flatten()
190	}
191
192	pub async fn wait_actor_registered_then_stopped(&self, actor_id: &str, generation: u32) {
193		let mut registered = false;
194		loop {
195			let notified = self.shared.actors_notify.notified();
196			if self.is_stopped() {
197				return;
198			}
199
200			let actor_is_registered = {
201				let guard = self
202					.shared
203					.actors
204					.lock()
205					.expect("shared actor registry poisoned");
206				guard
207					.get(actor_id)
208					.and_then(|generations| generations.get(&generation))
209					.is_some()
210			};
211
212			if registered && !actor_is_registered {
213				return;
214			}
215			if actor_is_registered {
216				registered = true;
217			}
218
219			tokio::select! {
220				_ = notified => {}
221				_ = self.wait_stopped() => return,
222			}
223		}
224	}
225
226	/// Resolve once the envoy has no active actors (or has stopped). Event-driven via
227	/// `actors_notify`, which `remove_actor` pings on deregister; armed with `enable`
228	/// before the count check so a drain-to-zero cannot race past the waiter.
229	pub async fn wait_actors_drained(&self) {
230		loop {
231			let notified = self.shared.actors_notify.notified();
232			tokio::pin!(notified);
233			// Register interest before reading the count so a `notify_waiters` that
234			// fires between the check and the await is not lost.
235			notified.as_mut().enable();
236
237			if self.is_stopped() || self.active_actor_count() == 0 {
238				return;
239			}
240
241			tokio::select! {
242				_ = notified => {}
243				_ = self.wait_stopped() => return,
244			}
245		}
246	}
247
248	pub fn http_request_counter(
249		&self,
250		actor_id: &str,
251		generation: Option<u32>,
252	) -> Option<Arc<AsyncCounter>> {
253		let guard = self
254			.shared
255			.actors
256			.lock()
257			.expect("shared actor registry poisoned");
258		let generations = guard.get(actor_id)?;
259
260		if let Some(generation) = generation {
261			return generations
262				.get(&generation)
263				.map(|actor| actor.active_http_request_count.clone());
264		}
265
266		generations
267			.iter()
268			.filter(|(_, actor)| !actor.handle.is_closed())
269			.max_by_key(|(generation, _)| *generation)
270			.map(|(_, actor)| actor.active_http_request_count.clone())
271	}
272
273	pub async fn get_active_http_request_count(
274		&self,
275		actor_id: &str,
276		generation: Option<u32>,
277	) -> Option<usize> {
278		self.http_request_counter(actor_id, generation)
279			.map(|counter| counter.load())
280	}
281
282	pub fn hibernatable_connection_is_live(
283		&self,
284		actor_id: &str,
285		_generation: Option<u32>,
286		gateway_id: protocol::GatewayId,
287		request_id: protocol::RequestId,
288	) -> bool {
289		let key = make_ws_key(&gateway_id, &request_id);
290		if self
291			.shared
292			.live_tunnel_requests
293			.lock()
294			.expect("shared live tunnel request registry poisoned")
295			.get(&key)
296			.is_some_and(|live_actor_id| live_actor_id == actor_id)
297		{
298			return true;
299		}
300
301		self.shared
302			.pending_hibernation_restores
303			.lock()
304			.expect("shared pending hibernation restore registry poisoned")
305			.get(actor_id)
306			.is_some_and(|entries| {
307				entries
308					.iter()
309					.any(|entry| entry.gateway_id == gateway_id && entry.request_id == request_id)
310			})
311	}
312
313	pub fn set_alarm(&self, actor_id: String, alarm_ts: Option<i64>, generation: Option<u32>) {
314		self.set_alarm_with_ack(actor_id, alarm_ts, generation, None);
315	}
316
317	pub fn set_alarm_with_ack(
318		&self,
319		actor_id: String,
320		alarm_ts: Option<i64>,
321		generation: Option<u32>,
322		ack_tx: Option<oneshot::Sender<()>>,
323	) {
324		let _ = crate::envoy::send_to_envoy_tx(
325			&self.shared,
326			ToEnvoyMessage::SetAlarm {
327				actor_id,
328				generation,
329				alarm_ts,
330				ack_tx,
331			},
332		);
333	}
334
335	pub async fn kv_get(
336		&self,
337		actor_id: String,
338		keys: Vec<Vec<u8>>,
339	) -> anyhow::Result<Vec<Option<Vec<u8>>>> {
340		let request_keys = keys.clone();
341		let response = self
342			.send_kv_request(
343				actor_id,
344				protocol::KvRequestData::KvGetRequest(protocol::KvGetRequest { keys }),
345			)
346			.await?;
347
348		match response {
349			protocol::KvResponseData::KvGetResponse(resp) => {
350				let mut result = Vec::with_capacity(request_keys.len());
351				for requested_key in &request_keys {
352					let mut found = false;
353					for (i, resp_key) in resp.keys.iter().enumerate() {
354						if requested_key == resp_key {
355							result.push(Some(resp.values[i].clone()));
356							found = true;
357							break;
358						}
359					}
360					if !found {
361						result.push(None);
362					}
363				}
364				Ok(result)
365			}
366			protocol::KvResponseData::KvErrorResponse(e) => {
367				anyhow::bail!("{}", e.message)
368			}
369			_ => anyhow::bail!("unexpected KV response type"),
370		}
371	}
372
373	pub async fn kv_list_all(
374		&self,
375		actor_id: String,
376		reverse: Option<bool>,
377		limit: Option<u64>,
378	) -> anyhow::Result<Vec<(Vec<u8>, Vec<u8>)>> {
379		let response = self
380			.send_kv_request(
381				actor_id,
382				protocol::KvRequestData::KvListRequest(protocol::KvListRequest {
383					query: protocol::KvListQuery::KvListAllQuery,
384					reverse,
385					limit,
386				}),
387			)
388			.await?;
389		parse_list_response(response)
390	}
391
392	pub async fn kv_list_range(
393		&self,
394		actor_id: String,
395		start: Vec<u8>,
396		end: Vec<u8>,
397		exclusive: bool,
398		reverse: Option<bool>,
399		limit: Option<u64>,
400	) -> anyhow::Result<Vec<(Vec<u8>, Vec<u8>)>> {
401		let response = self
402			.send_kv_request(
403				actor_id,
404				protocol::KvRequestData::KvListRequest(protocol::KvListRequest {
405					query: protocol::KvListQuery::KvListRangeQuery(protocol::KvListRangeQuery {
406						start,
407						end,
408						exclusive,
409					}),
410					reverse,
411					limit,
412				}),
413			)
414			.await?;
415		parse_list_response(response)
416	}
417
418	pub async fn kv_list_prefix(
419		&self,
420		actor_id: String,
421		prefix: Vec<u8>,
422		reverse: Option<bool>,
423		limit: Option<u64>,
424	) -> anyhow::Result<Vec<(Vec<u8>, Vec<u8>)>> {
425		let response = self
426			.send_kv_request(
427				actor_id,
428				protocol::KvRequestData::KvListRequest(protocol::KvListRequest {
429					query: protocol::KvListQuery::KvListPrefixQuery(protocol::KvListPrefixQuery {
430						key: prefix,
431					}),
432					reverse,
433					limit,
434				}),
435			)
436			.await?;
437		parse_list_response(response)
438	}
439
440	pub async fn kv_put(
441		&self,
442		actor_id: String,
443		entries: Vec<(Vec<u8>, Vec<u8>)>,
444	) -> anyhow::Result<()> {
445		let (keys, values): (Vec<_>, Vec<_>) = entries.into_iter().unzip();
446		let response = self
447			.send_kv_request(
448				actor_id,
449				protocol::KvRequestData::KvPutRequest(protocol::KvPutRequest { keys, values }),
450			)
451			.await?;
452		match response {
453			protocol::KvResponseData::KvPutResponse => Ok(()),
454			protocol::KvResponseData::KvErrorResponse(e) => anyhow::bail!("{}", e.message),
455			_ => anyhow::bail!("unexpected KV response type"),
456		}
457	}
458
459	pub async fn kv_delete(&self, actor_id: String, keys: Vec<Vec<u8>>) -> anyhow::Result<()> {
460		let response = self
461			.send_kv_request(
462				actor_id,
463				protocol::KvRequestData::KvDeleteRequest(protocol::KvDeleteRequest { keys }),
464			)
465			.await?;
466		match response {
467			protocol::KvResponseData::KvDeleteResponse => Ok(()),
468			protocol::KvResponseData::KvErrorResponse(e) => anyhow::bail!("{}", e.message),
469			_ => anyhow::bail!("unexpected KV response type"),
470		}
471	}
472
473	pub async fn kv_delete_range(
474		&self,
475		actor_id: String,
476		start: Vec<u8>,
477		end: Vec<u8>,
478	) -> anyhow::Result<()> {
479		let response = self
480			.send_kv_request(
481				actor_id,
482				protocol::KvRequestData::KvDeleteRangeRequest(protocol::KvDeleteRangeRequest {
483					start,
484					end,
485				}),
486			)
487			.await?;
488		match response {
489			protocol::KvResponseData::KvDeleteResponse => Ok(()),
490			protocol::KvResponseData::KvErrorResponse(e) => anyhow::bail!("{}", e.message),
491			_ => anyhow::bail!("unexpected KV response type"),
492		}
493	}
494
495	pub async fn kv_drop(&self, actor_id: String) -> anyhow::Result<()> {
496		let response = self
497			.send_kv_request(actor_id, protocol::KvRequestData::KvDropRequest)
498			.await?;
499		match response {
500			protocol::KvResponseData::KvDropResponse => Ok(()),
501			protocol::KvResponseData::KvErrorResponse(e) => anyhow::bail!("{}", e.message),
502			_ => anyhow::bail!("unexpected KV response type"),
503		}
504	}
505
506	pub async fn sqlite_get_pages(
507		&self,
508		request: protocol::SqliteGetPagesRequest,
509	) -> anyhow::Result<protocol::SqliteGetPagesResponse> {
510		match self
511			.send_sqlite_request(SqliteRequest::GetPages(request))
512			.await?
513		{
514			SqliteResponse::GetPages(response) => Ok(response),
515			_ => anyhow::bail!("unexpected sqlite get_pages response type"),
516		}
517	}
518
519	pub async fn sqlite_commit(
520		&self,
521		request: protocol::SqliteCommitRequest,
522	) -> anyhow::Result<protocol::SqliteCommitResponse> {
523		match self
524			.send_sqlite_request(SqliteRequest::Commit(request))
525			.await?
526		{
527			SqliteResponse::Commit(response) => Ok(response),
528			_ => anyhow::bail!("unexpected sqlite commit response type"),
529		}
530	}
531
532	pub async fn remote_sqlite_exec(
533		&self,
534		request: protocol::SqliteExecRequest,
535	) -> anyhow::Result<protocol::SqliteExecResponse> {
536		match self
537			.send_remote_sqlite_request(RemoteSqliteRequest::Exec(request), None)
538			.await?
539			.response
540		{
541			RemoteSqliteResponse::Exec(response) => Ok(response),
542			_ => anyhow::bail!("unexpected remote sqlite exec response type"),
543		}
544	}
545
546	pub async fn remote_sqlite_execute(
547		&self,
548		request: protocol::SqliteExecuteRequest,
549	) -> anyhow::Result<protocol::SqliteExecuteResponse> {
550		match self
551			.send_remote_sqlite_request(RemoteSqliteRequest::Execute(request), None)
552			.await?
553			.response
554		{
555			RemoteSqliteResponse::Execute(response) => Ok(response),
556			_ => anyhow::bail!("unexpected remote sqlite execute response type"),
557		}
558	}
559
560	pub async fn remote_sqlite_execute_batch(
561		&self,
562		request: protocol::SqliteExecuteBatchRequest,
563	) -> anyhow::Result<protocol::SqliteExecuteBatchResponse> {
564		match self
565			.send_remote_sqlite_request(RemoteSqliteRequest::ExecuteBatch(request), None)
566			.await?
567			.response
568		{
569			RemoteSqliteResponse::ExecuteBatch(response) => Ok(response),
570			_ => anyhow::bail!("unexpected remote sqlite execute batch response type"),
571		}
572	}
573
574	/// Executes remote SQLite on one exact WebSocket session and returns the
575	/// session that carried the response. Passing `None` allows an unsent request
576	/// to wait for the next connection; passing `Some` fails before sending if
577	/// that session has disconnected.
578	#[doc(hidden)]
579	pub async fn remote_sqlite_exec_with_session(
580		&self,
581		request: protocol::SqliteExecRequest,
582		expected_session: Option<u64>,
583	) -> anyhow::Result<(protocol::SqliteExecResponse, u64)> {
584		let envelope = self
585			.send_remote_sqlite_request(RemoteSqliteRequest::Exec(request), expected_session)
586			.await?;
587		match envelope.response {
588			RemoteSqliteResponse::Exec(response) => Ok((response, envelope.session)),
589			_ => anyhow::bail!("unexpected remote sqlite exec response type"),
590		}
591	}
592
593	#[doc(hidden)]
594	pub async fn remote_sqlite_execute_with_session(
595		&self,
596		request: protocol::SqliteExecuteRequest,
597		expected_session: Option<u64>,
598	) -> anyhow::Result<(protocol::SqliteExecuteResponse, u64)> {
599		let envelope = self
600			.send_remote_sqlite_request(RemoteSqliteRequest::Execute(request), expected_session)
601			.await?;
602		match envelope.response {
603			RemoteSqliteResponse::Execute(response) => Ok((response, envelope.session)),
604			_ => anyhow::bail!("unexpected remote sqlite execute response type"),
605		}
606	}
607
608	pub fn restore_hibernating_requests(
609		&self,
610		actor_id: String,
611		meta_entries: Vec<HibernatingWebSocketMetadata>,
612	) {
613		self.shared
614			.pending_hibernation_restores
615			.lock()
616			.expect("shared pending hibernation restore registry poisoned")
617			.insert(actor_id, meta_entries);
618	}
619
620	pub(crate) fn take_pending_hibernation_restore(
621		&self,
622		actor_id: &str,
623	) -> Option<Vec<HibernatingWebSocketMetadata>> {
624		self.shared
625			.pending_hibernation_restores
626			.lock()
627			.expect("shared pending hibernation restore registry poisoned")
628			.remove(actor_id)
629	}
630
631	pub fn send_hibernatable_ws_message_ack(
632		&self,
633		gateway_id: protocol::GatewayId,
634		request_id: protocol::RequestId,
635		client_message_index: u16,
636	) {
637		let _ = crate::envoy::send_to_envoy_tx(
638			&self.shared,
639			ToEnvoyMessage::HwsAck {
640				gateway_id,
641				request_id,
642				envoy_message_index: client_message_index,
643			},
644		);
645	}
646
647	pub(crate) async fn rebind_hibernating_websocket(
648		&self,
649		actor_id: String,
650		generation: u32,
651		gateway_id: protocol::GatewayId,
652		request_id: protocol::RequestId,
653	) -> bool {
654		let (response_tx, response_rx) = oneshot::channel();
655		if crate::envoy::send_to_envoy_tx(
656			&self.shared,
657			ToEnvoyMessage::RebindWebSocket {
658				actor_id,
659				generation,
660				gateway_id,
661				request_id,
662				response_tx,
663			},
664		)
665		.is_err()
666		{
667			return false;
668		}
669		response_rx.await.unwrap_or(false)
670	}
671
672	/// Inject a serverless start payload into the envoy.
673	/// The payload is a u16 LE protocol version followed by a serialized ToEnvoy message.
674	pub async fn start_serverless_actor(&self, payload: &[u8]) -> anyhow::Result<()> {
675		tracing::debug!(
676			envoy_key = %self.shared.envoy_key,
677			payload_len = payload.len(),
678			"received serverless start request"
679		);
680		let (message, _) = decode_serverless_actor_start_payload(payload)?;
681
682		// Wait for envoy to be started before injecting
683		self.started().await?;
684
685		tracing::debug!(
686			envoy_key = %self.shared.envoy_key,
687			data = crate::stringify::stringify_to_envoy(&message),
688			"received serverless start"
689		);
690		crate::envoy::send_to_envoy_tx(
691			&self.shared,
692			ToEnvoyMessage::ConnMessage {
693				message,
694				session: self.shared.connection_session.load(Ordering::Acquire),
695			},
696		)
697		.map_err(|_| anyhow::anyhow!("envoy channel closed"))?;
698
699		Ok(())
700	}
701
702	pub fn decode_serverless_actor_start(
703		&self,
704		payload: &[u8],
705	) -> anyhow::Result<ServerlessActorStart> {
706		let (_, actor_start) = decode_serverless_actor_start_payload(payload)?;
707		Ok(actor_start)
708	}
709}
710
711fn decode_serverless_actor_start_payload(
712	payload: &[u8],
713) -> anyhow::Result<(protocol::ToEnvoy, ServerlessActorStart)> {
714	use vbare::OwnedVersionedData;
715
716	if payload.len() < 2 {
717		anyhow::bail!("serverless start payload too short");
718	}
719
720	let version = u16::from_le_bytes([payload[0], payload[1]]);
721	if version != protocol::PROTOCOL_VERSION {
722		anyhow::bail!(
723			"serverless start payload does not match protocol version: {version} vs {}",
724			protocol::PROTOCOL_VERSION
725		);
726	}
727
728	let message = match crate::protocol::versioned::ToEnvoy::deserialize(&payload[2..], version) {
729		Ok(message) => message,
730		Err(err) if version == protocol::PROTOCOL_VERSION => {
731			tracing::debug!(
732				?err,
733				"serverless start payload failed current-version decode, retrying as v1-compatible body"
734			);
735			crate::protocol::versioned::ToEnvoy::deserialize(
736				&payload[2..],
737				protocol::PROTOCOL_VERSION - 1,
738			)?
739		}
740		Err(err) => return Err(err),
741	};
742
743	let protocol::ToEnvoy::ToEnvoyCommands(ref commands) = message else {
744		anyhow::bail!("invalid serverless payload: expected ToEnvoyCommands");
745	};
746	if commands.len() != 1 {
747		anyhow::bail!("invalid serverless payload: expected exactly 1 command");
748	}
749	if !matches!(commands[0].inner, protocol::Command::CommandStartActor(_)) {
750		anyhow::bail!("invalid serverless payload: expected CommandStartActor");
751	}
752
753	let actor_start = ServerlessActorStart {
754		actor_id: commands[0].checkpoint.actor_id.clone(),
755		generation: commands[0].checkpoint.generation,
756	};
757
758	Ok((message, actor_start))
759}
760
761impl EnvoyHandle {
762	async fn send_kv_request(
763		&self,
764		actor_id: String,
765		data: protocol::KvRequestData,
766	) -> anyhow::Result<protocol::KvResponseData> {
767		let (tx, rx) = tokio::sync::oneshot::channel();
768		crate::envoy::send_to_envoy_tx(
769			&self.shared,
770			ToEnvoyMessage::KvRequest {
771				actor_id,
772				data,
773				response_tx: tx,
774			},
775		)
776		.map_err(|_| anyhow::anyhow!("envoy channel closed"))?;
777		rx.await
778			.map_err(|_| anyhow::anyhow!("kv response channel closed"))?
779	}
780
781	async fn send_sqlite_request(&self, request: SqliteRequest) -> anyhow::Result<SqliteResponse> {
782		let kind = request.kind();
783		let total_start = crate::time::Instant::now();
784		let submit_start = crate::time::Instant::now();
785		let (tx, rx) = tokio::sync::oneshot::channel();
786		crate::envoy::send_to_envoy_tx(
787			&self.shared,
788			ToEnvoyMessage::SqliteRequest {
789				request,
790				response_tx: tx,
791			},
792		)
793		.map_err(|_| anyhow::anyhow!("envoy channel closed"))?;
794		let submit_elapsed = submit_start.elapsed();
795		METRICS
796			.sqlite_request_submit_duration_seconds
797			.with_label_values(&[kind])
798			.observe(submit_elapsed.as_secs_f64());
799
800		let wait_start = crate::time::Instant::now();
801		let result = rx
802			.await
803			.map_err(|_| anyhow::anyhow!("sqlite response channel closed"))?;
804		let wait_elapsed = wait_start.elapsed();
805		METRICS
806			.sqlite_request_wait_duration_seconds
807			.with_label_values(&[kind])
808			.observe(wait_elapsed.as_secs_f64());
809		METRICS
810			.sqlite_request_total_duration_seconds
811			.with_label_values(&[kind])
812			.observe(total_start.elapsed().as_secs_f64());
813		result
814	}
815
816	async fn send_remote_sqlite_request(
817		&self,
818		request: RemoteSqliteRequest,
819		expected_session: Option<u64>,
820	) -> anyhow::Result<crate::sqlite::RemoteSqliteResponseEnvelope> {
821		let kind = request.kind();
822		let total_start = crate::time::Instant::now();
823		let submit_start = crate::time::Instant::now();
824		let (tx, rx) = tokio::sync::oneshot::channel();
825		crate::envoy::send_to_envoy_tx(
826			&self.shared,
827			ToEnvoyMessage::RemoteSqliteRequest {
828				request,
829				expected_session,
830				response_tx: tx,
831			},
832		)
833		.map_err(|_| anyhow::anyhow!("envoy channel closed"))?;
834		let submit_elapsed = submit_start.elapsed();
835		METRICS
836			.sqlite_request_submit_duration_seconds
837			.with_label_values(&[kind])
838			.observe(submit_elapsed.as_secs_f64());
839
840		let wait_start = crate::time::Instant::now();
841		let result = rx
842			.await
843			.map_err(|_| anyhow::anyhow!("remote sqlite response channel closed"))?;
844		let wait_elapsed = wait_start.elapsed();
845		METRICS
846			.sqlite_request_wait_duration_seconds
847			.with_label_values(&[kind])
848			.observe(wait_elapsed.as_secs_f64());
849		METRICS
850			.sqlite_request_total_duration_seconds
851			.with_label_values(&[kind])
852			.observe(total_start.elapsed().as_secs_f64());
853		result
854	}
855}
856
857fn make_ws_key(gateway_id: &protocol::GatewayId, request_id: &protocol::RequestId) -> [u8; 8] {
858	let mut key = [0u8; 8];
859	key[..4].copy_from_slice(gateway_id);
860	key[4..].copy_from_slice(request_id);
861	key
862}
863
864fn parse_list_response(
865	response: protocol::KvResponseData,
866) -> anyhow::Result<Vec<(Vec<u8>, Vec<u8>)>> {
867	match response {
868		protocol::KvResponseData::KvListResponse(resp) => {
869			Ok(resp.keys.into_iter().zip(resp.values).collect())
870		}
871		protocol::KvResponseData::KvErrorResponse(e) => anyhow::bail!("{}", e.message),
872		_ => anyhow::bail!("unexpected KV response type"),
873	}
874}