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 http_ws_tx: Arc::new(tokio::sync::Mutex::new(None)),
589 connection_session: std::sync::atomic::AtomicU64::new(0),
590 next_connection_session: std::sync::atomic::AtomicU64::new(0),
591 connection_session_tx: tokio::sync::watch::channel(0).0,
592 protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)),
593 shutting_down: std::sync::atomic::AtomicBool::new(false),
594 last_ping_ts: std::sync::atomic::AtomicI64::new(0),
595 stopped_tx: tokio::sync::watch::channel(true).0,
596 });
597
598 EnvoyContext {
599 shared,
600 shutting_down: false,
601 actors: HashMap::new(),
602 buffered_actor_messages: HashMap::new(),
603 kv_requests: HashMap::new(),
604 next_kv_request_id: 0,
605 sqlite_requests: HashMap::new(),
606 next_sqlite_request_id: 0,
607 remote_sqlite_requests: HashMap::new(),
608 next_remote_sqlite_request_id: 0,
609 request_to_actor: BufferMap::new(),
610 http_request_routes: BufferMap::new(),
611 http_message_indices: BufferMap::new(),
612 http_request_cancellations: HashMap::new(),
613 buffered_messages: Vec::new(),
614 processed_command_idx: HashMap::new(),
615 }
616 }
617
618 fn exec_request() -> protocol::SqliteExecRequest {
619 protocol::SqliteExecRequest {
620 namespace_id: "ns".to_string(),
621 actor_id: "actor".to_string(),
622 generation: 1,
623 sql: "select 1".to_string(),
624 }
625 }
626
627 fn execute_request() -> protocol::SqliteExecuteRequest {
628 protocol::SqliteExecuteRequest {
629 namespace_id: "ns".to_string(),
630 actor_id: "actor".to_string(),
631 generation: 1,
632 sql: "select ?".to_string(),
633 params: Some(vec![protocol::SqliteBindParam::SqliteValueInteger(
634 protocol::SqliteValueInteger { value: 1 },
635 )]),
636 }
637 }
638
639 fn execute_batch_request() -> protocol::SqliteExecuteBatchRequest {
640 protocol::SqliteExecuteBatchRequest {
641 namespace_id: "ns".to_string(),
642 actor_id: "actor".to_string(),
643 generation: 1,
644 statements: vec![
645 protocol::SqliteBatchStatement {
646 sql: "insert into t values (?)".to_string(),
647 params: Some(vec![protocol::SqliteBindParam::SqliteValueInteger(
648 protocol::SqliteValueInteger { value: 1 },
649 )]),
650 },
651 protocol::SqliteBatchStatement {
652 sql: "insert into t values (?)".to_string(),
653 params: Some(vec![protocol::SqliteBindParam::SqliteValueInteger(
654 protocol::SqliteValueInteger { value: 2 },
655 )]),
656 },
657 ],
658 }
659 }
660
661 #[tokio::test]
662 async fn remote_sqlite_exec_response_matches_pending_request() {
663 let mut ctx = new_envoy_context();
664 let (ws_tx, mut ws_rx) = tokio::sync::mpsc::unbounded_channel();
665 let session = crate::connection::install_connection(&ctx.shared, ws_tx).await;
666 let (tx, rx) = oneshot::channel();
667
668 handle_remote_sqlite_request(
669 &mut ctx,
670 RemoteSqliteRequest::Exec(exec_request()),
671 None,
672 tx,
673 )
674 .await;
675 assert!(matches!(ws_rx.recv().await, Some(WsTxMessage::Send(_))));
676 assert!(ctx.remote_sqlite_requests.contains_key(&0));
677
678 handle_remote_sqlite_exec_response(
679 &mut ctx,
680 protocol::ToEnvoySqliteExecResponse {
681 request_id: 0,
682 data: protocol::SqliteExecResponse::SqliteExecOk(protocol::SqliteExecOk {
683 result: protocol::SqliteQueryResult {
684 columns: vec!["one".to_string()],
685 rows: vec![vec![protocol::SqliteColumnValue::SqliteValueInteger(
686 protocol::SqliteValueInteger { value: 1 },
687 )]],
688 },
689 }),
690 },
691 )
692 .await;
693
694 let response = rx
695 .await
696 .expect("response sender should complete")
697 .expect("response should succeed");
698 assert_eq!(response.session, session);
699 match response.response {
700 RemoteSqliteResponse::Exec(protocol::SqliteExecResponse::SqliteExecOk(ok)) => {
701 assert_eq!(ok.result.columns, vec!["one"]);
702 assert_eq!(ok.result.rows.len(), 1);
703 }
704 _ => panic!("unexpected response"),
705 }
706 assert!(ctx.remote_sqlite_requests.is_empty());
707 }
708
709 #[tokio::test]
710 async fn remote_sqlite_batch_uses_one_websocket_request() {
711 let mut ctx = new_envoy_context();
712 let (ws_tx, mut ws_rx) = tokio::sync::mpsc::unbounded_channel();
713 crate::connection::install_connection(&ctx.shared, ws_tx).await;
714 let (tx, _rx) = oneshot::channel();
715
716 handle_remote_sqlite_request(
717 &mut ctx,
718 RemoteSqliteRequest::ExecuteBatch(execute_batch_request()),
719 None,
720 tx,
721 )
722 .await;
723
724 assert!(matches!(ws_rx.recv().await, Some(WsTxMessage::Send(_))));
725 assert!(
726 ws_rx.try_recv().is_err(),
727 "a batch must serialize as one WebSocket message"
728 );
729 assert_eq!(ctx.remote_sqlite_requests.len(), 1);
730 }
731
732 #[test]
733 fn remote_sqlite_requests_reject_protocol_v3_serialization() {
734 let requests = vec![
735 RemoteSqliteRequest::Exec(exec_request()),
736 RemoteSqliteRequest::Execute(execute_request()),
737 ];
738
739 for request in requests {
740 let message = remote_sqlite_request_to_message(7, request);
741 let err = protocol::versioned::ToRivet::wrap_latest(message)
742 .serialize(3)
743 .expect_err("remote sqlite requests should require protocol v4");
744 let compatibility = err
745 .downcast_ref::<protocol::versioned::ProtocolCompatibilityError>()
746 .expect("error should be a protocol compatibility error");
747 assert_eq!(
748 compatibility.feature,
749 protocol::versioned::ProtocolCompatibilityFeature::RemoteSqliteExecution
750 );
751 assert_eq!(compatibility.required_version, 4);
752 assert_eq!(compatibility.target_version, 3);
753 }
754 }
755
756 #[tokio::test]
757 async fn remote_sqlite_shutdown_cleanup_fails_pending_requests() {
758 let mut ctx = new_envoy_context();
759 let (tx, rx) = oneshot::channel();
760
761 handle_remote_sqlite_request(
762 &mut ctx,
763 RemoteSqliteRequest::Execute(execute_request()),
764 None,
765 tx,
766 )
767 .await;
768 fail_remote_sqlite_requests_with_shutdown(&mut ctx);
769
770 let err = rx
771 .await
772 .expect("response sender should complete")
773 .expect_err("pending request should fail during shutdown");
774 assert!(err.downcast_ref::<EnvoyShutdownError>().is_some());
775 assert!(ctx.remote_sqlite_requests.is_empty());
776 }
777
778 #[tokio::test]
779 async fn sent_remote_sqlite_request_fails_indeterminate_on_disconnect() {
780 let mut ctx = new_envoy_context();
781 let (ws_tx, mut ws_rx) = tokio::sync::mpsc::unbounded_channel();
782 crate::connection::install_connection(&ctx.shared, ws_tx).await;
783 let (tx, rx) = oneshot::channel();
784
785 handle_remote_sqlite_request(
786 &mut ctx,
787 RemoteSqliteRequest::Execute(execute_request()),
788 None,
789 tx,
790 )
791 .await;
792 assert!(matches!(ws_rx.recv().await, Some(WsTxMessage::Send(_))));
793 assert!(
794 ctx.remote_sqlite_requests
795 .get(&0)
796 .expect("request should be pending")
797 .sent
798 );
799
800 fail_sent_remote_sqlite_requests_with_indeterminate_result(&mut ctx);
801
802 let err = rx
803 .await
804 .expect("response sender should complete")
805 .expect_err("sent write should fail indeterminate on disconnect");
806 let indeterminate = err
807 .downcast_ref::<RemoteSqliteIndeterminateResultError>()
808 .expect("error should describe indeterminate remote sqlite result");
809 assert_eq!(indeterminate.operation, "execute");
810 assert!(ctx.remote_sqlite_requests.is_empty());
811 }
812
813 #[tokio::test]
814 async fn unsent_remote_sqlite_request_survives_disconnect_and_sends_on_reconnect() {
815 let mut ctx = new_envoy_context();
816 let (tx, mut rx) = oneshot::channel();
817
818 handle_remote_sqlite_request(
819 &mut ctx,
820 RemoteSqliteRequest::Execute(execute_request()),
821 None,
822 tx,
823 )
824 .await;
825 assert!(
826 !ctx.remote_sqlite_requests
827 .get(&0)
828 .expect("request should be pending")
829 .sent
830 );
831
832 fail_sent_remote_sqlite_requests_with_indeterminate_result(&mut ctx);
833 assert!(matches!(
834 rx.try_recv(),
835 Err(tokio::sync::oneshot::error::TryRecvError::Empty)
836 ));
837 assert!(ctx.remote_sqlite_requests.contains_key(&0));
838
839 let (ws_tx, mut ws_rx) = tokio::sync::mpsc::unbounded_channel();
840 crate::connection::install_connection(&ctx.shared, ws_tx).await;
841 process_unsent_remote_sqlite_requests(&mut ctx).await;
842
843 assert!(matches!(ws_rx.recv().await, Some(WsTxMessage::Send(_))));
844 assert!(
845 ctx.remote_sqlite_requests
846 .get(&0)
847 .expect("request should still be pending")
848 .sent
849 );
850 }
851
852 #[tokio::test]
853 async fn transaction_request_never_crosses_connection_sessions() {
854 let mut ctx = new_envoy_context();
855 let (first_tx, _first_rx) = tokio::sync::mpsc::unbounded_channel();
856 let first_session = crate::connection::install_connection(&ctx.shared, first_tx).await;
857 crate::connection::remove_connection(&ctx.shared).await;
858 let (second_tx, mut second_rx) = tokio::sync::mpsc::unbounded_channel();
859 let second_session = crate::connection::install_connection(&ctx.shared, second_tx).await;
860 assert_ne!(first_session, second_session);
861
862 let (response_tx, response_rx) = oneshot::channel();
863 handle_remote_sqlite_request(
864 &mut ctx,
865 RemoteSqliteRequest::Execute(execute_request()),
866 Some(first_session),
867 response_tx,
868 )
869 .await;
870
871 let error = response_rx
872 .await
873 .expect("response sender should complete")
874 .expect_err("stale transaction session must fail before send");
875 let lost = error
876 .downcast_ref::<RemoteSqliteConnectionSessionLostError>()
877 .expect("error should identify the stale session");
878 assert_eq!(lost.expected, first_session);
879 assert_eq!(lost.current, Some(second_session));
880 assert!(second_rx.try_recv().is_err(), "stale request was sent");
881 assert!(ctx.remote_sqlite_requests.is_empty());
882 }
883}