@@ -3,8 +3,9 @@ use std::sync::Arc;
33use std:: { collections:: VecDeque , path:: PathBuf } ;
44
55use parking_lot:: Mutex ;
6- use rusqlite:: OpenFlags ;
6+ use rusqlite:: { ErrorCode , OpenFlags } ;
77use tokio:: task:: spawn_blocking;
8+ use tokio_util:: sync:: CancellationToken ;
89
910use crate :: frame:: { Frame , FrameNo } ;
1011use crate :: rpc:: replication:: Frame as RpcFrame ;
@@ -23,6 +24,32 @@ pub type FrameBuffer = Arc<Mutex<VecDeque<Frame>>>;
2324
2425pub struct SqliteInjector {
2526 pub ( in super :: super ) inner : Arc < Mutex < SqliteInjectorInner > > ,
27+ interrupt_handle : InjectorInterruptHandle ,
28+ }
29+
30+ #[ derive( Clone , Default ) ]
31+ struct InjectorInterruptHandle {
32+ current : Arc < Mutex < Option < rusqlite:: InterruptHandle > > > ,
33+ }
34+
35+ impl InjectorInterruptHandle {
36+ fn set ( & self , handle : rusqlite:: InterruptHandle ) {
37+ * self . current . lock ( ) = Some ( handle) ;
38+ }
39+
40+ fn interrupt ( & self ) {
41+ if let Some ( handle) = self . current . lock ( ) . as_ref ( ) {
42+ handle. interrupt ( ) ;
43+ }
44+ }
45+ }
46+
47+ fn is_cancellation_interrupt ( error : & Error ) -> bool {
48+ matches ! (
49+ error,
50+ Error :: Sqlite ( e)
51+ if e. sqlite_error_code( ) == Some ( ErrorCode :: OperationInterrupted )
52+ )
2653}
2754
2855impl Injector for SqliteInjector {
@@ -35,6 +62,41 @@ impl Injector for SqliteInjector {
3562 . unwrap ( )
3663 }
3764
65+ async fn inject_frame_with_cancellation (
66+ & mut self ,
67+ frame : RpcFrame ,
68+ token : & CancellationToken ,
69+ ) -> Result < Option < FrameNo > > {
70+ let inner = self . inner . clone ( ) ;
71+ let interrupt_handle = self . interrupt_handle . clone ( ) ;
72+ let frame =
73+ Frame :: try_from ( & frame. data [ ..] ) . map_err ( |e| Error :: FatalInjectError ( e. into ( ) ) ) ?;
74+ let mut join = spawn_blocking ( move || inner. lock ( ) . inject_frame ( frame) ) ;
75+
76+ tokio:: select! {
77+ biased;
78+
79+ result = & mut join => result. map_err( |e| Error :: FatalInjectError ( e. into( ) ) ) ?,
80+ _ = token. cancelled( ) => {
81+ interrupt_handle. interrupt( ) ;
82+ match join. await {
83+ Ok ( Ok ( result) ) => Ok ( result) ,
84+ Ok ( Err ( Error :: SyncCancelledForShutdown ) ) => {
85+ Err ( Error :: SyncCancelledForShutdown )
86+ }
87+ Ok ( Err ( e) ) if is_cancellation_interrupt( & e) => {
88+ Err ( Error :: SyncCancelledForShutdown )
89+ }
90+ Ok ( Err ( e) ) => {
91+ tracing:: warn!( error = %e, "injector_error_after_sync_cancellation" ) ;
92+ Err ( e)
93+ }
94+ Err ( e) => Err ( Error :: FatalInjectError ( e. into( ) ) ) ,
95+ }
96+ }
97+ }
98+ }
99+
38100 async fn rollback ( & mut self ) {
39101 let inner = self . inner . clone ( ) ;
40102 spawn_blocking ( move || inner. lock ( ) . rollback ( ) )
@@ -58,14 +120,23 @@ impl SqliteInjector {
58120 auto_checkpoint : u32 ,
59121 encryption_config : Option < libsql_sys:: EncryptionConfig > ,
60122 ) -> super :: Result < Self > {
123+ let interrupt_handle = InjectorInterruptHandle :: default ( ) ;
124+ let inner_interrupt_handle = interrupt_handle. clone ( ) ;
61125 let inner = spawn_blocking ( move || {
62- SqliteInjectorInner :: new ( path, capacity, auto_checkpoint, encryption_config)
126+ SqliteInjectorInner :: new (
127+ path,
128+ capacity,
129+ auto_checkpoint,
130+ encryption_config,
131+ inner_interrupt_handle,
132+ )
63133 } )
64134 . await
65135 . unwrap ( ) ?;
66136
67137 Ok ( Self {
68138 inner : Arc :: new ( Mutex :: new ( inner) ) ,
139+ interrupt_handle,
69140 } )
70141 }
71142}
@@ -86,6 +157,7 @@ pub(in super::super) struct SqliteInjectorInner {
86157 path : PathBuf ,
87158 encryption_config : Option < libsql_sys:: EncryptionConfig > ,
88159 auto_checkpoint : u32 ,
160+ interrupt_handle : InjectorInterruptHandle ,
89161}
90162
91163/// Methods from this trait are called before and after performing a frame injection.
@@ -98,6 +170,7 @@ impl SqliteInjectorInner {
98170 capacity : usize ,
99171 auto_checkpoint : u32 ,
100172 encryption_config : Option < libsql_sys:: EncryptionConfig > ,
173+ interrupt_handle : InjectorInterruptHandle ,
101174 ) -> Result < Self , Error > {
102175 let path = path. as_ref ( ) . to_path_buf ( ) ;
103176
@@ -113,6 +186,7 @@ impl SqliteInjectorInner {
113186 auto_checkpoint,
114187 encryption_config. clone ( ) ,
115188 ) ?;
189+ interrupt_handle. set ( connection. get_interrupt_handle ( ) ) ;
116190
117191 Ok ( Self {
118192 is_txn : false ,
@@ -124,6 +198,7 @@ impl SqliteInjectorInner {
124198 path,
125199 encryption_config,
126200 auto_checkpoint,
201+ interrupt_handle,
127202 } )
128203 }
129204
@@ -235,6 +310,7 @@ impl SqliteInjectorInner {
235310 self . auto_checkpoint ,
236311 self . encryption_config . clone ( ) ,
237312 ) ?;
313+ self . interrupt_handle . set ( new_conn. get_interrupt_handle ( ) ) ;
238314
239315 let _ = std:: mem:: replace ( & mut * conn, new_conn) ;
240316 }
@@ -266,6 +342,9 @@ impl SqliteInjectorInner {
266342mod test {
267343 use crate :: frame:: FrameBorrowed ;
268344 use std:: mem:: size_of;
345+ use std:: time:: Duration ;
346+
347+ use tokio:: sync:: oneshot;
269348
270349 use super :: * ;
271350 /// this this is generated by creating a table test, inserting 5 rows into it, and then
@@ -277,12 +356,87 @@ mod test {
277356 . map ( |b| Frame :: try_from ( b) . unwrap ( ) )
278357 }
279358
359+ #[ test]
360+ fn cancellation_interrupt_errors_are_classified_narrowly ( ) {
361+ let interrupted = Error :: Sqlite ( rusqlite:: Error :: SqliteFailure (
362+ rusqlite:: ffi:: Error {
363+ code : ErrorCode :: OperationInterrupted ,
364+ extended_code : rusqlite:: ffi:: SQLITE_INTERRUPT ,
365+ } ,
366+ None ,
367+ ) ) ;
368+ assert ! ( is_cancellation_interrupt( & interrupted) ) ;
369+
370+ let corrupt = Error :: Sqlite ( rusqlite:: Error :: SqliteFailure (
371+ rusqlite:: ffi:: Error {
372+ code : ErrorCode :: DatabaseCorrupt ,
373+ extended_code : rusqlite:: ffi:: SQLITE_CORRUPT ,
374+ } ,
375+ None ,
376+ ) ) ;
377+ assert ! ( !is_cancellation_interrupt( & corrupt) ) ;
378+ }
379+
380+ #[ tokio:: test]
381+ async fn cancellation_waits_for_real_blocking_injection_task_to_settle ( ) {
382+ let temp = tempfile:: tempdir ( ) . unwrap ( ) ;
383+ let mut injector = SqliteInjector :: new ( temp. path ( ) . join ( "data" ) , 10 , 10000 , None )
384+ . await
385+ . unwrap ( ) ;
386+ let frame = RpcFrame {
387+ data : wal_log ( ) . next ( ) . unwrap ( ) . bytes ( ) ,
388+ timestamp : None ,
389+ durable_frame_no : None ,
390+ } ;
391+
392+ let inner = injector. inner . clone ( ) ;
393+ let ( locked_tx, locked_rx) = oneshot:: channel ( ) ;
394+ let ( release_tx, release_rx) = std:: sync:: mpsc:: channel ( ) ;
395+ let blocker = tokio:: task:: spawn_blocking ( move || {
396+ let _guard = inner. lock ( ) ;
397+ let _ = locked_tx. send ( ( ) ) ;
398+ let _ = release_rx. recv ( ) ;
399+ } ) ;
400+ locked_rx. await . unwrap ( ) ;
401+
402+ let token = CancellationToken :: new ( ) ;
403+ let cancel = token. clone ( ) ;
404+ let injection = tokio:: spawn ( async move {
405+ let result = injector. inject_frame_with_cancellation ( frame, & token) . await ;
406+ ( injector, result)
407+ } ) ;
408+
409+ cancel. cancel ( ) ;
410+ tokio:: time:: sleep ( Duration :: from_millis ( 50 ) ) . await ;
411+ assert ! (
412+ !injection. is_finished( ) ,
413+ "cancelled injection returned before the blocking task settled"
414+ ) ;
415+
416+ release_tx. send ( ( ) ) . unwrap ( ) ;
417+ blocker. await . unwrap ( ) ;
418+ let ( _injector, result) = tokio:: time:: timeout ( Duration :: from_secs ( 5 ) , injection)
419+ . await
420+ . expect ( "cancelled injection did not settle after releasing the blocking task" )
421+ . unwrap ( ) ;
422+ assert ! ( matches!(
423+ result,
424+ Ok ( _) | Err ( Error :: SyncCancelledForShutdown )
425+ ) ) ;
426+ }
427+
280428 #[ test]
281429 fn test_simple_inject_frames ( ) {
282430 let temp = tempfile:: tempdir ( ) . unwrap ( ) ;
283431
284- let mut injector =
285- SqliteInjectorInner :: new ( temp. path ( ) . join ( "data" ) , 10 , 10000 , None ) . unwrap ( ) ;
432+ let mut injector = SqliteInjectorInner :: new (
433+ temp. path ( ) . join ( "data" ) ,
434+ 10 ,
435+ 10000 ,
436+ None ,
437+ InjectorInterruptHandle :: default ( ) ,
438+ )
439+ . unwrap ( ) ;
286440 let log = wal_log ( ) ;
287441 for frame in log {
288442 injector. inject_frame ( frame) . unwrap ( ) ;
@@ -302,8 +456,14 @@ mod test {
302456 let temp = tempfile:: tempdir ( ) . unwrap ( ) ;
303457
304458 // inject one frame at a time
305- let mut injector =
306- SqliteInjectorInner :: new ( temp. path ( ) . join ( "data" ) , 1 , 10000 , None ) . unwrap ( ) ;
459+ let mut injector = SqliteInjectorInner :: new (
460+ temp. path ( ) . join ( "data" ) ,
461+ 1 ,
462+ 10000 ,
463+ None ,
464+ InjectorInterruptHandle :: default ( ) ,
465+ )
466+ . unwrap ( ) ;
307467 let log = wal_log ( ) ;
308468 for frame in log {
309469 injector. inject_frame ( frame) . unwrap ( ) ;
@@ -323,8 +483,14 @@ mod test {
323483 let temp = tempfile:: tempdir ( ) . unwrap ( ) ;
324484
325485 // inject one frame at a time
326- let mut injector =
327- SqliteInjectorInner :: new ( temp. path ( ) . join ( "data" ) , 10 , 1000 , None ) . unwrap ( ) ;
486+ let mut injector = SqliteInjectorInner :: new (
487+ temp. path ( ) . join ( "data" ) ,
488+ 10 ,
489+ 1000 ,
490+ None ,
491+ InjectorInterruptHandle :: default ( ) ,
492+ )
493+ . unwrap ( ) ;
328494 let mut frames = wal_log ( ) ;
329495
330496 assert ! ( injector
0 commit comments