1use rivet_envoy_protocol as protocol;
2use tokio::sync::oneshot;
3
4use crate::connection::{WsSendResult, ws_send, ws_send_for_session};
5use crate::envoy::EnvoyContext;
6use crate::kv::KV_EXPIRE_MS;
7use crate::metrics::METRICS;
8use crate::utils::{
9 EnvoyShutdownError, RemoteSqliteConnectionSessionLostError,
10 RemoteSqliteIndeterminateResultError,
11};
12
13#[derive(Clone)]
14pub enum SqliteRequest {
15 GetPages(protocol::SqliteGetPagesRequest),
16 Commit(protocol::SqliteCommitRequest),
17}
18
19impl SqliteRequest {
20 pub fn kind(&self) -> &'static str {
21 match self {
22 SqliteRequest::GetPages(_) => "get_pages",
23 SqliteRequest::Commit(_) => "commit",
24 }
25 }
26}
27
28pub enum SqliteResponse {
29 GetPages(protocol::SqliteGetPagesResponse),
30 Commit(protocol::SqliteCommitResponse),
31}
32
33#[derive(Clone, Debug)]
34pub enum RemoteSqliteRequest {
35 Exec(protocol::SqliteExecRequest),
36 Execute(protocol::SqliteExecuteRequest),
37 ExecuteBatch(protocol::SqliteExecuteBatchRequest),
38}
39
40#[derive(Debug)]
41pub enum RemoteSqliteResponse {
42 Exec(protocol::SqliteExecResponse),
43 Execute(protocol::SqliteExecuteResponse),
44 ExecuteBatch(protocol::SqliteExecuteBatchResponse),
45}
46
47#[derive(Debug)]
48pub struct RemoteSqliteResponseEnvelope {
49 pub response: RemoteSqliteResponse,
50 pub session: u64,
51}
52
53impl RemoteSqliteRequest {
54 fn operation(&self) -> &'static str {
55 match self {
56 RemoteSqliteRequest::Exec(_) => "exec",
57 RemoteSqliteRequest::Execute(_) => "execute",
58 RemoteSqliteRequest::ExecuteBatch(_) => "execute_batch",
59 }
60 }
61
62 pub fn kind(&self) -> &'static str {
63 match self {
64 RemoteSqliteRequest::Exec(_) => "remote_exec",
65 RemoteSqliteRequest::Execute(_) => "remote_execute",
66 RemoteSqliteRequest::ExecuteBatch(_) => "remote_execute_batch",
67 }
68 }
69}
70
71pub struct SqliteRequestEntry {
72 pub request: SqliteRequest,
73 pub response_tx: oneshot::Sender<anyhow::Result<SqliteResponse>>,
74 pub sent: bool,
75 pub timestamp: crate::time::Instant,
76}
77
78pub struct RemoteSqliteRequestEntry {
79 pub request: RemoteSqliteRequest,
80 pub response_tx: oneshot::Sender<anyhow::Result<RemoteSqliteResponseEnvelope>>,
81 pub expected_session: Option<u64>,
82 pub sent_session: Option<u64>,
83 pub sent: bool,
84 pub timestamp: crate::time::Instant,
85}
86
87pub async fn handle_sqlite_request(
88 ctx: &mut EnvoyContext,
89 request: SqliteRequest,
90 response_tx: oneshot::Sender<anyhow::Result<SqliteResponse>>,
91) {
92 let request_id = ctx.next_sqlite_request_id;
93 ctx.next_sqlite_request_id += 1;
94
95 let entry = SqliteRequestEntry {
96 request,
97 response_tx,
98 sent: false,
99 timestamp: crate::time::Instant::now(),
100 };
101
102 ctx.sqlite_requests.insert(request_id, entry);
103 METRICS.sqlite_requests_inflight.inc();
104
105 let ws_available = {
106 let guard = ctx.shared.ws_tx.lock().await;
107 guard.is_some()
108 };
109
110 if ws_available {
111 send_single_sqlite_request(ctx, request_id).await;
112 }
113}
114
115pub async fn handle_remote_sqlite_request(
116 ctx: &mut EnvoyContext,
117 request: RemoteSqliteRequest,
118 expected_session: Option<u64>,
119 response_tx: oneshot::Sender<anyhow::Result<RemoteSqliteResponseEnvelope>>,
120) {
121 let request_id = ctx.next_remote_sqlite_request_id;
122 ctx.next_remote_sqlite_request_id += 1;
123
124 let entry = RemoteSqliteRequestEntry {
125 request,
126 response_tx,
127 expected_session,
128 sent_session: None,
129 sent: false,
130 timestamp: crate::time::Instant::now(),
131 };
132
133 ctx.remote_sqlite_requests.insert(request_id, entry);
134 METRICS.remote_sqlite_requests_inflight.inc();
135
136 send_single_remote_sqlite_request(ctx, request_id).await;
137}
138
139pub async fn handle_sqlite_get_pages_response(
140 ctx: &mut EnvoyContext,
141 response: protocol::ToEnvoySqliteGetPagesResponse,
142) {
143 handle_sqlite_response(
144 ctx,
145 response.request_id,
146 SqliteResponse::GetPages(response.data),
147 "sqlite_get_pages",
148 );
149}
150
151pub async fn handle_sqlite_commit_response(
152 ctx: &mut EnvoyContext,
153 response: protocol::ToEnvoySqliteCommitResponse,
154) {
155 handle_sqlite_response(
156 ctx,
157 response.request_id,
158 SqliteResponse::Commit(response.data),
159 "sqlite_commit",
160 );
161}
162
163pub async fn handle_remote_sqlite_exec_response(
164 ctx: &mut EnvoyContext,
165 response: protocol::ToEnvoySqliteExecResponse,
166) {
167 handle_remote_sqlite_response(
168 ctx,
169 response.request_id,
170 RemoteSqliteResponse::Exec(response.data),
171 "remote_sqlite_exec",
172 );
173}
174
175pub async fn handle_remote_sqlite_execute_response(
176 ctx: &mut EnvoyContext,
177 response: protocol::ToEnvoySqliteExecuteResponse,
178) {
179 handle_remote_sqlite_response(
180 ctx,
181 response.request_id,
182 RemoteSqliteResponse::Execute(response.data),
183 "remote_sqlite_execute",
184 );
185}
186
187pub async fn handle_remote_sqlite_execute_batch_response(
188 ctx: &mut EnvoyContext,
189 response: protocol::ToEnvoySqliteExecuteBatchResponse,
190) {
191 handle_remote_sqlite_response(
192 ctx,
193 response.request_id,
194 RemoteSqliteResponse::ExecuteBatch(response.data),
195 "remote_sqlite_execute_batch",
196 );
197}
198
199fn handle_sqlite_response(
200 ctx: &mut EnvoyContext,
201 request_id: u32,
202 response: SqliteResponse,
203 op: &str,
204) {
205 let request = ctx.sqlite_requests.remove(&request_id);
206
207 if let Some(request) = request {
208 METRICS.sqlite_requests_inflight.dec();
209 let _ = request.response_tx.send(Ok(response));
210 } else {
211 tracing::error!(
212 request_id,
213 op,
214 "received sqlite response for unknown request id"
215 );
216 }
217}
218
219fn handle_remote_sqlite_response(
220 ctx: &mut EnvoyContext,
221 request_id: u32,
222 response: RemoteSqliteResponse,
223 op: &str,
224) {
225 let request = ctx.remote_sqlite_requests.remove(&request_id);
226
227 if let Some(request) = request {
228 METRICS.remote_sqlite_requests_inflight.dec();
229 let session = request
230 .sent_session
231 .expect("a remote sqlite response must belong to a sent request");
232 let _ = request
233 .response_tx
234 .send(Ok(RemoteSqliteResponseEnvelope { response, session }));
235 } else {
236 tracing::error!(
237 request_id,
238 op,
239 "received remote sqlite response for unknown request id"
240 );
241 }
242}
243
244pub async fn send_single_sqlite_request(ctx: &mut EnvoyContext, request_id: u32) {
245 let request = ctx.sqlite_requests.get_mut(&request_id);
246 let Some(request) = request else { return };
247 if request.sent {
248 return;
249 }
250
251 let message =
252 match request.request.clone() {
253 SqliteRequest::GetPages(data) => protocol::ToRivet::ToRivetSqliteGetPagesRequest(
254 protocol::ToRivetSqliteGetPagesRequest { request_id, data },
255 ),
256 SqliteRequest::Commit(data) => protocol::ToRivet::ToRivetSqliteCommitRequest(
257 protocol::ToRivetSqliteCommitRequest { request_id, data },
258 ),
259 };
260
261 ws_send(&ctx.shared, message).await;
262
263 if let Some(request) = ctx.sqlite_requests.get_mut(&request_id) {
264 request.sent = true;
265 request.timestamp = crate::time::Instant::now();
266 }
267}
268
269pub async fn send_single_remote_sqlite_request(ctx: &mut EnvoyContext, request_id: u32) {
270 let request = ctx.remote_sqlite_requests.get_mut(&request_id);
271 let Some(request) = request else { return };
272 if request.sent {
273 return;
274 }
275
276 let expected_session = request.expected_session;
277 let message = remote_sqlite_request_to_message(request_id, request.request.clone());
278
279 match ws_send_for_session(&ctx.shared, message, expected_session).await {
280 WsSendResult::Sent { session } => {
281 if let Some(request) = ctx.remote_sqlite_requests.get_mut(&request_id) {
282 request.sent = true;
283 request.sent_session = Some(session);
284 request.timestamp = crate::time::Instant::now();
285 }
286 }
287 WsSendResult::Unavailable if expected_session.is_none() => {
288 }
291 WsSendResult::Unavailable => {
292 fail_remote_sqlite_session_request(ctx, request_id, None);
293 }
294 WsSendResult::StaleSession { current } => {
295 fail_remote_sqlite_session_request(ctx, request_id, current);
296 }
297 }
298}
299
300fn fail_remote_sqlite_session_request(
301 ctx: &mut EnvoyContext,
302 request_id: u32,
303 current: Option<u64>,
304) {
305 let Some(request) = ctx.remote_sqlite_requests.remove(&request_id) else {
306 return;
307 };
308 METRICS.remote_sqlite_requests_inflight.dec();
309 let expected = request
310 .expected_session
311 .expect("only session-affine requests fail session admission");
312 let _ = request.response_tx.send(Err(anyhow::anyhow!(
313 RemoteSqliteConnectionSessionLostError { expected, current }
314 )));
315}
316
317pub fn remote_sqlite_request_to_message(
318 request_id: u32,
319 request: RemoteSqliteRequest,
320) -> protocol::ToRivet {
321 match request {
322 RemoteSqliteRequest::Exec(data) => {
323 protocol::ToRivet::ToRivetSqliteExecRequest(protocol::ToRivetSqliteExecRequest {
324 request_id,
325 data,
326 })
327 }
328 RemoteSqliteRequest::Execute(data) => {
329 protocol::ToRivet::ToRivetSqliteExecuteRequest(protocol::ToRivetSqliteExecuteRequest {
330 request_id,
331 data,
332 })
333 }
334 RemoteSqliteRequest::ExecuteBatch(data) => {
335 protocol::ToRivet::ToRivetSqliteExecuteBatchRequest(
336 protocol::ToRivetSqliteExecuteBatchRequest { request_id, data },
337 )
338 }
339 }
340}
341
342pub async fn process_unsent_sqlite_requests(ctx: &mut EnvoyContext) {
343 let ws_available = {
344 let guard = ctx.shared.ws_tx.lock().await;
345 guard.is_some()
346 };
347
348 if !ws_available {
349 return;
350 }
351
352 let unsent: Vec<u32> = ctx
353 .sqlite_requests
354 .iter()
355 .filter(|(_, req)| !req.sent)
356 .map(|(id, _)| *id)
357 .collect();
358
359 for request_id in unsent {
360 send_single_sqlite_request(ctx, request_id).await;
361 }
362}
363
364pub async fn process_unsent_remote_sqlite_requests(ctx: &mut EnvoyContext) {
365 let ws_available = {
366 let guard = ctx.shared.ws_tx.lock().await;
367 guard.is_some()
368 };
369
370 if !ws_available {
371 return;
372 }
373
374 let unsent: Vec<u32> = ctx
375 .remote_sqlite_requests
376 .iter()
377 .filter(|(_, req)| !req.sent)
378 .map(|(id, _)| *id)
379 .collect();
380
381 for request_id in unsent {
382 send_single_remote_sqlite_request(ctx, request_id).await;
383 }
384}
385
386pub fn cleanup_old_sqlite_requests(ctx: &mut EnvoyContext) {
387 let now = crate::time::Instant::now();
388 let mut to_delete = Vec::new();
389
390 for (request_id, request) in &ctx.sqlite_requests {
391 if now.duration_since(request.timestamp).as_millis() > KV_EXPIRE_MS as u128 {
392 to_delete.push(*request_id);
393 }
394 }
395
396 for request_id in to_delete {
397 if let Some(request) = ctx.sqlite_requests.remove(&request_id) {
398 let kind = request.request.kind();
399 let was_sent = if request.sent { "true" } else { "false" };
400 METRICS
401 .sqlite_request_expired_total
402 .with_label_values(&[kind, was_sent])
403 .inc();
404 METRICS.sqlite_requests_inflight.dec();
405 tracing::warn!(
406 request_id,
407 kind,
408 was_sent = request.sent,
409 age_ms = now.duration_since(request.timestamp).as_millis() as u64,
410 "sqlite request expired by cleanup; if was_sent=true this indicates an abandoned in-flight request"
411 );
412 let _ = request
413 .response_tx
414 .send(Err(anyhow::anyhow!("sqlite request timed out")));
415 }
416 }
417}
418
419pub fn cleanup_old_remote_sqlite_requests(ctx: &mut EnvoyContext) {
420 let now = crate::time::Instant::now();
421 let mut to_delete = Vec::new();
422
423 for (request_id, request) in &ctx.remote_sqlite_requests {
424 if now.duration_since(request.timestamp).as_millis() > KV_EXPIRE_MS as u128 {
425 to_delete.push(*request_id);
426 }
427 }
428
429 for request_id in to_delete {
430 if let Some(request) = ctx.remote_sqlite_requests.remove(&request_id) {
431 let kind = request.request.kind();
432 let was_sent = if request.sent { "true" } else { "false" };
433 METRICS
434 .sqlite_request_expired_total
435 .with_label_values(&[kind, was_sent])
436 .inc();
437 METRICS.remote_sqlite_requests_inflight.dec();
438 tracing::warn!(
439 request_id,
440 kind,
441 was_sent = request.sent,
442 age_ms = now.duration_since(request.timestamp).as_millis() as u64,
443 "remote sqlite request expired by cleanup; if was_sent=true this indicates an abandoned in-flight request"
444 );
445 let _ = request
446 .response_tx
447 .send(Err(anyhow::anyhow!("remote sqlite request timed out")));
448 }
449 }
450}
451
452pub fn fail_sqlite_requests_with_shutdown(ctx: &mut EnvoyContext) {
453 for (_id, request) in ctx.sqlite_requests.drain() {
454 METRICS.sqlite_requests_inflight.dec();
455 let _ = request
456 .response_tx
457 .send(Err(anyhow::anyhow!(EnvoyShutdownError)));
458 }
459}
460
461pub fn fail_remote_sqlite_requests_with_shutdown(ctx: &mut EnvoyContext) {
462 for (_id, request) in ctx.remote_sqlite_requests.drain() {
463 METRICS.remote_sqlite_requests_inflight.dec();
464 let _ = request
465 .response_tx
466 .send(Err(anyhow::anyhow!(EnvoyShutdownError)));
467 }
468}
469
470pub fn fail_sent_remote_sqlite_requests_with_indeterminate_result(ctx: &mut EnvoyContext) {
471 let request_ids: Vec<u32> = ctx
472 .remote_sqlite_requests
473 .iter()
474 .filter(|(_, request)| request.sent)
475 .map(|(request_id, _)| *request_id)
476 .collect();
477
478 for request_id in request_ids {
479 if let Some(request) = ctx.remote_sqlite_requests.remove(&request_id) {
480 METRICS.remote_sqlite_requests_inflight.dec();
481 let operation = request.request.operation();
482 tracing::warn!(
483 request_id,
484 operation,
485 "remote sqlite response lost after websocket disconnect"
486 );
487 let _ = request.response_tx.send(Err(anyhow::anyhow!(
488 RemoteSqliteIndeterminateResultError { operation }
489 )));
490 }
491 }
492}
493
494#[cfg(test)]
495mod tests {
496 use std::collections::HashMap;
497 use std::sync::Arc;
498
499 use vbare::OwnedVersionedData;
500
501 use super::*;
502 use crate::config::{
503 BoxFuture, EnvoyCallbacks, EnvoyConfig, HttpRequest, HttpResponse, WebSocketHandler,
504 WebSocketSender,
505 };
506 use crate::context::{SharedContext, WsTxMessage};
507 use crate::handle::EnvoyHandle;
508 use crate::utils::{BufferMap, RemoteSqliteIndeterminateResultError};
509
510 struct IdleCallbacks;
511
512 impl EnvoyCallbacks for IdleCallbacks {
513 fn on_actor_start(
514 &self,
515 _handle: EnvoyHandle,
516 _actor_id: String,
517 _generation: u32,
518 _config: protocol::ActorConfig,
519 _preloaded_kv: Option<protocol::PreloadedKv>,
520 ) -> BoxFuture<anyhow::Result<()>> {
521 Box::pin(async { Ok(()) })
522 }
523
524 fn on_shutdown(&self) {}
525
526 fn fetch(
527 &self,
528 _handle: EnvoyHandle,
529 _actor_id: String,
530 _gateway_id: protocol::GatewayId,
531 _request_id: protocol::RequestId,
532 _request: HttpRequest,
533 ) -> BoxFuture<anyhow::Result<HttpResponse>> {
534 Box::pin(async { anyhow::bail!("fetch should not be called in sqlite tests") })
535 }
536
537 fn websocket(
538 &self,
539 _handle: EnvoyHandle,
540 _actor_id: String,
541 _gateway_id: protocol::GatewayId,
542 _request_id: protocol::RequestId,
543 _request: HttpRequest,
544 _path: String,
545 _headers: HashMap<String, String>,
546 _is_hibernatable: bool,
547 _is_restoring_hibernatable: bool,
548 _sender: WebSocketSender,
549 ) -> BoxFuture<anyhow::Result<WebSocketHandler>> {
550 Box::pin(async { anyhow::bail!("websocket should not be called in sqlite tests") })
551 }
552
553 fn can_hibernate(
554 &self,
555 _actor_id: &str,
556 _gateway_id: &protocol::GatewayId,
557 _request_id: &protocol::RequestId,
558 _request: &HttpRequest,
559 ) -> BoxFuture<anyhow::Result<bool>> {
560 Box::pin(async { Ok(false) })
561 }
562 }
563
564 fn new_envoy_context() -> EnvoyContext {
565 let (envoy_tx, _envoy_rx) = tokio::sync::mpsc::unbounded_channel();
566 let shared = Arc::new(SharedContext {
567 config: EnvoyConfig {
568 version: 1,
569 endpoint: "http://127.0.0.1:1".to_string(),
570 token: None,
571 namespace: "test".to_string(),
572 pool_name: "test".to_string(),
573 prepopulate_actor_names: HashMap::new(),
574 metadata: None,
575 not_global: true,
576 debug_latency_ms: None,
577 callbacks: Arc::new(IdleCallbacks),
578 },
579 envoy_key: "test-envoy".to_string(),
580 envoy_tx,
581 actors: Arc::new(std::sync::Mutex::new(HashMap::new())),
582 actors_notify: Arc::new(tokio::sync::Notify::new()),
583 live_tunnel_requests: Arc::new(std::sync::Mutex::new(HashMap::new())),
584 pending_hibernation_restores: Arc::new(std::sync::Mutex::new(HashMap::new())),
585 ws_tx: Arc::new(tokio::sync::Mutex::new(
586 None::<tokio::sync::mpsc::UnboundedSender<WsTxMessage>>,
587 )),
588 connection_session: std::sync::atomic::AtomicU64::new(0),
589 next_connection_session: std::sync::atomic::AtomicU64::new(0),
590 connection_session_tx: tokio::sync::watch::channel(0).0,
591 protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)),
592 shutting_down: std::sync::atomic::AtomicBool::new(false),
593 last_ping_ts: std::sync::atomic::AtomicI64::new(0),
594 stopped_tx: tokio::sync::watch::channel(true).0,
595 });
596
597 EnvoyContext {
598 shared,
599 shutting_down: false,
600 actors: HashMap::new(),
601 buffered_actor_messages: HashMap::new(),
602 kv_requests: HashMap::new(),
603 next_kv_request_id: 0,
604 sqlite_requests: HashMap::new(),
605 next_sqlite_request_id: 0,
606 remote_sqlite_requests: HashMap::new(),
607 next_remote_sqlite_request_id: 0,
608 request_to_actor: BufferMap::new(),
609 buffered_messages: Vec::new(),
610 processed_command_idx: HashMap::new(),
611 }
612 }
613
614 fn exec_request() -> protocol::SqliteExecRequest {
615 protocol::SqliteExecRequest {
616 namespace_id: "ns".to_string(),
617 actor_id: "actor".to_string(),
618 generation: 1,
619 sql: "select 1".to_string(),
620 }
621 }
622
623 fn execute_request() -> protocol::SqliteExecuteRequest {
624 protocol::SqliteExecuteRequest {
625 namespace_id: "ns".to_string(),
626 actor_id: "actor".to_string(),
627 generation: 1,
628 sql: "select ?".to_string(),
629 params: Some(vec![protocol::SqliteBindParam::SqliteValueInteger(
630 protocol::SqliteValueInteger { value: 1 },
631 )]),
632 }
633 }
634
635 fn execute_batch_request() -> protocol::SqliteExecuteBatchRequest {
636 protocol::SqliteExecuteBatchRequest {
637 namespace_id: "ns".to_string(),
638 actor_id: "actor".to_string(),
639 generation: 1,
640 statements: vec![
641 protocol::SqliteBatchStatement {
642 sql: "insert into t values (?)".to_string(),
643 params: Some(vec![protocol::SqliteBindParam::SqliteValueInteger(
644 protocol::SqliteValueInteger { value: 1 },
645 )]),
646 },
647 protocol::SqliteBatchStatement {
648 sql: "insert into t values (?)".to_string(),
649 params: Some(vec![protocol::SqliteBindParam::SqliteValueInteger(
650 protocol::SqliteValueInteger { value: 2 },
651 )]),
652 },
653 ],
654 }
655 }
656
657 #[tokio::test]
658 async fn remote_sqlite_exec_response_matches_pending_request() {
659 let mut ctx = new_envoy_context();
660 let (ws_tx, mut ws_rx) = tokio::sync::mpsc::unbounded_channel();
661 let session = crate::connection::install_connection(&ctx.shared, ws_tx).await;
662 let (tx, rx) = oneshot::channel();
663
664 handle_remote_sqlite_request(
665 &mut ctx,
666 RemoteSqliteRequest::Exec(exec_request()),
667 None,
668 tx,
669 )
670 .await;
671 assert!(matches!(ws_rx.recv().await, Some(WsTxMessage::Send(_))));
672 assert!(ctx.remote_sqlite_requests.contains_key(&0));
673
674 handle_remote_sqlite_exec_response(
675 &mut ctx,
676 protocol::ToEnvoySqliteExecResponse {
677 request_id: 0,
678 data: protocol::SqliteExecResponse::SqliteExecOk(protocol::SqliteExecOk {
679 result: protocol::SqliteQueryResult {
680 columns: vec!["one".to_string()],
681 rows: vec![vec![protocol::SqliteColumnValue::SqliteValueInteger(
682 protocol::SqliteValueInteger { value: 1 },
683 )]],
684 },
685 }),
686 },
687 )
688 .await;
689
690 let response = rx
691 .await
692 .expect("response sender should complete")
693 .expect("response should succeed");
694 assert_eq!(response.session, session);
695 match response.response {
696 RemoteSqliteResponse::Exec(protocol::SqliteExecResponse::SqliteExecOk(ok)) => {
697 assert_eq!(ok.result.columns, vec!["one"]);
698 assert_eq!(ok.result.rows.len(), 1);
699 }
700 _ => panic!("unexpected response"),
701 }
702 assert!(ctx.remote_sqlite_requests.is_empty());
703 }
704
705 #[tokio::test]
706 async fn remote_sqlite_batch_uses_one_websocket_request() {
707 let mut ctx = new_envoy_context();
708 let (ws_tx, mut ws_rx) = tokio::sync::mpsc::unbounded_channel();
709 crate::connection::install_connection(&ctx.shared, ws_tx).await;
710 let (tx, _rx) = oneshot::channel();
711
712 handle_remote_sqlite_request(
713 &mut ctx,
714 RemoteSqliteRequest::ExecuteBatch(execute_batch_request()),
715 None,
716 tx,
717 )
718 .await;
719
720 assert!(matches!(ws_rx.recv().await, Some(WsTxMessage::Send(_))));
721 assert!(
722 ws_rx.try_recv().is_err(),
723 "a batch must serialize as one WebSocket message"
724 );
725 assert_eq!(ctx.remote_sqlite_requests.len(), 1);
726 }
727
728 #[test]
729 fn remote_sqlite_requests_reject_protocol_v3_serialization() {
730 let requests = vec![
731 RemoteSqliteRequest::Exec(exec_request()),
732 RemoteSqliteRequest::Execute(execute_request()),
733 ];
734
735 for request in requests {
736 let message = remote_sqlite_request_to_message(7, request);
737 let err = protocol::versioned::ToRivet::wrap_latest(message)
738 .serialize(3)
739 .expect_err("remote sqlite requests should require protocol v4");
740 let compatibility = err
741 .downcast_ref::<protocol::versioned::ProtocolCompatibilityError>()
742 .expect("error should be a protocol compatibility error");
743 assert_eq!(
744 compatibility.feature,
745 protocol::versioned::ProtocolCompatibilityFeature::RemoteSqliteExecution
746 );
747 assert_eq!(compatibility.required_version, 4);
748 assert_eq!(compatibility.target_version, 3);
749 }
750 }
751
752 #[tokio::test]
753 async fn remote_sqlite_shutdown_cleanup_fails_pending_requests() {
754 let mut ctx = new_envoy_context();
755 let (tx, rx) = oneshot::channel();
756
757 handle_remote_sqlite_request(
758 &mut ctx,
759 RemoteSqliteRequest::Execute(execute_request()),
760 None,
761 tx,
762 )
763 .await;
764 fail_remote_sqlite_requests_with_shutdown(&mut ctx);
765
766 let err = rx
767 .await
768 .expect("response sender should complete")
769 .expect_err("pending request should fail during shutdown");
770 assert!(err.downcast_ref::<EnvoyShutdownError>().is_some());
771 assert!(ctx.remote_sqlite_requests.is_empty());
772 }
773
774 #[tokio::test]
775 async fn sent_remote_sqlite_request_fails_indeterminate_on_disconnect() {
776 let mut ctx = new_envoy_context();
777 let (ws_tx, mut ws_rx) = tokio::sync::mpsc::unbounded_channel();
778 crate::connection::install_connection(&ctx.shared, ws_tx).await;
779 let (tx, rx) = oneshot::channel();
780
781 handle_remote_sqlite_request(
782 &mut ctx,
783 RemoteSqliteRequest::Execute(execute_request()),
784 None,
785 tx,
786 )
787 .await;
788 assert!(matches!(ws_rx.recv().await, Some(WsTxMessage::Send(_))));
789 assert!(
790 ctx.remote_sqlite_requests
791 .get(&0)
792 .expect("request should be pending")
793 .sent
794 );
795
796 fail_sent_remote_sqlite_requests_with_indeterminate_result(&mut ctx);
797
798 let err = rx
799 .await
800 .expect("response sender should complete")
801 .expect_err("sent write should fail indeterminate on disconnect");
802 let indeterminate = err
803 .downcast_ref::<RemoteSqliteIndeterminateResultError>()
804 .expect("error should describe indeterminate remote sqlite result");
805 assert_eq!(indeterminate.operation, "execute");
806 assert!(ctx.remote_sqlite_requests.is_empty());
807 }
808
809 #[tokio::test]
810 async fn unsent_remote_sqlite_request_survives_disconnect_and_sends_on_reconnect() {
811 let mut ctx = new_envoy_context();
812 let (tx, mut rx) = oneshot::channel();
813
814 handle_remote_sqlite_request(
815 &mut ctx,
816 RemoteSqliteRequest::Execute(execute_request()),
817 None,
818 tx,
819 )
820 .await;
821 assert!(
822 !ctx.remote_sqlite_requests
823 .get(&0)
824 .expect("request should be pending")
825 .sent
826 );
827
828 fail_sent_remote_sqlite_requests_with_indeterminate_result(&mut ctx);
829 assert!(matches!(
830 rx.try_recv(),
831 Err(tokio::sync::oneshot::error::TryRecvError::Empty)
832 ));
833 assert!(ctx.remote_sqlite_requests.contains_key(&0));
834
835 let (ws_tx, mut ws_rx) = tokio::sync::mpsc::unbounded_channel();
836 crate::connection::install_connection(&ctx.shared, ws_tx).await;
837 process_unsent_remote_sqlite_requests(&mut ctx).await;
838
839 assert!(matches!(ws_rx.recv().await, Some(WsTxMessage::Send(_))));
840 assert!(
841 ctx.remote_sqlite_requests
842 .get(&0)
843 .expect("request should still be pending")
844 .sent
845 );
846 }
847
848 #[tokio::test]
849 async fn transaction_request_never_crosses_connection_sessions() {
850 let mut ctx = new_envoy_context();
851 let (first_tx, _first_rx) = tokio::sync::mpsc::unbounded_channel();
852 let first_session = crate::connection::install_connection(&ctx.shared, first_tx).await;
853 crate::connection::remove_connection(&ctx.shared).await;
854 let (second_tx, mut second_rx) = tokio::sync::mpsc::unbounded_channel();
855 let second_session = crate::connection::install_connection(&ctx.shared, second_tx).await;
856 assert_ne!(first_session, second_session);
857
858 let (response_tx, response_rx) = oneshot::channel();
859 handle_remote_sqlite_request(
860 &mut ctx,
861 RemoteSqliteRequest::Execute(execute_request()),
862 Some(first_session),
863 response_tx,
864 )
865 .await;
866
867 let error = response_rx
868 .await
869 .expect("response sender should complete")
870 .expect_err("stale transaction session must fail before send");
871 let lost = error
872 .downcast_ref::<RemoteSqliteConnectionSessionLostError>()
873 .expect("error should identify the stale session");
874 assert_eq!(lost.expected, first_session);
875 assert_eq!(lost.current, Some(second_session));
876 assert!(second_rx.try_recv().is_err(), "stale request was sent");
877 assert!(ctx.remote_sqlite_requests.is_empty());
878 }
879}