use std::{str::FromStr, sync::Arc}; use async_trait::async_trait; use librqbit::{ AddTorrent, AddTorrentOptions, ManagedTorrent, Session, SessionOptions, api::TorrentIdOrHash, }; use librqbit_core::Id20; use snafu::ResultExt; use tracing::instrument; use util::errors::AnyhowResultExt; use super::task::{RqbitCreation, RqbitHash, RqbitSelector, RqbitState, RqbitTask}; use crate::{ DownloaderError, bittorrent::{ downloader::TorrentDownloaderTrait, source::{HashTorrentSource, HashTorrentSourceTrait}, }, core::{DownloadIdSelector, DownloaderTrait}, errors::RqbitSnafu, }; #[derive(Debug)] pub struct RqbitDownloaderCreation { pub save_path: String, pub subscriber_id: i32, pub downloader_id: i32, } impl RqbitDownloaderCreation {} pub struct RqbitDownloader { pub save_path: String, pub subscriber_id: i32, pub downloader_id: i32, pub session: Arc, } impl RqbitDownloader { #[instrument(level = "debug")] pub async fn from_creation( creation: RqbitDownloaderCreation, ) -> Result, DownloaderError> { let session_opt = SessionOptions { ..Default::default() }; let session = Session::new_with_opts(creation.save_path.clone().into(), session_opt) .await .to_dyn_boxed() .context(RqbitSnafu {})?; Ok(Arc::new(Self { session, save_path: creation.save_path, subscriber_id: creation.subscriber_id, downloader_id: creation.downloader_id, })) } pub async fn add_torrent( &self, source: HashTorrentSource, opt: Option, ) -> Result { let hash = Id20::from_str(&source.hash_info() as &str) .to_dyn_boxed() .context(RqbitSnafu {})?; let source = match source { HashTorrentSource::TorrentFile(file) => AddTorrent::TorrentFileBytes(file.payload), HashTorrentSource::MagnetUrl(magnet) => AddTorrent::Url(magnet.url.into()), }; let response = self .session .add_torrent(source, opt) .await .to_dyn_boxed() .context(RqbitSnafu {})?; let handle = response .into_handle() .ok_or_else(|| anyhow::anyhow!("failed to get handle of add torrent task")) .to_dyn_boxed() .context(RqbitSnafu {})?; handle .wait_until_initialized() .await .to_dyn_boxed() .context(RqbitSnafu {})?; Ok(hash) } fn query_torrent_impl(&self, hash: RqbitHash) -> Result, DownloaderError> { let torrent = self .session .get(TorrentIdOrHash::Hash(hash)) .ok_or_else(|| anyhow::anyhow!("could not find torrent by hash {}", hash.as_string())) .to_dyn_boxed() .context(RqbitSnafu {})?; Ok(torrent) } pub fn query_torrent(&self, hash: RqbitHash) -> Result { let torrent = self.query_torrent_impl(hash)?; let task = RqbitTask::from_query(torrent)?; Ok(task) } pub async fn pause_torrent(&self, hash: RqbitHash) -> Result<(), DownloaderError> { let t = self.query_torrent_impl(hash)?; self.session .pause(&t) .await .to_dyn_boxed() .context(RqbitSnafu {})?; Ok(()) } pub async fn resume_torrent(&self, hash: RqbitHash) -> Result<(), DownloaderError> { let t = self.query_torrent_impl(hash)?; self.session .unpause(&t) .await .to_dyn_boxed() .context(RqbitSnafu {})?; Ok(()) } pub async fn delete_torrent(&self, hash: RqbitHash) -> Result<(), DownloaderError> { self.session .delete(TorrentIdOrHash::Hash(hash), true) .await .to_dyn_boxed() .context(RqbitSnafu {})?; Ok(()) } } #[async_trait] impl DownloaderTrait for RqbitDownloader { type State = RqbitState; type Id = RqbitHash; type Task = RqbitTask; type Creation = RqbitCreation; type Selector = RqbitSelector; #[instrument(level = "debug", skip(self))] async fn add_downloads( &self, creation: RqbitCreation, ) -> Result::Id>, DownloaderError> { let mut sources = creation.sources; if sources.len() == 1 { let hash = self .add_torrent( sources.pop().unwrap(), Some(AddTorrentOptions { paused: false, output_folder: Some(self.save_path.clone()), ..Default::default() }), ) .await?; Ok(vec![hash]) } else { let tasks = sources .into_iter() .map(|s| { self.add_torrent( s, Some(AddTorrentOptions { paused: false, output_folder: Some(self.save_path.clone()), ..Default::default() }), ) }) .collect::>(); let results = futures::future::try_join_all(tasks).await?; Ok(results) } } async fn pause_downloads( &self, selector: ::Selector, ) -> Result, DownloaderError> { ::pause_downloads(self, selector).await } async fn resume_downloads( &self, selector: ::Selector, ) -> Result, DownloaderError> { ::resume_downloads(self, selector).await } async fn remove_downloads( &self, selector: ::Selector, ) -> Result, DownloaderError> { ::remove_downloads(self, selector).await } #[instrument(level = "debug", skip(self))] async fn query_downloads( &self, selector: RqbitSelector, ) -> Result::Task>, DownloaderError> { let hashes = selector.into_iter(); let tasks = hashes .map(|h| self.query_torrent(h)) .collect::, DownloaderError>>()?; Ok(tasks) } } #[async_trait] impl TorrentDownloaderTrait for RqbitDownloader { type IdSelector = DownloadIdSelector; #[instrument(level = "debug", skip(self))] async fn pause_torrents( &self, selector: Self::IdSelector, ) -> Result { let mut hashes: Vec<_> = selector.clone(); if hashes.len() == 1 { self.pause_torrent(hashes.pop().unwrap()).await?; } else { futures::future::try_join_all(hashes.into_iter().map(|h| self.pause_torrent(h))) .await?; } Ok(selector) } #[instrument(level = "debug", skip(self))] async fn resume_torrents( &self, selector: Self::IdSelector, ) -> Result { let mut hashes: Vec<_> = selector.clone(); if hashes.len() == 1 { self.resume_torrent(hashes.pop().unwrap()).await?; } else { futures::future::try_join_all(hashes.into_iter().map(|h| self.resume_torrent(h))) .await?; } Ok(selector) } #[instrument(level = "debug", skip(self))] async fn remove_torrents( &self, selector: Self::IdSelector, ) -> Result { let mut hashes: Vec<_> = selector.clone(); if hashes.len() == 1 { self.delete_torrent(hashes.pop().unwrap()).await?; } else { futures::future::try_join_all(hashes.into_iter().map(|h| self.delete_torrent(h))) .await?; } Ok(selector) } }