diff --git a/dragonfly-client-storage/src/lib.rs b/dragonfly-client-storage/src/lib.rs index 5412ce8f4f2..dc0b9fb5818 100644 --- a/dragonfly-client-storage/src/lib.rs +++ b/dragonfly-client-storage/src/lib.rs @@ -400,6 +400,11 @@ impl Storage { self.metadata.get_piece(task_id, number) } + // get_pieces returns the piece metadatas. + pub fn get_pieces(&self, task_id: &str) -> Result> { + self.metadata.get_pieces(task_id) + } + // piece_id returns the piece id. #[instrument(skip_all)] pub fn piece_id(&self, task_id: &str, number: u32) -> String { diff --git a/dragonfly-client-storage/src/metadata.rs b/dragonfly-client-storage/src/metadata.rs index 2d13ffd576d..7cebc3da1dc 100644 --- a/dragonfly-client-storage/src/metadata.rs +++ b/dragonfly-client-storage/src/metadata.rs @@ -824,6 +824,23 @@ impl Metadata { self.db.get(self.piece_id(task_id, number).as_bytes()) } + // get_pieces gets the piece metadatas. + pub fn get_pieces(&self, task_id: &str) -> Result> { + let pieces = self + .db + .prefix_iter_raw::(task_id.as_bytes())? + .map(|ele| { + let (_, value) = ele?; + Ok(value) + }) + .collect::>>>()?; + + pieces + .par_iter() + .map(|piece| Piece::deserialize_from(piece)) + .collect() + } + // delete_piece deletes the piece metadata. #[instrument(skip_all)] pub fn delete_piece(&self, task_id: &str, number: u32) -> Result<()> { @@ -891,15 +908,18 @@ impl Metadata { #[cfg(test)] mod tests { use super::*; - use tempdir::TempDir; + const MOCK_TASK_ID_1: &str = "d3c4e940ad06c47fc36ac67801e6f8e36cb400e2391708620bc7e865b102062c"; + const MOCK_TASK_ID_2: &str = "a535b115f18d96870f0422ac891f91dd162f2f391e4778fb84279701fcd02dd1"; + #[test] fn should_create_metadata_db() { let dir = TempDir::new("metadata_db").unwrap(); let log_dir = dir.path().join("log"); let metadata = Metadata::new(Arc::new(Config::default()), dir.path(), &log_dir).unwrap(); assert!(metadata.get_tasks().unwrap().is_empty()); + assert!(metadata.get_pieces(MOCK_TASK_ID_1).unwrap().is_empty()); } #[test] @@ -908,7 +928,7 @@ mod tests { let log_dir = dir.path().join("log"); let metadata = Metadata::new(Arc::new(Config::default()), dir.path(), &log_dir).unwrap(); - let task_id = "task1"; + let task_id = MOCK_TASK_ID_1; // Test download_task_started. metadata @@ -967,7 +987,7 @@ mod tests { ); // Test get_tasks. - let task_id = "task2"; + let task_id = MOCK_TASK_ID_2; metadata .download_task_started(task_id, Some(1024), None, None) @@ -986,7 +1006,7 @@ mod tests { let dir = TempDir::new("metadata_db").unwrap(); let log_dir = dir.path().join("log"); let metadata = Metadata::new(Arc::new(Config::default()), dir.path(), &log_dir).unwrap(); - let task_id = "task3"; + let task_id = MOCK_TASK_ID_1; // Test download_piece_started. metadata.download_piece_started(task_id, 1).unwrap(); @@ -1010,6 +1030,12 @@ mod tests { "piece should be updated after download_piece_finished" ); + // Test get_pieces. + metadata.download_piece_started(task_id, 2).unwrap(); + metadata.download_piece_started(task_id, 3).unwrap(); + let pieces = metadata.get_pieces(task_id).unwrap(); + assert_eq!(pieces.len(), 3, "should get 3 pieces in total"); + // Test download_piece_failed. metadata.download_piece_started(task_id, 2).unwrap(); metadata.download_piece_started(task_id, 3).unwrap(); @@ -1048,5 +1074,10 @@ mod tests { piece.uploading_count, 0, "piece should be updated after upload_piece_failed" ); + + // Test delete_pieces. + metadata.delete_pieces(task_id).unwrap(); + let pieces = metadata.get_pieces(task_id).unwrap(); + assert!(pieces.is_empty(), "should get 0 pieces after delete_pieces"); } }