diff --git a/libsql-replication/src/injector/error.rs b/libsql-replication/src/injector/error.rs index ac8f1be711..bafd73dc47 100644 --- a/libsql-replication/src/injector/error.rs +++ b/libsql-replication/src/injector/error.rs @@ -9,4 +9,6 @@ pub enum Error { Sqlite(#[from] rusqlite::Error), #[error("A fatal error occured injecting frames: {0}")] FatalInjectError(BoxError), + #[error("sync cancelled for terminal shutdown")] + SyncCancelledForShutdown, } diff --git a/libsql-replication/src/injector/mod.rs b/libsql-replication/src/injector/mod.rs index 8e58e99a31..69728202a6 100644 --- a/libsql-replication/src/injector/mod.rs +++ b/libsql-replication/src/injector/mod.rs @@ -2,6 +2,7 @@ use std::future::Future; use super::rpc::replication::Frame as RpcFrame; pub use sqlite_injector::SqliteInjector; +use tokio_util::sync::CancellationToken; use crate::frame::FrameNo; @@ -18,6 +19,22 @@ pub trait Injector { frame: RpcFrame, ) -> impl Future>> + Send; + /// Inject a singular frame, cooperatively observing terminal sync cancellation. + /// + /// Cancellation is a best-effort capability for injectors. The default + /// implementation preserves the base `Injector` contract and does not + /// interrupt in-flight work; injectors that can safely interrupt their + /// backend should override this method and return + /// [`Error::SyncCancelledForShutdown`] only for cancellation-caused + /// interruptions. + fn inject_frame_with_cancellation( + &mut self, + frame: RpcFrame, + _token: &CancellationToken, + ) -> impl Future>> + Send { + self.inject_frame(frame) + } + /// Discard any uncommintted frames. fn rollback(&mut self) -> impl Future + Send; diff --git a/libsql-replication/src/injector/sqlite_injector/mod.rs b/libsql-replication/src/injector/sqlite_injector/mod.rs index 17fec1c553..a3c14b0be0 100644 --- a/libsql-replication/src/injector/sqlite_injector/mod.rs +++ b/libsql-replication/src/injector/sqlite_injector/mod.rs @@ -3,8 +3,9 @@ use std::sync::Arc; use std::{collections::VecDeque, path::PathBuf}; use parking_lot::Mutex; -use rusqlite::OpenFlags; +use rusqlite::{ErrorCode, OpenFlags}; use tokio::task::spawn_blocking; +use tokio_util::sync::CancellationToken; use crate::frame::{Frame, FrameNo}; use crate::rpc::replication::Frame as RpcFrame; @@ -23,6 +24,32 @@ pub type FrameBuffer = Arc>>; pub struct SqliteInjector { pub(in super::super) inner: Arc>, + interrupt_handle: InjectorInterruptHandle, +} + +#[derive(Clone, Default)] +struct InjectorInterruptHandle { + current: Arc>>, +} + +impl InjectorInterruptHandle { + fn set(&self, handle: rusqlite::InterruptHandle) { + *self.current.lock() = Some(handle); + } + + fn interrupt(&self) { + if let Some(handle) = self.current.lock().as_ref() { + handle.interrupt(); + } + } +} + +fn is_cancellation_interrupt(error: &Error) -> bool { + matches!( + error, + Error::Sqlite(e) + if e.sqlite_error_code() == Some(ErrorCode::OperationInterrupted) + ) } impl Injector for SqliteInjector { @@ -35,6 +62,41 @@ impl Injector for SqliteInjector { .unwrap() } + async fn inject_frame_with_cancellation( + &mut self, + frame: RpcFrame, + token: &CancellationToken, + ) -> Result> { + let inner = self.inner.clone(); + let interrupt_handle = self.interrupt_handle.clone(); + let frame = + Frame::try_from(&frame.data[..]).map_err(|e| Error::FatalInjectError(e.into()))?; + let mut join = spawn_blocking(move || inner.lock().inject_frame(frame)); + + tokio::select! { + biased; + + result = &mut join => result.map_err(|e| Error::FatalInjectError(e.into()))?, + _ = token.cancelled() => { + interrupt_handle.interrupt(); + match join.await { + Ok(Ok(result)) => Ok(result), + Ok(Err(Error::SyncCancelledForShutdown)) => { + Err(Error::SyncCancelledForShutdown) + } + Ok(Err(e)) if is_cancellation_interrupt(&e) => { + Err(Error::SyncCancelledForShutdown) + } + Ok(Err(e)) => { + tracing::warn!(error = %e, "injector_error_after_sync_cancellation"); + Err(e) + } + Err(e) => Err(Error::FatalInjectError(e.into())), + } + } + } + } + async fn rollback(&mut self) { let inner = self.inner.clone(); spawn_blocking(move || inner.lock().rollback()) @@ -58,14 +120,23 @@ impl SqliteInjector { auto_checkpoint: u32, encryption_config: Option, ) -> super::Result { + let interrupt_handle = InjectorInterruptHandle::default(); + let inner_interrupt_handle = interrupt_handle.clone(); let inner = spawn_blocking(move || { - SqliteInjectorInner::new(path, capacity, auto_checkpoint, encryption_config) + SqliteInjectorInner::new( + path, + capacity, + auto_checkpoint, + encryption_config, + inner_interrupt_handle, + ) }) .await .unwrap()?; Ok(Self { inner: Arc::new(Mutex::new(inner)), + interrupt_handle, }) } } @@ -86,6 +157,7 @@ pub(in super::super) struct SqliteInjectorInner { path: PathBuf, encryption_config: Option, auto_checkpoint: u32, + interrupt_handle: InjectorInterruptHandle, } /// Methods from this trait are called before and after performing a frame injection. @@ -98,6 +170,7 @@ impl SqliteInjectorInner { capacity: usize, auto_checkpoint: u32, encryption_config: Option, + interrupt_handle: InjectorInterruptHandle, ) -> Result { let path = path.as_ref().to_path_buf(); @@ -113,6 +186,7 @@ impl SqliteInjectorInner { auto_checkpoint, encryption_config.clone(), )?; + interrupt_handle.set(connection.get_interrupt_handle()); Ok(Self { is_txn: false, @@ -124,6 +198,7 @@ impl SqliteInjectorInner { path, encryption_config, auto_checkpoint, + interrupt_handle, }) } @@ -235,6 +310,7 @@ impl SqliteInjectorInner { self.auto_checkpoint, self.encryption_config.clone(), )?; + self.interrupt_handle.set(new_conn.get_interrupt_handle()); let _ = std::mem::replace(&mut *conn, new_conn); } @@ -266,6 +342,9 @@ impl SqliteInjectorInner { mod test { use crate::frame::FrameBorrowed; use std::mem::size_of; + use std::time::Duration; + + use tokio::sync::oneshot; use super::*; /// this this is generated by creating a table test, inserting 5 rows into it, and then @@ -277,12 +356,87 @@ mod test { .map(|b| Frame::try_from(b).unwrap()) } + #[test] + fn cancellation_interrupt_errors_are_classified_narrowly() { + let interrupted = Error::Sqlite(rusqlite::Error::SqliteFailure( + rusqlite::ffi::Error { + code: ErrorCode::OperationInterrupted, + extended_code: rusqlite::ffi::SQLITE_INTERRUPT, + }, + None, + )); + assert!(is_cancellation_interrupt(&interrupted)); + + let corrupt = Error::Sqlite(rusqlite::Error::SqliteFailure( + rusqlite::ffi::Error { + code: ErrorCode::DatabaseCorrupt, + extended_code: rusqlite::ffi::SQLITE_CORRUPT, + }, + None, + )); + assert!(!is_cancellation_interrupt(&corrupt)); + } + + #[tokio::test] + async fn cancellation_waits_for_real_blocking_injection_task_to_settle() { + let temp = tempfile::tempdir().unwrap(); + let mut injector = SqliteInjector::new(temp.path().join("data"), 10, 10000, None) + .await + .unwrap(); + let frame = RpcFrame { + data: wal_log().next().unwrap().bytes(), + timestamp: None, + durable_frame_no: None, + }; + + let inner = injector.inner.clone(); + let (locked_tx, locked_rx) = oneshot::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let blocker = tokio::task::spawn_blocking(move || { + let _guard = inner.lock(); + let _ = locked_tx.send(()); + let _ = release_rx.recv(); + }); + locked_rx.await.unwrap(); + + let token = CancellationToken::new(); + let cancel = token.clone(); + let injection = tokio::spawn(async move { + let result = injector.inject_frame_with_cancellation(frame, &token).await; + (injector, result) + }); + + cancel.cancel(); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!( + !injection.is_finished(), + "cancelled injection returned before the blocking task settled" + ); + + release_tx.send(()).unwrap(); + blocker.await.unwrap(); + let (_injector, result) = tokio::time::timeout(Duration::from_secs(5), injection) + .await + .expect("cancelled injection did not settle after releasing the blocking task") + .unwrap(); + assert!(matches!( + result, + Ok(_) | Err(Error::SyncCancelledForShutdown) + )); + } + #[test] fn test_simple_inject_frames() { let temp = tempfile::tempdir().unwrap(); - let mut injector = - SqliteInjectorInner::new(temp.path().join("data"), 10, 10000, None).unwrap(); + let mut injector = SqliteInjectorInner::new( + temp.path().join("data"), + 10, + 10000, + None, + InjectorInterruptHandle::default(), + ) + .unwrap(); let log = wal_log(); for frame in log { injector.inject_frame(frame).unwrap(); @@ -302,8 +456,14 @@ mod test { let temp = tempfile::tempdir().unwrap(); // inject one frame at a time - let mut injector = - SqliteInjectorInner::new(temp.path().join("data"), 1, 10000, None).unwrap(); + let mut injector = SqliteInjectorInner::new( + temp.path().join("data"), + 1, + 10000, + None, + InjectorInterruptHandle::default(), + ) + .unwrap(); let log = wal_log(); for frame in log { injector.inject_frame(frame).unwrap(); @@ -323,8 +483,14 @@ mod test { let temp = tempfile::tempdir().unwrap(); // inject one frame at a time - let mut injector = - SqliteInjectorInner::new(temp.path().join("data"), 10, 1000, None).unwrap(); + let mut injector = SqliteInjectorInner::new( + temp.path().join("data"), + 10, + 1000, + None, + InjectorInterruptHandle::default(), + ) + .unwrap(); let mut frames = wal_log(); assert!(injector diff --git a/libsql-replication/src/replicator.rs b/libsql-replication/src/replicator.rs index ca2de123ca..631e058962 100644 --- a/libsql-replication/src/replicator.rs +++ b/libsql-replication/src/replicator.rs @@ -2,6 +2,7 @@ use std::path::PathBuf; use tokio::time::Duration; use tokio_stream::{Stream, StreamExt}; +use tokio_util::sync::CancellationToken; use tonic::{Code, Status}; use crate::frame::{Frame, FrameNo}; @@ -38,6 +39,8 @@ pub enum Error { NoHandshake, #[error("Requested namespace doesn't exist")] NamespaceDoesntExist, + #[error("sync cancelled for terminal shutdown")] + SyncCancelledForShutdown, } impl From for Error { @@ -67,11 +70,41 @@ pub trait ReplicatorClient { /// Perform handshake with remote async fn handshake(&mut self) -> Result<(), Error>; + /// Perform handshake with remote, cooperatively observing terminal sync cancellation. + async fn handshake_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result<(), Error> { + tokio::select! { + result = self.handshake() => result, + _ = token.cancelled() => Err(Error::SyncCancelledForShutdown), + } + } /// Return a stream of frames to apply to the database async fn next_frames(&mut self) -> Result; + /// Return a stream of frames to apply to the database, cooperatively observing terminal sync cancellation. + async fn next_frames_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result { + tokio::select! { + result = self.next_frames() => result, + _ = token.cancelled() => Err(Error::SyncCancelledForShutdown), + } + } /// Return a snapshot for the current replication index. Called after next_frame has returned a /// NeedSnapshot error async fn snapshot(&mut self) -> Result; + /// Return a snapshot for the current replication index, cooperatively observing terminal sync cancellation. + async fn snapshot_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result { + tokio::select! { + result = self.snapshot() => result, + _ = token.cancelled() => Err(Error::SyncCancelledForShutdown), + } + } /// set the new commit frame_no async fn commit_frame_no(&mut self, frame_no: FrameNo) -> Result<(), Error>; /// Returns the currently committed replication index @@ -94,6 +127,17 @@ where Either::Right(b) => b.handshake().await, } } + + async fn handshake_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result<(), Error> { + match self { + Either::Left(a) => a.handshake_with_cancellation(token).await, + Either::Right(b) => b.handshake_with_cancellation(token).await, + } + } + /// Return a stream of frames to apply to the database async fn next_frames(&mut self) -> Result { match self { @@ -101,6 +145,23 @@ where Either::Right(b) => b.next_frames().await.map(Either::Right), } } + + async fn next_frames_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result { + match self { + Either::Left(a) => a + .next_frames_with_cancellation(token) + .await + .map(Either::Left), + Either::Right(b) => b + .next_frames_with_cancellation(token) + .await + .map(Either::Right), + } + } + /// Return a snapshot for the current replication index. Called after next_frame has returned a /// NeedSnapshot error async fn snapshot(&mut self) -> Result { @@ -109,6 +170,16 @@ where Either::Right(b) => b.snapshot().await.map(Either::Right), } } + + async fn snapshot_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result { + match self { + Either::Left(a) => a.snapshot_with_cancellation(token).await.map(Either::Left), + Either::Right(b) => b.snapshot_with_cancellation(token).await.map(Either::Right), + } + } /// set the new commit frame_no async fn commit_frame_no(&mut self, frame_no: FrameNo) -> Result<(), Error> { match self { @@ -154,7 +225,7 @@ enum ReplicatorState { impl Replicator where - C: ReplicatorClient, + C: ReplicatorClient + Send, { /// Creates a replicator for the db file pointed at by `db_path` pub async fn new_sqlite( @@ -177,7 +248,7 @@ where impl Replicator where - C: ReplicatorClient, + C: ReplicatorClient + Send, I: Injector, { pub fn new(client: C, injector: I) -> Self { @@ -250,19 +321,61 @@ where } } + pub async fn replicate_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result<(), Error> { + loop { + self.try_replicate_step_with_cancellation(token).await?; + if self.state == ReplicatorState::Exit { + self.state = ReplicatorState::NeedFrames; + return Ok(()); + } + } + } + async fn try_replicate_step(&mut self) -> Result<(), Error> { + self.try_replicate_step_inner(None).await + } + + async fn try_replicate_step_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result<(), Error> { + self.try_replicate_step_inner(Some(token)).await + } + + async fn try_replicate_step_inner( + &mut self, + token: Option<&CancellationToken>, + ) -> Result<(), Error> { let state = self.state; - let ret = match state { - ReplicatorState::NeedHandshake => self.try_perform_handshake().await, - ReplicatorState::NeedFrames => self.try_replicate().await, - ReplicatorState::NeedSnapshot => self.load_snapshot().await, - ReplicatorState::Exit => unreachable!("trying to step replicator on exit"), + let ret = match (state, token) { + (ReplicatorState::NeedHandshake, Some(token)) => { + self.try_perform_handshake_with_cancellation(token).await + } + (ReplicatorState::NeedHandshake, None) => self.try_perform_handshake().await, + (ReplicatorState::NeedFrames, Some(token)) => { + self.try_replicate_with_cancellation(token).await + } + (ReplicatorState::NeedFrames, None) => self.try_replicate().await, + (ReplicatorState::NeedSnapshot, Some(token)) => { + self.load_snapshot_with_cancellation(token).await + } + (ReplicatorState::NeedSnapshot, None) => self.load_snapshot().await, + (ReplicatorState::Exit, _) => unreachable!("trying to step replicator on exit"), }; // in case of error we rollback the current injector transaction, and start over. if ret.is_err() { + if matches!(ret, Err(Error::SyncCancelledForShutdown)) { + tracing::info!("sync_cancellation_rollback_started"); + } self.client.rollback(); self.injector.rollback().await; + if matches!(ret, Err(Error::SyncCancelledForShutdown)) { + tracing::info!("sync_cancellation_rollback_finished"); + } } self.state = match ret { @@ -294,6 +407,51 @@ where Ok(()) } + async fn try_perform_handshake_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result<(), Error> { + let mut error_printed = false; + for _ in 0..self.max_handshake_retries { + if token.is_cancelled() { + tracing::info!("sync_cancelled_during_handshake"); + return Err(Error::SyncCancelledForShutdown); + } + + tracing::debug!("Attempting to perform handshake with primary."); + match self.client.handshake_with_cancellation(token).await { + Ok(_) => { + self.state = ReplicatorState::NeedFrames; + return Ok(()); + } + Err(Error::Client(e)) if !error_printed => { + if e.downcast_ref::().is_some() { + tracing::warn!("error connecting to primary. retrying. Verify that the libsql server version is `>=0.22` error: {e}"); + } else { + tracing::warn!("error connecting to primary. retrying. error: {e}"); + } + + error_printed = true; + } + Err(Error::Client(_)) if error_printed => (), + Err(Error::SyncCancelledForShutdown) => { + tracing::info!("sync_cancelled_during_handshake"); + return Err(Error::SyncCancelledForShutdown); + } + Err(e) => return Err(e), + } + tokio::select! { + _ = tokio::time::sleep(Duration::from_secs(1)) => {}, + _ = token.cancelled() => { + tracing::info!("sync_cancelled_during_handshake"); + return Err(Error::SyncCancelledForShutdown); + } + } + } + + Err(Error::PrimaryHandshakeTimeout) + } + async fn try_replicate(&mut self) -> Result<(), Error> { let mut stream = self.client.next_frames().await?; @@ -304,6 +462,43 @@ where Ok(()) } + async fn try_replicate_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result<(), Error> { + if token.is_cancelled() { + tracing::info!("sync_cancelled_during_frame_fetch"); + return Err(Error::SyncCancelledForShutdown); + } + let mut stream = self.client.next_frames_with_cancellation(token).await?; + + loop { + let next = tokio::select! { + next = stream.next() => next, + _ = token.cancelled() => { + tracing::info!("sync_cancelled_during_frame_fetch"); + return Err(Error::SyncCancelledForShutdown); + } + }; + + let Some(frame) = next.transpose()? else { + if token.is_cancelled() { + tracing::info!("sync_cancelled_during_frame_fetch"); + return Err(Error::SyncCancelledForShutdown); + } + break; + }; + self.inject_frame_with_cancellation(frame, token).await?; + } + + if token.is_cancelled() { + tracing::info!("sync_cancelled_during_frame_fetch"); + return Err(Error::SyncCancelledForShutdown); + } + + Ok(()) + } + async fn load_snapshot(&mut self) -> Result<(), Error> { self.client.rollback(); self.injector.rollback().await; @@ -325,6 +520,52 @@ where } } + async fn load_snapshot_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result<(), Error> { + self.client.rollback(); + self.injector.rollback().await; + loop { + if token.is_cancelled() { + tracing::info!("sync_cancelled_during_snapshot_fetch"); + return Err(Error::SyncCancelledForShutdown); + } + + match self.client.snapshot_with_cancellation(token).await { + Ok(mut stream) => loop { + let next = tokio::select! { + next = stream.next() => next, + _ = token.cancelled() => { + tracing::info!("sync_cancelled_during_snapshot_stream"); + return Err(Error::SyncCancelledForShutdown); + } + }; + + let Some(frame) = next else { + if token.is_cancelled() { + tracing::info!("sync_cancelled_during_snapshot_stream"); + return Err(Error::SyncCancelledForShutdown); + } + return Ok(()); + }; + self.inject_frame_with_cancellation(frame?, token).await?; + }, + Err(Error::SnapshotPending) => { + tracing::info!("snapshot not ready yet, waiting 1s..."); + tokio::select! { + _ = tokio::time::sleep(Duration::from_secs(1)) => {}, + _ = token.cancelled() => { + tracing::info!("sync_cancelled_during_snapshot_fetch"); + return Err(Error::SyncCancelledForShutdown); + } + } + } + Err(e) => return Err(e), + } + } + } + async fn inject_frame(&mut self, frame: RpcFrame) -> Result<(), Error> { self.frames_synced += 1; @@ -332,22 +573,56 @@ where self.injector.durable_frame_no(frame_no); } - match self.injector.inject_frame(frame).await? { - Some(commit_fno) => { + if let Some(commit_fno) = self.injector.inject_frame(frame).await? { + self.client.commit_frame_no(commit_fno).await?; + } + + Ok(()) + } + + async fn inject_frame_with_cancellation( + &mut self, + frame: RpcFrame, + token: &CancellationToken, + ) -> Result<(), Error> { + if token.is_cancelled() { + tracing::info!("sync_cancelled_during_injection"); + return Err(Error::SyncCancelledForShutdown); + } + + self.frames_synced += 1; + + if let Some(frame_no) = frame.durable_frame_no { + self.injector.durable_frame_no(frame_no); + } + + match self + .injector + .inject_frame_with_cancellation(frame, token) + .await + { + Ok(Some(commit_fno)) => { self.client.commit_frame_no(commit_fno).await?; } - None => (), + Ok(None) => (), + Err(crate::injector::Error::SyncCancelledForShutdown) => { + tracing::info!("sync_cancelled_during_injection"); + return Err(Error::SyncCancelledForShutdown); + } + Err(e) => return Err(e.into()), + } + + if token.is_cancelled() { + tracing::info!("sync_cancelled_during_injection"); + return Err(Error::SyncCancelledForShutdown); } Ok(()) } pub async fn flush(&mut self) -> Result<(), Error> { - match self.injector.flush().await? { - Some(commit_fno) => { - self.client.commit_frame_no(commit_fno).await?; - } - None => (), + if let Some(commit_fno) = self.injector.flush().await? { + self.client.commit_frame_no(commit_fno).await?; } Ok(()) @@ -366,7 +641,14 @@ pub fn map_frame_err(f: Result) -> Result { #[cfg(test)] mod test { - use std::{mem::size_of, pin::Pin}; + use std::{ + mem::size_of, + pin::Pin, + sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + Arc, + }, + }; use async_stream::stream; @@ -375,6 +657,607 @@ mod test { use super::*; + struct BlockingInjector { + inject_called: Arc, + rolled_back: Arc, + } + + impl BlockingInjector { + fn new() -> Self { + Self { + inject_called: Arc::new(AtomicBool::new(false)), + rolled_back: Arc::new(AtomicBool::new(false)), + } + } + } + + impl Injector for BlockingInjector { + async fn inject_frame( + &mut self, + _frame: RpcFrame, + ) -> std::result::Result, crate::injector::Error> { + Ok(None) + } + + async fn inject_frame_with_cancellation( + &mut self, + _frame: RpcFrame, + token: &CancellationToken, + ) -> std::result::Result, crate::injector::Error> { + self.inject_called.store(true, Ordering::SeqCst); + token.cancelled().await; + Err(crate::injector::Error::SyncCancelledForShutdown) + } + + async fn rollback(&mut self) { + self.rolled_back.store(true, Ordering::SeqCst); + } + + async fn flush(&mut self) -> std::result::Result, crate::injector::Error> { + Ok(None) + } + + fn durable_frame_no(&mut self, _frame_no: u64) {} + } + + #[tokio::test] + async fn cancellation_while_waiting_for_handshake_rolls_back() { + struct Client { + rolled_back: Arc, + } + + #[async_trait::async_trait] + impl ReplicatorClient for Client { + type FrameStream = + Pin> + Send + 'static>>; + + async fn handshake(&mut self) -> Result<(), Error> { + std::future::pending().await + } + + async fn next_frames(&mut self) -> Result { + unreachable!() + } + + async fn snapshot(&mut self) -> Result { + unreachable!() + } + + async fn commit_frame_no(&mut self, _frame_no: FrameNo) -> Result<(), Error> { + unreachable!() + } + + fn committed_frame_no(&self) -> Option { + None + } + + fn rollback(&mut self) { + self.rolled_back.store(true, Ordering::SeqCst); + } + } + + let token = CancellationToken::new(); + let client_rolled_back = Arc::new(AtomicBool::new(false)); + let injector = BlockingInjector::new(); + let injector_rolled_back = injector.rolled_back.clone(); + let mut replicator = Replicator::new( + Client { + rolled_back: client_rolled_back.clone(), + }, + injector, + ); + replicator.state = ReplicatorState::NeedHandshake; + + let cancel = token.clone(); + let (result, _) = tokio::join!( + async { + replicator + .try_replicate_step_with_cancellation(&token) + .await + }, + async { + tokio::task::yield_now().await; + cancel.cancel(); + } + ); + + assert!(matches!( + result.unwrap_err(), + Error::SyncCancelledForShutdown + )); + assert!(client_rolled_back.load(Ordering::SeqCst)); + assert!(injector_rolled_back.load(Ordering::SeqCst)); + assert_eq!(replicator.state, ReplicatorState::NeedHandshake); + } + + #[tokio::test] + async fn cancellation_during_injection_waits_and_rolls_back() { + struct Client { + rolled_back: Arc, + } + + #[async_trait::async_trait] + impl ReplicatorClient for Client { + type FrameStream = + Pin> + Send + 'static>>; + + async fn handshake(&mut self) -> Result<(), Error> { + Ok(()) + } + + async fn next_frames(&mut self) -> Result { + Ok(Box::pin(tokio_stream::iter([Ok(RpcFrame { + data: Vec::new().into(), + timestamp: None, + durable_frame_no: None, + })]))) + } + + async fn snapshot(&mut self) -> Result { + unreachable!() + } + + async fn commit_frame_no(&mut self, _frame_no: FrameNo) -> Result<(), Error> { + Ok(()) + } + + fn committed_frame_no(&self) -> Option { + None + } + + fn rollback(&mut self) { + self.rolled_back.store(true, Ordering::SeqCst); + } + } + + let token = CancellationToken::new(); + let client_rolled_back = Arc::new(AtomicBool::new(false)); + let injector = BlockingInjector::new(); + let inject_called = injector.inject_called.clone(); + let injector_rolled_back = injector.rolled_back.clone(); + let mut replicator = Replicator::new( + Client { + rolled_back: client_rolled_back.clone(), + }, + injector, + ); + replicator.state = ReplicatorState::NeedFrames; + + let cancel = token.clone(); + let (result, _) = tokio::join!( + async { + replicator + .try_replicate_step_with_cancellation(&token) + .await + }, + async { + tokio::time::timeout(Duration::from_secs(5), async { + while !inject_called.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + }) + .await + .expect("inject_frame_with_cancellation was not called"); + cancel.cancel(); + } + ); + + assert!(matches!( + result.unwrap_err(), + Error::SyncCancelledForShutdown + )); + assert!(client_rolled_back.load(Ordering::SeqCst)); + assert!(injector_rolled_back.load(Ordering::SeqCst)); + assert_eq!(replicator.state, ReplicatorState::NeedHandshake); + } + + struct SuccessfulAfterCancelInjector { + inject_called: Arc, + rolled_back: Arc, + } + + impl SuccessfulAfterCancelInjector { + fn new() -> Self { + Self { + inject_called: Arc::new(AtomicBool::new(false)), + rolled_back: Arc::new(AtomicBool::new(false)), + } + } + } + + impl Injector for SuccessfulAfterCancelInjector { + async fn inject_frame( + &mut self, + _frame: RpcFrame, + ) -> std::result::Result, crate::injector::Error> { + Ok(None) + } + + async fn inject_frame_with_cancellation( + &mut self, + _frame: RpcFrame, + token: &CancellationToken, + ) -> std::result::Result, crate::injector::Error> { + self.inject_called.store(true, Ordering::SeqCst); + token.cancel(); + Ok(None) + } + + async fn rollback(&mut self) { + self.rolled_back.store(true, Ordering::SeqCst); + } + + async fn flush(&mut self) -> std::result::Result, crate::injector::Error> { + Ok(None) + } + + fn durable_frame_no(&mut self, _frame_no: u64) {} + } + + #[tokio::test] + async fn cancellation_after_successful_injection_is_not_swallowed() { + struct Client { + rolled_back: Arc, + } + + #[async_trait::async_trait] + impl ReplicatorClient for Client { + type FrameStream = + Pin> + Send + 'static>>; + + async fn handshake(&mut self) -> Result<(), Error> { + Ok(()) + } + + async fn next_frames(&mut self) -> Result { + Ok(Box::pin(tokio_stream::iter([Ok(RpcFrame { + data: Vec::new().into(), + timestamp: None, + durable_frame_no: None, + })]))) + } + + async fn snapshot(&mut self) -> Result { + unreachable!() + } + + async fn commit_frame_no(&mut self, _frame_no: FrameNo) -> Result<(), Error> { + Ok(()) + } + + fn committed_frame_no(&self) -> Option { + None + } + + fn rollback(&mut self) { + self.rolled_back.store(true, Ordering::SeqCst); + } + } + + let token = CancellationToken::new(); + let client_rolled_back = Arc::new(AtomicBool::new(false)); + let injector = SuccessfulAfterCancelInjector::new(); + let inject_called = injector.inject_called.clone(); + let injector_rolled_back = injector.rolled_back.clone(); + let mut replicator = Replicator::new( + Client { + rolled_back: client_rolled_back.clone(), + }, + injector, + ); + replicator.state = ReplicatorState::NeedFrames; + + let result = replicator + .try_replicate_step_with_cancellation(&token) + .await; + + assert!(inject_called.load(Ordering::SeqCst)); + assert!(matches!( + result.unwrap_err(), + Error::SyncCancelledForShutdown + )); + assert!(client_rolled_back.load(Ordering::SeqCst)); + assert!(injector_rolled_back.load(Ordering::SeqCst)); + assert_eq!(replicator.state, ReplicatorState::NeedHandshake); + } + + #[tokio::test] + async fn cancellation_while_waiting_for_frame_fetch_rolls_back() { + struct Client { + rolled_back: Arc, + } + + #[async_trait::async_trait] + impl ReplicatorClient for Client { + type FrameStream = + Pin> + Send + 'static>>; + + async fn handshake(&mut self) -> Result<(), Error> { + Ok(()) + } + + async fn next_frames(&mut self) -> Result { + std::future::pending().await + } + + async fn snapshot(&mut self) -> Result { + unreachable!() + } + + async fn commit_frame_no(&mut self, _frame_no: FrameNo) -> Result<(), Error> { + unreachable!() + } + + fn committed_frame_no(&self) -> Option { + None + } + + fn rollback(&mut self) { + self.rolled_back.store(true, Ordering::SeqCst); + } + } + + let token = CancellationToken::new(); + let client_rolled_back = Arc::new(AtomicBool::new(false)); + let injector = BlockingInjector::new(); + let injector_rolled_back = injector.rolled_back.clone(); + let mut replicator = Replicator::new( + Client { + rolled_back: client_rolled_back.clone(), + }, + injector, + ); + replicator.state = ReplicatorState::NeedFrames; + + let cancel = token.clone(); + let (result, _) = tokio::join!( + async { + replicator + .try_replicate_step_with_cancellation(&token) + .await + }, + async { + tokio::task::yield_now().await; + cancel.cancel(); + } + ); + + assert!(matches!( + result.unwrap_err(), + Error::SyncCancelledForShutdown + )); + assert!(client_rolled_back.load(Ordering::SeqCst)); + assert!(injector_rolled_back.load(Ordering::SeqCst)); + assert_eq!(replicator.state, ReplicatorState::NeedHandshake); + } + + #[tokio::test] + async fn cancellation_while_snapshot_request_pending_rolls_back() { + struct Client { + rolled_back: Arc, + } + + #[async_trait::async_trait] + impl ReplicatorClient for Client { + type FrameStream = + Pin> + Send + 'static>>; + + async fn handshake(&mut self) -> Result<(), Error> { + unreachable!() + } + + async fn next_frames(&mut self) -> Result { + unreachable!() + } + + async fn snapshot(&mut self) -> Result { + std::future::pending().await + } + + async fn commit_frame_no(&mut self, _frame_no: FrameNo) -> Result<(), Error> { + unreachable!() + } + + fn committed_frame_no(&self) -> Option { + None + } + + fn rollback(&mut self) { + self.rolled_back.store(true, Ordering::SeqCst); + } + } + + let token = CancellationToken::new(); + let client_rolled_back = Arc::new(AtomicBool::new(false)); + let injector = BlockingInjector::new(); + let injector_rolled_back = injector.rolled_back.clone(); + let mut replicator = Replicator::new( + Client { + rolled_back: client_rolled_back.clone(), + }, + injector, + ); + replicator.state = ReplicatorState::NeedSnapshot; + + let cancel = token.clone(); + let (result, _) = tokio::join!( + async { + replicator + .try_replicate_step_with_cancellation(&token) + .await + }, + async { + tokio::task::yield_now().await; + cancel.cancel(); + } + ); + + assert!(matches!( + result.unwrap_err(), + Error::SyncCancelledForShutdown + )); + assert!(client_rolled_back.load(Ordering::SeqCst)); + assert!(injector_rolled_back.load(Ordering::SeqCst)); + assert_eq!(replicator.state, ReplicatorState::NeedHandshake); + } + + #[tokio::test] + async fn cancellation_while_snapshot_pending_sleep_rolls_back() { + struct Client { + rolled_back: Arc, + calls: Arc, + } + + #[async_trait::async_trait] + impl ReplicatorClient for Client { + type FrameStream = + Pin> + Send + 'static>>; + + async fn handshake(&mut self) -> Result<(), Error> { + unreachable!() + } + + async fn next_frames(&mut self) -> Result { + unreachable!() + } + + async fn snapshot(&mut self) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(Error::SnapshotPending) + } + + async fn commit_frame_no(&mut self, _frame_no: FrameNo) -> Result<(), Error> { + unreachable!() + } + + fn committed_frame_no(&self) -> Option { + None + } + + fn rollback(&mut self) { + self.rolled_back.store(true, Ordering::SeqCst); + } + } + + let token = CancellationToken::new(); + let client_rolled_back = Arc::new(AtomicBool::new(false)); + let calls = Arc::new(AtomicUsize::new(0)); + let injector = BlockingInjector::new(); + let injector_rolled_back = injector.rolled_back.clone(); + let mut replicator = Replicator::new( + Client { + rolled_back: client_rolled_back.clone(), + calls: calls.clone(), + }, + injector, + ); + replicator.state = ReplicatorState::NeedSnapshot; + + let cancel = token.clone(); + let (result, _) = tokio::join!( + async { + replicator + .try_replicate_step_with_cancellation(&token) + .await + }, + async { + tokio::time::timeout(Duration::from_secs(5), async { + while calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("snapshot was not requested"); + cancel.cancel(); + } + ); + + assert!(matches!( + result.unwrap_err(), + Error::SyncCancelledForShutdown + )); + assert!(client_rolled_back.load(Ordering::SeqCst)); + assert!(injector_rolled_back.load(Ordering::SeqCst)); + assert_eq!(replicator.state, ReplicatorState::NeedHandshake); + } + + #[tokio::test] + async fn cancellation_while_snapshot_stream_pending_rolls_back() { + struct Client { + rolled_back: Arc, + } + + #[async_trait::async_trait] + impl ReplicatorClient for Client { + type FrameStream = + Pin> + Send + 'static>>; + + async fn handshake(&mut self) -> Result<(), Error> { + unreachable!() + } + + async fn next_frames(&mut self) -> Result { + unreachable!() + } + + async fn snapshot(&mut self) -> Result { + Ok(Box::pin(stream! { + std::future::pending::<()>().await; + yield Ok(RpcFrame { + data: Vec::new().into(), + timestamp: None, + durable_frame_no: None, + }); + })) + } + + async fn commit_frame_no(&mut self, _frame_no: FrameNo) -> Result<(), Error> { + unreachable!() + } + + fn committed_frame_no(&self) -> Option { + None + } + + fn rollback(&mut self) { + self.rolled_back.store(true, Ordering::SeqCst); + } + } + + let token = CancellationToken::new(); + let client_rolled_back = Arc::new(AtomicBool::new(false)); + let injector = BlockingInjector::new(); + let injector_rolled_back = injector.rolled_back.clone(); + let mut replicator = Replicator::new( + Client { + rolled_back: client_rolled_back.clone(), + }, + injector, + ); + replicator.state = ReplicatorState::NeedSnapshot; + + let cancel = token.clone(); + let (result, _) = tokio::join!( + async { + replicator + .try_replicate_step_with_cancellation(&token) + .await + }, + async { + tokio::task::yield_now().await; + cancel.cancel(); + } + ); + + assert!(matches!( + result.unwrap_err(), + Error::SyncCancelledForShutdown + )); + assert!(client_rolled_back.load(Ordering::SeqCst)); + assert!(injector_rolled_back.load(Ordering::SeqCst)); + assert_eq!(replicator.state, ReplicatorState::NeedHandshake); + } + #[tokio::test] async fn handshake_error_namespace_doesnt_exist() { let tmp = tempfile::NamedTempFile::new().unwrap(); diff --git a/libsql-server/src/namespace/configurator/replica.rs b/libsql-server/src/namespace/configurator/replica.rs index b1a108af73..0499100ce8 100644 --- a/libsql-server/src/namespace/configurator/replica.rs +++ b/libsql-server/src/namespace/configurator/replica.rs @@ -162,7 +162,8 @@ impl ConfigureNamespace for ReplicaConfigurator { e @ (Error::Internal(_) | Error::Client(_) | Error::PrimaryHandshakeTimeout - | Error::NeedSnapshot) => { + | Error::NeedSnapshot + | Error::SyncCancelledForShutdown) => { tracing::warn!("non-fatal replication error, retrying from last commit index: {e}"); }, Error::NoHandshake => { diff --git a/libsql-sqlite3/test/rust_suite/Cargo.toml b/libsql-sqlite3/test/rust_suite/Cargo.toml index 23c1735abe..bbd1b71bb3 100644 --- a/libsql-sqlite3/test/rust_suite/Cargo.toml +++ b/libsql-sqlite3/test/rust_suite/Cargo.toml @@ -2,6 +2,7 @@ name = "libsql_rust_suite" version = "0.2.0" edition = "2021" +rust-version = "1.85" [workspace] @@ -18,6 +19,17 @@ rustc-hash = "1" home = { version = "=0.5.9" } which = "=4.4.0" +# Keep transitive URL/idna dependencies compatible with the Rust toolchain used by +# extension-test CI. Newer releases require rustc 1.86+. +idna_adapter = "=1.2.1" +icu_collections = "=2.1.1" +icu_locale_core = "=2.1.1" +icu_normalizer = "=2.1.1" +icu_normalizer_data = "=2.1.1" +icu_properties = "=2.1.1" +icu_properties_data = "=2.1.1" +icu_provider = "=2.1.1" + [dev-dependencies.psm] version = "=0.1.21" diff --git a/libsql/Cargo.toml b/libsql/Cargo.toml index e8a7e0f909..f1cefea70a 100644 --- a/libsql/Cargo.toml +++ b/libsql/Cargo.toml @@ -88,6 +88,7 @@ replication = [ "dep:bytes", "dep:uuid", "dep:tokio-stream", + "dep:tokio-util", "dep:parking_lot", "dep:tokio", "dep:tonic", diff --git a/libsql/src/database.rs b/libsql/src/database.rs index b4da6171bf..520529fdbf 100644 --- a/libsql/src/database.rs +++ b/libsql/src/database.rs @@ -411,6 +411,15 @@ cfg_replication! { } } + /// Cancels the currently running v1 remote-replica sync, if any, for terminal shutdown/cleanup. + #[cfg(feature = "replication")] + pub fn cancel_current_sync_for_shutdown(&self) -> Result { + match &self.db_type { + DbType::Sync { db, encryption_config: _ } => db.cancel_current_sync_for_shutdown(), + _ => Err(Error::SyncNotSupported(format!("{:?}", self.db_type))), + } + } + /// Sync database from remote until it gets to a given replication_index or further, /// and returns the committed frame_no after syncing, if applicable. pub async fn sync_until(&self, replication_index: FrameNo) -> Result { diff --git a/libsql/src/errors.rs b/libsql/src/errors.rs index f612aa62f1..3ab73a2356 100644 --- a/libsql/src/errors.rs +++ b/libsql/src/errors.rs @@ -47,6 +47,8 @@ pub enum Error { RemoteSqliteFailure(i32, i32, String), #[error("replication error: {0}")] Replication(crate::BoxError), + #[error("sync cancelled for terminal shutdown")] + SyncCancelledForShutdown, #[error("path has invalid UTF-8")] InvalidUTF8Path, #[error("freeze is not supported in {0} mode.")] diff --git a/libsql/src/local/database.rs b/libsql/src/local/database.rs index 5d5eda0b06..74e1f36e37 100644 --- a/libsql/src/local/database.rs +++ b/libsql/src/local/database.rs @@ -393,6 +393,19 @@ impl Database { Ok(self.sync_oneshot().await?) } + #[cfg(feature = "replication")] + /// Cancels the currently running v1 remote-replica sync, if any, for terminal shutdown/cleanup. + pub fn cancel_current_sync_for_shutdown(&self) -> Result { + if let Some(ctx) = &self.replication_ctx { + Ok(ctx.replicator.cancel_current_sync_for_shutdown()) + } else { + Err(crate::errors::Error::Misuse( + "No replicator available. Use Database::with_replicator() to enable replication" + .to_string(), + )) + } + } + #[cfg(feature = "replication")] /// Return detailed logs about bytes synced with primary pub async fn get_sync_usage_stats(&self) -> Result { diff --git a/libsql/src/replication/mod.rs b/libsql/src/replication/mod.rs index 57ef8fc6a5..5baa249b68 100644 --- a/libsql/src/replication/mod.rs +++ b/libsql/src/replication/mod.rs @@ -1,7 +1,7 @@ //! Utilities used when using a replicated version of libsql. use std::path::PathBuf; -use std::sync::atomic::AtomicUsize; +use std::sync::atomic::{AtomicU64, AtomicUsize}; use std::sync::Arc; use std::time::Duration; @@ -14,8 +14,10 @@ use libsql_replication::rpc::proxy::{ query::Params, DescribeRequest, DescribeResult, ExecuteResults, Positional, Program, ProgramReq, Query, Step, }; +use parking_lot::Mutex as ParkingMutex; use tokio::sync::Mutex; use tokio::task::AbortHandle; +use tokio_util::sync::CancellationToken; use tracing::Instrument; use crate::database::EncryptionConfig; @@ -167,11 +169,82 @@ pub(crate) struct EmbeddedReplicator { replicator: Arc, SqliteInjector>>>, bg_abort: Option>, last_frames_synced: Arc, + current_sync_cancellation: CurrentSyncCancellation, +} + +#[derive(Clone, Default)] +struct CurrentSyncCancellation { + current: Arc>>, + next_id: Arc, +} + +#[derive(Clone)] +struct ActiveSyncToken { + id: u64, + token: CancellationToken, +} + +struct ActiveSyncCancellationGuard { + state: CurrentSyncCancellation, + active: ActiveSyncToken, +} + +impl CurrentSyncCancellation { + fn begin_sync(&self) -> ActiveSyncCancellationGuard { + let id = self + .next_id + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let active = ActiveSyncToken { + id, + token: CancellationToken::new(), + }; + *self.current.lock() = Some(active.clone()); + tracing::debug!("sync_cancellation_registered"); + ActiveSyncCancellationGuard { + state: self.clone(), + active, + } + } + + fn cancel_current_sync_for_shutdown(&self) -> bool { + let token = self.current.lock().as_ref().map(|active| active.token.clone()); + if let Some(token) = token { + tracing::info!("sync_cancellation_requested"); + token.cancel(); + true + } else { + false + } + } +} + +impl ActiveSyncCancellationGuard { + fn token(&self) -> &CancellationToken { + &self.active.token + } +} + +impl Drop for ActiveSyncCancellationGuard { + fn drop(&mut self) { + let mut current = self.state.current.lock(); + if current + .as_ref() + .is_some_and(|active| active.id == self.active.id) + { + *current = None; + tracing::debug!("sync_cancellation_settled"); + } + } } impl From for errors::Error { fn from(err: libsql_replication::replicator::Error) -> Self { - errors::Error::Replication(err.into()) + match err { + libsql_replication::replicator::Error::SyncCancelledForShutdown => { + errors::Error::SyncCancelledForShutdown + } + err => errors::Error::Replication(err.into()), + } } } @@ -197,6 +270,7 @@ impl EmbeddedReplicator { replicator, bg_abort: None, last_frames_synced: Arc::new(AtomicUsize::new(0)), + current_sync_cancellation: CurrentSyncCancellation::default(), }; if let Some(sync_duration) = perodic_sync { @@ -241,9 +315,15 @@ impl EmbeddedReplicator { replicator, bg_abort: None, last_frames_synced: Arc::new(AtomicUsize::new(0)), + current_sync_cancellation: CurrentSyncCancellation::default(), }) } + pub fn cancel_current_sync_for_shutdown(&self) -> bool { + self.current_sync_cancellation + .cancel_current_sync_for_shutdown() + } + pub async fn get_sync_usage_stats(&self) -> Result { let mut replicator = self.replicator.lock().await; match replicator.client_mut() { @@ -283,11 +363,14 @@ impl EmbeddedReplicator { )); } + let active_sync = self.current_sync_cancellation.begin_sync(); + let token = active_sync.token().clone(); + // we force a handshake to get the most up to date replication index from the primary. replicator.force_handshake(); loop { - match replicator.replicate().await { + match replicator.replicate_with_cancellation(&token).await { Err(libsql_replication::replicator::Error::Meta( libsql_replication::meta::Error::LogIncompatible, )) => { @@ -295,11 +378,11 @@ impl EmbeddedReplicator { // this time. tracing::debug!("re-replicating database after LogIncompatible error"); replicator - .replicate() + .replicate_with_cancellation(&token) .await - .map_err(|e| crate::Error::Replication(e.into()))?; + .map_err(crate::Error::from)?; } - Err(e) => return Err(crate::Error::Replication(e.into())), + Err(e) => return Err(crate::Error::from(e)), Ok(_) => { let Either::Left(client) = replicator.client_mut() else { unreachable!() diff --git a/libsql/src/replication/remote_client.rs b/libsql/src/replication/remote_client.rs index d55f68ec4d..d0eda382b4 100644 --- a/libsql/src/replication/remote_client.rs +++ b/libsql/src/replication/remote_client.rs @@ -14,6 +14,7 @@ use libsql_replication::rpc::replication::{ Frame as RpcFrame, verify_session_token, Frames, HelloRequest, HelloResponse, LogOffset, SESSION_TOKEN_KEY, }; use tokio_stream::Stream; +use tokio_util::sync::CancellationToken; use tonic::metadata::AsciiMetadataValue; use tonic::{Response, Status}; use zerocopy::FromBytes; @@ -111,11 +112,11 @@ impl RemoteClient { dirty: false, last_handshake_replication_index: None, prefetched_batch_log_entries: None, - handshake_latency_sum: Duration::default(), + handshake_latency_sum: Duration::ZERO, handshake_latency_count: 0, - frames_latency_sum: Duration::default(), + frames_latency_sum: Duration::ZERO, frames_latency_count: 0, - snapshot_latency_sum: Duration::default(), + snapshot_latency_sum: Duration::ZERO, snapshot_latency_count: 0, sync_stats: Arc::new(SyncStats::new()), }) @@ -172,7 +173,10 @@ impl RemoteClient { Ok(new_session) } - async fn do_handshake_with_prefetch(&mut self) -> (Result<(), Error>, Duration) { + async fn do_handshake_with_prefetch( + &mut self, + token: Option<&CancellationToken>, + ) -> (Result<(), Error>, Duration) { tracing::info!("Attempting to perform handshake with primary."); if let Some((Ok(frames), _)) = &self.prefetched_batch_log_entries { // TODO: check if it's ok to just do 4096 * frames.len() @@ -192,18 +196,40 @@ impl RemoteClient { wal_flavor: None, }); let mut client_clone = self.remote.clone(); - let hello_fut = time(async { - let res = self.remote.replication.hello(hello_req).await; - self.handle_handshake_response(res).await - }); let (hello, frames) = if prefetch { - let (hello, frames) = tokio::join!( - hello_fut, - time(client_clone.replication.batch_log_entries(log_offset_req)) - ); + let mut hello_client = self.remote.clone(); + let hello_rpc = hello_client.replication.hello(hello_req); + let frames_rpc = client_clone.replication.batch_log_entries(log_offset_req); + let hello = async { + let (hello_result, hello_time) = time(hello_rpc).await; + ( + self.handle_handshake_response(hello_result).await, + hello_time, + ) + }; + let (hello, frames) = if let Some(token) = token { + tokio::select! { + joined = async { tokio::join!(hello, time(frames_rpc)) } => joined, + _ = token.cancelled() => return (Err(Error::SyncCancelledForShutdown), Duration::ZERO), + } + } else { + tokio::join!(hello, time(frames_rpc)) + }; (hello, Some(frames)) } else { - (hello_fut.await, None) + let hello_rpc = self.remote.replication.hello(hello_req); + let (hello_result, hello_time) = if let Some(token) = token { + tokio::select! { + result = time(hello_rpc) => result, + _ = token.cancelled() => return (Err(Error::SyncCancelledForShutdown), Duration::ZERO), + } + } else { + time(hello_rpc).await + }; + ( + (self.handle_handshake_response(hello_result).await, hello_time), + None, + ) }; let mut prefetched_bytes = None; if let Some((Ok(frames), _)) = &frames { @@ -265,6 +291,7 @@ impl RemoteClient { async fn do_next_frames( &mut self, + token: Option<&CancellationToken>, ) -> ( Result<::FrameStream, Error>, Duration, @@ -276,7 +303,14 @@ impl RemoteClient { next_offset: self.next_offset(), wal_flavor: None, }); - let result = time(self.remote.replication.batch_log_entries(req)).await; + let result = if let Some(token) = token { + tokio::select! { + result = time(self.remote.replication.batch_log_entries(req)) => result, + _ = token.cancelled() => return (Err(Error::SyncCancelledForShutdown), Duration::ZERO), + } + } else { + time(self.remote.replication.batch_log_entries(req)).await + }; (result, false) } }; @@ -284,17 +318,24 @@ impl RemoteClient { (res, time) } - async fn do_snapshot(&mut self) -> Result<::FrameStream, Error> { + async fn do_snapshot( + &mut self, + token: Option<&CancellationToken>, + ) -> Result<::FrameStream, Error> { let req = self.make_request(LogOffset { next_offset: self.next_offset(), wal_flavor: None, }); let sync_stats = self.sync_stats.clone(); - let mut frames = self - .remote - .replication - .snapshot(req) - .await? + let response = if let Some(token) = token { + tokio::select! { + response = self.remote.replication.snapshot(req) => response, + _ = token.cancelled() => return Err(Error::SyncCancelledForShutdown), + } + } else { + self.remote.replication.snapshot(req).await + }?; + let mut frames = response .into_inner() .map_err(|e| e.into()) .map_ok(move |f| { @@ -307,7 +348,16 @@ impl RemoteClient { let frames = Pin::new(&mut frames); // the first frame is the one with the highest frame_no in the snapshot - if let Some(Ok(f)) = frames.peek().await { + let first = if let Some(token) = token { + tokio::select! { + first = frames.peek() => first, + _ = token.cancelled() => return Err(Error::SyncCancelledForShutdown), + } + } else { + frames.peek().await + }; + + if let Some(Ok(f)) = first { let header: FrameHeader = FrameHeader::read_from_prefix(&f.data[..]).unwrap(); self.last_received = Some(header.frame_no.get()); } @@ -349,7 +399,22 @@ impl ReplicatorClient for RemoteClient { /// Perform handshake with remote async fn handshake(&mut self) -> Result<(), Error> { - let (result, time) = self.do_handshake_with_prefetch().await; + let (result, time) = self.do_handshake_with_prefetch(None).await; + maybe_log( + time, + &mut self.handshake_latency_sum, + &mut self.handshake_latency_count, + &result, + "handshake", + ); + result + } + + async fn handshake_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result<(), Error> { + let (result, time) = self.do_handshake_with_prefetch(Some(token)).await; maybe_log( time, &mut self.handshake_latency_sum, @@ -362,7 +427,22 @@ impl ReplicatorClient for RemoteClient { /// Return a stream of frames to apply to the database async fn next_frames(&mut self) -> Result { - let (result, time) = self.do_next_frames().await; + let (result, time) = self.do_next_frames(None).await; + maybe_log( + time, + &mut self.frames_latency_sum, + &mut self.frames_latency_count, + &result, + "frames fetch", + ); + result + } + + async fn next_frames_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result { + let (result, time) = self.do_next_frames(Some(token)).await; maybe_log( time, &mut self.frames_latency_sum, @@ -376,7 +456,22 @@ impl ReplicatorClient for RemoteClient { /// Return a snapshot for the current replication index. Called after next_frame has returned a /// NeedSnapshot error async fn snapshot(&mut self) -> Result { - let (snapshot, time) = time(self.do_snapshot()).await; + let (snapshot, time) = time(self.do_snapshot(None)).await; + maybe_log( + time, + &mut self.snapshot_latency_sum, + &mut self.snapshot_latency_count, + &snapshot, + "snapshot fetch", + ); + snapshot + } + + async fn snapshot_with_cancellation( + &mut self, + token: &CancellationToken, + ) -> Result { + let (snapshot, time) = time(self.do_snapshot(Some(token))).await; maybe_log( time, &mut self.snapshot_latency_sum, diff --git a/libsql/tests/replication.rs b/libsql/tests/replication.rs index dff32b0fae..3d518d3a47 100644 --- a/libsql/tests/replication.rs +++ b/libsql/tests/replication.rs @@ -1,5 +1,7 @@ #![allow(deprecated)] +use std::sync::Arc; + use libsql::{replication::Frames, Database}; use libsql_replication::{ frame::{FrameBorrowed, FrameHeader, FrameMut}, @@ -8,6 +10,67 @@ use libsql_replication::{ const DB: &[u8] = include_bytes!("test.db"); +#[tokio::test] +async fn cancel_current_sync_for_shutdown_without_active_sync_returns_false() { + let tmp = tempfile::tempdir().unwrap(); + let db = Database::open_with_local_sync(tmp.path().join("data").to_str().unwrap(), None) + .await + .unwrap(); + + assert!(!db.cancel_current_sync_for_shutdown().unwrap()); +} + +#[tokio::test] +async fn cancel_current_sync_for_shutdown_cancels_active_remote_sync_and_clears_token() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + loop { + let Ok((_stream, _peer)) = listener.accept().await else { + break; + }; + } + }); + + let tmp = tempfile::tempdir().unwrap(); + let db = Arc::new( + Database::open_with_remote_sync( + tmp.path().join("data").to_str().unwrap(), + format!("http://{addr}"), + "token", + None, + ) + .await + .unwrap(), + ); + + let syncing = db.clone(); + let sync = tokio::spawn(async move { syncing.sync().await }); + + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if db.cancel_current_sync_for_shutdown().unwrap() { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("sync did not register a cancellable token"); + + let result = tokio::time::timeout(std::time::Duration::from_secs(5), sync) + .await + .expect("sync did not finish after cancellation") + .unwrap(); + assert!(matches!( + result.unwrap_err(), + libsql::Error::SyncCancelledForShutdown + )); + assert!(!db.cancel_current_sync_for_shutdown().unwrap()); + + server.abort(); +} + #[tokio::test] async fn inject_frames() { let tmp = tempfile::tempdir().unwrap();