@@ -5,6 +5,7 @@ use std::{collections::VecDeque, path::PathBuf};
55use parking_lot:: Mutex ;
66use rusqlite:: 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,24 @@ 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+ }
2645}
2746
2847impl Injector for SqliteInjector {
@@ -35,6 +54,32 @@ impl Injector for SqliteInjector {
3554 . unwrap ( )
3655 }
3756
57+ async fn inject_frame_with_cancellation (
58+ & mut self ,
59+ frame : RpcFrame ,
60+ token : & CancellationToken ,
61+ ) -> Result < Option < FrameNo > > {
62+ let inner = self . inner . clone ( ) ;
63+ let interrupt_handle = self . interrupt_handle . clone ( ) ;
64+ let frame =
65+ Frame :: try_from ( & frame. data [ ..] ) . map_err ( |e| Error :: FatalInjectError ( e. into ( ) ) ) ?;
66+ let mut join = spawn_blocking ( move || inner. lock ( ) . inject_frame ( frame) ) ;
67+
68+ tokio:: select! {
69+ biased;
70+
71+ result = & mut join => result. unwrap( ) ,
72+ _ = token. cancelled( ) => {
73+ interrupt_handle. interrupt( ) ;
74+ match join. await {
75+ Ok ( Ok ( result) ) => Ok ( result) ,
76+ Ok ( Err ( _) ) => Err ( Error :: SyncCancelledForShutdown ) ,
77+ Err ( e) => Err ( Error :: FatalInjectError ( e. into( ) ) ) ,
78+ }
79+ }
80+ }
81+ }
82+
3883 async fn rollback ( & mut self ) {
3984 let inner = self . inner . clone ( ) ;
4085 spawn_blocking ( move || inner. lock ( ) . rollback ( ) )
@@ -58,14 +103,23 @@ impl SqliteInjector {
58103 auto_checkpoint : u32 ,
59104 encryption_config : Option < libsql_sys:: EncryptionConfig > ,
60105 ) -> super :: Result < Self > {
106+ let interrupt_handle = InjectorInterruptHandle :: default ( ) ;
107+ let inner_interrupt_handle = interrupt_handle. clone ( ) ;
61108 let inner = spawn_blocking ( move || {
62- SqliteInjectorInner :: new ( path, capacity, auto_checkpoint, encryption_config)
109+ SqliteInjectorInner :: new (
110+ path,
111+ capacity,
112+ auto_checkpoint,
113+ encryption_config,
114+ inner_interrupt_handle,
115+ )
63116 } )
64117 . await
65118 . unwrap ( ) ?;
66119
67120 Ok ( Self {
68121 inner : Arc :: new ( Mutex :: new ( inner) ) ,
122+ interrupt_handle,
69123 } )
70124 }
71125}
@@ -86,6 +140,7 @@ pub(in super::super) struct SqliteInjectorInner {
86140 path : PathBuf ,
87141 encryption_config : Option < libsql_sys:: EncryptionConfig > ,
88142 auto_checkpoint : u32 ,
143+ interrupt_handle : InjectorInterruptHandle ,
89144}
90145
91146/// Methods from this trait are called before and after performing a frame injection.
@@ -98,6 +153,7 @@ impl SqliteInjectorInner {
98153 capacity : usize ,
99154 auto_checkpoint : u32 ,
100155 encryption_config : Option < libsql_sys:: EncryptionConfig > ,
156+ interrupt_handle : InjectorInterruptHandle ,
101157 ) -> Result < Self , Error > {
102158 let path = path. as_ref ( ) . to_path_buf ( ) ;
103159
@@ -113,6 +169,7 @@ impl SqliteInjectorInner {
113169 auto_checkpoint,
114170 encryption_config. clone ( ) ,
115171 ) ?;
172+ interrupt_handle. set ( connection. get_interrupt_handle ( ) ) ;
116173
117174 Ok ( Self {
118175 is_txn : false ,
@@ -124,6 +181,7 @@ impl SqliteInjectorInner {
124181 path,
125182 encryption_config,
126183 auto_checkpoint,
184+ interrupt_handle,
127185 } )
128186 }
129187
@@ -235,6 +293,7 @@ impl SqliteInjectorInner {
235293 self . auto_checkpoint ,
236294 self . encryption_config . clone ( ) ,
237295 ) ?;
296+ self . interrupt_handle . set ( new_conn. get_interrupt_handle ( ) ) ;
238297
239298 let _ = std:: mem:: replace ( & mut * conn, new_conn) ;
240299 }
@@ -281,8 +340,14 @@ mod test {
281340 fn test_simple_inject_frames ( ) {
282341 let temp = tempfile:: tempdir ( ) . unwrap ( ) ;
283342
284- let mut injector =
285- SqliteInjectorInner :: new ( temp. path ( ) . join ( "data" ) , 10 , 10000 , None ) . unwrap ( ) ;
343+ let mut injector = SqliteInjectorInner :: new (
344+ temp. path ( ) . join ( "data" ) ,
345+ 10 ,
346+ 10000 ,
347+ None ,
348+ InjectorInterruptHandle :: default ( ) ,
349+ )
350+ . unwrap ( ) ;
286351 let log = wal_log ( ) ;
287352 for frame in log {
288353 injector. inject_frame ( frame) . unwrap ( ) ;
@@ -302,8 +367,14 @@ mod test {
302367 let temp = tempfile:: tempdir ( ) . unwrap ( ) ;
303368
304369 // inject one frame at a time
305- let mut injector =
306- SqliteInjectorInner :: new ( temp. path ( ) . join ( "data" ) , 1 , 10000 , None ) . unwrap ( ) ;
370+ let mut injector = SqliteInjectorInner :: new (
371+ temp. path ( ) . join ( "data" ) ,
372+ 1 ,
373+ 10000 ,
374+ None ,
375+ InjectorInterruptHandle :: default ( ) ,
376+ )
377+ . unwrap ( ) ;
307378 let log = wal_log ( ) ;
308379 for frame in log {
309380 injector. inject_frame ( frame) . unwrap ( ) ;
@@ -323,8 +394,14 @@ mod test {
323394 let temp = tempfile:: tempdir ( ) . unwrap ( ) ;
324395
325396 // inject one frame at a time
326- let mut injector =
327- SqliteInjectorInner :: new ( temp. path ( ) . join ( "data" ) , 10 , 1000 , None ) . unwrap ( ) ;
397+ let mut injector = SqliteInjectorInner :: new (
398+ temp. path ( ) . join ( "data" ) ,
399+ 10 ,
400+ 1000 ,
401+ None ,
402+ InjectorInterruptHandle :: default ( ) ,
403+ )
404+ . unwrap ( ) ;
328405 let mut frames = wal_log ( ) ;
329406
330407 assert ! ( injector
0 commit comments