Skip to content

Commit

Permalink
test: add tests for calculate_piece_range
Browse files Browse the repository at this point in the history 
Signed-off-by: Gaius <[email protected]>
  • Loading branch information
@gaius-qi
gaius-qi committed Dec 20, 2024
1 parent 6ace03a commit 637d455
Show file tree
Hide file tree
Showing 5 changed files with 199 additions and 44 deletions.
1 change: 1 addition & 0 deletions Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions dragonfly-client-storage/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,3 +32,4 @@ rayon = "1.10.0"

[dev-dependencies]
tempdir = "0.3"
tempfile.workspace = true
116 changes: 98 additions & 18 deletions dragonfly-client-storage/src/content.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -252,22 +252,16 @@ impl Content {
range: Option<Range>,
) -> Result<impl AsyncRead> {
let task_path = self.get_task_path(task_id);

// Calculate the target offset and length based on the range.
let (target_offset, target_length) = calculate_piece_range(offset, length, range);

let f = File::open(task_path.as_path()).await.map_err(|err| {
error!("open {:?} failed: {}", task_path, err);
err
})?;
let mut f_reader = BufReader::with_capacity(self.config.storage.read_buffer_size, f);

// Calculate the target offset and length based on the range.
let (target_offset, target_length) = if let Some(range) = range {
let target_offset = max(offset, range.start);
let target_length =
min(offset + length - 1, range.start + range.length - 1) - target_offset + 1;
(target_offset, target_length)
} else {
(offset, length)
};

f_reader
.seek(SeekFrom::Start(target_offset))
.await
Expand All @@ -292,14 +286,7 @@ impl Content {
let task_path = self.get_task_path(task_id);

// Calculate the target offset and length based on the range.
let (target_offset, target_length) = if let Some(range) = range {
let target_offset = max(offset, range.start);
let target_length =
min(offset + length - 1, range.start + range.length - 1) - target_offset + 1;
(target_offset, target_length)
} else {
(offset, length)
};
let (target_offset, target_length) = calculate_piece_range(offset, length, range);

let f = File::open(task_path.as_path()).await.map_err(|err| {
error!("open {:?} failed: {}", task_path, err);
Expand Down Expand Up @@ -605,3 +592,96 @@ impl Content {
Ok(task_dir.join(task_id))
}
}

/// calculate_piece_range calculates the target offset and length based on the piece range and
/// request range.
pub fn calculate_piece_range(offset: u64, length: u64, range: Option<Range>) -> (u64, u64) {
if let Some(range) = range {
let target_offset = max(offset, range.start);
let target_length =
min(offset + length - 1, range.start + range.length - 1) - target_offset + 1;
(target_offset, target_length)
} else {
(offset, length)
}
}

#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;

#[tokio::test]
async fn should_calculate_piece_range() {
let test_cases = vec![
(1, 4, None, 1, 4),
(
1,
4,
Some(Range {
start: 1,
length: 4,
}),
1,
4,
),
(
1,
4,
Some(Range {
start: 2,
length: 1,
}),
2,
1,
),
(
1,
4,
Some(Range {
start: 1,
length: 1,
}),
1,
1,
),
(
1,
4,
Some(Range {
start: 4,
length: 1,
}),
4,
1,
),
(
1,
4,
Some(Range {
start: 0,
length: 2,
}),
1,
1,
),
(
1,
4,
Some(Range {
start: 4,
length: 3,
}),
4,
1,
),
];

for (piece_offset, piece_length, range, expected_offset, expected_length) in test_cases {
let (target_offset, target_length) =
calculate_piece_range(piece_offset, piece_length, range);
assert_eq!(target_offset, expected_offset);
assert_eq!(target_length, expected_length);
}
}
}
123 changes: 97 additions & 26 deletions dragonfly-client/src/proxy/cache.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,11 +93,11 @@ impl Cache {
};

// Calculate the target offset and length based on the range.
let (target_offset, target_length) =
self.calculate_piece_range(interested_piece.offset, interested_piece.length, range);
let (piece_target_offset, piece_target_length) =
calculate_piece_range(interested_piece.offset, interested_piece.length, range);

let begin = target_offset;
let end = target_offset + target_length;
let begin = piece_target_offset;
let end = piece_target_offset + piece_target_length;
if begin >= piece_content.len() || end > piece_content.len() {
return Err(Error::InvalidParameter);
}
Expand All @@ -109,28 +109,6 @@ impl Cache {
Ok(Some(content.freeze()))
}

/// calculate_piece_range calculates the target offset and length based on the piece range and
/// request range.
fn calculate_piece_range(
&self,
piece_offset: u64,
piece_length: u64,
range: Option<Range>,
) -> (usize, usize) {
if let Some(range) = range {
let target_offset = max(piece_offset, range.start) - piece_offset;

let interested_piece_end = piece_offset + piece_length - 1;
let range_end = range.start + range.length - 1;
let target_length =
min(interested_piece_end, range_end) - target_offset - piece_offset + 1;

(target_offset as usize, target_length as usize)
} else {
(0, piece_length as usize)
}
}

/// get_piece gets the piece content from the cache.
pub fn get_piece(&self, id: &str) -> Option<bytes::Bytes> {
let mut pieces = self.pieces.lock().unwrap();
Expand All @@ -154,3 +132,96 @@ impl Cache {
pieces.contains(id)
}
}

/// calculate_piece_range calculates the target offset and length based on the piece range and
/// request range.
pub fn calculate_piece_range(offset: u64, length: u64, range: Option<Range>) -> (usize, usize) {
if let Some(range) = range {
let target_offset = max(offset, range.start) - offset;
let target_length =
min(offset + length - 1, range.start + range.length - 1) - target_offset - offset + 1;

(target_offset as usize, target_length as usize)
} else {
(0, length as usize)
}
}

#[cfg(test)]
mod tests {
use super::*;

#[tokio::test]
async fn should_calculate_piece_range() {
let test_cases = vec![
(1, 4, None, 0, 4),
(
1,
4,
Some(Range {
start: 1,
length: 4,
}),
0,
4,
),
(
1,
4,
Some(Range {
start: 2,
length: 1,
}),
1,
1,
),
(
1,
4,
Some(Range {
start: 1,
length: 1,
}),
0,
1,
),
(
1,
4,
Some(Range {
start: 4,
length: 1,
}),
3,
1,
),
(
1,
4,
Some(Range {
start: 0,
length: 2,
}),
0,
1,
),
(
1,
4,
Some(Range {
start: 4,
length: 3,
}),
3,
1,
),
];

for (piece_offset, piece_length, range, expected_offset, expected_length) in test_cases {
let (target_offset, target_length) =
calculate_piece_range(piece_offset, piece_length, range);
assert_eq!(target_offset, expected_offset);
assert_eq!(target_length, expected_length);
}
}
}
2 changes: 2 additions & 0 deletions dragonfly-client/src/proxy/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -887,6 +887,8 @@ async fn proxy_via_dfdaemon(
}
};

// When the piece reader readsto the end, add the piece
// to the cache.
if n == 0 {
cache.add_piece(&piece_id, content.freeze());
break;
Expand Down

0 comments on commit 637d455

Please sign in to comment.