Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
242 changes: 149 additions & 93 deletions dash-spv/src/storage/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -394,134 +394,190 @@ impl masternode::MasternodeStateStorage for DiskStorageManager {

#[cfg(test)]
mod tests {
use crate::ChainState;

use super::*;
use dashcore::Header as BlockHeader;
use tempfile::{tempdir, TempDir};

#[tokio::test]
async fn test_load_headers() -> Result<(), Box<dyn std::error::Error>> {
async fn test_store_load_headers() -> Result<(), Box<dyn std::error::Error>> {
// Create a temporary directory for the test
let temp_dir = TempDir::new()?;
let mut storage = DiskStorageManager::new(temp_dir.path().to_path_buf())
.await
.expect("Unable to create storage");
let mut storage =
DiskStorageManager::new(temp_dir.path()).await.expect("Unable to create storage");

// Create a test header
let test_header = BlockHeader::dummy(1);
let headers = BlockHeader::dummy_batch(0..60_000);

// Store just one header
storage.store_headers(&[test_header]).await?;
storage.store_headers(&headers[0..0]).await.expect("Should handle empty header batch");
assert_eq!(storage.get_tip_height().await, None);

storage.store_headers(&headers[0..1]).await.expect("Failed to store headers");
let loaded_headers = storage.load_headers(0..1).await?;

// Should only get back the one header we stored
assert_eq!(loaded_headers.len(), 1);
assert_eq!(loaded_headers[0], test_header);
assert_eq!(loaded_headers[0], headers[0]);

storage.store_headers(&headers[1..100]).await.expect("Failed to store headers");
let loaded_headers = storage.load_headers(50..60).await.unwrap();
assert_eq!(loaded_headers.len(), 10);
assert_eq!(&loaded_headers, &headers[50..60]);

storage.store_headers(&headers[100..headers.len()]).await.expect("Failed to store headers");

let tip_height = storage.get_tip_height().await.unwrap();
let tip_header = storage.get_header(tip_height).await.unwrap().unwrap();
let expected_header = &headers[headers.len() - 1];
assert_eq!(tip_header, *expected_header);

let non_existing_height = tip_height + 1;
let non_existing_header = storage.get_header(non_existing_height).await.unwrap();
assert!(non_existing_header.is_none());

storage.shutdown().await;
drop(storage);
let storage =
DiskStorageManager::new(temp_dir.path()).await.expect("Unable to open storage");

let loaded_headers = storage.load_headers(49_999..50_002).await.unwrap();
assert_eq!(loaded_headers.len(), 3);
assert_eq!(&loaded_headers, &headers[49_999..50_002]);

Ok(())
Comment on lines +434 to 443
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Stop the background worker before dropping storage instances.

Dropping DiskStorageManager without shutdown() leaves the background task running, which can keep writing to a temp path after the test finishes and cause flaky cleanup. Consider shutting down before drop, even in tests.

🔧 Suggested fix
@@
-        let storage =
-            DiskStorageManager::new(temp_dir.path()).await.expect("Unable to open storage");
+        let mut storage =
+            DiskStorageManager::new(temp_dir.path()).await.expect("Unable to open storage");
@@
         let loaded_headers = storage.load_headers(49_999..50_002).await.unwrap();
         assert_eq!(loaded_headers.len(), 3);
         assert_eq!(&loaded_headers, &headers[49_999..50_002]);
+        storage.shutdown().await;
@@
-            let storage = DiskStorageManager::new(&temp_dir.path()).await.unwrap();
+            let mut storage = DiskStorageManager::new(temp_dir.path()).await.unwrap();
@@
             for i in 0..10 {
                 let stored_header = storage.get_header(i).await.unwrap().unwrap();
                 let hash = stored_header.block_hash();
                 let height = storage.get_header_height_by_hash(&hash).await.unwrap();
                 assert_eq!(height, Some(i), "Height mismatch after reload for header {}", i);
             }
+            storage.shutdown().await;
@@
-        drop(storage1);
+        storage1.shutdown().await;
+        drop(storage1);

Also applies to: 523-535, 575-577

🤖 Prompt for AI Agents
In `@dash-spv/src/storage/mod.rs` around lines 434 - 443, The test drops
DiskStorageManager while its background worker may still be running; ensure you
call storage.shutdown().await before drop(storage) to stop the worker and avoid
post-test writes. Update the instances around DiskStorageManager::new(...) and
subsequent drop(storage) to invoke the async shutdown() on the storage handle
(e.g., call storage.shutdown().await) prior to dropping it; apply the same
change to the other similar blocks referenced (around the other occurrences
~lines 523-535 and 575-577).

}

#[tokio::test]
async fn test_checkpoint_storage_indexing() -> StorageResult<()> {
let temp_dir = tempdir().expect("Failed to create temp dir");
let mut storage = DiskStorageManager::new(temp_dir.path().to_path_buf()).await?;
let mut storage = DiskStorageManager::new(temp_dir.path()).await?;

// Create test headers starting from checkpoint height
let checkpoint_height = 1_100_000;
let headers = BlockHeader::dummy_batch(checkpoint_height..checkpoint_height + 100);

let mut base_state = ChainState::new();
base_state.sync_base_height = checkpoint_height;
storage.store_chain_state(&base_state).await?;

storage.store_headers_at_height(&headers, checkpoint_height).await?;
assert_eq!(storage.get_stored_headers_len().await, headers.len() as u32);

// Verify headers are stored at correct blockchain heights
let header_at_base = storage.get_header(checkpoint_height).await?;
assert_eq!(
header_at_base.expect("Header at base blockchain height should exist"),
headers[0]
);

let header_at_ending = storage.get_header(checkpoint_height + 99).await?;
assert_eq!(
header_at_ending.expect("Header at ending blockchain height should exist"),
headers[99]
);

// Test the reverse index (hash -> blockchain height)
let hash_0 = headers[0].block_hash();
let height_0 = storage.get_header_height_by_hash(&hash_0).await?;
assert_eq!(
height_0,
Some(checkpoint_height),
"Hash should map to blockchain height 1,100,000"
);

let hash_99 = headers[99].block_hash();
let height_99 = storage.get_header_height_by_hash(&hash_99).await?;
assert_eq!(
height_99,
Some(checkpoint_height + 99),
"Hash should map to blockchain height 1,100,099"
);

// Store chain state to persist sync_base_height
let mut chain_state = ChainState::new();
chain_state.sync_base_height = checkpoint_height;
storage.store_chain_state(&chain_state).await?;

// Force save to disk
storage.persist().await;
const CHECKPOINT_HEIGHT: u32 = 1_100_000;
let headers: Vec<BlockHeader> = BlockHeader::dummy_batch(0..100);

storage.store_headers_at_height(&headers, CHECKPOINT_HEIGHT).await?;

check_storage(&storage, &headers).await?;

storage.shutdown().await;
drop(storage);

// Create a new storage instance to test index rebuilding
let storage2 = DiskStorageManager::new(temp_dir.path().to_path_buf()).await?;

// Verify the index was rebuilt correctly
let height_after_rebuild = storage2.get_header_height_by_hash(&hash_0).await?;
assert_eq!(
height_after_rebuild,
Some(checkpoint_height),
"After index rebuild, hash should still map to blockchain height 1,100,000"
);

// Verify header can still be retrieved by blockchain height after reload
let header_after_reload = storage2.get_header(checkpoint_height).await?;
assert!(
header_after_reload.is_some(),
"Header at base blockchain height should exist after reload"
);
assert_eq!(header_after_reload.unwrap(), headers[0]);
let storage = DiskStorageManager::new(temp_dir.path()).await?;

Ok(())
check_storage(&storage, &headers).await?;

return Ok(());

async fn check_storage(
storage: &DiskStorageManager,
headers: &[BlockHeader],
) -> StorageResult<()> {
assert_eq!(storage.get_stored_headers_len().await, headers.len() as u32);

let header_at_base = storage.get_header(CHECKPOINT_HEIGHT).await?;
assert_eq!(header_at_base, Some(headers[0]));

let header_at_ending = storage.get_header(CHECKPOINT_HEIGHT + 99).await?;
assert_eq!(header_at_ending, Some(headers[99]));

// Test the reverse index (hash -> blockchain height)
let hash_0 = headers[0].block_hash();
let height_0 = storage.get_header_height_by_hash(&hash_0).await?;
assert_eq!(
height_0,
Some(CHECKPOINT_HEIGHT),
"Hash should map to blockchain height 1,100,000"
);

let hash_99 = headers[99].block_hash();
let height_99 = storage.get_header_height_by_hash(&hash_99).await?;
assert_eq!(
height_99,
Some(CHECKPOINT_HEIGHT + 99),
"Hash should map to blockchain height 1,100,099"
);

Ok(())
}
}

#[tokio::test]
async fn test_shutdown_flushes_index() -> Result<(), Box<dyn std::error::Error>> {
let temp_dir = TempDir::new()?;
let base_path = temp_dir.path().to_path_buf();
let headers = BlockHeader::dummy_batch(0..11_000);
let last_hash = headers.last().unwrap().block_hash();
async fn test_reverse_index_disk_storage() {
let temp_dir = tempfile::tempdir().unwrap();

{
let mut storage = DiskStorageManager::new(base_path.clone()).await?;
let mut storage = DiskStorageManager::new(temp_dir.path()).await.unwrap();

// Create and store headers
let headers = BlockHeader::dummy_batch(0..10);

storage.store_headers(&headers[..10_000]).await?;
storage.persist().await;
storage.store_headers(&headers).await.unwrap();

// Test reverse lookups
for (i, header) in headers.iter().enumerate() {
let hash = header.block_hash();
let height = storage.get_header_height_by_hash(&hash).await.unwrap();
assert_eq!(height, Some(i as u32), "Height mismatch for header {}", i);
}

storage.store_headers(&headers[10_000..]).await?;
storage.shutdown().await;
}

let storage = DiskStorageManager::new(base_path).await?;
let height = storage.get_header_height_by_hash(&last_hash).await?;
assert_eq!(height, Some(10_999));
// Test persistence - reload storage and verify index still works
{
let storage = DiskStorageManager::new(&temp_dir.path()).await.unwrap();

// The index should have been rebuilt from the loaded headers
// We need to get the actual headers that were stored to test properly
for i in 0..10 {
let stored_header = storage.get_header(i).await.unwrap().unwrap();
let hash = stored_header.block_hash();
let height = storage.get_header_height_by_hash(&hash).await.unwrap();
assert_eq!(height, Some(i), "Height mismatch after reload for header {}", i);
}
}
}

Ok(())
#[tokio::test]
async fn test_clear_clears_index() {
let mut storage =
DiskStorageManager::with_temp_dir().await.expect("Failed to create tmp storage");

// Store some headers
let header = BlockHeader::dummy_batch(0..1);
storage.store_headers(&header).await.unwrap();

let hash = header[0].block_hash();
assert!(storage.get_header_height_by_hash(&hash).await.unwrap().is_some());

// Clear storage
storage.clear().await.unwrap();

Comment on lines +540 to +552
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Keep the TempDir alive for the storage lifetime.

DiskStorageManager::with_temp_dir() drops the TempDir immediately, which can delete the backing directory while the storage is still active. Use an explicit TempDir in the test to keep the directory alive.

🔧 Suggested fix
@@
-        let mut storage =
-            DiskStorageManager::with_temp_dir().await.expect("Failed to create tmp storage");
+        let temp_dir = TempDir::new().expect("Failed to create temp directory");
+        let mut storage =
+            DiskStorageManager::new(temp_dir.path()).await.expect("Failed to create tmp storage");
@@
         // Verify index is cleared
         assert!(storage.get_header_height_by_hash(&hash).await.unwrap().is_none());
+        storage.shutdown().await;
🤖 Prompt for AI Agents
In `@dash-spv/src/storage/mod.rs` around lines 540 - 552,
DiskStorageManager::with_temp_dir() currently drops the TempDir immediately,
risking deletion of the backing directory while storage is still used; fix the
test by creating and retaining an explicit TempDir (e.g., via tempfile::TempDir)
and passing its path into DiskStorageManager (or into DiskStorageManager::new)
so the TempDir lives for the lifetime of storage; update the test around
DiskStorageManager::with_temp_dir(), BlockHeader::dummy_batch,
storage.store_headers, storage.get_header_height_by_hash, and storage.clear to
hold the TempDir variable until after storage.clear() completes.

// Verify index is cleared
assert!(storage.get_header_height_by_hash(&hash).await.unwrap().is_none());
}

#[tokio::test]
async fn test_lock_lifecycle() {
let temp_dir = TempDir::new().expect("Failed to create temp directory");
let path = temp_dir.path().to_path_buf();
let lock_path = {
let mut lock_file = path.clone();
lock_file.set_extension("lock");
lock_file
};

let mut storage1 = DiskStorageManager::new(&path).await.unwrap();
assert!(lock_path.exists(), "Lock file should exist while storage is open");
storage1.clear().await.expect("Failed to clear the storage");
assert!(lock_path.exists(), "Lock file should exist after storage is cleared");

let storage2 = DiskStorageManager::new(&path).await;
assert!(storage2.is_err(), "Second storage manager should fail");

// Lock file removed when storage drops
drop(storage1);
assert!(!lock_path.exists(), "Lock file should be removed after storage drops");

// Can reopen storage after previous one dropped
let storage3 = DiskStorageManager::new(&path).await;
assert!(storage3.is_ok(), "Should reopen after previous storage dropped");
}
}
Loading
Loading