use CancellationToken

Author: clitic

Cargo.lock

@@ -46,6 +46,7 @@ tokio = { version = "1", features = [
     "signal",
 ] }
 tokio-stream = { version = "0.1", optional = true }
+tokio-util = "0.7"
 url = "2"
 vsd-mp4 = { version = "0.2.1", path = "../vsd-mp4", features = ["full"] }
 widevine = { version = "0.1.0", optional = true }

vsd/Cargo.toml

@@ -3,16 +3,14 @@
 // [dependencies]
 // vsd = { version = "0.5", default-features = false, features = ["rustls-tls"]}
 
-use std::{
-    path::PathBuf,
-    sync::{Arc, atomic::AtomicBool},
-};
+use std::{path::PathBuf, sync::Arc};
 use vsd::{
     Downloader, Error, Muxer, Result,
     playlist::MediaType,
     progress::{ProgressCallback, ProgressState},
     reqwest::Client,
     tokio,
+    tokio_util::sync::CancellationToken,
 };
 
 struct Progress;
@@ -46,8 +44,8 @@ async fn main() -> Result<()> {
         )
         .await?;
 
-    // You can clone this var and pause download by setting its value to false.
-    let running = Arc::new(AtomicBool::new(true));
+    // You can clone this token and call .cancel() to pause a download.
+    let token = CancellationToken::new();
     let mut muxer = Muxer(Vec::new());
 
     // Download first subtitle stream.
@@ -59,7 +57,7 @@ async fn main() -> Result<()> {
             );
 
             // If stream is already downloaded then no progress updates will be triggered.
-            let dl_info = match stream.download(&config, &running, Arc::new(Progress)).await {
+            let dl_info = match stream.download(&config, Arc::new(Progress), &token).await {
                 Ok(info) => info,
                 Err(Error::UnsupportedEncryption(e)) => {
                     println!("Unsupported encryption {}", e);
@@ -69,8 +67,8 @@ async fn main() -> Result<()> {
                     println!("Stream has no segments");
                     continue;
                 }
-                Err(Error::MissingKey(kid)) => {
-                    println!("Missing decryption key for {kid}");
+                Err(Error::MissingKey(key_id)) => {
+                    println!("Missing decryption key for {key_id}");
                     continue;
                 }
                 Err(Error::DownloadInterrupted) => {

vsd/examples/download.rs

@@ -6,13 +6,8 @@ use crate::{
 };
 use colored::Colorize;
 use log::{info, warn};
-use std::{
-    collections::HashSet,
-    sync::{
-        Arc,
-        atomic::{AtomicBool, Ordering},
-    },
-};
+use std::collections::HashSet;
+use tokio_util::sync::CancellationToken;
 use vsd_mp4::{boxes::TencBox, pssh::PsshBox};
 
 pub async fn download_streams(
@@ -23,13 +18,13 @@ pub async fn download_streams(
         dump_pssh_info(config, &streams).await?;
     }
 
-    let running = Arc::new(AtomicBool::new(true));
-    let ctrlc = running.clone();
+    let token = CancellationToken::new();
+    let ctrlc_token = token.clone();
 
     tokio::spawn(async move {
-        if tokio::signal::ctrl_c().await.is_ok() && ctrlc.load(Ordering::SeqCst) {
+        if tokio::signal::ctrl_c().await.is_ok() && !ctrlc_token.is_cancelled() {
             warn!("Aborting download due to Ctrl+C.");
-            ctrlc.store(false, Ordering::SeqCst);
+            ctrlc_token.cancel();
         }
 
         if tokio::signal::ctrl_c().await.is_ok() {
@@ -39,7 +34,6 @@ pub async fn download_streams(
     });
 
     let mut muxer = Muxer(Vec::new());
-    let running = running.clone();
     let total = streams.len();
 
     for (i, stream) in streams.iter().enumerate() {
@@ -60,11 +54,11 @@ pub async fn download_streams(
         if stream.media_type == MediaType::Subtitles {
             muxer
                 .0
-                .push(sub::download(config, &running, pb, stream).await?);
+                .push(sub::download(config, pb, &token, stream).await?);
         } else {
             muxer
                 .0
-                .push(vid::download(config, &running, pb, stream).await?);
+                .push(vid::download(config, pb, &token, stream).await?);
         }
     }
 

vsd/src/core/dl.rs

@@ -8,8 +8,8 @@ use crate::{
 use colored::Colorize;
 use log::{debug, info, warn};
 use reqwest::{Url, header};
-use std::sync::atomic::{AtomicBool, Ordering};
 use tokio::{fs::File, io::AsyncWriteExt, task::JoinSet};
+use tokio_util::sync::CancellationToken;
 use vsd_mp4::sub::{StppSubsParser, WvttSubsParser, ttml};
 
 enum SubtitleType {
@@ -49,8 +49,8 @@ fn detect_codec(codecs: Option<&str>, data: &[u8], ext: &str) -> (&'static str,
 
 pub async fn download(
     config: &DownloadConfig,
-    running: &AtomicBool,
-    pb: Progress,
+    progress: Progress,
+    token: &CancellationToken,
     stream: &MediaPlaylist,
 ) -> Result<Stream> {
     let base_url = stream.uri.parse::<Url>()?;
@@ -99,18 +99,18 @@ pub async fn download(
         );
     }
 
-    pb.update(size);
+    progress.update(size);
 
     let remaining = &stream.segments[1..];
 
     if !remaining.is_empty() {
-        let pb_handle = pb.spawn();
+        let progress_handle = progress.spawn();
         let max_threads = config.max_threads as usize;
         let mut set: JoinSet<Result<(usize, Vec<u8>)>> = JoinSet::new();
         let mut results = vec![None; remaining.len()];
 
         for (i, segment) in remaining.iter().enumerate() {
-            if !running.load(Ordering::SeqCst) {
+            if token.is_cancelled() {
                 break;
             }
 
@@ -123,7 +123,7 @@ pub async fn download(
                             return Err(e);
                         }
                     };
-                    pb.update(bytes.len());
+                    progress.update(bytes.len());
                     results[i] = Some(bytes);
                 }
             }
@@ -150,20 +150,20 @@ pub async fn download(
                     return Err(e);
                 }
             };
-            pb.update(bytes.len());
+            progress.update(bytes.len());
             results[i] = Some(bytes);
         }
 
         for mut bytes in results.into_iter().flatten() {
             data.append(&mut bytes);
         }
 
-        pb_handle.abort();
+        progress_handle.abort();
     }
 
-    pb.finish();
+    progress.finish();
 
-    if !running.load(Ordering::SeqCst) {
+    if token.is_cancelled() {
         return Err(Error::DownloadInterrupted);
     }
 

vsd/src/core/sub.rs

@@ -7,15 +7,13 @@ use crate::{
 use colored::Colorize;
 use log::{debug, info, trace, warn};
 use reqwest::{StatusCode, Url, header};
-use std::sync::{
-    Arc,
-    atomic::{AtomicBool, Ordering},
-};
+use std::sync::Arc;
 use tokio::{
     fs::{self, File},
     io::{self, AsyncWriteExt},
     task::JoinSet,
 };
+use tokio_util::sync::CancellationToken;
 use vsd_mp4::{
     boxes::TencBox,
     decrypt::{CencDecrypter, HlsAes128Decrypter, HlsSampleAesDecrypter},
@@ -26,8 +24,8 @@ const PNG_HEADER: [u8; 8] = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
 
 pub async fn download(
     config: &DownloadConfig,
-    running: &AtomicBool,
-    pb: Progress,
+    progress: Progress,
+    token: &CancellationToken,
     stream: &MediaPlaylist,
 ) -> Result<Stream> {
     if let Some(Segment {
@@ -69,7 +67,7 @@ pub async fn download(
 
     let base_url = stream.uri.parse::<Url>()?;
     let ext = stream.extension();
-    let pb_handle = pb.spawn();
+    let progress_handle = progress.spawn();
     let temp_dir = temp_file.with_extension("");
     let mut auto_increment_iv = false;
     let mut decrypter = Decrypter::None;
@@ -97,14 +95,14 @@ pub async fn download(
     let mut set: JoinSet<Result<usize>> = JoinSet::new();
 
     for (i, segment) in stream.segments.iter().enumerate() {
-        if !running.load(Ordering::SeqCst) {
+        if token.is_cancelled() {
             break;
         }
 
         while set.len() >= max_threads {
             if let Some(Ok(result)) = set.join_next().await {
                 match result {
-                    Ok(bytes) => pb.update(bytes),
+                    Ok(bytes) => progress.update(bytes),
                     Err(e) => {
                         set.abort_all();
                         return Err(e);
@@ -190,7 +188,7 @@ pub async fn download(
 
         if out_file.exists() {
             let size = fs::metadata(&out_file).await?.len();
-            pb.skip(size as usize);
+            progress.skip(size as usize);
             continue;
         }
 
@@ -261,18 +259,18 @@ pub async fn download(
 
     while let Some(Ok(result)) = set.join_next().await {
         match result {
-            Ok(bytes) => pb.update(bytes),
+            Ok(bytes) => progress.update(bytes),
             Err(e) => {
                 set.abort_all();
                 return Err(e);
             }
         }
     }
 
-    pb_handle.abort();
-    pb.finish();
+    progress_handle.abort();
+    progress.finish();
 
-    if !running.load(Ordering::SeqCst) {
+    if token.is_cancelled() {
         return Err(Error::DownloadInterrupted);
     }
 

vsd/src/core/vid.rs

@@ -21,5 +21,6 @@ pub use core::{DownloadConfig, Downloader, Muxer, Stream};
 pub use error::{Error, Result};
 pub use reqwest;
 pub use tokio;
+pub use tokio_util;
 pub use utils::find_ffmpeg;
 pub use vsd_mp4;

vsd/src/lib.rs

@@ -12,13 +12,8 @@ use reqwest::{
     header::{self, HeaderValue},
 };
 use serde::Serialize;
-use std::{
-    cmp::Reverse,
-    collections::HashSet,
-    fmt::Display,
-    path::PathBuf,
-    sync::{Arc, atomic::AtomicBool},
-};
+use std::{cmp::Reverse, collections::HashSet, fmt::Display, path::PathBuf, sync::Arc};
+use tokio_util::sync::CancellationToken;
 use vsd_mp4::{boxes::TencBox, pssh::PsshBox};
 
 #[derive(Debug, Serialize)]
@@ -277,18 +272,18 @@ impl MediaPlaylist {
     pub async fn download(
         &self,
         config: &DownloadConfig,
-        running: &AtomicBool,
-        callback: Arc<dyn ProgressCallback>,
+        progress: Arc<dyn ProgressCallback>,
+        token: &CancellationToken,
     ) -> Result<Stream> {
         if self.segments.is_empty() {
             return Err(Error::MissingSegments);
         }
 
-        let pb = Progress::new(&self.id, self.segments.len(), Some(callback));
+        let progress = Progress::new(&self.id, self.segments.len(), Some(progress));
         let temp_file = if self.media_type == MediaType::Subtitles {
-            core::sub::download(config, running, pb, self).await?
+            core::sub::download(config, progress, token, self).await?
         } else {
-            core::vid::download(config, running, pb, self).await?
+            core::vid::download(config, progress, token, self).await?
         };
 
         Ok(temp_file)