use std::ops::Deref; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::thread; use std::time::Duration; use anyhow::{anyhow, Result}; use signal_hook::consts::{SIGINT, SIGTERM}; use tokio::fs::File; use tokio::io::AsyncWriteExt; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::sync::mpsc::Sender; use tokio::task::JoinHandle; use signal_hook::iterator::Signals; use log::{debug, error, info}; use tokio::time::{sleep, timeout}; fn setup_graceful_shutdown(cancelled: Arc) { thread::spawn(move || { let signals = Signals::new([SIGINT, SIGTERM]); match signals { Ok(mut signal_info) => { if signal_info.forever().next().is_some() { info!("user request cancellation"); cancelled.store(true, Ordering::Relaxed); } } Err(error) => { error!("Failed to setup signal handler: {error}") } } }); } async fn download_image(tx: Sender>, url: &str) { info!("Downloading image for {}", url); let resp = reqwest::get(url) .await; if let Ok(resp) = resp { let image_bytes = resp.bytes().await; match image_bytes { Ok(data) => { tx.send(Ok(data)).await.expect("failed to send message"); } Err(err) => { tx.send(Err(anyhow!("failed to get image data: {}", err))).await.expect("failed to send message"); } } } else { tx.send(Err(anyhow!("failed to get response"))).await.expect("failed to send message"); } } #[tokio::main] async fn main() -> Result<()> { env_logger::init(); let task_timeout = Duration::from_millis(500); let (tx, mut rx) = mpsc::channel(32); let cancelled = Arc::from(AtomicBool::new(false)); setup_graceful_shutdown(cancelled.clone()); let urls: Vec<&str> = vec![ "https://images.unsplash.com/photo-1739117956532-8a37c76093dd?fm=jpg", "https://images.unsplash.com/photo-1735930371721-d1243e8122ab?fm=jpg", "https://images.unsplash.com/photo-1732204662871-f853329ce638?fm=jpg", "https://images.unsplash.com/photo-1712677925320-5f93d7f853f5?fm=jpg", "https://images.unsplash.com/photo-1712677925900-12574cb3cf15?fm=jpg", "https://images.unsplash.com/photo-1710436000845-bb707af976a6?fm=jpg", "https://images.unsplash.com/photo-1700386277812-5ff74eb03650?fm=jpg", "https://images.unsplash.com/photo-1700386277415-c9c5270ee8b8?fm=jpg", "https://images.unsplash.com/photo-1691947563165-28011f40d786?fm=jpg", "https://images.unsplash.com/photo-1691782359338-97fcbcce8ce8?fm=jpg" ]; let mut handles: Vec> = Vec::new(); for url in urls { let tx = tx.clone(); let handle = tokio::spawn(async move { let timeout_status = timeout(task_timeout, download_image(tx, url)).await; debug!("timeout status = {:?} for url {}", timeout_status, url) }); handles.push(handle); } let wait_cancelled = cancelled.clone(); let wait_task = tokio::spawn(async move { for handle in handles { if wait_cancelled.load(Ordering::Relaxed) { handle.abort(); } else { let _ = handle.await; } } drop(tx); }); let mut counter = 1; while let Some(message) = rx.recv().await { if cancelled.load(Ordering::Relaxed) { break } if let Ok(image_bytes) = message { let file = File::create(format!("{}.jpg", counter)).await; if let Ok(mut file) = file { let _ = file.write_all(&image_bytes).await; info!("Wrote image {}", counter); } else { error!("Failed to write image {}", counter) } counter += 1; } } let _ = wait_task.await; Ok(()) }