use std::collections::HashMap;
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum TransferDirection {
Upload,
Download,
}
#[derive(Clone, Debug)]
pub struct TransferChunk {
pub chunk_index: u32,
pub data: Vec<u8>,
pub checksum: u32,
}
impl TransferChunk {
pub fn compute_checksum(data: &[u8]) -> u32 {
data.iter().fold(0u32, |acc, &b| acc ^ u32::from(b))
}
pub fn is_valid(&self) -> bool {
self.checksum == Self::compute_checksum(&self.data)
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum TransferState {
Pending,
InProgress {
chunks_done: u32,
total_chunks: u32,
},
Paused {
chunks_done: u32,
},
Completed,
Failed {
reason: String,
},
}
#[derive(Clone, Debug)]
pub struct BlockTransfer {
pub transfer_id: u64,
pub cid: String,
pub peer_id: String,
pub direction: TransferDirection,
pub total_bytes: u64,
pub chunk_size: u64,
pub state: TransferState,
pub received_chunks: Vec<u32>,
pub started_at_tick: u64,
}
impl BlockTransfer {
pub fn total_chunks(&self) -> u32 {
if self.chunk_size == 0 {
return 0;
}
self.total_bytes.div_ceil(self.chunk_size) as u32
}
pub fn progress(&self) -> f64 {
let total = self.total_chunks();
if total == 0 {
return 1.0;
}
let done = match &self.state {
TransferState::Pending => 0,
TransferState::InProgress { chunks_done, .. } => *chunks_done,
TransferState::Paused { chunks_done } => *chunks_done,
TransferState::Completed => total,
TransferState::Failed { .. } => self.received_chunks.len() as u32,
};
f64::from(done) / f64::from(total)
}
pub fn is_complete(&self) -> bool {
self.state == TransferState::Completed
}
}
#[derive(Clone, Debug, Default)]
pub struct TransferManagerStats {
pub active_transfers: usize,
pub completed_transfers: u64,
pub failed_transfers: u64,
pub total_bytes_transferred: u64,
}
pub struct StreamingBlockTransfer {
pub transfers: HashMap<u64, BlockTransfer>,
pub next_id: u64,
pub current_tick: u64,
}
impl StreamingBlockTransfer {
pub fn new() -> Self {
Self {
transfers: HashMap::new(),
next_id: 1,
current_tick: 0,
}
}
pub fn start_transfer(
&mut self,
cid: String,
peer_id: String,
direction: TransferDirection,
total_bytes: u64,
chunk_size: u64,
) -> u64 {
let id = self.next_id;
self.next_id += 1;
let effective_chunk_size = if chunk_size == 0 { 65536 } else { chunk_size };
let transfer = BlockTransfer {
transfer_id: id,
cid,
peer_id,
direction,
total_bytes,
chunk_size: effective_chunk_size,
state: TransferState::Pending,
received_chunks: Vec::new(),
started_at_tick: self.current_tick,
};
self.transfers.insert(id, transfer);
id
}
pub fn receive_chunk(&mut self, transfer_id: u64, chunk: TransferChunk) -> Result<(), String> {
let transfer = self
.transfers
.get_mut(&transfer_id)
.ok_or_else(|| format!("unknown transfer id: {transfer_id}"))?;
match &transfer.state {
TransferState::Paused { .. } => {
return Err(format!("transfer {transfer_id} is paused"));
}
TransferState::Completed => {
return Err(format!("transfer {transfer_id} is already completed"));
}
TransferState::Failed { reason } => {
return Err(format!("transfer {transfer_id} has failed: {reason}"));
}
TransferState::Pending | TransferState::InProgress { .. } => {}
}
if !chunk.is_valid() {
return Err(format!(
"checksum mismatch for chunk {} of transfer {transfer_id}: expected {}, got {}",
chunk.chunk_index,
TransferChunk::compute_checksum(&chunk.data),
chunk.checksum,
));
}
if transfer.received_chunks.contains(&chunk.chunk_index) {
return Err(format!(
"chunk {} already received for transfer {transfer_id}",
chunk.chunk_index
));
}
transfer.received_chunks.push(chunk.chunk_index);
let chunks_done = transfer.received_chunks.len() as u32;
let total = transfer.total_chunks();
if chunks_done >= total {
transfer.state = TransferState::Completed;
} else {
transfer.state = TransferState::InProgress {
chunks_done,
total_chunks: total,
};
}
Ok(())
}
pub fn pause(&mut self, transfer_id: u64) -> bool {
let Some(transfer) = self.transfers.get_mut(&transfer_id) else {
return false;
};
let chunks_done = match &transfer.state {
TransferState::Pending => 0,
TransferState::InProgress { chunks_done, .. } => *chunks_done,
_ => return false,
};
transfer.state = TransferState::Paused { chunks_done };
true
}
pub fn resume(&mut self, transfer_id: u64) -> bool {
let Some(transfer) = self.transfers.get_mut(&transfer_id) else {
return false;
};
let chunks_done = match &transfer.state {
TransferState::Paused { chunks_done } => *chunks_done,
_ => return false,
};
let total_chunks = transfer.total_chunks();
transfer.state = TransferState::InProgress {
chunks_done,
total_chunks,
};
true
}
pub fn fail(&mut self, transfer_id: u64, reason: String) -> bool {
let Some(transfer) = self.transfers.get_mut(&transfer_id) else {
return false;
};
match transfer.state {
TransferState::Completed | TransferState::Failed { .. } => return false,
_ => {}
}
transfer.state = TransferState::Failed { reason };
true
}
pub fn stats(&self) -> TransferManagerStats {
let mut stats = TransferManagerStats::default();
for transfer in self.transfers.values() {
match &transfer.state {
TransferState::Pending
| TransferState::InProgress { .. }
| TransferState::Paused { .. } => {
stats.active_transfers += 1;
}
TransferState::Completed => {
stats.completed_transfers += 1;
stats.total_bytes_transferred += transfer.total_bytes;
}
TransferState::Failed { .. } => {
stats.failed_transfers += 1;
}
}
}
stats
}
pub fn advance_tick(&mut self) {
self.current_tick += 1;
}
}
impl Default for StreamingBlockTransfer {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_chunk(index: u32, data: Vec<u8>) -> TransferChunk {
let checksum = TransferChunk::compute_checksum(&data);
TransferChunk {
chunk_index: index,
data,
checksum,
}
}
#[test]
fn test_start_transfer_unique_ids() {
let mut mgr = StreamingBlockTransfer::new();
let id1 = mgr.start_transfer(
"cid1".into(),
"peer1".into(),
TransferDirection::Download,
1024,
512,
);
let id2 = mgr.start_transfer(
"cid2".into(),
"peer2".into(),
TransferDirection::Upload,
2048,
512,
);
assert_ne!(id1, id2);
assert!(id2 > id1);
}
#[test]
fn test_start_transfer_pending_state() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer(
"cid".into(),
"peer".into(),
TransferDirection::Download,
100,
64,
);
let t = &mgr.transfers[&id];
assert_eq!(t.state, TransferState::Pending);
}
#[test]
fn test_total_chunks_exact_multiple() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer(
"c".into(),
"p".into(),
TransferDirection::Download,
65536,
65536,
);
assert_eq!(mgr.transfers[&id].total_chunks(), 1);
}
#[test]
fn test_total_chunks_partial() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer(
"c".into(),
"p".into(),
TransferDirection::Download,
65537,
65536,
);
assert_eq!(mgr.transfers[&id].total_chunks(), 2);
}
#[test]
fn test_total_chunks_zero_bytes() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer(
"c".into(),
"p".into(),
TransferDirection::Download,
0,
65536,
);
assert_eq!(mgr.transfers[&id].total_chunks(), 0);
}
#[test]
fn test_receive_chunk_updates_state() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 20, 10);
let chunk = make_chunk(0, vec![1u8; 10]);
mgr.receive_chunk(id, chunk)
.expect("receive_chunk should succeed");
match &mgr.transfers[&id].state {
TransferState::InProgress {
chunks_done,
total_chunks,
} => {
assert_eq!(*chunks_done, 1);
assert_eq!(*total_chunks, 2);
}
other => panic!("expected InProgress, got {other:?}"),
}
}
#[test]
fn test_receive_chunk_checksum_mismatch() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 10, 10);
let chunk = TransferChunk {
chunk_index: 0,
data: vec![42u8; 10],
checksum: 0xDEAD_BEEF, };
let result = mgr.receive_chunk(id, chunk);
assert!(result.is_err(), "expected error for checksum mismatch");
assert!(result.unwrap_err().contains("checksum mismatch"));
}
#[test]
fn test_all_chunks_received_completes_transfer() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 30, 10);
for i in 0..3u32 {
let chunk = make_chunk(i, vec![i as u8; 10]);
mgr.receive_chunk(id, chunk)
.expect("chunk should be accepted");
}
assert_eq!(mgr.transfers[&id].state, TransferState::Completed);
assert!(mgr.transfers[&id].is_complete());
}
#[test]
fn test_duplicate_chunk_rejected() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 20, 10);
let chunk0 = make_chunk(0, vec![1u8; 10]);
mgr.receive_chunk(id, chunk0.clone())
.expect("first receive should succeed");
let dup = make_chunk(0, vec![1u8; 10]);
let result = mgr.receive_chunk(id, dup);
assert!(result.is_err());
assert!(result.unwrap_err().contains("already received"));
}
#[test]
fn test_pause_in_progress() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 30, 10);
let chunk = make_chunk(0, vec![0u8; 10]);
mgr.receive_chunk(id, chunk)
.expect("test: receive_chunk should accept valid chunk");
assert!(mgr.pause(id));
match &mgr.transfers[&id].state {
TransferState::Paused { chunks_done } => assert_eq!(*chunks_done, 1),
other => panic!("expected Paused, got {other:?}"),
}
}
#[test]
fn test_pause_pending() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 100, 10);
assert!(mgr.pause(id));
assert_eq!(
mgr.transfers[&id].state,
TransferState::Paused { chunks_done: 0 }
);
}
#[test]
fn test_resume_paused() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 30, 10);
let chunk = make_chunk(0, vec![7u8; 10]);
mgr.receive_chunk(id, chunk)
.expect("test: receive_chunk should accept valid chunk");
mgr.pause(id);
assert!(mgr.resume(id));
match &mgr.transfers[&id].state {
TransferState::InProgress {
chunks_done,
total_chunks,
} => {
assert_eq!(*chunks_done, 1);
assert_eq!(*total_chunks, 3);
}
other => panic!("expected InProgress, got {other:?}"),
}
}
#[test]
fn test_resume_non_paused_returns_false() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 10, 10);
assert!(!mgr.resume(id)); }
#[test]
fn test_fail_sets_failed_state() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 100, 10);
assert!(mgr.fail(id, "network timeout".into()));
match &mgr.transfers[&id].state {
TransferState::Failed { reason } => assert_eq!(reason, "network timeout"),
other => panic!("expected Failed, got {other:?}"),
}
}
#[test]
fn test_fail_on_completed_returns_false() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 10, 10);
let chunk = make_chunk(0, vec![0u8; 10]);
mgr.receive_chunk(id, chunk)
.expect("test: receive_chunk should accept valid chunk");
assert!(mgr.transfers[&id].is_complete());
assert!(!mgr.fail(id, "too late".into()));
}
#[test]
fn test_progress_values() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 40, 10);
assert_eq!(mgr.transfers[&id].progress(), 0.0);
let chunk0 = make_chunk(0, vec![1u8; 10]);
mgr.receive_chunk(id, chunk0)
.expect("test: receive_chunk should accept chunk 0");
let p = mgr.transfers[&id].progress();
assert!((p - 0.25).abs() < f64::EPSILON, "expected 0.25 got {p}");
let chunk1 = make_chunk(1, vec![2u8; 10]);
mgr.receive_chunk(id, chunk1)
.expect("test: receive_chunk should accept chunk 1");
let p = mgr.transfers[&id].progress();
assert!((p - 0.5).abs() < f64::EPSILON, "expected 0.5 got {p}");
}
#[test]
fn test_stats_counts() {
let mut mgr = StreamingBlockTransfer::new();
mgr.start_transfer(
"c1".into(),
"p1".into(),
TransferDirection::Download,
10,
10,
);
mgr.start_transfer("c2".into(), "p2".into(), TransferDirection::Upload, 10, 10);
let id_done = mgr.start_transfer(
"c3".into(),
"p3".into(),
TransferDirection::Download,
10,
10,
);
let chunk = make_chunk(0, vec![5u8; 10]);
mgr.receive_chunk(id_done, chunk)
.expect("test: receive_chunk should accept valid chunk for completed transfer");
let id_fail = mgr.start_transfer(
"c4".into(),
"p4".into(),
TransferDirection::Download,
10,
10,
);
mgr.fail(id_fail, "oops".into());
let stats = mgr.stats();
assert_eq!(stats.active_transfers, 2);
assert_eq!(stats.completed_transfers, 1);
assert_eq!(stats.failed_transfers, 1);
assert_eq!(stats.total_bytes_transferred, 10);
}
#[test]
fn test_advance_tick() {
let mut mgr = StreamingBlockTransfer::new();
assert_eq!(mgr.current_tick, 0);
mgr.advance_tick();
mgr.advance_tick();
assert_eq!(mgr.current_tick, 2);
}
#[test]
fn test_started_at_tick_recorded() {
let mut mgr = StreamingBlockTransfer::new();
mgr.advance_tick();
mgr.advance_tick();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Upload, 10, 10);
assert_eq!(mgr.transfers[&id].started_at_tick, 2);
}
#[test]
fn test_receive_chunk_rejected_when_paused() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 20, 10);
mgr.pause(id);
let chunk = make_chunk(0, vec![3u8; 10]);
let result = mgr.receive_chunk(id, chunk);
assert!(result.is_err());
assert!(result.unwrap_err().contains("paused"));
}
#[test]
fn test_is_complete_false_for_in_progress() {
let mut mgr = StreamingBlockTransfer::new();
let id = mgr.start_transfer("c".into(), "p".into(), TransferDirection::Download, 20, 10);
let chunk = make_chunk(0, vec![9u8; 10]);
mgr.receive_chunk(id, chunk)
.expect("test: receive_chunk should accept valid chunk");
assert!(!mgr.transfers[&id].is_complete());
}
#[test]
fn test_receive_chunk_unknown_id() {
let mut mgr = StreamingBlockTransfer::new();
let chunk = make_chunk(0, vec![1u8]);
let result = mgr.receive_chunk(9999, chunk);
assert!(result.is_err());
assert!(result.unwrap_err().contains("unknown transfer id"));
}
}