1- use std:: {
2- sync:: {
3- atomic:: { AtomicU64 , Ordering } ,
4- Arc ,
5- } ,
6- time:: Duration ,
7- } ;
1+ use std:: { sync:: Arc , time:: Duration } ;
82
9- use log:: { error, info, warn} ;
3+ use futures:: TryFutureExt as _;
4+ use log:: { error, info} ;
105use spacetimedb_commitlog:: payload:: {
116 txdata:: { Mutations , Ops } ,
127 Txdata ,
@@ -18,27 +13,32 @@ use spacetimedb_lib::Identity;
1813use spacetimedb_primitives:: TableId ;
1914use tokio:: {
2015 runtime,
21- sync:: mpsc:: { channel, unbounded_channel, Receiver , Sender , UnboundedReceiver , UnboundedSender } ,
22- time:: { timeout, Instant } ,
16+ sync:: {
17+ futures:: OwnedNotified ,
18+ mpsc:: { channel, unbounded_channel, Receiver , Sender , UnboundedReceiver , UnboundedSender } ,
19+ oneshot, Notify ,
20+ } ,
21+ time:: timeout,
2322} ;
2423
25- use crate :: db:: { lock_file :: LockFile , persistence:: Durability } ;
24+ use crate :: db:: persistence:: Durability ;
2625
2726/// A request to persist a transaction or to terminate the actor.
2827pub struct DurabilityRequest {
2928 reducer_context : Option < ReducerContext > ,
3029 tx_data : Arc < TxData > ,
3130}
3231
32+ type ShutdownReply = oneshot:: Sender < OwnedNotified > ;
33+
3334/// Represents a handle to a background task that persists transactions
3435/// according to the [`Durability`] policy provided.
3536///
3637/// This exists to avoid holding a transaction lock while
3738/// preparing the [TxData] for processing by the [Durability] layer.
3839pub struct DurabilityWorker {
3940 request_tx : UnboundedSender < DurabilityRequest > ,
40- requested_tx_offset : AtomicU64 ,
41- shutdown : Sender < ( ) > ,
41+ shutdown : Sender < ShutdownReply > ,
4242 durability : Arc < Durability > ,
4343 runtime : runtime:: Handle ,
4444}
@@ -64,7 +64,6 @@ impl DurabilityWorker {
6464
6565 Self {
6666 request_tx,
67- requested_tx_offset : AtomicU64 :: new ( 0 ) ,
6867 shutdown : shutdown_tx,
6968 durability,
7069 runtime,
@@ -97,14 +96,6 @@ impl DurabilityWorker {
9796 reducer_context,
9897 tx_data : tx_data. clone ( ) ,
9998 } )
100- . inspect ( |( ) | {
101- // If `tx_data` has a `None` tx offset, the actor will ignore it.
102- // Otherwise, record the offset as requested, so that
103- // [Self::shutdown] can determine when the queue is drained.
104- if let Some ( tx_offset) = tx_data. tx_offset ( ) {
105- self . requested_tx_offset . fetch_max ( tx_offset, Ordering :: SeqCst ) ;
106- }
107- } )
10899 . expect ( HUNG_UP ) ;
109100 }
110101
@@ -121,22 +112,26 @@ impl DurabilityWorker {
121112 ///
122113 /// # Panics
123114 ///
124- /// If [Self::request_durability] is called after [Self::shutdown], the
125- /// former will panic.
126- pub async fn shutdown ( & self ) -> anyhow:: Result < TxOffset > {
127- // Request the actor to shutdown.
128- // Ignore send errors -- in that case a shutdown is already in progress.
129- let _ = self . shutdown . try_send ( ( ) ) ;
130- // Wait for the request channel to be closed.
131- self . request_tx . closed ( ) . await ;
132- // Load the latest tx offset and wait for it to become durable.
133- let latest_tx_offset = self . requested_tx_offset . load ( Ordering :: SeqCst ) ;
134- let durable_offset = self . durable_tx_offset ( ) . wait_for ( latest_tx_offset) . await ?;
135-
136- Ok ( durable_offset)
115+ /// After this method was called, calling [Self::request_durability]
116+ /// will panic.
117+ pub async fn close ( & self ) -> Option < TxOffset > {
118+ let ( done_tx, done_rx) = oneshot:: channel ( ) ;
119+ // Channel errors can be ignored.
120+ // It just means that the actor already exited.
121+ let _ = self
122+ . shutdown
123+ . send ( done_tx)
124+ . map_err ( drop)
125+ . and_then ( |( ) | done_rx. map_err ( drop) )
126+ . and_then ( |done| async move {
127+ done. await ;
128+ Ok ( ( ) )
129+ } )
130+ . await ;
131+ self . durability . close ( ) . await
137132 }
138133
139- /// Consume `self` and run [Self::shutdown ].
134+ /// Consume `self` and run [Self::close ].
140135 ///
141136 /// The `lock_file` is not dropped until the shutdown is complete (either
142137 /// successfully or unsuccessfully). This is to prevent the database to be
@@ -151,59 +146,44 @@ impl DurabilityWorker {
151146 /// owning this durability worker.
152147 ///
153148 /// This method is used in the `Drop` impl for [crate::db::relational_db::RelationalDB].
154- pub ( super ) fn spawn_shutdown ( self , database_identity : Identity , lock_file : LockFile ) {
149+ pub ( super ) fn spawn_close ( self , database_identity : Identity ) {
155150 let rt = self . runtime . clone ( ) ;
156- let mut shutdown = rt. spawn ( async move { self . shutdown ( ) . await } ) ;
157151 rt. spawn ( async move {
158152 let label = format ! ( "[{database_identity}]" ) ;
159- let start = Instant :: now ( ) ;
160- loop {
161- // Warn every 5s if the shutdown doesn't appear to make progress.
162- // The backing durability could still be writing to disk,
163- // but we can't cancel it from here,
164- // so dropping the lock file would be unsafe.
165- match timeout ( Duration :: from_secs ( 5 ) , & mut shutdown) . await {
166- Err ( _elapsed) => {
167- let since = start. elapsed ( ) . as_secs_f32 ( ) ;
168- error ! ( "{label} waiting for durability worker shutdown since {since}s" , ) ;
169- continue ;
170- }
171- Ok ( res) => {
172- let Ok ( done) = res else {
173- warn ! ( "{label} durability worker shutdown cancelled" ) ;
174- break ;
175- } ;
176- match done {
177- Ok ( offset) => info ! ( "{label} durability worker shut down at tx offset: {offset}" ) ,
178- Err ( e) => warn ! ( "{label} error shutting down durability worker: {e:#}" ) ,
179- }
180- break ;
181- }
153+ // Apply a timeout, in case `Durability::close` doesn't terminate
154+ // as advertised. This is a bug, but panicking here would not
155+ // unwind at the call site.
156+ match timeout ( Duration :: from_secs ( 10 ) , self . close ( ) ) . await {
157+ Err ( _elapsed) => {
158+ error ! ( "{label} timeout waiting for durability worker shutdown" ) ;
159+ }
160+ Ok ( offset) => {
161+ info ! ( "{label} durability worker shut down at tx offset: {offset:?}" ) ;
182162 }
183163 }
184- drop ( lock_file) ;
185164 } ) ;
186165 }
187166}
188167
189168pub struct DurabilityWorkerActor {
190169 request_rx : UnboundedReceiver < DurabilityRequest > ,
191- shutdown : Receiver < ( ) > ,
170+ shutdown : Receiver < ShutdownReply > ,
192171 durability : Arc < Durability > ,
193172}
194173
195174impl DurabilityWorkerActor {
196175 /// Processes requests to do durability.
197176 async fn run ( mut self ) {
177+ let done = scopeguard:: guard ( Arc :: new ( Notify :: new ( ) ) , |done| done. notify_waiters ( ) ) ;
198178 loop {
199179 tokio:: select! {
200180 // Biased towards the shutdown channel,
201181 // so that adding new requests is prevented promptly.
202182 biased;
203183
204- Some ( ( ) ) = self . shutdown. recv( ) => {
184+ Some ( reply ) = self . shutdown. recv( ) => {
205185 self . request_rx. close( ) ;
206- self . shutdown . close ( ) ;
186+ let _ = reply . send ( done . clone ( ) . notified_owned ( ) ) ;
207187 } ,
208188
209189 req = self . request_rx. recv( ) => {
@@ -214,6 +194,8 @@ impl DurabilityWorkerActor {
214194 }
215195 }
216196 }
197+
198+ info ! ( "durability worker actor done" ) ;
217199 }
218200
219201 pub fn do_durability ( durability : & Durability , reducer_context : Option < ReducerContext > , tx_data : & TxData ) {
@@ -280,6 +262,7 @@ impl DurabilityWorkerActor {
280262mod tests {
281263 use std:: { pin:: pin, task:: Poll } ;
282264
265+ use futures:: FutureExt as _;
283266 use pretty_assertions:: assert_matches;
284267 use spacetimedb_sats:: product;
285268 use tokio:: sync:: watch;
@@ -318,6 +301,22 @@ mod tests {
318301 fn durable_tx_offset ( & self ) -> DurableOffset {
319302 self . durable . subscribe ( ) . into ( )
320303 }
304+
305+ fn close ( & self ) -> spacetimedb_durability:: Close {
306+ let mut durable = self . durable . subscribe ( ) ;
307+ let appended = self . appended . subscribe ( ) ;
308+ async move {
309+ let durable_offset = durable
310+ . wait_for ( |durable| match ( * durable) . zip ( * appended. borrow ( ) ) {
311+ Some ( ( durable_offset, appended_offset) ) => durable_offset >= appended_offset,
312+ None => false ,
313+ } )
314+ . await
315+ . unwrap ( ) ;
316+ * durable_offset
317+ }
318+ . boxed ( )
319+ }
321320 }
322321
323322 #[ tokio:: test]
@@ -333,13 +332,8 @@ mod tests {
333332
334333 worker. request_durability ( None , & Arc :: new ( txdata) ) ;
335334 }
336- assert_eq ! (
337- 10 ,
338- worker. requested_tx_offset. load( Ordering :: Relaxed ) ,
339- "worker should have requested up to tx offset 10"
340- ) ;
341335
342- let shutdown = worker. shutdown ( ) ;
336+ let shutdown = worker. close ( ) ;
343337 let mut shutdown_fut = pin ! ( shutdown) ;
344338 assert_matches ! (
345339 futures:: poll!( & mut shutdown_fut) ,
@@ -357,7 +351,7 @@ mod tests {
357351 durability. mark_durable ( 10 ) . await ;
358352 assert_matches ! (
359353 futures:: poll!( & mut shutdown_fut) ,
360- Poll :: Ready ( Ok ( 10 ) ) ,
354+ Poll :: Ready ( Some ( 10 ) ) ,
361355 "shutdown returns, reporting durable offset at 10"
362356 ) ;
363357 assert_eq ! (
0 commit comments