stormdl_segment/
manager.rs

1use crate::splitter::{initial_segments, split_range};
2use parking_lot::RwLock;
3use std::sync::Arc;
4use stormdl_core::{ByteRange, SegmentState, SegmentStatus};
5
6pub struct SegmentManager {
7    segments: Arc<RwLock<Vec<SegmentState>>>,
8    total_size: u64,
9    min_segment_size: u64,
10    max_segments: usize,
11}
12
13impl SegmentManager {
14    pub fn new(total_size: u64) -> Self {
15        Self {
16            segments: Arc::new(RwLock::new(Vec::new())),
17            total_size,
18            min_segment_size: 256 * 1024,
19            max_segments: 32,
20        }
21    }
22
23    pub fn with_config(total_size: u64, min_segment_size: u64, max_segments: usize) -> Self {
24        Self {
25            segments: Arc::new(RwLock::new(Vec::new())),
26            total_size,
27            min_segment_size,
28            max_segments,
29        }
30    }
31
32    pub fn with_segments(total_size: u64, num_segments: usize) -> Self {
33        let manager = Self::new(total_size);
34        let ranges = split_range(total_size, num_segments);
35        let segments: Vec<SegmentState> = ranges
36            .into_iter()
37            .enumerate()
38            .map(|(id, range)| SegmentState::new(id, range))
39            .collect();
40        *manager.segments.write() = segments;
41        manager
42    }
43
44    pub fn initialize(&self) -> Vec<SegmentState> {
45        let num_segments = initial_segments(self.total_size);
46        let ranges = split_range(self.total_size, num_segments);
47
48        let segments: Vec<SegmentState> = ranges
49            .into_iter()
50            .enumerate()
51            .map(|(id, range)| SegmentState::new(id, range))
52            .collect();
53
54        *self.segments.write() = segments.clone();
55        segments
56    }
57
58    pub fn get_segments(&self) -> Vec<SegmentState> {
59        self.segments.read().clone()
60    }
61
62    pub fn update_segment(&self, id: usize, downloaded: u64, speed: f64) {
63        let mut segments = self.segments.write();
64        if let Some(segment) = segments.get_mut(id) {
65            segment.downloaded = downloaded;
66            segment.speed = speed;
67        }
68    }
69
70    pub fn mark_complete(&self, id: usize) {
71        let mut segments = self.segments.write();
72        if let Some(segment) = segments.get_mut(id) {
73            segment.status = SegmentStatus::Complete;
74            segment.downloaded = segment.range.len();
75        }
76    }
77
78    pub fn mark_error(&self, id: usize) {
79        let mut segments = self.segments.write();
80        if let Some(segment) = segments.get_mut(id) {
81            segment.status = SegmentStatus::Error;
82        }
83    }
84
85    pub fn split_segment(&self, id: usize) -> Option<SegmentState> {
86        let mut segments = self.segments.write();
87
88        if segments.len() >= self.max_segments {
89            return None;
90        }
91
92        let segment = segments.get(id)?;
93        let remaining = segment.remaining();
94
95        if remaining < self.min_segment_size * 2 {
96            return None;
97        }
98
99        let current_offset = segment.range.start + segment.downloaded;
100        let split_point = current_offset + remaining / 2;
101
102        let new_id = segments.len();
103        let new_range = ByteRange::new(split_point, segment.range.end);
104
105        segments.get_mut(id)?.range.end = split_point;
106
107        let new_segment = SegmentState::new(new_id, new_range);
108        segments.push(new_segment.clone());
109
110        Some(new_segment)
111    }
112
113    pub fn total_downloaded(&self) -> u64 {
114        self.segments.read().iter().map(|s| s.downloaded).sum()
115    }
116
117    pub fn all_complete(&self) -> bool {
118        self.segments
119            .read()
120            .iter()
121            .all(|s| s.status == SegmentStatus::Complete)
122    }
123
124    pub fn average_speed(&self) -> f64 {
125        let segments = self.segments.read();
126        let active: Vec<_> = segments
127            .iter()
128            .filter(|s| s.status == SegmentStatus::Active)
129            .collect();
130
131        if active.is_empty() {
132            return 0.0;
133        }
134
135        active.iter().map(|s| s.speed).sum::<f64>() / active.len() as f64
136    }
137}
stormdl_segment/manager.rs

stormdl_segment/
manager.rs
