ic_bn_lib/
lib.rs

1// Needed for certain macros
2#![recursion_limit = "256"]
3#![warn(clippy::nursery)]
4#![warn(tail_expr_drop_order)]
5#![allow(clippy::cognitive_complexity)]
6#![allow(clippy::field_reassign_with_default)]
7
8#[cfg(feature = "custom-domains")]
9pub mod custom_domains;
10pub mod http;
11pub mod pubsub;
12pub mod tasks;
13pub mod tests;
14pub mod tls;
15pub mod utils;
16#[cfg(feature = "vector")]
17pub mod vector;
18
19use std::{fs::File, path::Path};
20
21use anyhow::{Context, anyhow};
22use bytes::Bytes;
23use futures::StreamExt;
24use ic_bn_lib_common::Error;
25use tokio::io::AsyncWriteExt;
26
27pub use hyper;
28pub use hyper_util;
29pub use ic_agent;
30pub use prometheus;
31pub use reqwest;
32pub use rustls;
33pub use uuid;
34
35/// Error to be used with `retry_async` macro
36/// which indicates whether it should be retried or not.
37#[derive(thiserror::Error, Debug)]
38pub enum RetryError {
39    #[error("Permanent error: {0:?}")]
40    Permanent(anyhow::Error),
41    #[error("Transient error: {0:?}")]
42    Transient(anyhow::Error),
43}
44
45/// Downloads the given url to given path.
46/// Destination folder must exist.
47pub fn download_url_to(url: &str, path: &Path) -> Result<u64, Error> {
48    let mut r = reqwest::blocking::get(url).context("unable to perform HTTP request")?;
49    if !r.status().is_success() {
50        return Err(anyhow!("incorrect HTTP code: {}", r.status()).into());
51    }
52
53    let mut file = File::create(path).context("could not create file")?;
54    Ok(r.copy_to(&mut file)
55        .context("unable to write body to file")?)
56}
57
58/// Downloads the given url and returns it as Bytes
59pub fn download_url(url: &str) -> Result<Bytes, Error> {
60    let r = reqwest::blocking::get(url).context("unable to perform HTTP request")?;
61    if !r.status().is_success() {
62        return Err(anyhow!("incorrect HTTP code: {}", r.status()).into());
63    }
64
65    Ok(r.bytes().context("unable to fetch file")?)
66}
67
68/// Downloads the given url to given path.
69/// Destination folder must exist.
70pub async fn download_url_to_async(url: &str, path: &Path) -> Result<(), Error> {
71    let r = reqwest::get(url)
72        .await
73        .context("unable to perform HTTP request")?;
74    if !r.status().is_success() {
75        return Err(anyhow!("incorrect HTTP code: {}", r.status()).into());
76    }
77
78    let mut file = tokio::fs::File::create(path)
79        .await
80        .context("could not create file")?;
81
82    let mut stream = r.bytes_stream();
83    while let Some(v) = stream.next().await {
84        file.write(&v.context("unable to read chunk")?)
85            .await
86            .context("unable to write chunk")?;
87    }
88
89    Ok(())
90}
91
92/// Downloads the given url and returns it as Bytes
93pub async fn download_url_async(url: &str) -> Result<Bytes, Error> {
94    let r = reqwest::get(url)
95        .await
96        .context("unable to perform HTTP request")?;
97
98    if !r.status().is_success() {
99        return Err(anyhow!("incorrect HTTP code: {}", r.status()).into());
100    }
101
102    Ok(r.bytes().await.context("unable to fetch file")?)
103}
104
105/// Retrying async closures/functions holding mutable references is a pain in Rust.
106/// So, for now, we'll have to use a macro to work that around.
107#[macro_export]
108macro_rules! retry_async {
109    ($f:expr, $timeout:expr, $delay:expr) => {{
110        use rand::{Rng, SeedableRng};
111        // SmallRng is Send which we require
112        let mut rng = rand::rngs::SmallRng::from_entropy();
113
114        let start = std::time::Instant::now();
115        let mut delay = $delay;
116
117        let result = loop {
118            // Run the function wrapping it into Tokio timeout future so
119            // its execution time doesn't exceed our configured limit
120            let Ok(res) = tokio::time::timeout($timeout, $f).await else {
121                break Err(anyhow::anyhow!("Timed out"));
122            };
123
124            let err = match res {
125                Ok(v) => break Ok(v),
126                Err($crate::RetryError::Permanent(e)) => break Err(e),
127                Err($crate::RetryError::Transient(e)) => e,
128            };
129
130            let left = $timeout.saturating_sub(start.elapsed());
131            if left == std::time::Duration::ZERO {
132                break Err(err);
133            }
134
135            delay = left.min(delay * 2);
136            // Generate a random jitter in 0.0..0.1 range
137            let jitter: f64 = (rng.r#gen::<f64>() / 10.0);
138            let d64 = delay.as_secs_f64();
139            delay = Duration::from_secs_f64(d64.mul_add(0.95, d64 * jitter));
140            tokio::time::sleep(delay).await;
141        };
142
143        result
144    }};
145
146    ($f:expr, $timeout:expr) => {
147        retry_async!($f, $timeout, Duration::from_millis(500))
148    };
149
150    ($f:expr) => {
151        retry_async!($f, Duration::from_secs(60), Duration::from_millis(500))
152    };
153}
154
155#[macro_export]
156macro_rules! dyn_event {
157    ($lvl:ident, $($arg:tt)+) => {
158        match $lvl {
159            ::tracing::Level::TRACE => ::tracing::trace!($($arg)+),
160            ::tracing::Level::DEBUG => ::tracing::debug!($($arg)+),
161            ::tracing::Level::INFO => ::tracing::info!($($arg)+),
162            ::tracing::Level::WARN => ::tracing::warn!($($arg)+),
163            ::tracing::Level::ERROR => ::tracing::error!($($arg)+),
164        }
165    };
166}