diff --git a/Cargo.lock b/Cargo.lock index 2fad33f..10d2ad4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2870,6 +2870,7 @@ dependencies = [ "serde", "serde-wasm-bindgen 0.6.5", "serde_json", + "sha2", "sqlite-wasm-rs", "uuid", "wasm-bindgen", diff --git a/Cargo.toml b/Cargo.toml index aea03e5..a829c8c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,10 +37,18 @@ web-sys = { version = "0.3", features = [ "FileSystemHandle", "FileSystemRemoveOptions", "FileSystemGetDirectoryOptions", - "DomException" + "DomException", + "Response", + "AbortController", + "AbortSignal", + "Request", + "RequestInit", + "ReadableStream", + "ReadableStreamDefaultReader" ]} serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" +sha2 = "0.10" serde-wasm-bindgen = "0.6" uuid = { version = "1.0", features = ["v4", "js"] } console_error_panic_hook = "0.1" diff --git a/packages/sqlite-web-core/Cargo.toml b/packages/sqlite-web-core/Cargo.toml index 2cf1380..f4b412e 100644 --- a/packages/sqlite-web-core/Cargo.toml +++ b/packages/sqlite-web-core/Cargo.toml @@ -13,6 +13,7 @@ js-sys = { workspace = true } web-sys = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +sha2 = { workspace = true } serde-wasm-bindgen = { workspace = true } uuid = { workspace = true } console_error_panic_hook = { workspace = true } diff --git a/packages/sqlite-web-core/src/coordination.rs b/packages/sqlite-web-core/src/coordination.rs index 0756ec4..ade4d8d 100644 --- a/packages/sqlite-web-core/src/coordination.rs +++ b/packages/sqlite-web-core/src/coordination.rs @@ -1,6 +1,6 @@ use js_sys::{Function, Object, Promise, Reflect}; use std::cell::{Cell, RefCell}; -use std::collections::{HashMap, VecDeque}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::future::Future; use std::pin::Pin; use std::rc::Rc; @@ -14,10 +14,10 @@ use web_sys::{ Blob, BlobPropertyBag, BroadcastChannel, DedicatedWorkerGlobalScope, MessageEvent, Url, Worker, }; -use crate::database::SQLiteDatabase; +use crate::database::{SQLiteDatabase, SnapshotCancellation, MAX_SNAPSHOT_SIZE}; use crate::messages::{ - ChannelMessage, MainThreadMessage, SqlBatchStatement, WorkerErrorPayload, WorkerMessage, - WORKER_ERROR_TYPE_INITIALIZATION_PENDING, + ChannelMessage, MainThreadMessage, SnapshotCompression, SqlBatchStatement, WorkerErrorPayload, + WorkerMessage, WORKER_ERROR_TYPE_INITIALIZATION_PENDING, }; use crate::util::{js_value_to_string, sanitize_identifier, set_js_property}; @@ -28,6 +28,8 @@ pub enum LeadershipRole { } const MAX_DB_WORKER_RESPAWNS: u32 = 3; +const SNAPSHOT_CANCELLATION_TOMBSTONE_TTL_MS: f64 = 30_000.0; +const MAX_SNAPSHOT_CANCELLATION_TOMBSTONES: usize = 1024; pub struct WorkerConfig { pub db_name: String, pub follower_timeout_ms: f64, @@ -81,9 +83,71 @@ pub fn worker_config_from_global() -> Result { }) } +fn worker_id_from_global() -> String { + Reflect::get(&js_sys::global(), &JsValue::from_str("__SQLITE_CLIENT_ID")) + .ok() + .and_then(|value| value.as_string()) + .filter(|value| !value.is_empty()) + .unwrap_or_else(|| Uuid::new_v4().to_string()) +} + +#[derive(Clone)] enum DbRequestOrigin { - Local { request_id: u32 }, - Forwarded { query_id: String }, + Local { + request_id: u32, + }, + Forwarded { + query_id: String, + }, + Snapshot { + key: SnapshotKey, + waiters: Vec, + cancelling: Vec, + phase: SnapshotJobPhase, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct SnapshotKey { + url: String, + compression: SnapshotCompression, + sha256: String, + uncompressed_size: u64, +} + +#[derive(Clone)] +enum SnapshotWaiter { + Local { + request_id: u32, + requester_id: String, + query_id: String, + }, + Forwarded { + requester_id: String, + query_id: String, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum SnapshotJobPhase { + Active, + CancellingLastWaiter, +} + +impl SnapshotWaiter { + fn snapshot_identity(&self) -> (&str, &str) { + match self { + Self::Local { + requester_id, + query_id, + .. + } + | Self::Forwarded { + requester_id, + query_id, + } => (requester_id, query_id), + } + } } struct DbJob { @@ -99,6 +163,14 @@ pub(crate) enum DbJobPayload { Batch { statements: Vec, }, + Snapshot { + db_name: String, + url: String, + compression: SnapshotCompression, + sha256: String, + uncompressed_size: f64, + cancellation: SnapshotCancellation, + }, } type DbExecFuture = Pin> + 'static>>; @@ -143,6 +215,10 @@ pub struct CoordinatorState { pub follower_pending: Rc>>, pub next_db_request_id: Rc>, db_worker_restart_attempts: Rc>, + cancelled_snapshot_requests: Rc>>, + snapshot_cancellation_pending: Rc>>, + leader_draining: Rc>, + shutdown_requests: Rc>>, } pub struct DbWorkerState { @@ -150,6 +226,7 @@ pub struct DbWorkerState { pub db_name: String, db_queue: Rc>>, db_processing: Rc>, + snapshot_cancellations: Rc>>, hooks: DbWorkerHooks, } @@ -161,7 +238,7 @@ pub fn create_broadcast_channel(db_name: &str) -> Result Result, JsValue> { Ok(Rc::new(CoordinatorState { - worker_id: Uuid::new_v4().to_string(), + worker_id: worker_id_from_global(), role: Rc::new(RefCell::new(LeadershipRole::Follower)), leader_id: Rc::new(RefCell::new(None)), leader_ready: Rc::new(RefCell::new(false)), @@ -176,6 +253,10 @@ impl CoordinatorState { follower_pending: Rc::new(RefCell::new(HashMap::new())), next_db_request_id: Rc::new(RefCell::new(1)), db_worker_restart_attempts: Rc::new(Cell::new(0)), + cancelled_snapshot_requests: Rc::new(RefCell::new(HashMap::new())), + snapshot_cancellation_pending: Rc::new(RefCell::new(HashSet::new())), + leader_draining: Rc::new(Cell::new(false)), + shutdown_requests: Rc::new(RefCell::new(Vec::new())), })) } @@ -201,11 +282,109 @@ impl CoordinatorState { } } + fn attach_to_inflight_snapshot(&self, key: &SnapshotKey, waiter: SnapshotWaiter) -> bool { + let mut pending = self.db_pending.borrow_mut(); + if let Some(DbRequestOrigin::Snapshot { waiters, .. }) = pending + .values_mut() + .find(|origin| matches!(origin, DbRequestOrigin::Snapshot { key: pending_key, phase: SnapshotJobPhase::Active, .. } if pending_key == key)) + { + waiters.push(waiter); + true + } else { + false + } + } + + fn prune_snapshot_cancellation_tombstones(&self, now_ms: f64) { + let mut tombstones = self.cancelled_snapshot_requests.borrow_mut(); + tombstones.retain(|_, expiry| *expiry > now_ms); + while tombstones.len() > MAX_SNAPSHOT_CANCELLATION_TOMBSTONES { + let Some(oldest) = tombstones + .iter() + .min_by(|left, right| left.1.total_cmp(right.1)) + .map(|(key, _)| key.clone()) + else { + break; + }; + tombstones.remove(&oldest); + } + } + + fn remember_snapshot_cancellation(self: &Rc, requester_id: String, query_id: String) { + let now_ms = js_sys::Date::now(); + self.prune_snapshot_cancellation_tombstones(now_ms); + let key = (requester_id, query_id); + let expiry = now_ms + SNAPSHOT_CANCELLATION_TOMBSTONE_TTL_MS; + { + let mut tombstones = self.cancelled_snapshot_requests.borrow_mut(); + if tombstones.len() >= MAX_SNAPSHOT_CANCELLATION_TOMBSTONES { + if let Some(oldest) = tombstones + .iter() + .min_by(|left, right| left.1.total_cmp(right.1)) + .map(|(key, _)| key.clone()) + { + tombstones.remove(&oldest); + } + } + tombstones.insert(key.clone(), expiry); + } + let state = Rc::clone(self); + spawn_local(async move { + sleep_ms(SNAPSHOT_CANCELLATION_TOMBSTONE_TTL_MS as i32).await; + let mut tombstones = state.cancelled_snapshot_requests.borrow_mut(); + if tombstones.get(&key).is_some_and(|stored| *stored <= expiry) { + tombstones.remove(&key); + } + }); + } + + fn take_snapshot_cancellation(&self, requester_id: &str, query_id: &str) -> bool { + self.prune_snapshot_cancellation_tombstones(js_sys::Date::now()); + self.cancelled_snapshot_requests + .borrow_mut() + .remove(&(requester_id.to_string(), query_id.to_string())) + .is_some() + } + fn handle_follower_forward_failure(&self, query_id: &str, request_id: u32, err: &str) { self.follower_pending.borrow_mut().remove(query_id); let _ = send_query_result_to_main(request_id, Err(err.to_string())); } + fn begin_snapshot_timeout_cancellation(&self, query_id: &str) -> bool { + if !self.follower_pending.borrow().contains_key(query_id) { + return false; + } + self.snapshot_cancellation_pending + .borrow_mut() + .insert(query_id.to_string()); + if let Err(error) = send_channel_message( + &self.channel, + &ChannelMessage::CancelSnapshotRequest { + requester_id: self.worker_id.clone(), + query_id: query_id.to_string(), + }, + ) { + let _ = send_worker_error_message(&error); + } + true + } + + fn reject_if_leader_draining(&self, query_id: &str) -> bool { + if !self.leader_draining.get() { + return false; + } + let _ = send_channel_message( + &self.channel, + &ChannelMessage::QueryResponse { + query_id: query_id.to_string(), + result: None, + error: Some("Leader is shutting down; retry on the next leader".to_string()), + }, + ); + true + } + pub fn start_leader_probe(self: &Rc) { if matches!(*self.role.borrow(), LeadershipRole::Leader) { return; @@ -304,6 +483,9 @@ impl CoordinatorState { } fn on_lock_granted(self: &Rc) { + if self.leader_draining.get() { + return; + } *self.role.borrow_mut() = LeadershipRole::Leader; self.mark_leader_known(self.worker_id.clone()); @@ -398,6 +580,14 @@ impl CoordinatorState { }) => { self.handle_db_query_result(request_id, result, error); } + Ok(MainThreadMessage::SnapshotCancelled { + request_id, + cancelled, + }) => { + if cancelled { + self.handle_snapshot_cancelled(request_id); + } + } Err(_) => { if let Some(err) = parse_worker_error_payload(&data) { self.handle_db_worker_failure(err); @@ -420,6 +610,10 @@ impl CoordinatorState { for (_, origin) in pending { self.fail_origin(origin, error.clone()); } + self.finish_shutdown_if_drained(); + if self.leader_draining.get() { + return; + } if attempts > MAX_DB_WORKER_RESPAWNS { let message = format!( "DB worker restart limit reached (max {MAX_DB_WORKER_RESPAWNS}); leaving worker failed" @@ -440,6 +634,13 @@ impl CoordinatorState { params, } => match *self.role.borrow() { LeadershipRole::Leader => { + if self.leader_draining.get() { + let _ = send_query_result_to_main( + request_id, + Err("Coordinator is shutting down".to_string()), + ); + return; + } if !*self.db_worker_ready.borrow() { let _ = send_query_result_to_main( request_id, @@ -490,6 +691,13 @@ impl CoordinatorState { statements, } => match *self.role.borrow() { LeadershipRole::Leader => { + if self.leader_draining.get() { + let _ = send_query_result_to_main( + request_id, + Err("Coordinator is shutting down".to_string()), + ); + return; + } if !*self.db_worker_ready.borrow() { let _ = send_query_result_to_main( request_id, @@ -534,6 +742,111 @@ impl CoordinatorState { } } }, + WorkerMessage::InstallSnapshot { + request_id, + url, + compression, + sha256, + uncompressed_size, + } => { + let query_id = format!("{}:{request_id}", self.worker_id); + if self.take_snapshot_cancellation(&self.worker_id, &query_id) { + let _ = send_query_result_to_main( + request_id, + Err("Snapshot installation cancelled".to_string()), + ); + return; + } + match *self.role.borrow() { + LeadershipRole::Leader => { + if self.leader_draining.get() { + let _ = send_query_result_to_main( + request_id, + Err("Coordinator is shutting down".to_string()), + ); + return; + } + if !*self.db_worker_ready.borrow() { + let _ = send_query_result_to_main( + request_id, + Err(WORKER_ERROR_TYPE_INITIALIZATION_PENDING.to_string()), + ); + return; + } + self.forward_snapshot_to_db( + DbRequestOrigin::Local { request_id }, + SnapshotWaiter::Local { + request_id, + requester_id: self.worker_id.clone(), + query_id, + }, + url, + compression, + sha256, + uncompressed_size, + ); + } + LeadershipRole::Follower => { + if !*self.leader_ready.borrow() { + let _ = send_query_result_to_main( + request_id, + Err(WORKER_ERROR_TYPE_INITIALIZATION_PENDING.to_string()), + ); + return; + } + self.follower_pending + .borrow_mut() + .insert(query_id.clone(), request_id); + let timeout_query_id = query_id.clone(); + let timeout_state = Rc::clone(self); + // Snapshot downloads may legitimately outlive ordinary + // query timeouts, but must still settle if no leader reply + // ever arrives. + let timeout = self.query_timeout_ms.max(600_000.0); + spawn_local(async move { + sleep_ms(timeout.ceil().min(i32::MAX as f64) as i32).await; + timeout_state.begin_snapshot_timeout_cancellation(&timeout_query_id); + }); + let request = ChannelMessage::InstallSnapshotRequest { + query_id: query_id.clone(), + requester_id: self.worker_id.clone(), + url, + compression, + sha256, + uncompressed_size, + }; + if let Err(err) = send_channel_message(&self.channel, &request) { + self.handle_follower_forward_failure(&query_id, request_id, &err); + let _ = send_worker_error_message(&err); + } + } + } + } + WorkerMessage::CancelSnapshot { .. } => {} + WorkerMessage::CancelForwardedSnapshot { request_id } => { + let query_id = format!("{}:{request_id}", self.worker_id); + if let Err(error) = send_channel_message( + &self.channel, + &ChannelMessage::CancelSnapshotRequest { + requester_id: self.worker_id.clone(), + query_id, + }, + ) { + let _ = send_worker_error_message(&error); + } + } + WorkerMessage::Shutdown { request_id } => { + self.leader_draining.set(true); + if matches!(*self.role.borrow(), LeadershipRole::Leader) { + self.shutdown_requests.borrow_mut().push(request_id); + self.finish_shutdown_if_drained(); + } else { + let _ = send_query_result_to_main( + request_id, + Ok("Coordinator drained".to_string()), + ); + } + } } } @@ -577,6 +890,9 @@ impl CoordinatorState { params, } => { if matches!(*self.role.borrow(), LeadershipRole::Leader) { + if self.reject_if_leader_draining(&query_id) { + return; + } if !*self.db_worker_ready.borrow() { let _ = send_channel_message( &self.channel, @@ -596,6 +912,9 @@ impl CoordinatorState { statements, } => { if matches!(*self.role.borrow(), LeadershipRole::Leader) { + if self.reject_if_leader_draining(&query_id) { + return; + } if !*self.db_worker_ready.borrow() { let _ = send_channel_message( &self.channel, @@ -610,12 +929,80 @@ impl CoordinatorState { self.forward_batch_to_db(DbRequestOrigin::Forwarded { query_id }, statements); } } + ChannelMessage::InstallSnapshotRequest { + query_id, + requester_id, + url, + compression, + sha256, + uncompressed_size, + } => { + if matches!(*self.role.borrow(), LeadershipRole::Leader) { + if self.reject_if_leader_draining(&query_id) { + return; + } + if self.take_snapshot_cancellation(&requester_id, &query_id) { + return; + } + if !*self.db_worker_ready.borrow() { + let _ = send_channel_message( + &self.channel, + &ChannelMessage::QueryResponse { + query_id, + result: None, + error: Some(WORKER_ERROR_TYPE_INITIALIZATION_PENDING.to_string()), + }, + ); + return; + } + self.forward_snapshot_to_db( + DbRequestOrigin::Forwarded { + query_id: query_id.clone(), + }, + SnapshotWaiter::Forwarded { + requester_id, + query_id, + }, + url, + compression, + sha256, + uncompressed_size, + ); + } + } + ChannelMessage::CancelSnapshotRequest { + requester_id, + query_id, + } => { + if matches!(*self.role.borrow(), LeadershipRole::Leader) { + if !self.detach_snapshot_waiter(&requester_id, &query_id) { + self.remember_snapshot_cancellation(requester_id.clone(), query_id.clone()); + let _ = send_channel_message( + &self.channel, + &ChannelMessage::SnapshotCancellationResponse { + query_id, + result: None, + error: Some("Snapshot installation cancelled".to_string()), + }, + ); + } + } + } + ChannelMessage::SnapshotCancellationResponse { query_id, .. } => { + self.follower_pending.borrow_mut().remove(&query_id); + self.snapshot_cancellation_pending + .borrow_mut() + .remove(&query_id); + } ChannelMessage::QueryResponse { query_id, result, error, } => { if let Some(request_id) = self.follower_pending.borrow_mut().remove(&query_id) { + self.snapshot_cancellation_pending + .borrow_mut() + .remove(&query_id); let outcome = match (result, error) { (Some(res), _) => Ok(res), (_, Some(err)) => Err(err), @@ -653,6 +1040,11 @@ impl CoordinatorState { }, ); } + DbRequestOrigin::Snapshot { .. } => { + unreachable!( + "snapshot origins are only dispatched by forward_snapshot_to_db" + ) + } } return; }; @@ -700,6 +1092,11 @@ impl CoordinatorState { }, ); } + DbRequestOrigin::Snapshot { .. } => { + unreachable!( + "snapshot origins are only dispatched by forward_snapshot_to_db" + ) + } } return; }; @@ -721,6 +1118,180 @@ impl CoordinatorState { self.post_db_worker_message(worker, db_request_id, msg); } + fn forward_snapshot_to_db( + self: &Rc, + origin: DbRequestOrigin, + waiter: SnapshotWaiter, + url: String, + compression: SnapshotCompression, + sha256: String, + uncompressed_size: f64, + ) { + if !uncompressed_size.is_finite() + || uncompressed_size < 512.0 + || uncompressed_size.fract() != 0.0 + || uncompressed_size > MAX_SNAPSHOT_SIZE as f64 + { + self.fail_origin( + origin, + format!( + "Snapshot uncompressed size must be an integer between 512 and {MAX_SNAPSHOT_SIZE}" + ), + ); + return; + } + let key = SnapshotKey { + url: url.clone(), + compression, + sha256: sha256.to_ascii_lowercase(), + uncompressed_size: uncompressed_size as u64, + }; + if self.attach_to_inflight_snapshot(&key, waiter.clone()) { + return; + } + let worker = { + let borrow = self.db_worker.borrow(); + let Some(worker) = borrow.as_ref() else { + self.fail_origin(origin, WORKER_ERROR_TYPE_INITIALIZATION_PENDING.to_string()); + return; + }; + worker.clone() + }; + let db_request_id = { + let mut next = self.next_db_request_id.borrow_mut(); + let id = *next; + *next = next.wrapping_add(1).max(1); + id + }; + self.db_pending.borrow_mut().insert( + db_request_id, + DbRequestOrigin::Snapshot { + key, + waiters: vec![waiter], + cancelling: Vec::new(), + phase: SnapshotJobPhase::Active, + }, + ); + self.post_db_worker_message( + worker, + db_request_id, + WorkerMessage::InstallSnapshot { + request_id: db_request_id, + url, + compression, + sha256, + uncompressed_size, + }, + ); + } + + fn detach_snapshot_waiter(&self, requester_id: &str, query_id: &str) -> bool { + let mut cancel_db_request = None; + let mut acknowledge_now = Vec::new(); + let mut detached = false; + { + let mut pending = self.db_pending.borrow_mut(); + for (request_id, origin) in pending.iter_mut() { + let DbRequestOrigin::Snapshot { + waiters, + cancelling, + phase, + .. + } = origin + else { + continue; + }; + if cancelling + .iter() + .any(|waiter| waiter.snapshot_identity() == (requester_id, query_id)) + { + return true; + } + let Some(index) = waiters + .iter() + .position(|waiter| waiter.snapshot_identity() == (requester_id, query_id)) + else { + continue; + }; + let waiter = waiters.remove(index); + detached = true; + if waiters.is_empty() { + cancelling.push(waiter); + *phase = SnapshotJobPhase::CancellingLastWaiter; + cancel_db_request = Some(*request_id); + } else if matches!(waiter, SnapshotWaiter::Local { .. }) { + // The leader coordinator owns the DB worker, so its local + // close cannot complete until the shared job has delivered + // the forwarded waiters that still depend on it. + cancelling.push(waiter); + } else { + acknowledge_now.push(waiter); + } + break; + } + } + + for waiter in acknowledge_now { + self.deliver_snapshot_cancellation( + waiter, + Err("Snapshot installation cancelled".to_string()), + ); + } + let worker = self.db_worker.borrow().clone(); + if let (Some(worker), Some(request_id)) = (worker, cancel_db_request) { + self.post_db_worker_cancellation(&worker, request_id); + } + detached + } + + fn deliver_snapshot_cancellation( + &self, + waiter: SnapshotWaiter, + outcome: Result, + ) { + let (_, query_id) = waiter.snapshot_identity(); + let (result, error) = match outcome { + Ok(result) => (Some(result), None), + Err(error) => (None, Some(error)), + }; + let _ = send_channel_message( + &self.channel, + &ChannelMessage::SnapshotCancellationResponse { + query_id: query_id.to_string(), + result, + error, + }, + ); + } + + fn finish_shutdown_if_drained(&self) { + if !self.leader_draining.get() || !self.db_pending.borrow().is_empty() { + return; + } + for request_id in self.shutdown_requests.borrow_mut().drain(..) { + let _ = send_query_result_to_main(request_id, Ok("Coordinator drained".to_string())); + } + } + + fn post_db_worker_cancellation(&self, worker: &Worker, request_id: u32) { + let message = WorkerMessage::CancelSnapshot { request_id }; + match serde_wasm_bindgen::to_value(&message) { + Ok(value) => { + if let Err(error) = worker.post_message(&value) { + let _ = send_worker_error_message(&format!( + "Failed to cancel snapshot installation: {}", + js_value_to_string(&error) + )); + } + } + Err(error) => { + let _ = send_worker_error_message(&format!( + "Failed to serialize snapshot cancellation: {error:?}" + )); + } + } + } + fn post_db_worker_message(&self, worker: Worker, db_request_id: u32, msg: WorkerMessage) { match serde_wasm_bindgen::to_value(&msg) { Ok(val) => { @@ -731,6 +1302,7 @@ impl CoordinatorState { origin, "Failed to dispatch query to DB worker".to_string(), ); + self.finish_shutdown_if_drained(); } } } @@ -738,6 +1310,7 @@ impl CoordinatorState { let _ = send_worker_error_message(&format!("{err:?}")); if let Some(origin) = self.db_pending.borrow_mut().remove(&db_request_id) { self.fail_origin(origin, "Failed to serialize query".to_string()); + self.finish_shutdown_if_drained(); } } } @@ -758,6 +1331,40 @@ impl CoordinatorState { }, ); } + DbRequestOrigin::Snapshot { + waiters, + cancelling, + .. + } => { + for waiter in waiters { + self.deliver_snapshot_waiter(waiter, Err(error.clone())); + } + for waiter in cancelling { + self.deliver_snapshot_cancellation(waiter, Err(error.clone())); + } + } + } + } + + fn deliver_snapshot_waiter(&self, waiter: SnapshotWaiter, outcome: Result) { + match waiter { + SnapshotWaiter::Local { request_id, .. } => { + let _ = send_query_result_to_main(request_id, outcome); + } + SnapshotWaiter::Forwarded { query_id, .. } => { + let (result, error) = match outcome { + Ok(result) => (Some(result), None), + Err(error) => (None, Some(error)), + }; + let _ = send_channel_message( + &self.channel, + &ChannelMessage::QueryResponse { + query_id, + result, + error, + }, + ); + } } } @@ -801,11 +1408,67 @@ impl CoordinatorState { ); } }, + DbRequestOrigin::Snapshot { + waiters, + cancelling, + phase, + .. + } => { + for waiter in waiters { + self.deliver_snapshot_waiter(waiter, outcome.clone()); + } + for waiter in cancelling { + let cancellation_outcome = match phase { + SnapshotJobPhase::Active => { + Err("Snapshot installation cancelled".to_string()) + } + SnapshotJobPhase::CancellingLastWaiter => outcome.clone(), + }; + self.deliver_snapshot_cancellation(waiter, cancellation_outcome); + } + } } + self.finish_shutdown_if_drained(); + } + + fn handle_snapshot_cancelled(&self, db_request_id: u32) { + let Some(DbRequestOrigin::Snapshot { + waiters, + cancelling, + .. + }) = self.db_pending.borrow_mut().remove(&db_request_id) + else { + return; + }; + debug_assert!(waiters.is_empty()); + for waiter in cancelling { + self.deliver_snapshot_cancellation( + waiter, + Err("Snapshot installation cancelled".to_string()), + ); + } + self.finish_shutdown_if_drained(); } fn mark_leader_known(&self, leader_id: String) { - *self.leader_id.borrow_mut() = Some(leader_id); + let previous = self.leader_id.borrow_mut().replace(leader_id.clone()); + if previous + .as_ref() + .is_some_and(|current| current != &leader_id) + { + self.snapshot_cancellation_pending.borrow_mut().clear(); + let pending = self + .follower_pending + .borrow_mut() + .drain() + .collect::>(); + for (_, request_id) in pending { + let _ = send_query_result_to_main( + request_id, + Err("Leader changed while request was pending".to_string()), + ); + } + } } fn signal_ready_once(&self) { @@ -830,6 +1493,7 @@ impl DbWorkerState { db_name: config.db_name, db_queue: Rc::new(RefCell::new(VecDeque::new())), db_processing: Rc::new(Cell::new(false)), + snapshot_cancellations: Rc::new(RefCell::new(HashMap::new())), hooks, }) } @@ -864,6 +1528,47 @@ impl DbWorkerState { } => { self.enqueue_job(request_id, DbJobPayload::Batch { statements }); } + WorkerMessage::InstallSnapshot { + request_id, + url, + compression, + sha256, + uncompressed_size, + } => { + let cancellation = SnapshotCancellation::default(); + self.snapshot_cancellations + .borrow_mut() + .insert(request_id, cancellation.clone()); + self.enqueue_job( + request_id, + DbJobPayload::Snapshot { + db_name: self.db_name.clone(), + url, + compression, + sha256, + uncompressed_size, + cancellation, + }, + ); + } + WorkerMessage::CancelSnapshot { request_id } => { + let cancelled = if let Some(cancellation) = + self.snapshot_cancellations.borrow().get(&request_id) + { + cancellation.cancel(); + true + } else { + false + }; + match make_snapshot_cancelled_message(request_id, cancelled) { + Ok(response) => self.hooks.deliver.as_ref()(&response), + Err(error) => { + let _ = send_worker_error(error); + } + } + } + WorkerMessage::CancelForwardedSnapshot { .. } => {} + WorkerMessage::Shutdown { .. } => {} } } @@ -892,6 +1597,10 @@ impl DbWorkerState { let exec = Rc::clone(&hooks.exec); let deliver = Rc::clone(&hooks.deliver); let result = exec.as_ref()(db, job.payload).await; + state + .snapshot_cancellations + .borrow_mut() + .remove(&job.request_id); match make_query_result_message(job.request_id, result) { Ok(resp) => deliver.as_ref()(&resp), Err(err) => { @@ -1024,6 +1733,21 @@ pub fn make_query_result_message( Ok(response) } +fn make_snapshot_cancelled_message( + request_id: u32, + cancelled: bool, +) -> Result { + let response = js_sys::Object::new(); + set_js_property(&response, "type", &JsValue::from_str("snapshot-cancelled"))?; + set_js_property( + &response, + "requestId", + &JsValue::from_f64(request_id as f64), + )?; + set_js_property(&response, "cancelled", &JsValue::from_bool(cancelled))?; + Ok(response) +} + pub fn send_query_result_to_main( request_id: u32, result: Result, @@ -1043,17 +1767,53 @@ async fn exec_on_db( ) -> Result { let db_opt = db.borrow_mut().take(); let result = match db_opt { - Some(mut database) => { - let result = match payload { - DbJobPayload::Query { sql, params } => match params { + Some(mut database) => match payload { + DbJobPayload::Snapshot { + db_name, + url, + compression, + sha256, + uncompressed_size, + cancellation, + } => { + drop(database); + match SQLiteDatabase::install_snapshot_cancellable( + &db_name, + &url, + compression, + &sha256, + uncompressed_size, + cancellation, + ) + .await + { + Ok((database, stats)) => { + *db.borrow_mut() = Some(database); + serde_json::to_string(&stats) + .map_err(|e| format!("Failed to encode snapshot result: {e}")) + } + Err(err) => { + if let Ok(database) = SQLiteDatabase::initialize_opfs(&db_name).await { + *db.borrow_mut() = Some(database); + } + Err(err) + } + } + } + DbJobPayload::Query { sql, params } => { + let result = match params { Some(p) => database.exec_with_params(&sql, p).await, None => database.exec(&sql).await, - }, - DbJobPayload::Batch { statements } => database.exec_batch(statements).await, - }; - *db.borrow_mut() = Some(database); - result - } + }; + *db.borrow_mut() = Some(database); + result + } + DbJobPayload::Batch { statements } => { + let result = database.exec_batch(statements).await; + *db.borrow_mut() = Some(database); + result + } + }, None => Err(WORKER_ERROR_TYPE_INITIALIZATION_PENDING.to_string()), }; result @@ -1210,19 +1970,27 @@ mod tests { listener.forget(); state.on_lock_granted(); - sleep_ms(50).await; + for _ in 0..40 { + let ready_received = received + .borrow() + .iter() + .any(|message| matches!(message, ChannelMessage::LeaderReady { .. })); + if ready_received { + break; + } + sleep_ms(25).await; + } - let msgs = received.borrow(); - assert!( - msgs.iter() - .any(|m| matches!(m, ChannelMessage::NewLeader { .. })), - "should announce new-leader" - ); - assert!( - msgs.iter() - .any(|m| matches!(m, ChannelMessage::LeaderReady { .. })), - "should announce leader-ready" - ); + let new_leader_received = received + .borrow() + .iter() + .any(|message| matches!(message, ChannelMessage::NewLeader { .. })); + let ready_received = received + .borrow() + .iter() + .any(|message| matches!(message, ChannelMessage::LeaderReady { .. })); + assert!(new_leader_received, "should announce new-leader"); + assert!(ready_received, "should announce leader-ready"); } #[wasm_bindgen_test(async)] @@ -1461,7 +2229,12 @@ mod tests { params: None, }); - sleep_ms(30).await; + for _ in 0..40 { + if results.length() == 2 { + break; + } + sleep_ms(5).await; + } assert_eq!(results.length(), 2, "both queued queries should complete"); for entry in results.iter() { @@ -1477,4 +2250,396 @@ mod tests { assert!(error.is_none(), "no error expected"); } } + + #[wasm_bindgen_test] + fn identical_snapshot_requests_coalesce_to_one_db_job() { + let state = CoordinatorState::new(WorkerConfig { + db_name: "testdb-snapshot-coalesce".to_string(), + follower_timeout_ms: 10.0, + query_timeout_ms: 10.0, + }) + .expect("state"); + let key = SnapshotKey { + url: "https://example.invalid/snapshot.db".to_string(), + compression: SnapshotCompression::None, + sha256: "a".repeat(64), + uncompressed_size: 1024, + }; + state.db_pending.borrow_mut().insert( + 1, + DbRequestOrigin::Snapshot { + key: key.clone(), + waiters: vec![SnapshotWaiter::Local { + request_id: 10, + requester_id: "local".to_string(), + query_id: "local:10".to_string(), + }], + cancelling: Vec::new(), + phase: SnapshotJobPhase::Active, + }, + ); + + assert!(state.attach_to_inflight_snapshot( + &key, + SnapshotWaiter::Forwarded { + requester_id: "follower-a".to_string(), + query_id: "remote-1".to_string(), + }, + )); + assert_eq!(state.db_pending.borrow().len(), 1); + let pending = state.db_pending.borrow(); + let DbRequestOrigin::Snapshot { waiters, .. } = pending.get(&1).expect("job") else { + panic!("expected snapshot job"); + }; + assert_eq!(waiters.len(), 2); + } + + #[wasm_bindgen_test] + fn snapshot_cancellation_detaches_only_matching_forwarded_waiters() { + let state = CoordinatorState::new(WorkerConfig { + db_name: "testdb-snapshot-detach".to_string(), + follower_timeout_ms: 10.0, + query_timeout_ms: 10.0, + }) + .expect("state"); + let key = SnapshotKey { + url: "https://example.invalid/snapshot.db".to_string(), + compression: SnapshotCompression::None, + sha256: "a".repeat(64), + uncompressed_size: 1024, + }; + state.db_pending.borrow_mut().insert( + 1, + DbRequestOrigin::Snapshot { + key, + waiters: vec![ + SnapshotWaiter::Forwarded { + requester_id: "closing-follower".to_string(), + query_id: "closing-query".to_string(), + }, + SnapshotWaiter::Forwarded { + requester_id: "live-follower".to_string(), + query_id: "live-query".to_string(), + }, + ], + cancelling: Vec::new(), + phase: SnapshotJobPhase::Active, + }, + ); + + state.detach_snapshot_waiter("closing-follower", "closing-query"); + + let pending = state.db_pending.borrow(); + let DbRequestOrigin::Snapshot { waiters, .. } = pending.get(&1).expect("live job") else { + panic!("expected snapshot job"); + }; + assert_eq!(waiters.len(), 1); + assert!(matches!( + &waiters[0], + SnapshotWaiter::Forwarded { requester_id, query_id } + if requester_id == "live-follower" && query_id == "live-query" + )); + } + + #[wasm_bindgen_test] + fn cancelling_last_waiter_does_not_accept_a_new_coalesced_waiter() { + let state = CoordinatorState::new(WorkerConfig { + db_name: "testdb-snapshot-cancelling-successor".to_string(), + follower_timeout_ms: 10.0, + query_timeout_ms: 10.0, + }) + .expect("state"); + let key = SnapshotKey { + url: "https://example.invalid/snapshot.db".to_string(), + compression: SnapshotCompression::None, + sha256: "a".repeat(64), + uncompressed_size: 1024, + }; + state.db_pending.borrow_mut().insert( + 1, + DbRequestOrigin::Snapshot { + key: key.clone(), + waiters: vec![SnapshotWaiter::Forwarded { + requester_id: "follower-a".to_string(), + query_id: "a:1".to_string(), + }], + cancelling: Vec::new(), + phase: SnapshotJobPhase::Active, + }, + ); + assert!(state.detach_snapshot_waiter("follower-a", "a:1")); + assert!(!state.attach_to_inflight_snapshot( + &key, + SnapshotWaiter::Forwarded { + requester_id: "follower-b".to_string(), + query_id: "b:1".to_string(), + }, + )); + let pending = state.db_pending.borrow(); + let DbRequestOrigin::Snapshot { + waiters, + cancelling, + phase, + .. + } = pending.get(&1).expect("cancelling job") + else { + panic!("expected snapshot job"); + }; + assert!(waiters.is_empty()); + assert_eq!(cancelling.len(), 1); + assert_eq!(*phase, SnapshotJobPhase::CancellingLastWaiter); + } + + #[wasm_bindgen_test] + fn leader_local_cancellation_waits_for_forwarded_coalesced_waiter() { + let state = CoordinatorState::new(WorkerConfig { + db_name: "testdb-snapshot-leader-close".to_string(), + follower_timeout_ms: 10.0, + query_timeout_ms: 10.0, + }) + .expect("state"); + state.db_pending.borrow_mut().insert( + 1, + DbRequestOrigin::Snapshot { + key: SnapshotKey { + url: "https://example.invalid/snapshot.db".to_string(), + compression: SnapshotCompression::None, + sha256: "a".repeat(64), + uncompressed_size: 1024, + }, + waiters: vec![ + SnapshotWaiter::Local { + request_id: 1, + requester_id: "leader".to_string(), + query_id: "leader:1".to_string(), + }, + SnapshotWaiter::Forwarded { + requester_id: "follower".to_string(), + query_id: "follower:1".to_string(), + }, + ], + cancelling: Vec::new(), + phase: SnapshotJobPhase::Active, + }, + ); + assert!(state.detach_snapshot_waiter("leader", "leader:1")); + let pending = state.db_pending.borrow(); + let DbRequestOrigin::Snapshot { + waiters, + cancelling, + phase, + .. + } = pending.get(&1).expect("shared job") + else { + panic!("expected snapshot job"); + }; + assert_eq!(waiters.len(), 1); + assert_eq!(cancelling.len(), 1); + assert_eq!(*phase, SnapshotJobPhase::Active); + } + + #[wasm_bindgen_test] + fn leader_shutdown_rejects_successor_before_acknowledging_local_cancel() { + let state = CoordinatorState::new(WorkerConfig { + db_name: "testdb-snapshot-draining-successor".to_string(), + follower_timeout_ms: 10.0, + query_timeout_ms: 10.0, + }) + .expect("state"); + *state.role.borrow_mut() = LeadershipRole::Leader; + *state.db_worker_ready.borrow_mut() = true; + state.db_pending.borrow_mut().insert( + 1, + DbRequestOrigin::Snapshot { + key: SnapshotKey { + url: "https://example.invalid/snapshot.db".to_string(), + compression: SnapshotCompression::None, + sha256: "a".repeat(64), + uncompressed_size: 1024, + }, + waiters: vec![SnapshotWaiter::Local { + request_id: 1, + requester_id: "leader".to_string(), + query_id: "leader:1".to_string(), + }], + cancelling: Vec::new(), + phase: SnapshotJobPhase::Active, + }, + ); + assert!(state.detach_snapshot_waiter("leader", "leader:1")); + state.handle_main_message(WorkerMessage::Shutdown { request_id: 99 }); + assert!(state.leader_draining.get()); + assert_eq!(state.shutdown_requests.borrow().as_slice(), &[99]); + + state.handle_channel_message(ChannelMessage::InstallSnapshotRequest { + requester_id: "follower".to_string(), + query_id: "follower:2".to_string(), + url: "https://example.invalid/snapshot.db".to_string(), + compression: SnapshotCompression::None, + sha256: "a".repeat(64), + uncompressed_size: 1024.0, + }); + assert_eq!(state.db_pending.borrow().len(), 1); + + state.handle_snapshot_cancelled(1); + assert!(state.db_pending.borrow().is_empty()); + assert!(state.shutdown_requests.borrow().is_empty()); + } + + #[wasm_bindgen_test] + fn local_cancel_before_request_still_drains_unrelated_forwarded_work() { + let state = CoordinatorState::new(WorkerConfig { + db_name: "testdb-local-cancel-before-request-drain".to_string(), + follower_timeout_ms: 10.0, + query_timeout_ms: 10.0, + }) + .expect("state"); + *state.role.borrow_mut() = LeadershipRole::Leader; + *state.db_worker_ready.borrow_mut() = true; + state.db_pending.borrow_mut().insert( + 7, + DbRequestOrigin::Forwarded { + query_id: "follower:7".to_string(), + }, + ); + + let local_query_id = format!("{}:5", state.worker_id); + state.handle_channel_message(ChannelMessage::CancelSnapshotRequest { + requester_id: state.worker_id.clone(), + query_id: local_query_id, + }); + assert_eq!(state.cancelled_snapshot_requests.borrow().len(), 1); + + state.handle_main_message(WorkerMessage::InstallSnapshot { + request_id: 5, + url: "https://example.invalid/snapshot.db".to_string(), + compression: SnapshotCompression::None, + sha256: "a".repeat(64), + uncompressed_size: 1024.0, + }); + assert!(state.cancelled_snapshot_requests.borrow().is_empty()); + assert_eq!(state.db_pending.borrow().len(), 1); + + state.handle_main_message(WorkerMessage::Shutdown { request_id: 99 }); + assert!(state.leader_draining.get()); + assert_eq!(state.shutdown_requests.borrow().as_slice(), &[99]); + + state.handle_channel_message(ChannelMessage::QueryRequest { + query_id: "new-follower:8".to_string(), + sql: "SELECT 1".to_string(), + params: None, + }); + assert_eq!(state.db_pending.borrow().len(), 1); + + state.handle_db_query_result(7, Some("[]".to_string()), None); + assert!(state.db_pending.borrow().is_empty()); + assert!(state.shutdown_requests.borrow().is_empty()); + } + + #[wasm_bindgen_test] + fn snapshot_cancellation_arriving_before_request_prevents_dispatch() { + let state = CoordinatorState::new(WorkerConfig { + db_name: "testdb-snapshot-cancel-before-request".to_string(), + follower_timeout_ms: 10.0, + query_timeout_ms: 10.0, + }) + .expect("state"); + *state.role.borrow_mut() = LeadershipRole::Leader; + state.handle_channel_message(ChannelMessage::CancelSnapshotRequest { + requester_id: "closing-follower".to_string(), + query_id: "closing-follower:7".to_string(), + }); + assert_eq!(state.cancelled_snapshot_requests.borrow().len(), 1); + + state.handle_channel_message(ChannelMessage::InstallSnapshotRequest { + requester_id: "closing-follower".to_string(), + query_id: "closing-follower:7".to_string(), + url: "https://example.invalid/snapshot.db".to_string(), + compression: SnapshotCompression::None, + sha256: "a".repeat(64), + uncompressed_size: 1024.0, + }); + + assert!(state.cancelled_snapshot_requests.borrow().is_empty()); + assert!(state.db_pending.borrow().is_empty()); + } + + #[wasm_bindgen_test] + fn unmatched_snapshot_cancellation_tombstones_expire_and_remain_bounded() { + let state = CoordinatorState::new(WorkerConfig { + db_name: "testdb-snapshot-cancel-expiry".to_string(), + follower_timeout_ms: 10.0, + query_timeout_ms: 10.0, + }) + .expect("state"); + state + .cancelled_snapshot_requests + .borrow_mut() + .insert(("late".to_string(), "completed".to_string()), 10.0); + state.prune_snapshot_cancellation_tombstones(10.0); + assert!(state.cancelled_snapshot_requests.borrow().is_empty()); + + for index in 0..(MAX_SNAPSHOT_CANCELLATION_TOMBSTONES + 20) { + state.cancelled_snapshot_requests.borrow_mut().insert( + ("orphan".to_string(), format!("query-{index}")), + 1000.0 + index as f64, + ); + } + state.prune_snapshot_cancellation_tombstones(0.0); + assert_eq!( + state.cancelled_snapshot_requests.borrow().len(), + MAX_SNAPSHOT_CANCELLATION_TOMBSTONES + ); + } + + #[wasm_bindgen_test] + fn follower_timeout_waits_for_authoritative_activation_outcome() { + let state = CoordinatorState::new(WorkerConfig { + db_name: "testdb-snapshot-timeout-outcome".to_string(), + follower_timeout_ms: 10.0, + query_timeout_ms: 10.0, + }) + .expect("state"); + let query_id = "follower:9".to_string(); + state + .follower_pending + .borrow_mut() + .insert(query_id.clone(), 9); + + assert!(state.begin_snapshot_timeout_cancellation(&query_id)); + assert_eq!(state.follower_pending.borrow().get(&query_id), Some(&9)); + assert!(state + .snapshot_cancellation_pending + .borrow() + .contains(&query_id)); + + state.handle_channel_message(ChannelMessage::SnapshotCancellationResponse { + query_id: query_id.clone(), + result: Some("activation-won".to_string()), + error: None, + }); + assert!(!state.follower_pending.borrow().contains_key(&query_id)); + assert!(!state + .snapshot_cancellation_pending + .borrow() + .contains(&query_id)); + } + + #[wasm_bindgen_test] + fn leader_change_cleans_up_forwarded_requests() { + let state = CoordinatorState::new(WorkerConfig { + db_name: "testdb-leader-change".to_string(), + follower_timeout_ms: 10.0, + query_timeout_ms: 10.0, + }) + .expect("state"); + state.mark_leader_known("leader-a".to_string()); + state + .follower_pending + .borrow_mut() + .insert("snapshot-1".to_string(), 99); + + state.mark_leader_known("leader-b".to_string()); + assert!(state.follower_pending.borrow().is_empty()); + } } diff --git a/packages/sqlite-web-core/src/database.rs b/packages/sqlite-web-core/src/database.rs index 32a13c7..fb6e085 100644 --- a/packages/sqlite-web-core/src/database.rs +++ b/packages/sqlite-web-core/src/database.rs @@ -1,11 +1,67 @@ use crate::database_functions::register_custom_functions; -use crate::messages::SqlBatchStatement; +use crate::messages::{SnapshotCompression, SqlBatchStatement}; use crate::util::sanitize_db_filename; use base64::Engine; +use js_sys::{Array, Date, Function, Reflect, Uint8Array}; +use sha2::{Digest, Sha256}; use sqlite_wasm_rs::export::{install_opfs_sahpool, *}; +use std::cell::{Cell, RefCell}; use std::ffi::{CStr, CString}; use std::os::raw::c_void; +use std::rc::Rc; use wasm_bindgen::prelude::*; +use wasm_bindgen::JsCast; +use wasm_bindgen_futures::JsFuture; +use web_sys::{ + AbortController, DedicatedWorkerGlobalScope, ReadableStream, ReadableStreamDefaultReader, + Request, RequestInit, Response, +}; + +const JS_MAX_SAFE_INTEGER: u64 = 9_007_199_254_740_991; +const OPFS_SAH_HEADER_SIZE: u64 = 4096; +pub(crate) const MAX_SNAPSHOT_SIZE: u64 = JS_MAX_SAFE_INTEGER - OPFS_SAH_HEADER_SIZE; + +#[derive(Clone, Default)] +pub(crate) struct SnapshotCancellation { + cancelled: Rc>, + abort_controller: Rc>>, +} + +impl SnapshotCancellation { + pub(crate) fn cancel(&self) { + self.cancelled.set(true); + if let Some(controller) = self.abort_controller.borrow().as_ref() { + controller.abort(); + } + } + + fn register(&self, controller: &AbortController) { + if self.cancelled.get() { + controller.abort(); + } else { + self.abort_controller + .borrow_mut() + .replace(controller.clone()); + } + } + + fn ensure_active(&self) -> Result<(), String> { + if self.cancelled.get() { + Err("Snapshot installation cancelled".to_string()) + } else { + Ok(()) + } + } +} + +#[derive(Debug, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct SnapshotInstallResult { + pub bytes_written: u64, + pub compression: SnapshotCompression, + pub elapsed_ms: f64, + pub sha256: String, +} // Real SQLite database using sqlite-wasm-rs FFI pub struct SQLiteDatabase { @@ -698,16 +754,9 @@ impl SQLiteDatabase { } } - pub async fn initialize_opfs(db_name: &str) -> Result { - // Install OPFS VFS and set as default - install_opfs_sahpool(None, true) - .await - .map_err(|e| JsValue::from_str(&format!("Failed to install OPFS VFS: {e:?}")))?; - - // Open database with OPFS + fn open_opfs_path(db_path: &str) -> Result { let mut db: *mut sqlite3 = std::ptr::null_mut(); - let sanitized = sanitize_db_filename(db_name); - let open_uri = format!("opfs-sahpool:{}", sanitized); + let open_uri = format!("opfs-sahpool:{db_path}"); let db_name = CString::new(open_uri.clone()).map_err(|e| { JsValue::from_str(&format!( "Invalid database URI (NUL found): {open_uri} ({e})" @@ -756,6 +805,401 @@ impl SQLiteDatabase { }) } + fn open_opfs(db_name: &str) -> Result { + Self::open_opfs_path(&sanitize_db_filename(db_name)) + } + + fn snapshot_paths(db_name: &str) -> Result<(String, String, String), String> { + let target = sanitize_db_filename(db_name); + let staging = format!("{target}.snapshot-staging"); + let backup = format!("{target}.snapshot-backup"); + // sqlite-wasm-rs publishes utility names as rooted paths and requires + // their complete UTF-8 representation to be strictly shorter than 512. + // SQLite also passes the VFS-prefixed spelling through xFullPathname, + // including its trailing NUL, before staged validation can open it. + let fits_snapshot_path = |path: &str| { + format!("/{path}").len() < 512 && "opfs-sahpool:".len() + path.len() + 1 <= 512 + }; + if !fits_snapshot_path(&staging) || !fits_snapshot_path(&backup) { + return Err("Database name is too long for atomic snapshot installation".to_string()); + } + Ok((target, staging, backup)) + } + + /// Complete or roll back an interrupted snapshot activation. The commit + /// protocol first moves the old target to `backup`, then moves the fully + /// validated staging database to `target`. Therefore a missing target plus + /// a backup always means activation was interrupted before commit. + async fn validate_snapshot_path(path: &str, label: &str) -> Result<(), String> { + let mut database = + Self::open_opfs_path(path).map_err(|e| format!("Failed to open {label}: {e:?}"))?; + let raw = database + .exec("PRAGMA quick_check") + .await + .map_err(|e| format!("{label} failed SQLite integrity validation: {e}"))?; + let rows: Vec = serde_json::from_str(&raw) + .map_err(|e| format!("{label} returned an invalid integrity result: {e}"))?; + let valid = rows.len() == 1 + && rows[0] + .get("quick_check") + .and_then(serde_json::Value::as_str) + == Some("ok"); + if !valid { + return Err(format!("{label} failed SQLite integrity validation: {raw}")); + } + Ok(()) + } + + async fn recover_snapshot_activation( + util: &OpfsSAHPoolUtil, + target: &str, + staging: &str, + backup: &str, + ) -> Result<(), String> { + if !util.has_path(target) && util.has_path(backup) { + util.rename_path(backup, target) + .map_err(|e| format!("Failed to restore snapshot backup: {e:?}"))?; + } else if util.has_path(target) && util.has_path(backup) { + // Both names exist between activation and verified commit. Never + // discard the known-good backup merely because a target exists. + match Self::validate_snapshot_path(target, "recovered snapshot").await { + Ok(()) => { + util.unlink(backup).map_err(|e| { + format!("Recovered snapshot is valid, but its backup could not be removed: {e:?}") + })?; + } + Err(validation_error) => { + Self::restore_snapshot_backup(util, target, backup).map_err(|rollback| { + format!("{validation_error}; rollback also failed: {rollback}") + })?; + } + } + } + + if util.has_path(staging) { + let _ = util.unlink(staging); + } + Ok(()) + } + + fn activate_snapshot( + util: &OpfsSAHPoolUtil, + target: &str, + staging: &str, + backup: &str, + ) -> Result { + if util.has_path(backup) { + util.unlink(backup) + .map_err(|e| format!("Failed to clear previous snapshot backup: {e:?}"))?; + } + + let had_previous = util.has_path(target); + if had_previous { + util.rename_path(target, backup) + .map_err(|e| format!("Failed to preserve current database: {e:?}"))?; + } + + if let Err(error) = util.rename_path(staging, target) { + if had_previous && util.has_path(backup) && !util.has_path(target) { + if let Err(rollback) = util.rename_path(backup, target) { + return Err(format!( + "Failed to activate snapshot: {error:?}; rollback also failed: {rollback:?}" + )); + } + } + return Err(format!("Failed to activate snapshot: {error:?}")); + } + Ok(had_previous) + } + + fn restore_snapshot_backup( + util: &OpfsSAHPoolUtil, + target: &str, + backup: &str, + ) -> Result<(), String> { + if util.has_path(target) { + util.unlink(target) + .map_err(|e| format!("Failed to discard invalid activated snapshot: {e:?}"))?; + } + if util.has_path(backup) { + util.rename_path(backup, target) + .map_err(|e| format!("Failed to restore previous database: {e:?}"))?; + } + Ok(()) + } + + pub async fn initialize_opfs(db_name: &str) -> Result { + let util = install_opfs_sahpool(None, true) + .await + .map_err(|e| JsValue::from_str(&format!("Failed to install OPFS VFS: {e:?}")))?; + // Snapshot staging names are longer than ordinary database names. + // A database which is valid for the VFS must remain openable even if + // it is too long to opt into snapshot installation. + if let Ok((target, staging, backup)) = Self::snapshot_paths(db_name) { + Self::recover_snapshot_activation(&util, &target, &staging, &backup) + .await + .map_err(|error| { + JsValue::from_str(&format!("Failed to recover snapshot activation: {error}")) + })?; + } + Self::open_opfs(db_name) + } + + /// Download a SQLite file and install it directly into OPFS without + /// materializing the full response in memory. Compressed snapshots are + /// decompressed as a stream in this database worker. + /// The caller must close the target database before invoking this method. + pub async fn install_snapshot( + db_name: &str, + url: &str, + compression: SnapshotCompression, + expected_sha256: &str, + expected_uncompressed_size: f64, + ) -> Result<(Self, SnapshotInstallResult), String> { + Self::install_snapshot_cancellable( + db_name, + url, + compression, + expected_sha256, + expected_uncompressed_size, + SnapshotCancellation::default(), + ) + .await + } + + pub(crate) async fn install_snapshot_cancellable( + db_name: &str, + url: &str, + compression: SnapshotCompression, + expected_sha256: &str, + expected_uncompressed_size: f64, + cancellation: SnapshotCancellation, + ) -> Result<(Self, SnapshotInstallResult), String> { + cancellation.ensure_active()?; + if !expected_uncompressed_size.is_finite() + || expected_uncompressed_size < 512.0 + || expected_uncompressed_size.fract() != 0.0 + || expected_uncompressed_size > MAX_SNAPSHOT_SIZE as f64 + { + return Err(format!( + "Expected uncompressed size must be an integer between 512 and {MAX_SNAPSHOT_SIZE}" + )); + } + if url.trim().is_empty() { + return Err("Snapshot URL is required".to_string()); + } + if expected_sha256.len() != 64 + || !expected_sha256.bytes().all(|byte| byte.is_ascii_hexdigit()) + { + return Err( + "Snapshot SHA-256 must contain exactly 64 hexadecimal characters".to_string(), + ); + } + let expected_uncompressed_size = expected_uncompressed_size as u64; + let started_at = Date::now(); + let util = install_opfs_sahpool(None, true) + .await + .map_err(|e| format!("Failed to install OPFS VFS: {e:?}"))?; + let (target, staging, backup) = Self::snapshot_paths(db_name)?; + Self::recover_snapshot_activation(&util, &target, &staging, &backup).await?; + if util.get_capacity() == util.get_file_count() { + util.add_capacity(1) + .await + .map_err(|e| format!("Failed to allocate snapshot staging space: {e:?}"))?; + } + let mut importer = util + .begin_import_db(&staging) + .map_err(|e| format!("Failed to begin snapshot import: {e:?}"))?; + let mut hasher = Sha256::new(); + let mut streamed_bytes = 0u64; + + let global: DedicatedWorkerGlobalScope = js_sys::global() + .dyn_into() + .map_err(|_| "Snapshot installation must run in a dedicated worker".to_string())?; + let abort_controller = AbortController::new() + .map_err(|e| format!("Failed to create snapshot abort controller: {e:?}"))?; + cancellation.register(&abort_controller); + cancellation.ensure_active()?; + let request_init = RequestInit::new(); + request_init.set_signal(Some(&abort_controller.signal())); + let request = Request::new_with_str_and_init(url, &request_init) + .map_err(|e| format!("Failed to create snapshot request: {e:?}"))?; + let response: Response = JsFuture::from(global.fetch_with_request(&request)) + .await + .map_err(|e| format!("Snapshot fetch failed: {e:?}"))? + .dyn_into() + .map_err(|_| "Snapshot fetch returned an invalid response".to_string())?; + if !response.ok() { + if let Some(body) = response.body() { + let _ = JsFuture::from(body.cancel()).await; + } + abort_controller.abort(); + return Err(format!( + "Snapshot fetch returned HTTP {} {}", + response.status(), + response.status_text() + )); + } + let body = match response.body() { + Some(body) => body, + None => { + abort_controller.abort(); + return Err("Snapshot response did not contain a body".to_string()); + } + }; + + let content_result: Result = match compression { + SnapshotCompression::None => Ok(body.clone()), + SnapshotCompression::Gzip => (|| { + // web-sys still marks DecompressionStream as unstable, so + // construct this worker API through reflection. + let ctor: Function = + Reflect::get(&js_sys::global(), &JsValue::from_str("DecompressionStream")) + .map_err(|e| format!("DecompressionStream is unavailable: {e:?}"))? + .dyn_into() + .map_err(|_| "DecompressionStream is unavailable".to_string())?; + let args = Array::new(); + args.push(&JsValue::from_str("gzip")); + let decompressor = Reflect::construct(&ctor, &args) + .map_err(|e| format!("Failed to create gzip decompressor: {e:?}"))?; + let pipe_through: Function = + Reflect::get(body.as_ref(), &JsValue::from_str("pipeThrough")) + .map_err(|e| format!("ReadableStream.pipeThrough is unavailable: {e:?}"))? + .dyn_into() + .map_err(|_| "ReadableStream.pipeThrough is unavailable".to_string())?; + pipe_through + .call1(body.as_ref(), &decompressor) + .map_err(|e| format!("Failed to start gzip decompression: {e:?}"))? + .dyn_into::() + .map_err(|_| "Gzip decompressor returned an invalid stream".to_string()) + })(), + }; + let content = match content_result { + Ok(content) => content, + Err(error) => { + let _ = JsFuture::from(body.cancel()).await; + abort_controller.abort(); + return Err(error); + } + }; + let reader = match ReadableStreamDefaultReader::new(&content) { + Ok(reader) => reader, + Err(error) => { + let _ = JsFuture::from(content.cancel()).await; + abort_controller.abort(); + return Err(format!("Failed to open snapshot stream: {error:?}")); + } + }; + + let stream_result: Result<(), String> = async { + loop { + cancellation.ensure_active()?; + let item = JsFuture::from(reader.read()) + .await + .map_err(|e| format!("Failed while reading snapshot: {e:?}"))?; + cancellation.ensure_active()?; + let done = Reflect::get(&item, &JsValue::from_str("done")) + .map_err(|e| format!("Invalid stream result: {e:?}"))? + .as_bool() + .unwrap_or(false); + if done { + break; + } + let value = Reflect::get(&item, &JsValue::from_str("value")) + .map_err(|e| format!("Invalid stream chunk: {e:?}"))?; + let chunk = Uint8Array::new(&value).to_vec(); + hasher.update(&chunk); + let chunk_len = u64::try_from(chunk.len()) + .map_err(|_| "Snapshot chunk size overflow".to_string())?; + streamed_bytes = streamed_bytes + .checked_add(chunk_len) + .ok_or_else(|| "Snapshot size overflow".to_string())?; + if streamed_bytes > expected_uncompressed_size { + return Err(format!( + "Snapshot exceeded expected size of {:.0} bytes", + expected_uncompressed_size + )); + } + importer + .write_chunk(&chunk) + .map_err(|e| format!("Failed to write snapshot chunk: {e:?}"))?; + } + Ok(()) + } + .await; + if let Err(error) = stream_result { + let _ = JsFuture::from(reader.cancel()).await; + reader.release_lock(); + abort_controller.abort(); + return Err(error); + } + reader.release_lock(); + + let actual_sha256 = format!("{:x}", hasher.finalize()); + cancellation.ensure_active()?; + if streamed_bytes != expected_uncompressed_size { + return Err(format!( + "Snapshot size mismatch: expected {:.0} bytes but received {}", + expected_uncompressed_size, streamed_bytes + )); + } + if !actual_sha256.eq_ignore_ascii_case(expected_sha256) { + return Err(format!( + "Snapshot SHA-256 mismatch: expected {} but received {}", + expected_sha256, actual_sha256 + )); + } + + let bytes_written = importer + .finish() + .map_err(|e| format!("Failed to finish snapshot import: {e:?}"))?; + + let staging_validation = Self::validate_snapshot_path(&staging, "staged snapshot").await; + if let Err(error) = staging_validation { + let _ = util.unlink(&staging); + return Err(error); + } + + // Give the DB worker one task turn to process a cancellation posted + // while synchronous importer finalization or quick_check was running. + crate::coordination::sleep_ms(0).await; + cancellation.ensure_active()?; + + let had_previous = Self::activate_snapshot(&util, &target, &staging, &backup)?; + let database_result = async { + let database = Self::open_opfs(db_name) + .map_err(|e| format!("Failed to open activated snapshot: {e:?}"))?; + drop(database); + Self::validate_snapshot_path(&target, "activated snapshot").await?; + let database = Self::open_opfs(db_name) + .map_err(|e| format!("Failed to reopen activated snapshot: {e:?}"))?; + Ok::(database) + } + .await; + let database = match database_result { + Ok(database) => database, + Err(error) => { + Self::restore_snapshot_backup(&util, &target, &backup)?; + return Err(error); + } + }; + + if had_previous && util.has_path(&backup) { + // Activation is already committed and verified. A stale backup is + // harmless and startup recovery will retry this cleanup. + let _ = util.unlink(&backup); + } + Ok(( + database, + SnapshotInstallResult { + bytes_written, + compression, + elapsed_ms: Date::now() - started_at, + sha256: actual_sha256, + }, + )) + } + /// Execute a prepared statement, collecting any result rows and the affected row count. /// Returns Some(rows) for queries (column count > 0), even if zero rows; None otherwise. fn exec_prepared_statement( @@ -2448,4 +2892,28 @@ mod tests { assert_eq!(array[0]["msg"].as_str().unwrap(), "insert; happened"); assert_eq!(array[1]["msg"].as_str().unwrap(), "second; line"); } + + #[wasm_bindgen_test] + fn snapshot_paths_honor_rooted_opfs_limit() { + let accepted = "a".repeat(478); + let (_, staging, _) = SQLiteDatabase::snapshot_paths(&accepted).expect("boundary path"); + assert_eq!("opfs-sahpool:".len() + staging.len() + 1, 512); + + let rejected = "a".repeat(479); + assert!(SQLiteDatabase::snapshot_paths(&rejected).is_err()); + } + + #[wasm_bindgen_test] + fn snapshot_result_bytes_are_exact_js_integers() { + let result = SnapshotInstallResult { + bytes_written: MAX_SNAPSHOT_SIZE, + compression: SnapshotCompression::None, + elapsed_ms: 1.0, + sha256: "0".repeat(64), + }; + let json = serde_json::to_string(&result).expect("serialize snapshot result"); + let value: serde_json::Value = serde_json::from_str(&json).expect("parse snapshot result"); + assert_eq!(value["bytesWritten"].as_u64(), Some(MAX_SNAPSHOT_SIZE)); + assert!(MAX_SNAPSHOT_SIZE <= JS_MAX_SAFE_INTEGER); + } } diff --git a/packages/sqlite-web-core/src/messages.rs b/packages/sqlite-web-core/src/messages.rs index 6fba81b..bc25974 100644 --- a/packages/sqlite-web-core/src/messages.rs +++ b/packages/sqlite-web-core/src/messages.rs @@ -21,6 +21,13 @@ pub struct SqlBatchStatement { pub params: Option>, } +#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, Hash)] +#[serde(rename_all = "lowercase")] +pub enum SnapshotCompression { + None, + Gzip, +} + // Message types for BroadcastChannel communication #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] #[serde(tag = "type")] @@ -50,6 +57,33 @@ pub enum ChannelMessage { query_id: String, statements: Vec, }, + #[serde(rename = "install-snapshot-request")] + InstallSnapshotRequest { + #[serde(rename = "queryId")] + query_id: String, + #[serde(rename = "requesterId")] + requester_id: String, + url: String, + compression: SnapshotCompression, + #[serde(rename = "sha256")] + sha256: String, + #[serde(rename = "uncompressedSize")] + uncompressed_size: f64, + }, + #[serde(rename = "cancel-snapshot-request")] + CancelSnapshotRequest { + #[serde(rename = "requesterId")] + requester_id: String, + #[serde(rename = "queryId")] + query_id: String, + }, + #[serde(rename = "snapshot-cancellation-response")] + SnapshotCancellationResponse { + #[serde(rename = "queryId")] + query_id: String, + result: Option, + error: Option, + }, #[serde(rename = "query-response")] QueryResponse { #[serde(rename = "queryId")] @@ -83,6 +117,32 @@ pub enum WorkerMessage { request_id: u32, statements: Vec, }, + #[serde(rename = "install-snapshot")] + InstallSnapshot { + #[serde(rename = "requestId")] + request_id: u32, + url: String, + compression: SnapshotCompression, + #[serde(rename = "sha256")] + sha256: String, + #[serde(rename = "uncompressedSize")] + uncompressed_size: f64, + }, + #[serde(rename = "cancel-snapshot")] + CancelSnapshot { + #[serde(rename = "requestId")] + request_id: u32, + }, + #[serde(rename = "cancel-forwarded-snapshot")] + CancelForwardedSnapshot { + #[serde(rename = "requestId")] + request_id: u32, + }, + #[serde(rename = "shutdown")] + Shutdown { + #[serde(rename = "requestId")] + request_id: u32, + }, } // Messages to main thread @@ -96,6 +156,12 @@ pub enum MainThreadMessage { result: Option, error: Option, }, + #[serde(rename = "snapshot-cancelled")] + SnapshotCancelled { + #[serde(rename = "requestId")] + request_id: u32, + cancelled: bool, + }, #[serde(rename = "worker-ready")] WorkerReady, } @@ -159,6 +225,29 @@ mod tests { assert!(json.contains("\"params\"")); }); + let snapshot_request = ChannelMessage::InstallSnapshotRequest { + query_id: "snapshot-456".to_string(), + requester_id: "worker-123".to_string(), + url: "https://example.com/database.sqlite.gz".to_string(), + compression: SnapshotCompression::Gzip, + sha256: "a".repeat(64), + uncompressed_size: 4096.0, + }; + assert_serialization_roundtrip(snapshot_request, "install-snapshot-request", |json| { + assert!(json.contains("\"requesterId\":\"worker-123\"")); + assert!(json.contains("\"compression\":\"gzip\"")); + assert!(json.contains("\"uncompressedSize\":4096.0")); + }); + + let snapshot_cancel = ChannelMessage::CancelSnapshotRequest { + requester_id: "worker-123".to_string(), + query_id: "snapshot-456".to_string(), + }; + assert_serialization_roundtrip(snapshot_cancel, "cancel-snapshot-request", |json| { + assert!(json.contains("\"requesterId\":\"worker-123\"")); + assert!(json.contains("\"queryId\":\"snapshot-456\"")); + }); + let query_success = ChannelMessage::QueryResponse { query_id: "query-789".to_string(), result: Some("[{\"id\": 1, \"name\": \"test\"}]".to_string()), @@ -242,6 +331,24 @@ mod tests { } } + #[wasm_bindgen_test] + fn test_worker_message_install_snapshot_serialization() { + let msg = WorkerMessage::InstallSnapshot { + request_id: 44, + url: "https://example.com/database.sqlite".to_string(), + compression: SnapshotCompression::None, + sha256: "b".repeat(64), + uncompressed_size: 8192.0, + }; + + let json = serde_json::to_string(&msg).expect("Should serialize"); + assert!(json.contains("\"type\":\"install-snapshot\"")); + assert!(json.contains("\"compression\":\"none\"")); + + let deserialized: WorkerMessage = serde_json::from_str(&json).expect("Should deserialize"); + assert_eq!(msg, deserialized); + } + #[wasm_bindgen_test] fn test_main_thread_messages_serialization() { let success_result = MainThreadMessage::QueryResult { diff --git a/packages/sqlite-web/src/db.rs b/packages/sqlite-web/src/db.rs index f3aa0ee..13149fd 100644 --- a/packages/sqlite-web/src/db.rs +++ b/packages/sqlite-web/src/db.rs @@ -1,5 +1,5 @@ -use std::cell::RefCell; -use std::collections::HashMap; +use std::cell::{Cell, RefCell}; +use std::collections::{HashMap, HashSet}; use std::rc::Rc; use js_sys::{Array, Reflect}; @@ -7,8 +7,10 @@ use serde::Serialize; use wasm_bindgen::prelude::*; use wasm_bindgen::JsCast; use wasm_bindgen_futures::JsFuture; + +const MAX_SNAPSHOT_SIZE: f64 = 9_007_199_254_736_895.0; use wasm_bindgen_utils::prelude::*; -use web_sys::Worker; +use web_sys::{BroadcastChannel, MessageEvent, Worker}; use crate::errors::SQLiteWasmDatabaseError; use crate::messages::WORKER_ERROR_TYPE_INITIALIZATION_PENDING; @@ -19,6 +21,97 @@ use crate::utils::describe_js_value; use crate::worker::{create_worker_from_code, install_onmessage_handler}; use crate::worker_template::generate_self_contained_worker; +fn sanitize_identifier(name: &str) -> String { + let sanitized: String = name + .trim() + .chars() + .map(|character| match character { + 'a'..='z' | 'A'..='Z' | '0'..='9' | '.' | '_' | '-' => character, + _ => '_', + }) + .collect(); + if sanitized.is_empty() { + "db".to_string() + } else { + sanitized + } +} + +fn install_snapshot_cancellation_listener( + channel: &BroadcastChannel, + pending_queries: Rc>>, + pending_snapshots: Rc>>, + client_id: Rc>, +) { + let onmessage = Closure::wrap(Box::new(move |event: MessageEvent| { + let data = event.data(); + let message_type = Reflect::get(&data, &JsValue::from_str("type")) + .ok() + .and_then(|value| value.as_string()); + if message_type.as_deref() != Some("snapshot-cancellation-response") { + return; + } + let Some(query_id) = Reflect::get(&data, &JsValue::from_str("queryId")) + .ok() + .and_then(|value| value.as_string()) + else { + return; + }; + let prefix = format!("{}:", client_id.borrow()); + let Some(request_id) = query_id + .strip_prefix(&prefix) + .and_then(|value| value.parse::().ok()) + else { + return; + }; + if !pending_snapshots.borrow_mut().remove(&request_id) { + return; + } + let Some((resolve, reject)) = pending_queries.borrow_mut().remove(&request_id) else { + return; + }; + let error = Reflect::get(&data, &JsValue::from_str("error")) + .ok() + .filter(|value| !value.is_null() && !value.is_undefined()); + if let Some(error) = error { + let _ = reject.call1(&JsValue::NULL, &error); + return; + } + let result = Reflect::get(&data, &JsValue::from_str("result")) + .ok() + .filter(|value| !value.is_null() && !value.is_undefined()) + .unwrap_or(JsValue::UNDEFINED); + let _ = resolve.call1(&JsValue::NULL, &result); + }) as Box); + channel.set_onmessage(Some(onmessage.as_ref().unchecked_ref())); + onmessage.forget(); +} + +async fn yield_main_thread() { + let promise = js_sys::Promise::new(&mut |resolve, _| { + let callback = Closure::once(move || { + let _ = resolve.call0(&JsValue::NULL); + }); + if let Some(window) = web_sys::window() { + if window + .set_timeout_with_callback_and_timeout_and_arguments_0( + callback.as_ref().unchecked_ref(), + 0, + ) + .is_ok() + { + callback.forget(); + return; + } + } + let _ = callback + .as_ref() + .unchecked_ref::() + .call0(&JsValue::NULL); + }); + let _ = JsFuture::from(promise).await; +} + #[wasm_bindgen] pub struct SQLiteWasmDatabase { worker: Rc>, @@ -26,6 +119,13 @@ pub struct SQLiteWasmDatabase { pending_queries: Rc>>, next_request_id: Rc>, ready_signal: ReadySignal, + worker_stopped: Rc>, + shutdown_complete: Rc>, + permanently_closed: Rc>, + lifecycle_generation: Rc>, + client_id: Rc>, + coordination_channel: BroadcastChannel, + pending_snapshot_requests: Rc>>, } impl Serialize for SQLiteWasmDatabase { @@ -56,24 +156,209 @@ impl SQLiteWasmDatabase { } fn construct(db_name: &str) -> Result { - let worker_code = generate_self_contained_worker(db_name); + let client_id = format!( + "{:.0}-{:016x}", + js_sys::Date::now(), + (js_sys::Math::random() * u64::MAX as f64) as u64 + ); + let worker_code = generate_self_contained_worker(db_name, &client_id); let worker = create_worker_from_code(&worker_code)?; + let channel_name = format!("sqlite-queries-{}", sanitize_identifier(db_name)); + let coordination_channel = BroadcastChannel::new(&channel_name)?; let pending_queries: Rc>> = Rc::new(RefCell::new(HashMap::new())); + let pending_snapshot_requests = Rc::new(RefCell::new(HashSet::new())); + let client_id_state = Rc::new(RefCell::new(client_id.clone())); let ready_signal = ReadySignal::new(); install_onmessage_handler(&worker, Rc::clone(&pending_queries), ready_signal.clone()); + let worker = Rc::new(RefCell::new(worker)); + let worker_stopped = Rc::new(Cell::new(false)); let next_request_id = Rc::new(RefCell::new(1u32)); + install_snapshot_cancellation_listener( + &coordination_channel, + Rc::clone(&pending_queries), + Rc::clone(&pending_snapshot_requests), + Rc::clone(&client_id_state), + ); Ok(SQLiteWasmDatabase { - worker: Rc::new(RefCell::new(worker)), + worker, db_name: db_name.to_string(), pending_queries, next_request_id, ready_signal, + worker_stopped, + shutdown_complete: Rc::new(Cell::new(false)), + permanently_closed: Rc::new(Cell::new(false)), + lifecycle_generation: Rc::new(Cell::new(0)), + client_id: client_id_state, + coordination_channel, + pending_snapshot_requests, }) } + fn ensure_open(&self) -> Result<(), SQLiteWasmDatabaseError> { + if self.permanently_closed.get() || self.worker_stopped.get() { + Err(SQLiteWasmDatabaseError::JsError(JsValue::from_str( + "Database is closed", + ))) + } else { + Ok(()) + } + } + + fn advance_lifecycle_generation(&self) -> u64 { + let generation = self.lifecycle_generation.get().wrapping_add(1); + self.lifecycle_generation.set(generation); + generation + } + + fn begin_worker_shutdown(&self, reason: &str) { + if !self.worker_stopped.replace(true) { + self.cancel_forwarded_snapshots(); + } + let error = JsValue::from_str(reason); + let snapshot_requests = self.pending_snapshot_requests.borrow(); + self.pending_queries + .borrow_mut() + .retain(|request_id, (_, reject)| { + if snapshot_requests.contains(request_id) { + true + } else { + let _ = reject.call1(&JsValue::NULL, &error); + false + } + }); + self.ready_signal.mark_failed(reason.to_string()); + } + + async fn request_coordinator_shutdown(&self) -> Result<(), SQLiteWasmDatabaseError> { + let request_id = { + let mut next = self.next_request_id.borrow_mut(); + let id = *next; + *next = next.wrapping_add(1).max(1); + id + }; + let message = js_sys::Object::new(); + Reflect::set( + &message, + &JsValue::from_str("type"), + &JsValue::from_str("shutdown"), + ) + .map_err(SQLiteWasmDatabaseError::JsError)?; + Reflect::set( + &message, + &JsValue::from_str("requestId"), + &JsValue::from_f64(request_id as f64), + ) + .map_err(SQLiteWasmDatabaseError::JsError)?; + + let worker = Rc::clone(&self.worker); + let pending_queries = Rc::clone(&self.pending_queries); + let promise = js_sys::Promise::new(&mut |resolve, reject| { + pending_queries + .borrow_mut() + .insert(request_id, (resolve, reject)); + if let Err(error) = worker.borrow().post_message(&message) { + if let Some((_, reject)) = pending_queries.borrow_mut().remove(&request_id) { + let _ = reject.call1(&JsValue::NULL, &error); + } + } + }); + JsFuture::from(promise) + .await + .map(|_| ()) + .map_err(SQLiteWasmDatabaseError::JsError) + } + + fn terminate_worker(&self, reason: &str) { + self.worker.borrow().terminate(); + let error = JsValue::from_str(reason); + for (_, (_, reject)) in self.pending_queries.borrow_mut().drain() { + let _ = reject.call1(&JsValue::NULL, &error); + } + self.pending_snapshot_requests.borrow_mut().clear(); + self.shutdown_complete.set(true); + } + + async fn shutdown_worker(&self, reason: &str) -> Result<(), SQLiteWasmDatabaseError> { + if self.worker_stopped.get() { + while !self.shutdown_complete.get() { + yield_main_thread().await; + } + return Ok(()); + } + self.shutdown_complete.set(false); + self.begin_worker_shutdown(reason); + let shutdown_result = self.request_coordinator_shutdown().await; + if shutdown_result.is_err() { + self.terminate_worker(reason); + yield_main_thread().await; + return shutdown_result; + } + while !self.pending_snapshot_requests.borrow().is_empty() { + yield_main_thread().await; + } + self.terminate_worker(reason); + yield_main_thread().await; + shutdown_result + } + + fn cancel_forwarded_snapshots(&self) { + let client_id = self.client_id.borrow().clone(); + for request_id in self.pending_snapshot_requests.borrow().iter() { + let message = js_sys::Object::new(); + if Reflect::set( + &message, + &JsValue::from_str("type"), + &JsValue::from_str("cancel-snapshot-request"), + ) + .is_err() + || Reflect::set( + &message, + &JsValue::from_str("requesterId"), + &JsValue::from_str(&client_id), + ) + .is_err() + || Reflect::set( + &message, + &JsValue::from_str("queryId"), + &JsValue::from_str(&format!("{client_id}:{request_id}")), + ) + .is_err() + { + continue; + } + if self.coordination_channel.post_message(&message).is_err() { + let fallback = js_sys::Object::new(); + let configured = Reflect::set( + &fallback, + &JsValue::from_str("type"), + &JsValue::from_str("cancel-forwarded-snapshot"), + ) + .and_then(|_| { + Reflect::set( + &fallback, + &JsValue::from_str("requestId"), + &JsValue::from_f64(*request_id as f64), + ) + }); + if configured.is_ok() { + let _ = self.worker.borrow().post_message(&fallback); + } + } + } + } + + fn close_permanently(&self, reason: &str) { + if !self.permanently_closed.replace(true) { + self.advance_lifecycle_generation(); + } + self.begin_worker_shutdown(reason); + self.terminate_worker(reason); + } + fn normalize_params(params: Option) -> Result { let params_js = params.map(JsValue::from).unwrap_or(JsValue::UNDEFINED); normalize_params_js(¶ms_js) @@ -136,6 +421,20 @@ impl SQLiteWasmDatabase { }) } + fn validate_snapshot_size(uncompressed_size: f64) -> Result<(), SQLiteWasmDatabaseError> { + if !uncompressed_size.is_finite() + || uncompressed_size < 512.0 + || uncompressed_size.fract() != 0.0 + || uncompressed_size > MAX_SNAPSHOT_SIZE + { + Err(SQLiteWasmDatabaseError::JsError(JsValue::from_str( + "Snapshot uncompressed size must be a JavaScript-safe integer supported by OPFS", + ))) + } else { + Ok(()) + } + } + async fn wait_until_ready(&self) -> Result<(), SQLiteWasmDatabaseError> { match self.ready_signal.current_state() { InitializationState::Ready => return Ok(()), @@ -173,6 +472,7 @@ impl SQLiteWasmDatabase { sql: &str, params: Option, ) -> Result { + self.ensure_open()?; let worker = Rc::clone(&self.worker); let pending_queries = Rc::clone(&self.pending_queries); let sql = sql.to_string(); @@ -254,6 +554,7 @@ impl SQLiteWasmDatabase { &self, statements: Array, ) -> Result { + self.ensure_open()?; let worker = Rc::clone(&self.worker); let pending_queries = Rc::clone(&self.pending_queries); let statements = Self::normalize_batch_statements(statements)?; @@ -316,20 +617,138 @@ impl SQLiteWasmDatabase { Ok(result.as_string().unwrap_or_else(|| format!("{result:?}"))) } + /// Stream a SQLite snapshot into OPFS in the database worker. Only the + /// source URL and generic transport/validation metadata cross worker + /// boundaries. Supported compression values are `none` and `gzip`. + #[wasm_export(js_name = "installSnapshot", unchecked_return_type = "string")] + pub async fn install_snapshot( + &self, + url: &str, + #[wasm_export(unchecked_param_type = "\"none\" | \"gzip\"")] compression: &str, + sha256: &str, + uncompressed_size: f64, + ) -> Result { + self.ensure_open()?; + let url = url.trim(); + let compression = match compression.trim().to_ascii_lowercase().as_str() { + "none" => "none", + "gzip" => "gzip", + _ => { + return Err(SQLiteWasmDatabaseError::JsError(JsValue::from_str( + "Snapshot compression must be either 'none' or 'gzip'", + ))) + } + }; + let sha256 = sha256.trim(); + if url.is_empty() { + return Err(SQLiteWasmDatabaseError::JsError(JsValue::from_str( + "Snapshot URL is required", + ))); + } + if sha256.len() != 64 || !sha256.bytes().all(|byte| byte.is_ascii_hexdigit()) { + return Err(SQLiteWasmDatabaseError::JsError(JsValue::from_str( + "Snapshot SHA-256 must contain exactly 64 hexadecimal characters", + ))); + } + Self::validate_snapshot_size(uncompressed_size)?; + + let worker = Rc::clone(&self.worker); + let pending_queries = Rc::clone(&self.pending_queries); + let request_id = { + let mut n = self.next_request_id.borrow_mut(); + let id = *n; + *n = n.wrapping_add(1).max(1); + id + }; + let message = js_sys::Object::new(); + Reflect::set( + &message, + &JsValue::from_str("type"), + &JsValue::from_str("install-snapshot"), + ) + .map_err(SQLiteWasmDatabaseError::JsError)?; + Reflect::set( + &message, + &JsValue::from_str("requestId"), + &JsValue::from_f64(request_id as f64), + ) + .map_err(SQLiteWasmDatabaseError::JsError)?; + Reflect::set(&message, &JsValue::from_str("url"), &JsValue::from_str(url)) + .map_err(SQLiteWasmDatabaseError::JsError)?; + Reflect::set( + &message, + &JsValue::from_str("compression"), + &JsValue::from_str(compression), + ) + .map_err(SQLiteWasmDatabaseError::JsError)?; + Reflect::set( + &message, + &JsValue::from_str("sha256"), + &JsValue::from_str(sha256), + ) + .map_err(SQLiteWasmDatabaseError::JsError)?; + Reflect::set( + &message, + &JsValue::from_str("uncompressedSize"), + &JsValue::from_f64(uncompressed_size), + ) + .map_err(SQLiteWasmDatabaseError::JsError)?; + + let promise = js_sys::Promise::new(&mut |resolve, reject| match worker + .borrow() + .post_message(&message) + { + Ok(()) => { + pending_queries + .borrow_mut() + .insert(request_id, (resolve, reject)); + self.pending_snapshot_requests + .borrow_mut() + .insert(request_id); + } + Err(err) => { + let _ = reject.call1(&JsValue::NULL, &err); + } + }); + let result = JsFuture::from(promise).await; + self.pending_snapshot_requests + .borrow_mut() + .remove(&request_id); + let result = result.map_err(SQLiteWasmDatabaseError::JsError)?; + Ok(result.as_string().unwrap_or_else(|| format!("{result:?}"))) + } + + /// Cancel outstanding work and terminate this database's worker. + /// This operation is idempotent. Callers must await it before calling + /// wasm-bindgen's generated `free()` or releasing any external ownership + /// lock so forwarded snapshot cancellation is acknowledged first. + #[wasm_export(js_name = "close", unchecked_return_type = "void")] + pub async fn close(&self) -> Result<(), SQLiteWasmDatabaseError> { + if self.permanently_closed.replace(true) { + return Ok(()); + } + self.advance_lifecycle_generation(); + self.shutdown_worker("Database closed").await + } + #[wasm_export(js_name = "wipeAndRecreate", unchecked_return_type = "void")] pub async fn wipe_and_recreate(&self) -> Result<(), SQLiteWasmDatabaseError> { - self.worker.borrow().terminate(); + self.ensure_open()?; + let wipe_generation = self.advance_lifecycle_generation(); + self.shutdown_worker("Database wipe in progress").await?; - for (_, (_, reject)) in self.pending_queries.borrow_mut().drain() { - let err = JsValue::from_str("Database wipe in progress"); - let _ = reject.call1(&JsValue::NULL, &err); + let deletion_result = delete_opfs_sahpool_directory_with_retries().await; + + if self.permanently_closed.get() || self.lifecycle_generation.get() != wipe_generation { + return Err(SQLiteWasmDatabaseError::JsError(JsValue::from_str( + "Database is closed", + ))); } self.ready_signal.reset(); - let deletion_result = delete_opfs_sahpool_directory_with_retries().await; - - let worker_code = generate_self_contained_worker(&self.db_name); + let client_id = self.client_id.borrow().clone(); + let worker_code = generate_self_contained_worker(&self.db_name, &client_id); let new_worker = create_worker_from_code(&worker_code).map_err(SQLiteWasmDatabaseError::JsError)?; @@ -340,13 +759,29 @@ impl SQLiteWasmDatabase { ); *self.worker.borrow_mut() = new_worker; + self.worker_stopped.set(false); + self.shutdown_complete.set(false); - self.wait_until_ready().await?; + let ready_result = self.wait_until_ready().await; + if self.permanently_closed.get() || self.lifecycle_generation.get() != wipe_generation { + self.begin_worker_shutdown("Database closed"); + self.terminate_worker("Database closed"); + return Err(SQLiteWasmDatabaseError::JsError(JsValue::from_str( + "Database is closed", + ))); + } + ready_result?; deletion_result } } +impl Drop for SQLiteWasmDatabase { + fn drop(&mut self) { + self.close_permanently("Database released"); + } +} + fn is_initialization_pending_error(err: &JsValue) -> bool { let error_type = Reflect::get(err, &JsValue::from_str("type")) .ok() @@ -494,6 +929,14 @@ mod tests { ); } + #[wasm_bindgen_test] + fn snapshot_size_enforces_exact_opfs_and_js_boundary() { + assert!(SQLiteWasmDatabase::validate_snapshot_size(MAX_SNAPSHOT_SIZE).is_ok()); + assert!(SQLiteWasmDatabase::validate_snapshot_size(MAX_SNAPSHOT_SIZE + 1.0).is_err()); + assert!(SQLiteWasmDatabase::validate_snapshot_size(511.0).is_err()); + assert!(SQLiteWasmDatabase::validate_snapshot_size(512.5).is_err()); + } + #[wasm_bindgen_test] fn normalize_batch_statements_propagates_params_getter_errors() { let statements = Array::new(); @@ -687,4 +1130,18 @@ mod tests { let res = db.query("SELECT ?", Some(arr)).await; assert!(res.is_err(), "-Infinity should be rejected"); } + + #[wasm_bindgen_test(async)] + async fn close_is_idempotent_and_prevents_new_work() { + let db = SQLiteWasmDatabase::new("test_close").await.unwrap(); + db.close().await.unwrap(); + db.close().await.unwrap(); + let error = db.query("SELECT 1", None).await.expect_err("closed DB"); + match error { + SQLiteWasmDatabaseError::JsError(value) => { + assert_eq!(value.as_string().as_deref(), Some("Database is closed")); + } + other => panic!("expected closed JS error, got {other:?}"), + } + } } diff --git a/packages/sqlite-web/src/tests.rs b/packages/sqlite-web/src/tests.rs index afefb90..27f4736 100644 --- a/packages/sqlite-web/src/tests.rs +++ b/packages/sqlite-web/src/tests.rs @@ -79,7 +79,7 @@ fn test_error_propagation_chain() { #[wasm_bindgen_test] fn test_worker_template_generation() { - let worker_code = generate_self_contained_worker("testdb"); + let worker_code = generate_self_contained_worker("testdb", "test-client"); assert!(!worker_code.is_empty()); assert!( diff --git a/packages/sqlite-web/src/worker_template.rs b/packages/sqlite-web/src/worker_template.rs index dc04b4a..abef6ec 100644 --- a/packages/sqlite-web/src/worker_template.rs +++ b/packages/sqlite-web/src/worker_template.rs @@ -1,15 +1,17 @@ /// Generate self-contained worker with embedded WASM and JS glue code /// and inject the database name into the worker global scope so core /// can read it during initialization. -pub fn generate_self_contained_worker(db_name: &str) -> String { +pub fn generate_self_contained_worker(db_name: &str, client_id: &str) -> String { // Safely JSON-encode the db name for JS embedding let encoded = serde_json::to_string(db_name).unwrap_or_else(|_| "\"unknown\"".to_string()); + let client_id = + serde_json::to_string(client_id).unwrap_or_else(|_| "\"unknown-client\"".to_string()); let embedded_body = serde_json::to_string(include_str!("embedded_worker.js")) .unwrap_or_else(|_| "\"\"".to_string()); // __SQLITE_EMBEDDED_WORKER stores the JSON-encoded embedded worker body (embedded_body) so the coordinator can spawn a separate DB worker (see coordination.rs:301-313); set when embedded-worker mode is used and consumers must JSON-decode before instantiating the worker. let prefix = format!( - "self.__SQLITE_DB_NAME = {};\nself.__SQLITE_FOLLOWER_TIMEOUT_MS = 5000.0;\nself.__SQLITE_QUERY_TIMEOUT_MS = 30000.0;\nself.__SQLITE_EMBEDDED_WORKER = {};\n", - encoded, embedded_body + "self.__SQLITE_DB_NAME = {};\nself.__SQLITE_CLIENT_ID = {};\nself.__SQLITE_FOLLOWER_TIMEOUT_MS = 5000.0;\nself.__SQLITE_QUERY_TIMEOUT_MS = 30000.0;\nself.__SQLITE_EMBEDDED_WORKER = {};\n", + encoded, client_id, embedded_body ); // Use the bundled worker template with embedded WASM let body = include_str!("embedded_worker.js"); @@ -25,7 +27,11 @@ mod tests { #[wasm_bindgen_test] fn embeds_db_name_and_timeout_configuration() { - let output = generate_self_contained_worker("my_db"); + let output = generate_self_contained_worker("my_db", "client-1"); + assert!( + output.contains("self.__SQLITE_CLIENT_ID = \"client-1\";"), + "client id should be JSON encoded in prefix" + ); assert!( output.contains("self.__SQLITE_DB_NAME = \"my_db\";"), "db name should be JSON encoded in prefix" @@ -46,7 +52,7 @@ mod tests { #[wasm_bindgen_test] fn appends_embedded_worker_body() { - let output = generate_self_contained_worker("whatever"); + let output = generate_self_contained_worker("whatever", "client-2"); let body = include_str!("embedded_worker.js"); assert!( output.ends_with(body), diff --git a/svelte-test/tests/integration/shutdown-handoff.test.ts b/svelte-test/tests/integration/shutdown-handoff.test.ts new file mode 100644 index 0000000..18126fe --- /dev/null +++ b/svelte-test/tests/integration/shutdown-handoff.test.ts @@ -0,0 +1,414 @@ +import { describe, expect, it } from 'vitest'; +import type { SQLiteWasmDatabase } from '@rainlanguage/sqlite-web'; +import { createTestDatabase } from '../fixtures/test-helpers.js'; + +const SNAPSHOT_SHA256 = '5d5408c09e537d8af271fff0eab22ab9ef624259476010ada25de6a3732dec08'; +const DRAIN_REGRESSION_DATABASE = 'shutdown-drain-regressions'; + +async function closeAfterLeadershipHandoff(db: SQLiteWasmDatabase): Promise { + for (let attempt = 0; attempt < 100; attempt += 1) { + const ready = await db.query('SELECT 1'); + if (!ready.error) { + const closed = await db.close(); + expect(closed.error).toBeUndefined(); + return; + } + await new Promise((resolve) => setTimeout(resolve, 25)); + } + throw new Error('surviving client did not complete leadership handoff'); +} + +describe('database worker shutdown handoff', () => { + it('drains follower-only work before an otherwise-idle leader closes', async () => { + const databaseName = DRAIN_REGRESSION_DATABASE; + const leader = await createTestDatabase(databaseName); + const follower = await createTestDatabase(databaseName); + try { + await fetch('/snapshot-stats?reset=1'); + const followerInstall = follower.installSnapshot( + new URL('/snapshot.slow.db', window.location.href).href, + 'none', + SNAPSHOT_SHA256, + 1024 + ); + for (let attempt = 0; attempt < 40; attempt += 1) { + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + if (stats.snapshotRequestCount === 1) break; + await new Promise((resolve) => setTimeout(resolve, 25)); + } + + const close = leader.close(); + const followerResult = await Promise.race([ + followerInstall, + new Promise((_, reject) => + setTimeout(() => reject(new Error('follower-only install did not settle')), 3000) + ) + ]); + const closeResult = await Promise.race([ + close, + new Promise((_, reject) => + setTimeout(() => reject(new Error('follower-only leader close did not settle')), 3000) + ) + ]); + expect(followerResult.error).toBeUndefined(); + expect(closeResult.error).toBeUndefined(); + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + expect(stats.snapshotRequestCount).toBe(1); + } finally { + leader.free(); + await closeAfterLeadershipHandoff(follower); + follower.free(); + await new Promise((resolve) => setTimeout(resolve, 100)); + } + }); + + it('drains unrelated follower work during an immediate local install and close', async () => { + const databaseName = DRAIN_REGRESSION_DATABASE; + const leader = await createTestDatabase(databaseName); + const follower = await createTestDatabase(databaseName); + try { + await fetch('/snapshot-stats?reset=1'); + const followerInstall = follower.installSnapshot( + new URL('/snapshot.slow.db', window.location.href).href, + 'none', + SNAPSHOT_SHA256, + 1024 + ); + for (let attempt = 0; attempt < 40; attempt += 1) { + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + if (stats.snapshotRequestCount === 1) break; + await new Promise((resolve) => setTimeout(resolve, 25)); + } + + const localInstall = leader.installSnapshot( + new URL('/snapshot.delayed.db?local=1', window.location.href).href, + 'none', + SNAPSHOT_SHA256, + 1024 + ); + const [localResult, followerResult, closeResult] = await Promise.race([ + Promise.all([localInstall, followerInstall, leader.close()]), + new Promise((_, reject) => + setTimeout(() => reject(new Error('cancel-before-request drain did not settle')), 3000) + ) + ]); + expect(localResult.error?.msg).toContain('Snapshot installation cancelled'); + expect(followerResult.error).toBeUndefined(); + expect(closeResult.error).toBeUndefined(); + } finally { + leader.free(); + await closeAfterLeadershipHandoff(follower); + follower.free(); + await new Promise((resolve) => setTimeout(resolve, 100)); + } + }); + + it('settles a follower arriving while the sole leader install is cancelling', async () => { + const databaseName = DRAIN_REGRESSION_DATABASE; + const leader = await createTestDatabase(databaseName); + const follower = await createTestDatabase(databaseName); + try { + await fetch('/snapshot-stats?reset=1'); + const url = new URL('/snapshot.slow.db', window.location.href).href; + const leaderInstall = leader.installSnapshot( + url, + 'none', + SNAPSHOT_SHA256, + 1024 + ); + for (let attempt = 0; attempt < 40; attempt += 1) { + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + if (stats.snapshotRequestCount === 1) break; + await new Promise((resolve) => setTimeout(resolve, 25)); + } + + const close = leader.close(); + await Promise.resolve(); + const followerInstall = follower.installSnapshot( + url, + 'none', + SNAPSHOT_SHA256, + 1024 + ); + const [leaderResult, followerResult, closeResult] = await Promise.race([ + Promise.all([leaderInstall, followerInstall, close]), + new Promise((_, reject) => + setTimeout(() => reject(new Error('successor ordering did not settle')), 3000) + ) + ]); + expect(leaderResult.error?.msg).toContain('Snapshot installation cancelled'); + expect(closeResult.error).toBeUndefined(); + expect(followerResult.error?.msg).toContain( + 'Leader is shutting down; retry on the next leader' + ); + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + expect(stats.snapshotRequestCount).toBe(1); + } finally { + leader.free(); + await closeAfterLeadershipHandoff(follower); + follower.free(); + await new Promise((resolve) => setTimeout(resolve, 100)); + } + }); + + it('settles install and close when cancellation arrives before forwarding', async () => { + const databaseName = `shutdown-immediate-close-${Date.now()}`; + const leader = await createTestDatabase(databaseName); + const follower = await createTestDatabase(databaseName); + try { + const write = await leader.query(` + CREATE TABLE immediate_close_sentinel (value TEXT NOT NULL); + INSERT INTO immediate_close_sentinel VALUES ('preserved'); + `); + expect(write.error).toBeUndefined(); + const install = follower.installSnapshot( + new URL('/snapshot.slow.db', window.location.href).href, + 'none', + SNAPSHOT_SHA256, + 1024 + ); + const close = follower.close(); + const [installResult, closeResult] = await Promise.race([ + Promise.all([install, close]), + new Promise((_, reject) => + setTimeout(() => reject(new Error('immediate close did not settle')), 2000) + ) + ]); + expect(installResult.error?.msg).toContain('Snapshot installation cancelled'); + expect(closeResult.error).toBeUndefined(); + const read = await leader.query('SELECT value FROM immediate_close_sentinel'); + expect(JSON.parse(read.value || '[]')).toEqual([{ value: 'preserved' }]); + } finally { + follower.free(); + await leader.close(); + leader.free(); + } + }); + + it('keeps a closing leader alive for a coalesced follower install', async () => { + const databaseName = `shutdown-leader-coalesced-${Date.now()}`; + const leader = await createTestDatabase(databaseName); + const follower = await createTestDatabase(databaseName); + try { + await fetch('/snapshot-stats?reset=1'); + const url = new URL('/snapshot.slow.db', window.location.href).href; + const leaderInstall = leader.installSnapshot( + url, + 'none', + SNAPSHOT_SHA256, + 1024 + ); + const followerInstall = follower.installSnapshot( + url, + 'none', + SNAPSHOT_SHA256, + 1024 + ); + for (let attempt = 0; attempt < 40; attempt += 1) { + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + if (stats.snapshotRequestCount === 1) break; + await new Promise((resolve) => setTimeout(resolve, 25)); + } + const close = leader.close(); + const [leaderResult, followerResult, closeResult] = await Promise.all([ + leaderInstall, + followerInstall, + close + ]); + expect(leaderResult.error?.msg).toContain('Snapshot installation cancelled'); + expect(followerResult.error).toBeUndefined(); + expect(closeResult.error).toBeUndefined(); + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + expect(stats.snapshotRequestCount).toBe(1); + let read = await follower.query('SELECT label FROM snapshot_items ORDER BY id'); + for (let attempt = 0; read.error && attempt < 80; attempt += 1) { + await new Promise((resolve) => setTimeout(resolve, 25)); + read = await follower.query('SELECT label FROM snapshot_items ORDER BY id'); + } + expect(read.error).toBeUndefined(); + } finally { + leader.free(); + await closeAfterLeadershipHandoff(follower); + follower.free(); + await new Promise((resolve) => setTimeout(resolve, 100)); + } + }); + + it('fences follower cancellation during snapshot finalization', async () => { + const databaseName = `shutdown-follower-finalize-${Date.now()}`; + const leader = await createTestDatabase(databaseName); + let follower: SQLiteWasmDatabase | undefined; + try { + const write = await leader.query(` + CREATE TABLE finalization_sentinel (value TEXT NOT NULL); + INSERT INTO finalization_sentinel VALUES ('must survive finalization'); + `); + expect(write.error).toBeUndefined(); + follower = await createTestDatabase(databaseName); + await fetch('/snapshot-stats?reset=1'); + const metadata = await fetch('/snapshot-finalization-meta').then((response) => + response.json() + ); + + const install = follower.installSnapshot( + new URL('/snapshot.finalization.db', window.location.href).href, + 'none', + metadata.sha256, + metadata.size + ); + for (let attempt = 0; attempt < 400; attempt += 1) { + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + if (stats.completedSnapshotRequestCount === 1) break; + await new Promise((resolve) => setTimeout(resolve, 2)); + } + const completed = await fetch('/snapshot-stats').then((response) => response.json()); + expect(completed.completedSnapshotRequestCount).toBe(1); + + const closed = await follower.close(); + expect(closed.error).toBeUndefined(); + const installResult = await install; + expect(installResult.error?.msg).toContain('Snapshot installation cancelled'); + follower.free(); + follower = undefined; + + const read = await leader.query('SELECT value FROM finalization_sentinel'); + expect(read.error).toBeUndefined(); + expect(JSON.parse(read.value || '[]')).toEqual([ + { value: 'must survive finalization' } + ]); + } finally { + if (follower) { + await follower.close(); + follower.free(); + } + await leader.close(); + leader.free(); + } + }, 30000); + + it('detaches a closing follower and cancels an orphaned leader install', async () => { + const databaseName = `shutdown-follower-snapshot-${Date.now()}`; + const leader = await createTestDatabase(databaseName); + let follower: SQLiteWasmDatabase | undefined; + try { + const write = await leader.query(` + CREATE TABLE follower_close_sentinel (value TEXT NOT NULL); + INSERT INTO follower_close_sentinel VALUES ('must survive'); + `); + expect(write.error).toBeUndefined(); + follower = await createTestDatabase(databaseName); + await fetch('/snapshot-stats?reset=1'); + + const install = follower.installSnapshot( + new URL('/snapshot.slow.db', window.location.href).href, + 'none', + SNAPSHOT_SHA256, + 1024 + ); + for (let attempt = 0; attempt < 40; attempt += 1) { + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + if (stats.snapshotRequestCount === 1) break; + await new Promise((resolve) => setTimeout(resolve, 25)); + } + const started = await fetch('/snapshot-stats').then((response) => response.json()); + expect(started.snapshotRequestCount).toBe(1); + + const closed = await follower.close(); + expect(closed.error).toBeUndefined(); + const installResult = await install; + expect(installResult.error?.msg).toContain('Snapshot installation cancelled'); + follower.free(); + follower = undefined; + + for (let attempt = 0; attempt < 40; attempt += 1) { + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + if (stats.cancelledSnapshotRequestCount > 0) break; + await new Promise((resolve) => setTimeout(resolve, 25)); + } + const cancelled = await fetch('/snapshot-stats').then((response) => response.json()); + expect(cancelled.cancelledSnapshotRequestCount).toBeGreaterThan(0); + + const read = await leader.query('SELECT value FROM follower_close_sentinel'); + expect(read.error).toBeUndefined(); + expect(JSON.parse(read.value || '[]')).toEqual([{ value: 'must survive' }]); + } finally { + if (follower) { + await follower.close(); + follower.free(); + } + await leader.close(); + leader.free(); + } + }); + + it('cannot resurrect a worker when close races wipeAndRecreate', async () => { + const database = await createTestDatabase(`shutdown-wipe-race-${Date.now()}`); + try { + const wipe = database.wipeAndRecreate(); + // wipeAndRecreate has stopped the old worker and yielded while deleting + // OPFS. Closing now permanently invalidates that wipe generation. + await Promise.resolve(); + const closed = await database.close(); + expect(closed.error).toBeUndefined(); + + const wipeResult = await wipe; + expect(wipeResult.error?.msg).toContain('Database is closed'); + const query = await database.query('SELECT 1 AS value'); + expect(query.error?.msg).toContain('Database is closed'); + } finally { + await database.close(); + database.free(); + } + }); + + it('closes a live worker before a second client reopens the same OPFS database', async () => { + const databaseName = `shutdown-handoff-${Date.now()}`; + let first: SQLiteWasmDatabase | undefined = await createTestDatabase(databaseName); + let second: SQLiteWasmDatabase | undefined; + + try { + const write = await first.query(` + CREATE TABLE shutdown_handoff_sentinel (value TEXT NOT NULL); + INSERT INTO shutdown_handoff_sentinel VALUES ('persisted before close'); + `); + expect(write.error).toBeUndefined(); + + // Leave a request outstanding to prove close settles public promises, + // rather than merely making future calls fail. + const pendingInstall = first.installSnapshot( + new URL('/snapshot.delayed.db', window.location.href).href, + 'none', + SNAPSHOT_SHA256, + 1024 + ); + await Promise.resolve(); + + const closed = await first.close(); + expect(closed.error).toBeUndefined(); + + const pendingResult = await pendingInstall; + expect(pendingResult.error?.msg).toMatch( + /Database closed|Snapshot installation cancelled/ + ); + first.free(); + first = undefined; + + // The document remains alive. Only the first SDK worker is gone. + second = await createTestDatabase(databaseName); + const read = await second.query('SELECT value FROM shutdown_handoff_sentinel'); + expect(read.error).toBeUndefined(); + expect(JSON.parse(read.value || '[]')).toEqual([ + { value: 'persisted before close' } + ]); + } finally { + if (first) { + await first.close(); + first.free(); + } + if (second) { + await second.close(); + second.free(); + } + } + }); +}); diff --git a/svelte-test/tests/integration/snapshot-import.test.ts b/svelte-test/tests/integration/snapshot-import.test.ts new file mode 100644 index 0000000..1cb2b49 --- /dev/null +++ b/svelte-test/tests/integration/snapshot-import.test.ts @@ -0,0 +1,184 @@ +import { afterEach, describe, expect, it } from 'vitest'; +import { createTestDatabase } from '../fixtures/test-helpers.js'; + +const SNAPSHOT_SHA256 = '5d5408c09e537d8af271fff0eab22ab9ef624259476010ada25de6a3732dec08'; +const SNAPSHOT_SIZE = 1024; +const CORRUPT_SNAPSHOT_SHA256 = 'ba23dd3c57c8f955a8abe95509623b444b68b2018c36730153e50a28e352ebd2'; +type TestDatabase = Awaited>; +const openDatabases: TestDatabase[] = []; + +async function createSnapshotDatabase(name: string): Promise { + const database = await createTestDatabase(name); + openDatabases.push(database); + return database; +} + +afterEach(async () => { + for (const database of openDatabases.splice(0).reverse()) { + await database.close(); + database.free(); + } +}); + +describe('streaming SQLite snapshot installation', () => { + it('opens a long VFS-valid name while rejecting snapshot staging paths', async () => { + const db = await createSnapshotDatabase('l'.repeat(480)); + const write = await db.query(` + CREATE TABLE long_name_sentinel (value TEXT NOT NULL); + INSERT INTO long_name_sentinel VALUES ('ordinary open works'); + `); + expect(write.error).toBeUndefined(); + + const install = await db.installSnapshot( + new URL('/snapshot.raw.db', window.location.href).href, + 'none', + SNAPSHOT_SHA256, + SNAPSHOT_SIZE + ); + expect(install.error?.msg).toContain('too long for atomic snapshot installation'); + + const read = await db.query('SELECT value FROM long_name_sentinel'); + expect(JSON.parse(read.value || '[]')).toEqual([{ value: 'ordinary open works' }]); + }); + + it('installs uncompressed and gzip-compressed SQLite snapshots', async () => { + const db = await createSnapshotDatabase(`snapshot-raw-poc-${Date.now()}`); + const rawInstall = await db.installSnapshot( + new URL('/snapshot.raw.db', window.location.href).href, + 'none', + SNAPSHOT_SHA256, + SNAPSHOT_SIZE + ); + + expect(rawInstall.error).toBeUndefined(); + const rawStats = JSON.parse(rawInstall.value || '{}'); + expect(rawStats).toMatchObject({ + bytesWritten: SNAPSHOT_SIZE, + compression: 'none', + sha256: SNAPSHOT_SHA256 + }); + + const rawResult = await db.query( + "SELECT COUNT(*) AS table_count FROM sqlite_schema WHERE type = 'table'" + ); + expect(rawResult.error).toBeUndefined(); + const [rawRow] = JSON.parse(rawResult.value || '[]'); + expect(rawRow.table_count).toBe(1); + const rawItems = await db.query('SELECT id, label FROM snapshot_items ORDER BY id'); + expect(rawItems.error).toBeUndefined(); + expect(JSON.parse(rawItems.value || '[]')).toEqual([ + { id: 1, label: 'alpha' }, + { id: 2, label: 'beta' }, + { id: 3, label: 'gamma' } + ]); + + const sentinelResult = await db.query(` + CREATE TABLE snapshot_install_sentinel (value TEXT NOT NULL); + INSERT INTO snapshot_install_sentinel VALUES ('original database'); + `); + expect(sentinelResult.error).toBeUndefined(); + + const rejectedInstall = await db.installSnapshot( + new URL('/snapshot.raw.db', window.location.href).href, + 'none', + '0'.repeat(64), + SNAPSHOT_SIZE + ); + expect(rejectedInstall.error?.msg).toContain('SHA-256 mismatch'); + + const preservedResult = await db.query( + 'SELECT value FROM snapshot_install_sentinel' + ); + expect(preservedResult.error).toBeUndefined(); + expect(JSON.parse(preservedResult.value || '[]')).toEqual([ + { value: 'original database' } + ]); + + const corruptInstall = await db.installSnapshot( + new URL('/snapshot.corrupt.db', window.location.href).href, + 'none', + CORRUPT_SNAPSHOT_SHA256, + SNAPSHOT_SIZE + ); + expect(corruptInstall.error?.msg).toContain('integrity validation'); + const preservedAfterCorruption = await db.query( + 'SELECT value FROM snapshot_install_sentinel' + ); + expect(JSON.parse(preservedAfterCorruption.value || '[]')).toEqual([ + { value: 'original database' } + ]); + + const install = await db.installSnapshot( + new URL('/snapshot.db.gz', window.location.href).href, + 'gzip', + SNAPSHOT_SHA256, + SNAPSHOT_SIZE + ); + expect(install.error).toBeUndefined(); + const stats = JSON.parse(install.value || '{}'); + expect(stats.bytesWritten).toBe(SNAPSHOT_SIZE); + expect(stats.compression).toBe('gzip'); + expect(stats.sha256).toBe(SNAPSHOT_SHA256); + + const result = await db.query(` + SELECT + (SELECT quick_check FROM pragma_quick_check LIMIT 1) AS quick_check, + (SELECT page_count FROM pragma_page_count) AS page_count, + (SELECT page_size FROM pragma_page_size) AS page_size, + (SELECT COUNT(*) FROM sqlite_schema WHERE type = 'table') AS table_count, + (SELECT COUNT(*) FROM sqlite_schema WHERE name = 'snapshot_install_sentinel') AS sentinel_count + `); + expect(result.error).toBeUndefined(); + const [row] = JSON.parse(result.value || '[]'); + expect(row.quick_check).toBe('ok'); + expect(row.page_count * row.page_size).toBe(SNAPSHOT_SIZE); + expect(row.table_count).toBe(1); + expect(row.sentinel_count).toBe(0); + }); + + it('cancels an oversized response and remains usable', async () => { + await fetch('/snapshot-stats?reset=1'); + const db = await createSnapshotDatabase(`snapshot-cancel-${Date.now()}`); + const rejected = await db.installSnapshot( + new URL('/snapshot.oversize.db', window.location.href).href, + 'none', + '0'.repeat(64), + SNAPSHOT_SIZE + ); + expect(rejected.error?.msg).toContain('exceeded expected size'); + + for (let i = 0; i < 20; i += 1) { + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + if (stats.cancelledSnapshotRequestCount > 0) break; + await new Promise((resolve) => setTimeout(resolve, 25)); + } + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + expect(stats.cancelledSnapshotRequestCount).toBeGreaterThan(0); + + const successful = await db.installSnapshot( + new URL('/snapshot.raw.db', window.location.href).href, + 'none', + SNAPSHOT_SHA256, + SNAPSHOT_SIZE + ); + expect(successful.error).toBeUndefined(); + }); + + it('coalesces matching installs from two clients into one download', async () => { + await fetch('/snapshot-stats?reset=1'); + const name = `snapshot-coalesce-${Date.now()}`; + const [leader, follower] = await Promise.all([ + createSnapshotDatabase(name), + createSnapshotDatabase(name) + ]); + const url = new URL('/snapshot.delayed.db', window.location.href).href; + const [first, second] = await Promise.all([ + leader.installSnapshot(url, 'none', SNAPSHOT_SHA256, SNAPSHOT_SIZE), + follower.installSnapshot(url, 'none', SNAPSHOT_SHA256, SNAPSHOT_SIZE) + ]); + expect(first.error).toBeUndefined(); + expect(second.error).toBeUndefined(); + const stats = await fetch('/snapshot-stats').then((response) => response.json()); + expect(stats.snapshotRequestCount).toBe(1); + }); +}); diff --git a/svelte-test/vitest.config.js b/svelte-test/vitest.config.js index 34f9c91..d17f302 100644 --- a/svelte-test/vitest.config.js +++ b/svelte-test/vitest.config.js @@ -1,12 +1,155 @@ import { defineConfig } from 'vitest/config'; import path from 'path'; import fs from 'fs'; +import os from 'node:os'; +import { gzipSync } from 'node:zlib'; +import { createHash } from 'node:crypto'; +import { DatabaseSync } from 'node:sqlite'; + +// A deterministic, two-page SQLite fixture with one `snapshot_items` table. +// Keeping it inline makes the browser integration test hermetic and avoids +// coupling the SDK suite to any consumer's production database dump. +const snapshotFixture = Buffer.from( + [ + 'U1FMaXRlIGZvcm1hdCAzAAIAAQEMQCAgAAAABAAAAAIAAAAAAAAAAAAAAAIAAAAEAAAAAAAAAAAAAAAB', + 'AAAABwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAC6N+A0AAAABAYEAAYEAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAHEBBxcpKQGBHXRhYmxlc25hcHNob3RfaXRlbXNzbmFwc2hv', + 'dF9pdGVtcwJDUkVBVEUgVEFCTEUgc25hcHNob3RfaXRlbXMoaWQgSU5URUdFUiBQUklNQVJZIEtFWSwg', + 'bGFiZWwgVEVYVCBOT1QgTlVMTCkAAAAAAAAAAAAAAAANAAAAAwHXAAHqAeEB1wAAAAAAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAIAwMAF2dhbW1hBwIDABViZXRhCAEDABdhbHBoYQAAAAAAAAAA', + 'AAAAAA==' + ].join(''), + 'base64' +); +const compressedSnapshotFixture = gzipSync(snapshotFixture, { mtime: 0 }); +const corruptSnapshotFixture = Buffer.from(snapshotFixture); +corruptSnapshotFixture[512] = 0; +const corruptSnapshotSha256 = createHash('sha256').update(corruptSnapshotFixture).digest('hex'); +const finalizationFixtureDirectory = fs.mkdtempSync(path.join(os.tmpdir(), 'sqlite-web-finalize-')); +const finalizationFixturePath = path.join(finalizationFixtureDirectory, 'snapshot.db'); +const finalizationDatabase = new DatabaseSync(finalizationFixturePath); +finalizationDatabase.exec(` + PRAGMA journal_mode = OFF; + CREATE TABLE snapshot_items(id INTEGER PRIMARY KEY, label TEXT NOT NULL); + INSERT INTO snapshot_items VALUES (1, 'replacement snapshot'); + CREATE TABLE validation_padding(payload BLOB NOT NULL); + INSERT INTO validation_padding VALUES (zeroblob(33554432)); +`); +finalizationDatabase.close(); +const finalizationSnapshotFixture = fs.readFileSync(finalizationFixturePath); +fs.rmSync(finalizationFixtureDirectory, { recursive: true }); +const finalizationSnapshotSha256 = createHash('sha256') + .update(finalizationSnapshotFixture) + .digest('hex'); +let snapshotRequestCount = 0; +let cancelledSnapshotRequestCount = 0; +let completedSnapshotRequestCount = 0; export default defineConfig({ plugins: [ { name: 'rainlanguage-sqlite-web-serve', configureServer(server) { + server.middlewares.use('/snapshot-stats', (req, res) => { + if (req.url?.includes('reset')) { + snapshotRequestCount = 0; + cancelledSnapshotRequestCount = 0; + completedSnapshotRequestCount = 0; + } + res.setHeader('Content-Type', 'application/json'); + res.end( + JSON.stringify({ + snapshotRequestCount, + cancelledSnapshotRequestCount, + completedSnapshotRequestCount + }) + ); + }); + server.middlewares.use('/snapshot-finalization-meta', (_req, res) => { + res.setHeader('Content-Type', 'application/json'); + res.end( + JSON.stringify({ + sha256: finalizationSnapshotSha256, + size: finalizationSnapshotFixture.length + }) + ); + }); + server.middlewares.use('/snapshot.finalization.db', (_req, res) => { + snapshotRequestCount += 1; + res.setHeader('Content-Type', 'application/vnd.sqlite3'); + res.setHeader('Content-Length', finalizationSnapshotFixture.length); + res.end(finalizationSnapshotFixture, () => { + completedSnapshotRequestCount += 1; + }); + }); + server.middlewares.use('/snapshot.corrupt.db', (_req, res) => { + res.setHeader('Content-Type', 'application/vnd.sqlite3'); + res.setHeader('X-Snapshot-Sha256', corruptSnapshotSha256); + res.end(corruptSnapshotFixture); + }); + server.middlewares.use('/snapshot.oversize.db', (req, res) => { + let complete = false; + const chunk = Buffer.alloc(64 * 1024, 1); + res.setHeader('Content-Type', 'application/vnd.sqlite3'); + const interval = setInterval(() => res.write(chunk), 10); + const finish = setTimeout(() => { + complete = true; + clearInterval(interval); + res.end(); + }, 2000); + req.on('close', () => { + clearInterval(interval); + clearTimeout(finish); + if (!complete) cancelledSnapshotRequestCount += 1; + }); + }); + server.middlewares.use('/snapshot.raw.db', (_req, res) => { + snapshotRequestCount += 1; + res.setHeader('Content-Type', 'application/vnd.sqlite3'); + res.setHeader('Content-Length', snapshotFixture.length); + res.end(snapshotFixture); + }); + server.middlewares.use('/snapshot.delayed.db', (_req, res) => { + snapshotRequestCount += 1; + res.setHeader('Content-Type', 'application/vnd.sqlite3'); + setTimeout(() => res.end(snapshotFixture), 150); + }); + server.middlewares.use('/snapshot.slow.db', (req, res) => { + snapshotRequestCount += 1; + let complete = false; + let offset = 0; + res.setHeader('Content-Type', 'application/vnd.sqlite3'); + const interval = setInterval(() => { + const next = Math.min(offset + 128, snapshotFixture.length); + res.write(snapshotFixture.subarray(offset, next)); + offset = next; + if (offset === snapshotFixture.length) { + complete = true; + clearInterval(interval); + res.end(); + } + }, 100); + req.on('close', () => { + clearInterval(interval); + if (!complete) cancelledSnapshotRequestCount += 1; + }); + }); + server.middlewares.use('/snapshot.db.gz', (_req, res) => { + res.setHeader('Content-Type', 'application/gzip'); + res.setHeader('Content-Length', compressedSnapshotFixture.length); + res.end(compressedSnapshotFixture); + }); server.middlewares.use('/pkg', (req, res, next) => { const filePath = req.url?.substring(1); const fullPath = path.join(process.cwd(), 'node_modules/@rainlanguage/sqlite-web', filePath || '');