Skip to main content

llama_server/
http.rs

1use crate::Error;
2
3use sipper::{Straw, sipper};
4use tokio::io::AsyncWrite;
5
6use std::sync::LazyLock;
7use std::time::Instant;
8
9pub fn client() -> reqwest::Client {
10    static CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
11        reqwest::Client::builder()
12            .user_agent(format!(
13                "{} {}",
14                env!("CARGO_PKG_NAME"),
15                env!("CARGO_PKG_VERSION")
16            ))
17            .build()
18            .expect("should be a valid client")
19    });
20
21    CLIENT.clone()
22}
23
24pub fn download<'a, W: AsyncWrite + Unpin>(
25    url: impl reqwest::IntoUrl + Send + 'a,
26    writer: &'a mut W,
27) -> impl Straw<(), Progress, Error> + 'a {
28    use tokio::io::AsyncWriteExt;
29
30    sipper(move |mut progress| async move {
31        let mut download = client().get(url).send().await?;
32        let start = Instant::now();
33        let total = download.content_length().unwrap_or_default();
34
35        let mut downloaded = 0;
36
37        progress
38            .send(Progress {
39                total,
40                downloaded,
41                speed: 0,
42            })
43            .await;
44
45        while let Some(chunk) = download.chunk().await? {
46            downloaded += chunk.len() as u64;
47            let speed = (downloaded as f32 / start.elapsed().as_secs_f32()) as u64;
48
49            progress
50                .send(Progress {
51                    total,
52                    downloaded,
53                    speed,
54                })
55                .await;
56
57            writer.write_all(&chunk).await?;
58        }
59
60        writer.flush().await?;
61
62        Ok(())
63    })
64}
65
66/// The progress of an HTTP download.
67#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
68pub struct Progress {
69    /// The amount of bytes downloaded.
70    pub downloaded: u64,
71    /// The total of bytes to download.
72    pub total: u64,
73    /// The current download speed, in bytes per second.
74    pub speed: u64,
75}