feat: 为下载加入 flush 与 content-length 检查

This commit is contained in:
amtoaer
2025-01-22 00:18:04 +08:00
parent 18c966a0f9
commit 6ae87364b4
+22 -9
View File
@@ -1,10 +1,11 @@
use core::str;
use std::path::Path; use std::path::Path;
use anyhow::{anyhow, Result}; use anyhow::{bail, ensure, Result};
use futures::StreamExt; use futures::StreamExt;
use reqwest::Method; use reqwest::Method;
use tokio::fs::{self, File}; use tokio::fs::{self, File};
use tokio::io; use tokio::io::{self, AsyncWriteExt};
use crate::bilibili::Client; use crate::bilibili::Client;
pub struct Downloader { pub struct Downloader {
@@ -24,10 +25,25 @@ impl Downloader {
fs::create_dir_all(parent).await?; fs::create_dir_all(parent).await?;
} }
let mut file = File::create(path).await?; let mut file = File::create(path).await?;
let mut res = self.client.request(Method::GET, url, None).send().await?.bytes_stream(); let resp = self.client.request(Method::GET, url, None).send().await?;
while let Some(item) = res.next().await { let expected = resp.content_length().unwrap_or_else(|| {
io::copy(&mut item?.as_ref(), &mut file).await?; warn!("content length is missing, fallback to 0");
0
});
let mut received = 0u64;
let mut stream = resp.bytes_stream();
while let Some(bytes) = stream.next().await {
let bytes = bytes?;
received += bytes.len() as u64;
io::copy(&mut bytes.as_ref(), &mut file).await?;
} }
file.flush().await?;
ensure!(
received >= expected,
"received {} bytes, expected {} bytes",
received,
expected
);
Ok(()) Ok(())
} }
@@ -46,10 +62,7 @@ impl Downloader {
.output() .output()
.await?; .await?;
if !output.status.success() { if !output.status.success() {
return match String::from_utf8(output.stderr) { bail!("ffmpeg error: {}", str::from_utf8(&output.stderr).unwrap_or("unknown"));
Ok(err) => Err(anyhow!(err)),
_ => Err(anyhow!("ffmpeg error")),
};
} }
Ok(()) Ok(())
} }