Skip to main content

rivet_envoy_client/
sqlite.rs

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			// Ordinary requests that have never crossed the socket remain eligible
289			// for replay. A transaction-affine request takes the error branch below.
290		}
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}