chore(storage): add test cases for Transactions and Receipts truncation

crates/storage/provider/src/providers/static_file/mod.rs

@@ -64,11 +64,18 @@ mod tests {
     };
     use reth_db_api::transaction::DbTxMut;
     use reth_primitives::{
-        static_file::{find_fixed_range, DEFAULT_BLOCKS_PER_STATIC_FILE},
-        BlockHash, Header,
+        static_file::{find_fixed_range, SegmentRangeInclusive, DEFAULT_BLOCKS_PER_STATIC_FILE},
+        BlockHash, Header, Receipt, TransactionSignedNoHash,
     };
     use reth_testing_utils::generators::{self, random_header_range};
-    use std::{fs, path::Path};
+    use std::{fmt::Debug, fs, ops::Range, path::Path};
+
+    fn assert_eyre<T: PartialEq + Debug>(got: T, expected: T, msg: &str) -> eyre::Result<()> {
+        if got != expected {
+            eyre::bail!("{msg} | got: {got:?} expected: {expected:?})");
+        }
+        Ok(())
+    }
 
     #[test]
     fn test_snap() {
@@ -186,18 +193,25 @@ mod tests {
             prune_count: u64,
             expected_tip: Option<u64>,
             expected_file_count: u64,
-        ) {
-            writer.prune_headers(prune_count).unwrap();
-            writer.commit().unwrap();
+        ) -> eyre::Result<()> {
+            writer.prune_headers(prune_count)?;
+            writer.commit()?;
 
             // Validate the highest block after pruning
-            assert_eq!(
+            assert_eyre(
                 sf_rw.get_highest_static_file_block(StaticFileSegment::Headers),
-                expected_tip
-            );
+                expected_tip,
+                "block mismatch",
+            )?;
 
             // Validate the number of files remaining in the directory
-            assert_eq!(fs::read_dir(static_dir).unwrap().count(), expected_file_count as usize);
+            assert_eyre(
+                fs::read_dir(static_dir)?.count(),
+                expected_file_count as usize,
+                "file count mismatch",
+            )?;
+
+            Ok(())
         }
 
         // [ Test Cases ]
@@ -206,31 +220,31 @@ mod tests {
         type ExpectedFileCount = u64;
         let mut tmp_tip = tip;
         let test_cases: Vec<(PruneCount, Option<ExpectedTip>, ExpectedFileCount)> = vec![
-            // Case 1: Pruning 1 header
+            // Case 0: Pruning 1 header
             {
                 tmp_tip -= 1;
                 (1, Some(tmp_tip), initial_file_count)
             },
-            // Case 2: Pruning remaining rows from file should result in its deletion
+            // Case 1: Pruning remaining rows from file should result in its deletion
             {
                 tmp_tip -= blocks_per_file - 1;
                 (blocks_per_file - 1, Some(tmp_tip), initial_file_count - files_per_range)
             },
-            // Case 3: Pruning more headers than a single file has (tip reduced by
+            // Case 2: Pruning more headers than a single file has (tip reduced by
             // blocks_per_file + 1) should result in a file set deletion
             {
                 tmp_tip -= blocks_per_file + 1;
                 (blocks_per_file + 1, Some(tmp_tip), initial_file_count - files_per_range * 2)
             },
-            // Case 4: Pruning all remaining headers from the file except the genesis header
+            // Case 3: Pruning all remaining headers from the file except the genesis header
             {
                 (
                     tmp_tip,
                     Some(0),             // Only genesis block remains
                     files_per_range + 1, // The file set with block 0 should remain
                 )
             },
-            // Case 5: Pruning the genesis header (should not delete the file set with block 0)
+            // Case 4: Pruning the genesis header (should not delete the file set with block 0)
             {
                 (
                     1,
@@ -254,15 +268,212 @@ mod tests {
 
             let mut header_writer = sf_rw.latest_writer(StaticFileSegment::Headers).unwrap();
 
-            for (prune_count, expected_tip, expected_file_count) in test_cases {
+            for (case, (prune_count, expected_tip, expected_file_count)) in
+                test_cases.into_iter().enumerate()
+            {
                 prune_and_validate(
                     &mut header_writer,
                     &sf_rw,
                     &static_dir,
                     prune_count,
                     expected_tip,
                     expected_file_count,
-                );
+                )
+                .map_err(|err| eyre::eyre!("Test case {case}: {err}"))
+                .unwrap();
+            }
+        }
+    }
+
+    /// 3 block ranges are built
+    ///
+    /// for `blocks_per_file = 10`:
+    /// * `0..=9` : except genesis, every block has a tx/receipt
+    /// * `10..=19`: no txs/receipts
+    /// * `20..=29`: only one tx/receipt
+    fn setup_tx_based_scenario(
+        sf_rw: &StaticFileProvider,
+        segment: StaticFileSegment,
+        blocks_per_file: u64,
+    ) {
+        fn setup_block_ranges(
+            writer: &mut StaticFileProviderRWRefMut<'_>,
+            sf_rw: &StaticFileProvider,
+            segment: StaticFileSegment,
+            block_range: &Range<u64>,
+            mut tx_count: u64,
+            next_tx_num: &mut u64,
+        ) {
+            for block in block_range.clone() {
+                writer.increment_block(block).unwrap();
+
+                // Append transaction/receipt if there's still a transaction count to append
+                if tx_count > 0 {
+                    if segment.is_receipts() {
+                        writer.append_receipt(*next_tx_num, &Receipt::default()).unwrap();
+                    } else {
+                        writer
+                            .append_transaction(*next_tx_num, &TransactionSignedNoHash::default())
+                            .unwrap();
+                    }
+                    *next_tx_num += 1;
+                    tx_count -= 1;
+                }
+            }
+            writer.commit().unwrap();
+
+            // Calculate expected values based on the range and transactions
+            let expected_block = block_range.end - 1;
+            let expected_tx = if tx_count == 0 { *next_tx_num - 1 } else { *next_tx_num };
+
+            // Perform assertions after processing the blocks
+            assert_eq!(sf_rw.get_highest_static_file_block(segment), Some(expected_block),);
+            assert_eq!(sf_rw.get_highest_static_file_tx(segment), Some(expected_tx),);
+        }
+
+        // Define the block ranges and transaction counts as vectors
+        let block_ranges = [
+            0..blocks_per_file,
+            blocks_per_file..blocks_per_file * 2,
+            blocks_per_file * 2..blocks_per_file * 3,
+        ];
+
+        let tx_counts = [
+            blocks_per_file - 1, // First range: tx per block except genesis
+            0,                   // Second range: no transactions
+            1,                   // Third range: 1 transaction in the second block
+        ];
+
+        let mut writer = sf_rw.latest_writer(segment).unwrap();
+        let mut next_tx_num = 0;
+
+        // Loop through setup scenarios
+        for (block_range, tx_count) in block_ranges.iter().zip(tx_counts.iter()) {
+            setup_block_ranges(
+                &mut writer,
+                sf_rw,
+                segment,
+                block_range,
+                *tx_count,
+                &mut next_tx_num,
+            );
+        }
+
+        // Ensure that scenario was properly setup
+        let expected_tx_ranges = vec![
+            Some(SegmentRangeInclusive::new(0, 8)),
+            None,
+            Some(SegmentRangeInclusive::new(9, 9)),
+        ];
+
+        block_ranges.iter().zip(expected_tx_ranges).for_each(|(block_range, expected_tx_range)| {
+            assert_eq!(
+                sf_rw
+                    .get_segment_provider_from_block(segment, block_range.start, None)
+                    .unwrap()
+                    .user_header()
+                    .tx_range(),
+                expected_tx_range.as_ref()
+            );
+        });
+    }
+
+    #[test]
+    #[ignore]
+    fn test_tx_based_truncation() {
+        let segments = [StaticFileSegment::Transactions, StaticFileSegment::Receipts];
+        let blocks_per_file = 10; // Number of blocks per file
+        let files_per_range = 3; // Number of files per range (data/conf/offset files)
+        let file_set_count = 3; // Number of sets of files to create
+        let initial_file_count = files_per_range * file_set_count + 1; // Includes lockfile
+
+        fn prune_and_validate(
+            sf_rw: &StaticFileProvider,
+            static_dir: impl AsRef<Path>,
+            segment: StaticFileSegment,
+            prune_count: u64,
+            last_block: u64,
+            expected_tx_tip: u64,
+            expected_file_count: i32,
+        ) -> eyre::Result<()> {
+            let mut writer = sf_rw.latest_writer(segment)?;
+
+            // Prune transactions or receipts based on the segment type
+            if segment.is_receipts() {
+                writer.prune_receipts(prune_count, last_block)?;
+            } else {
+                writer.prune_transactions(prune_count, last_block)?;
+            }
+            writer.commit()?;
+
+            // Verify the highest block and transaction tips
+            assert_eyre(
+                sf_rw.get_highest_static_file_block(segment),
+                Some(last_block),
+                "block mismatch",
+            )?;
+            assert_eyre(
+                sf_rw.get_highest_static_file_tx(segment),
+                Some(expected_tx_tip),
+                "tx mismatch",
+            )?;
+
+            // Ensure the file count has reduced as expected
+            assert_eyre(
+                fs::read_dir(static_dir)?.count(),
+                expected_file_count as usize,
+                "file count mismatch",
+            )?;
+            Ok(())
+        }
+
+        for segment in segments {
+            let (static_dir, _) = create_test_static_files_dir();
+
+            let sf_rw = StaticFileProvider::read_write(&static_dir)
+                .expect("Failed to create static file provider")
+                .with_custom_blocks_per_file(blocks_per_file);
+
+            setup_tx_based_scenario(&sf_rw, segment, blocks_per_file);
+
+            let sf_rw = StaticFileProvider::read_write(&static_dir)
+                .expect("Failed to create static file provider")
+                .with_custom_blocks_per_file(blocks_per_file);
+            let highest_tx = sf_rw.get_highest_static_file_tx(segment).unwrap();
+
+            // Test cases
+            // [prune_count, last_block, expected_tx_tip, expected_file_count)
+            let test_cases = vec![
+                // Case 0: 20..=29 has only one tx. Prune the only tx of the block range.
+                // It ensures that the file is not deleted even though there are no rows, since the
+                // `last_block` which is passed to the prune method is the first
+                // block of the range.
+                (1, blocks_per_file * 2, highest_tx - 1, initial_file_count),
+                // Case 1: 10..=19 has no txs. There are no txes in the whole block range, but want
+                // to unwind to block 9. Ensures that the 20..=29 and 10..=19 files
+                // are deleted.
+                (0, blocks_per_file - 1, highest_tx - 1, files_per_range + 1), // includes lockfile
+                // Case 2: Prune most txs up to block 1.
+                (7, 1, 1, files_per_range + 1),
+                // Case 3: Prune remaining tx and ensure that file is not deleted.
+                (1, 0, 0, files_per_range + 1),
+            ];
+
+            // Loop through test cases
+            for (case, (prune_count, last_block, expected_tx_tip, expected_file_count)) in
+                test_cases.into_iter().enumerate()
+            {
+                prune_and_validate(
+                    &sf_rw,
+                    &static_dir,
+                    segment,
+                    prune_count,
+                    last_block,
+                    expected_tx_tip,
+                    expected_file_count,
+                )
+                .map_err(|err| eyre::eyre!("Test case {case}: {err}"))
+                .unwrap();
             }
         }
     }