use std::collections::BTreeMap;
use std::fs::create_dir;
use std::path::PathBuf;
use crate::commands::*;
use crate::dependencies::*;
use crate::options::*;
use di::{Ref, inject, injectable};
use flat_db::{Hash, Table};
use futures::stream::{StreamExt, iter};
use log::error;
use rogue_logging::Error;
#[injectable]
pub(crate) struct Queue {
table: Table<20, 1, QueueItem>,
}
impl Queue {
pub(crate) fn from_path(path: PathBuf) -> Self {
Self {
table: Table::new(path),
}
}
#[inject]
pub(crate) fn from_options(options: Ref<CacheOptions>) -> Self {
let path = options.cache.clone().expect("queue path should be set");
let path = path.join("queue");
if !path.exists() {
create_dir(&path)
.expect("should be able to create queue directory if it does not exist");
}
Self::from_path(path)
}
pub(crate) fn get(&self, hash: Hash<20>) -> Result<Option<QueueItem>, Error> {
self.table.get(hash)
}
pub(crate) async fn get_unprocessed(
&mut self,
indexer: String,
transcode_enabled: bool,
upload_enabled: bool,
retry_failed_transcodes: bool,
) -> Result<Vec<Hash<20>>, Error> {
let is_red = indexer == "red";
let mut items = self
.get_unprocessed_internal(
indexer,
transcode_enabled,
upload_enabled,
retry_failed_transcodes,
)
.await?;
if is_red {
let mut pth_items = self
.get_unprocessed_internal(
"pth".to_owned(),
transcode_enabled,
upload_enabled,
retry_failed_transcodes,
)
.await?;
items.append(&mut pth_items);
}
Ok(items)
}
async fn get_unprocessed_internal(
&mut self,
indexer: String,
transcode_enabled: bool,
upload_enabled: bool,
retry_failed_transcodes: bool,
) -> Result<Vec<Hash<20>>, Error> {
let items = self.table.get_all().await?;
let mut items: Vec<&QueueItem> = items
.values()
.filter(|item| {
item.indexer == indexer
&& exclude_verified_if_transcode_disabled(item, transcode_enabled)
&& exclude_transcoded_if_upload_disabled(item, upload_enabled)
&& exclude_verify_failures(item)
&& exclude_transcode_failures(item, retry_failed_transcodes)
&& item.upload.is_none()
})
.collect();
items.sort_by_key(|x| &x.name);
let hashes = items.iter().map(|x| x.hash).collect();
Ok(hashes)
}
pub(crate) async fn get_all(&mut self) -> Result<BTreeMap<Hash<20>, QueueItem>, Error> {
self.table.get_all().await
}
pub(crate) async fn set(&mut self, item: QueueItem) -> Result<(), Error> {
self.table.set(item.hash, item).await
}
pub(crate) async fn set_many(
&self,
items: BTreeMap<Hash<20>, QueueItem>,
replace: bool,
) -> Result<usize, Error> {
self.table.set_many(items, replace).await
}
pub(crate) async fn remove(&mut self, hash: Hash<20>) -> Result<Option<QueueItem>, Error> {
self.table.remove(hash).await
}
pub(crate) async fn insert_new_torrent_files(
&mut self,
paths: Vec<PathBuf>,
) -> Result<usize, Error> {
let stream = iter(paths.into_iter());
let items: BTreeMap<_, _> = stream
.filter_map(|path| async {
let torrent = match ImdlCommand::show(&path).await {
Ok(torrent) => Some(torrent),
Err(error) => {
error!("Failed to read torrent: {}\n{error}", path.display());
None
}
};
let item = QueueItem::from_torrent(path, torrent?);
Some((item.hash, item))
})
.collect()
.await;
self.table.set_many(items, false).await
}
}
fn exclude_verify_failures(item: &QueueItem) -> bool {
!matches!(
item.verify,
Some(VerifyStatus {
verified: false,
..
})
)
}
fn exclude_transcode_failures(item: &QueueItem, retry_failed_transcodes: bool) -> bool {
retry_failed_transcodes
|| !matches!(item.transcode, Some(TranscodeStatus { success: false, .. }))
}
fn exclude_verified_if_transcode_disabled(item: &QueueItem, transcode_enabled: bool) -> bool {
transcode_enabled || item.verify.is_none()
}
fn exclude_transcoded_if_upload_disabled(item: &QueueItem, upload_enabled: bool) -> bool {
upload_enabled || !matches!(item.transcode, Some(TranscodeStatus { success: true, .. }))
}