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#[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 pub fn is_stopped(&self) -> bool {
49 *self.shared.stopped_tx.borrow()
50 }
51
52 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 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 pub const PING_HEALTHY_THRESHOLD_MS: i64 = 20_000;
81
82 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 #[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 pub async fn wait_actors_drained(&self) {
230 loop {
231 let notified = self.shared.actors_notify.notified();
232 tokio::pin!(notified);
233 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 #[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 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 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}