feat(stream-gears): 支持直播流下载的多条件分割 (#1400)

* remove: blive hack feature

* feat(stream-gears): 支持直播流下载的多条件分割

biliup/engine/download.py

@@ -473,6 +473,8 @@ class DownloadBase(ABC):
                 r = await __client_get(url, stream=True)
                 if r.headers.get('Location'):
                     url = r.headers['Location']
+                    # logger.info(f'{self.plugin_msg}: stream url: {url}')
+                    # url = self.replace_url(url)
                     logger.info(f'{self.plugin_msg}: stream url: {url}')
                     r = await __client_get(url, stream=True)
             if r.status_code == 200:
@@ -534,6 +536,7 @@ def stream_gears_download(url, headers, file_name, segment_time=None, file_size=
         segment.size = file_size
     if file_size is None and segment_time is None:
         segment.size = 8 * 1024 * 1024 * 1024
+    print('Python segment.__dict__:', segment.__dict__)
     # FIXME: 下载时如出现403，这里不会回到上层方法获取新链接
     if file_name_callback:
         stream_gears.download_with_callback(

crates/biliup/src/downloader/httpflv.rs

@@ -30,7 +30,7 @@ pub(crate) async fn parse_flv(
     mut segment: Segmentable,
 ) -> crate::downloader::error::Result<()> {
     let mut flv_tags_cache: Vec<(TagHeader, Bytes, Bytes)> = Vec::new();
-
+    // println!("parse_flv Segment: {:?}", segment);
     let _previous_tag_size = connection.read_frame(4).await?;
 
     let mut out = FlvFile::new(file)?;

crates/biliup/src/downloader/util.rs

@@ -13,17 +13,20 @@ pub enum Segment {
     Size(u64, u64),
     Never,
 }
+
 #[derive(Debug)]
 pub struct Segmentable {
     time: Time,
     size: Size,
 }
+
 #[derive(Debug)]
 struct Time {
     expected: Option<Duration>,
     start: Duration,
     current: Duration,
 }
+
 #[derive(Debug)]
 struct Size {
     expected: Option<u64>,
@@ -45,14 +48,82 @@ impl Segmentable {
         }
     }
 
+    /// 检查是否需要分割 - 只要时间或大小任一条件满足就返回 true
     pub fn needed(&self) -> bool {
+        // 检查时间条件
+        let time_exceeded = if let Some(expected_time) = self.time.expected {
+            (self.time.current - self.time.start) >= expected_time
+        } else {
+            false
+        };
+
+        // 检查大小条件
+        let size_exceeded = if let Some(expected_size) = self.size.expected {
+            self.size.current >= expected_size
+        } else {
+            false
+        };
+
+        // 任一条件满足就需要分割
+        let result = time_exceeded || size_exceeded;
+
+        // 添加调试信息
+        if result {
+            if time_exceeded && size_exceeded {
+                tracing::debug!(
+                    "Segmentation needed: Both time ({:?} >= {:?}) and size ({} >= {}) conditions met",
+                    self.time.current - self.time.start,
+                    self.time.expected.unwrap(),
+                    self.size.current,
+                    self.size.expected.unwrap()
+                );
+            } else if time_exceeded {
+                tracing::debug!(
+                    "Segmentation needed: Time condition met ({:?} >= {:?})",
+                    self.time.current - self.time.start,
+                    self.time.expected.unwrap()
+                );
+            } else if size_exceeded {
+                tracing::debug!(
+                    "Segmentation needed: Size condition met ({} >= {})",
+                    self.size.current,
+                    self.size.expected.unwrap()
+                );
+            }
+        }
+
+        result
+    }
+
+    /// 检查单独的时间条件
+    pub fn time_needed(&self) -> bool {
         if let Some(expected_time) = self.time.expected {
-            return (self.time.current - self.time.start) >= expected_time;
+            (self.time.current - self.time.start) >= expected_time
+        } else {
+            false
         }
+    }
+
+    /// 检查单独的大小条件
+    pub fn size_needed(&self) -> bool {
         if let Some(expected_size) = self.size.expected {
-            return self.size.current > expected_size;
+            self.size.current >= expected_size
+        } else {
+            false
+        }
+    }
+
+    /// 获取分割原因的描述
+    pub fn get_segment_reason(&self) -> String {
+        let time_exceeded = self.time_needed();
+        let size_exceeded = self.size_needed();
+
+        match (time_exceeded, size_exceeded) {
+            (true, true) => "Time and size limits reached".to_string(),
+            (true, false) => "Time limit reached".to_string(),
+            (false, true) => "Size limit reached".to_string(),
+            (false, false) => "No segmentation needed".to_string(),
         }
-        false
     }
 
     pub fn increase_time(&mut self, number: Duration) {
@@ -75,9 +146,44 @@ impl Segmentable {
         self.size.current = number
     }
 
+    /// 重置计数器，通常在创建新分割后调用
     pub fn reset(&mut self) {
+        self.size.current = 0;
+        self.time.start = self.time.current; // 保持当前时间位置，但重置起始点
+    }
+
+    /// 完全重置所有状态
+    pub fn full_reset(&mut self) {
         self.size.current = 0;
         self.time.current = Duration::ZERO;
+        self.time.start = Duration::ZERO;
+    }
+
+    /// 获取当前状态信息
+    pub fn get_status(&self) -> String {
+        let time_info = if let Some(expected) = self.time.expected {
+            format!(
+                "Time: {:.1}s/{:.1}s ({:.1}%)",
+                (self.time.current - self.time.start).as_secs_f64(),
+                expected.as_secs_f64(),
+                ((self.time.current - self.time.start).as_secs_f64() / expected.as_secs_f64() * 100.0).min(100.0)
+            )
+        } else {
+            "Time: No limit".to_string()
+        };
+
+        let size_info = if let Some(expected) = self.size.expected {
+            format!(
+                "Size: {}/{} bytes ({:.1}%)",
+                self.size.current,
+                expected,
+                (self.size.current as f64 / expected as f64 * 100.0).min(100.0)
+            )
+        } else {
+            "Size: No limit".to_string()
+        };
+
+        format!("{}, {}", time_info, size_info)
     }
 }
 
@@ -156,6 +262,7 @@ pub fn format_filename(file_name: &str) -> String {
 
 #[cfg(test)]
 mod tests {
+    use super::*;
     use anyhow::Result;
     use std::path::{Path, PathBuf};
 
@@ -171,4 +278,46 @@ mod tests {
 
         Ok(())
     }
-}
+
+    #[test]
+    fn test_segmentation_logic() -> Result<()> {
+        // 测试时间分割
+        let mut seg = Segmentable::new(Some(Duration::from_secs(10)), None);
+        assert!(!seg.needed());
+
+        seg.increase_time(Duration::from_secs(15));
+        assert!(seg.needed());
+        assert!(seg.time_needed());
+        assert!(!seg.size_needed());
+
+        // 测试大小分割
+        let mut seg = Segmentable::new(None, Some(1024));
+        assert!(!seg.needed());
+
+        seg.increase_size(2048);
+        assert!(seg.needed());
+        assert!(!seg.time_needed());
+        assert!(seg.size_needed());
+
+        // 测试双重条件
+        let mut seg = Segmentable::new(Some(Duration::from_secs(10)), Some(1024));
+        assert!(!seg.needed());
+
+        // 只满足时间条件
+        seg.increase_time(Duration::from_secs(15));
+        assert!(seg.needed());
+
+        // 重置并只满足大小条件
+        seg.full_reset();
+        seg.increase_size(2048);
+        assert!(seg.needed());
+
+        // 同时满足两个条件
+        seg.increase_time(Duration::from_secs(15));
+        assert!(seg.needed());
+        assert!(seg.time_needed());
+        assert!(seg.size_needed());
+
+        Ok(())
+    }
+}

crates/stream-gears/src/lib.rs

@@ -16,16 +16,12 @@ use biliup::downloader::extractor::CallbackFn;
 use biliup::downloader::util::Segmentable;
 use tracing_subscriber::layer::SubscriberExt;
 
-#[derive(FromPyObject)]
-pub enum PySegment {
-    Time {
-        #[pyo3(attribute("time"))]
-        time: u64,
-    },
-    Size {
-        #[pyo3(attribute("size"))]
-        size: u64,
-    },
+#[derive(FromPyObject, Debug)]
+pub struct PySegment {
+    #[pyo3(attribute("time"))]
+    time: Option<u64>,
+    #[pyo3(attribute("size"))]
+    size: Option<u64>,
 }
 
 #[pyfunction]
@@ -74,11 +70,29 @@ fn download_with_callback(
             .with_timer(local_time)
             .with_writer(non_blocking);
 
-        let segment = match segment {
-            PySegment::Time { time } => Segmentable::new(Some(Duration::from_secs(time)), None),
-            PySegment::Size { size } => Segmentable::new(None, Some(size)),
+        // println!("Input segment: {:?}", segment);
+        // println!("Input segment time: {:?}, size: {:?}", segment.time, segment.size);
+
+        let segmentable = match (segment.time, segment.size) {
+            (Some(time), Some(size)) => {
+                // 如果同时有时间和大小，可以选择优先使用哪个，或者创建支持两者的逻辑
+                // 这里假设优先使用时间分割
+                Segmentable::new(Some(Duration::from_secs(time)), Some(size))
+            }
+            (Some(time), None) => {
+                Segmentable::new(Some(Duration::from_secs(time)), None)
+            }
+            (None, Some(size)) => {
+                Segmentable::new(None, Some(size))
+            }
+            (None, None) => {
+                // 如果都没有，使用默认值
+                Segmentable::default()
+            }
         };
 
+        tracing::debug!("Segmentable: {:?}", segmentable);
+
         let file_name_hook = file_name_callback_fn.map(|callback_fn| -> CallbackFn {
             Box::new(move |fmt_file_name| {
                 Python::with_gil(|py| match callback_fn.call1(py, (fmt_file_name,)) {
@@ -96,7 +110,7 @@ fn download_with_callback(
                 url,
                 map,
                 file_name,
-                segment,
+                segmentable,
                 file_name_hook,
                 proxy.as_deref(),
             ) {

pyproject.toml

@@ -23,15 +23,15 @@ dependencies = [
     "tomli-w>=1.0.0",
     "psutil >= 5.4.6",
     "aiofiles >= 22.1.0",
-    "yt-dlp >= 2025.3.31",
+    "yt-dlp >= 2025.7.21",
     "Pillow >= 5.2.0",
-    "aiohttp[speedups] < 3.10.0", # aio-libs/aiohttp#10519
+    "aiohttp[speedups] >= 3.9.5", # aio-libs/aiohttp#10519
     "aiohttp_cors >= 0.7.0",
     "requests >= 2.32.3",
     "httpx[http2] >= 0.28.1",
     "httpcore >= 1.0.7",
     "PyYAML >= 4.2b1",
-    "streamlink >= 7.2.0",
+    "streamlink >= 7.5.0",
     "ykdl >= 1.8.0",
     "rsa >= 4.6",
     "tomli >= 1.1.0; python_version < '3.11'",