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