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			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}