@@ -23,8 +23,9 @@ use firewood_storage::{
2323 CheckOpt , CheckerReport , Committed , FileBacked , FileIoError , HashedNodeReader ,
2424 ImmutableProposal , NodeHashAlgorithm , NodeStore , Parentable , ReadableStorage , TrieReader ,
2525} ;
26+ use nonzero_ext:: nonzero;
2627use std:: io:: Write ;
27- use std:: num:: NonZeroUsize ;
28+ use std:: num:: { NonZeroU64 , NonZeroUsize } ;
2829use std:: path:: Path ;
2930use std:: sync:: Arc ;
3031use thiserror:: Error ;
@@ -135,6 +136,9 @@ pub struct DbConfig {
135136 /// Whether to enable `RootStore`.
136137 #[ builder( default = false ) ]
137138 pub root_store : bool ,
139+ /// The maximum number of unpersisted revisions that can exist at a given time.
140+ #[ builder( default = nonzero!( 1u64 ) ) ]
141+ pub deferred_persistence_commit_count : NonZeroU64 ,
138142}
139143
140144#[ derive( Debug ) ]
@@ -182,6 +186,7 @@ impl Db {
182186 . create ( cfg. create_if_missing )
183187 . truncate ( cfg. truncate )
184188 . root_store ( cfg. root_store )
189+ . deferred_persistence_commit_count ( cfg. deferred_persistence_commit_count )
185190 . manager ( cfg. manager )
186191 . build ( ) ;
187192 let manager = RevisionManager :: new ( config_manager) ?;
@@ -337,9 +342,9 @@ impl Db {
337342
338343 /// Closes the database gracefully.
339344 ///
340- /// This method shuts down the background persistence worker. If not called
341- /// explicitly, `Drop` will attempt a best-effort shutdown but cannot report
342- /// errors.
345+ /// This method shuts down the background persistence worker and persists
346+ /// the last committed revision. If not called explicitly, `Drop` will
347+ /// attempt a best-effort shutdown but cannot report errors.
343348 pub fn close ( mut self ) -> Result < ( ) , api:: Error > {
344349 self . manager . close ( ) . map_err ( Into :: into)
345350 }
@@ -423,18 +428,19 @@ mod test {
423428 use core:: iter:: Take ;
424429 use std:: collections:: HashMap ;
425430 use std:: iter:: Peekable ;
426- use std:: num:: NonZeroUsize ;
431+ use std:: num:: { NonZeroU64 , NonZeroUsize } ;
427432 use std:: ops:: { Deref , DerefMut } ;
428433 use std:: path:: Path ;
429434
430435 use firewood_storage:: {
431436 CheckOpt , CheckerError , HashedNodeReader , IntoHashType , LinearAddress , MaybePersistedNode ,
432437 NodeStore , RootReader , TrieHash ,
433438 } ;
439+ use nonzero_ext:: nonzero;
434440
435441 use crate :: db:: { Db , Proposal , UseParallel } ;
436442 use crate :: manager:: RevisionManagerConfig ;
437- use crate :: v2:: api:: { Db as _, DbView , Proposal as _} ;
443+ use crate :: v2:: api:: { Db as _, DbView , HashKeyExt , Proposal as _} ;
438444
439445 use super :: { BatchOp , DbConfig } ;
440446
@@ -1392,6 +1398,194 @@ mod test {
13921398 assert_eq ! ( value, new_value. as_ref( ) ) ;
13931399 }
13941400
1401+ #[ test]
1402+ fn test_deferred_persist_close_with_high_commit_count ( ) {
1403+ const HIGH_COMMIT_COUNT : NonZeroU64 = nonzero ! ( 1_000_000u64 ) ;
1404+
1405+ // Set commit count to an arbitrarily high number so persist happens
1406+ // only on shutdown
1407+ let dbcfg = DbConfig :: builder ( )
1408+ . deferred_persistence_commit_count ( HIGH_COMMIT_COUNT )
1409+ . build ( ) ;
1410+
1411+ let db = TestDb :: new_with_config ( dbcfg) ;
1412+
1413+ // Then, commit once and see what the latest revision is
1414+ let key = b"foo" ;
1415+ let value = b"bar" ;
1416+ let batch = vec ! [ BatchOp :: Put { key, value } ] ;
1417+ let proposal = db. propose ( batch) . unwrap ( ) ;
1418+ let root_hash = proposal. root_hash ( ) . unwrap ( ) . unwrap ( ) ;
1419+
1420+ proposal. commit ( ) . unwrap ( ) ;
1421+ let db = db. reopen ( ) ;
1422+
1423+ let revision = db. view ( root_hash) . unwrap ( ) ;
1424+ let new_value = revision. val ( b"foo" ) . unwrap ( ) . unwrap ( ) ;
1425+
1426+ assert_eq ! ( value, new_value. as_ref( ) ) ;
1427+ }
1428+
1429+ #[ test]
1430+ fn test_deferred_persist_after_reaching_sub_interval ( ) {
1431+ const COMMIT_COUNT : NonZeroU64 = nonzero ! ( 10u64 ) ;
1432+ const SUB_INTERVAL : u64 = COMMIT_COUNT . get ( ) / 2 ;
1433+
1434+ let dbcfg = DbConfig :: builder ( )
1435+ . deferred_persistence_commit_count ( COMMIT_COUNT )
1436+ . build ( ) ;
1437+
1438+ let db = TestDb :: new_with_config ( dbcfg) ;
1439+
1440+ // Commit SUB_INTERVAL proposals to trigger the first persist
1441+ for i in 0 ..SUB_INTERVAL {
1442+ let batch = vec ! [ BatchOp :: Put {
1443+ key: format!( "key{i}" ) . as_bytes( ) . to_vec( ) ,
1444+ value: format!( "value{i}" ) . as_bytes( ) . to_vec( ) ,
1445+ } ] ;
1446+ let proposal = db. propose ( batch) . unwrap ( ) ;
1447+ proposal. commit ( ) . unwrap ( ) ;
1448+ }
1449+
1450+ // Wait for the background thread to finish persisting
1451+ db. wait_persisted ( ) ;
1452+
1453+ // Verify the root is now persisted
1454+ let revision = db. manager . current_revision ( ) ;
1455+ let is_persisted = revision
1456+ . root_as_maybe_persisted_node ( )
1457+ . is_some_and ( |node| node. unpersisted ( ) . is_none ( ) ) ;
1458+
1459+ assert ! (
1460+ is_persisted,
1461+ "Root should be persisted after hitting commit count"
1462+ ) ;
1463+ }
1464+
1465+ /// Verifies that an unpersisted revision which wipes the database is
1466+ /// persisted when the database closes.
1467+ #[ test]
1468+ fn test_deferred_persistence_closing_on_empty_trie ( ) {
1469+ const COMMIT_COUNT : NonZeroU64 = nonzero ! ( 10u64 ) ;
1470+ const SUB_INTERVAL : u64 = COMMIT_COUNT . get ( ) / 2 ;
1471+
1472+ let dbcfg = DbConfig :: builder ( )
1473+ . deferred_persistence_commit_count ( COMMIT_COUNT )
1474+ . build ( ) ;
1475+
1476+ let db = TestDb :: new_with_config ( dbcfg) ;
1477+
1478+ // Commit SUB_INTERVAL proposals to trigger the first persist
1479+ for i in 0 ..SUB_INTERVAL {
1480+ let batch = vec ! [ BatchOp :: Put {
1481+ key: format!( "key{i}" ) . as_bytes( ) . to_vec( ) ,
1482+ value: format!( "value{i}" ) . as_bytes( ) . to_vec( ) ,
1483+ } ] ;
1484+ let proposal = db. propose ( batch) . unwrap ( ) ;
1485+ proposal. commit ( ) . unwrap ( ) ;
1486+ }
1487+
1488+ // Empty the trie
1489+ let batch: Vec < BatchOp < Vec < u8 > , Vec < u8 > > > = vec ! [ BatchOp :: DeleteRange { prefix: vec![ ] } ] ;
1490+ let proposal = db. propose ( batch) . unwrap ( ) ;
1491+ proposal. commit ( ) . unwrap ( ) ;
1492+
1493+ let db = db. reopen ( ) ;
1494+
1495+ // Verify that the latest committed revision is empty.
1496+ let last_committed_hash = db. root_hash ( ) . unwrap ( ) ;
1497+ assert_eq ! ( last_committed_hash, TrieHash :: default_root_hash( ) ) ;
1498+ }
1499+
1500+ #[ test]
1501+ fn test_deferred_persistence_root_store ( ) {
1502+ const NUM_COMMITS : usize = 20 ;
1503+ const COMMIT_COUNT : NonZeroU64 = nonzero ! ( 10u64 ) ;
1504+ const MAX_IN_MEMORY_REVISIONS : usize = 5 ;
1505+
1506+ // Revisions to verify after reopening (1-indexed commit numbers)
1507+ const CHECKPOINTS : [ usize ; 4 ] = [ 5 , 10 , 15 , 20 ] ;
1508+
1509+ let dbcfg = DbConfig :: builder ( )
1510+ . manager (
1511+ RevisionManagerConfig :: builder ( )
1512+ . max_revisions ( MAX_IN_MEMORY_REVISIONS )
1513+ . build ( ) ,
1514+ )
1515+ . deferred_persistence_commit_count ( COMMIT_COUNT )
1516+ . root_store ( true )
1517+ . build ( ) ;
1518+
1519+ let db = TestDb :: new_with_config ( dbcfg) ;
1520+
1521+ // Track root hashes at checkpoint commits
1522+ let mut checkpoint_roots: Vec < TrieHash > = Vec :: new ( ) ;
1523+
1524+ // Commit NUM_COMMITS proposals
1525+ let key = b"key" ;
1526+ for i in 1 ..=NUM_COMMITS {
1527+ let batch = vec ! [ BatchOp :: Put {
1528+ key,
1529+ value: format!( "{i}" ) . as_bytes( ) . to_vec( ) ,
1530+ } ] ;
1531+ let proposal = db. propose ( batch) . unwrap ( ) ;
1532+
1533+ if CHECKPOINTS . contains ( & i) {
1534+ checkpoint_roots. push ( proposal. root_hash ( ) . unwrap ( ) . unwrap ( ) ) ;
1535+ }
1536+
1537+ proposal. commit ( ) . unwrap ( ) ;
1538+ }
1539+
1540+ let db = db. reopen ( ) ;
1541+
1542+ // Verify that checkpoint revisions are accessible after reopening
1543+ // and contain the expected values
1544+ for ( i, root) in checkpoint_roots. into_iter ( ) . enumerate ( ) {
1545+ let view = db. view ( root) . unwrap ( ) ;
1546+ let checkpoint = CHECKPOINTS [ i] ;
1547+ let expected = format ! ( "{checkpoint}" ) ;
1548+ let actual = view. val ( key) . unwrap ( ) . unwrap ( ) ;
1549+ assert_eq ! ( expected. as_bytes( ) , actual. as_ref( ) ) ;
1550+ }
1551+ }
1552+
1553+ /// Verifies that non-persisted revisions are lost after reopening the database.
1554+ #[ test]
1555+ fn test_deferred_persistence_unpersisted_revisions ( ) {
1556+ const NUM_COMMITS : usize = 10 ;
1557+ const COMMIT_COUNT : NonZeroU64 = nonzero ! ( 10u64 ) ;
1558+
1559+ let dbcfg = DbConfig :: builder ( )
1560+ . deferred_persistence_commit_count ( COMMIT_COUNT )
1561+ . root_store ( true )
1562+ . build ( ) ;
1563+
1564+ let db = TestDb :: new_with_config ( dbcfg) ;
1565+
1566+ // Track root hashes for every commit
1567+ let mut root_hashes: Vec < TrieHash > = Vec :: new ( ) ;
1568+
1569+ let key = b"key" ;
1570+ for i in 1 ..=NUM_COMMITS {
1571+ let batch = vec ! [ BatchOp :: Put {
1572+ key,
1573+ value: format!( "{i}" ) . as_bytes( ) . to_vec( ) ,
1574+ } ] ;
1575+ let proposal = db. propose ( batch) . unwrap ( ) ;
1576+ root_hashes. push ( proposal. root_hash ( ) . unwrap ( ) . unwrap ( ) ) ;
1577+ proposal. commit ( ) . unwrap ( ) ;
1578+ }
1579+
1580+ let db = db. reopen ( ) ;
1581+
1582+ // Commits 1-4 and 6-9 were not persisted and should not be accessible
1583+ let unpersisted: [ usize ; 8 ] = [ 1 , 2 , 3 , 4 , 6 , 7 , 8 , 9 ] ;
1584+ for & i in & unpersisted {
1585+ assert ! ( db. view( root_hashes[ i - 1 ] . clone( ) ) . is_err( ) ) ;
1586+ }
1587+ }
1588+
13951589 // Testdb is a helper struct for testing the Db. Once it's dropped, the directory and file disappear
13961590 pub ( super ) struct TestDb {
13971591 db : Db ,
0 commit comments