333 lines
11 KiB
Rust
333 lines
11 KiB
Rust
//! Drains `AppState`'s three command channels — the daemon-side consumer
|
|
//! DESIGN.md §6.2 and §9.3 assume exists (issues #107 and #132). A `Search`
|
|
//! resets backoff and re-runs the targeted-search + grab lane immediately;
|
|
//! `Grab` sends an already-chosen release straight to Transmission, skipping
|
|
//! search. Movies, episodes and seasons share one lane because they share
|
|
//! one operator waiting on a 202.
|
|
|
|
use arr_api::{AppState, EpisodeCommand, MovieCommand, SeasonCommand};
|
|
use arr_db::Db;
|
|
use tokio::sync::watch;
|
|
|
|
use crate::grab::{GrabAction, GrabError};
|
|
use crate::tv_grab::TvGrabAction;
|
|
|
|
/// Which channel a drained command came off.
|
|
enum Command {
|
|
Movie(MovieCommand),
|
|
Episode(EpisodeCommand),
|
|
Season(SeasonCommand),
|
|
}
|
|
|
|
/// Run until every sender is dropped or `shutdown` fires. Either action is
|
|
/// `None` when Prowlarr (or TMDB, for movies) is not configured (`main.rs`
|
|
/// warns about this already for the reconcile lanes) — commands are still
|
|
/// drained so the channels never fill, they just cannot be acted on.
|
|
pub async fn run(
|
|
state: AppState,
|
|
database: Db,
|
|
movies: Option<GrabAction>,
|
|
tv: Option<TvGrabAction>,
|
|
mut shutdown: watch::Receiver<bool>,
|
|
) {
|
|
loop {
|
|
let command = tokio::select! {
|
|
biased;
|
|
changed = shutdown.changed() => {
|
|
if changed.is_err() || *shutdown.borrow() {
|
|
return;
|
|
}
|
|
continue;
|
|
}
|
|
command = state.next_movie_command() => command.map(Command::Movie),
|
|
command = state.next_episode_command() => command.map(Command::Episode),
|
|
command = state.next_season_command() => command.map(Command::Season),
|
|
};
|
|
let Some(command) = command else {
|
|
return;
|
|
};
|
|
match command {
|
|
Command::Movie(command) => {
|
|
let Some(movies) = &movies else {
|
|
tracing::warn!(
|
|
"manual movie command received but Prowlarr or TMDB is not configured"
|
|
);
|
|
continue;
|
|
};
|
|
if let Err(error) = handle_movie(movies, &database, command).await {
|
|
tracing::error!(%error, "manual movie command failed");
|
|
}
|
|
}
|
|
Command::Episode(command) => {
|
|
let Some(tv) = &tv else {
|
|
tracing::warn!(
|
|
"manual episode command received but Prowlarr is not configured"
|
|
);
|
|
continue;
|
|
};
|
|
if let Err(error) = handle_episode(tv, &database, command).await {
|
|
tracing::error!(%error, "manual episode command failed");
|
|
}
|
|
}
|
|
Command::Season(command) => {
|
|
let Some(tv) = &tv else {
|
|
tracing::warn!("manual season command received but Prowlarr is not configured");
|
|
continue;
|
|
};
|
|
if let Err(error) = handle_season(tv, &database, command).await {
|
|
tracing::error!(%error, "manual season command failed");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn handle_movie(
|
|
grab: &GrabAction,
|
|
database: &Db,
|
|
command: MovieCommand,
|
|
) -> Result<(), GrabError> {
|
|
match command {
|
|
MovieCommand::Search { movie_id } => {
|
|
grab.search_now(database, movie_id).await?;
|
|
}
|
|
MovieCommand::Grab {
|
|
movie_id,
|
|
release_id,
|
|
} => {
|
|
grab.grab_release_now(database, movie_id, release_id)
|
|
.await?;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn handle_episode(
|
|
tv: &TvGrabAction,
|
|
database: &Db,
|
|
command: EpisodeCommand,
|
|
) -> Result<(), GrabError> {
|
|
match command {
|
|
EpisodeCommand::Search { episode_id } => {
|
|
tv.search_episode_now(database, episode_id).await?;
|
|
}
|
|
EpisodeCommand::Grab {
|
|
episode_id,
|
|
release_id,
|
|
} => {
|
|
tv.grab_episode_release_now(database, episode_id, release_id)
|
|
.await?;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn handle_season(
|
|
tv: &TvGrabAction,
|
|
database: &Db,
|
|
command: SeasonCommand,
|
|
) -> Result<(), GrabError> {
|
|
match command {
|
|
SeasonCommand::Search { season_id } => {
|
|
tv.search_season_now(database, season_id).await?;
|
|
}
|
|
SeasonCommand::Grab {
|
|
season_id,
|
|
release_id,
|
|
} => {
|
|
tv.grab_season_release_now(database, season_id, release_id)
|
|
.await?;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[allow(clippy::unwrap_used)]
|
|
mod tests {
|
|
use std::collections::HashMap;
|
|
use std::path::PathBuf;
|
|
|
|
use arr_db::Db;
|
|
use arr_dl::TransmissionClient;
|
|
use arr_indexer::ProwlarrClient;
|
|
use wiremock::matchers::{method, path, path_regex, query_param};
|
|
use wiremock::{Mock, MockServer, ResponseTemplate};
|
|
|
|
use super::*;
|
|
use crate::grab::{SeedingLimits, SeedingRules};
|
|
|
|
async fn database_with_wanted_movie() -> (tempfile::TempDir, Db) {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let database = Db::connect(dir.path().join("arr.db")).await.unwrap();
|
|
database.migrate().await.unwrap();
|
|
sqlx::query(
|
|
"INSERT INTO movies (tmdb_id, title, year, original_language, root_id)
|
|
SELECT 693134, 'Dune Part Two', 2024, 'en', id
|
|
FROM roots WHERE kind = 'movie' AND audience = 'main'",
|
|
)
|
|
.execute(database.pool())
|
|
.await
|
|
.unwrap();
|
|
(dir, database)
|
|
}
|
|
|
|
async fn transmission() -> MockServer {
|
|
let server = MockServer::start().await;
|
|
Mock::given(method("POST"))
|
|
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
|
"result": "success",
|
|
"arguments": {"torrent-added": {"id": 1, "name": "x", "hashString": "aaaa"}}
|
|
})))
|
|
.mount(&server)
|
|
.await;
|
|
server
|
|
}
|
|
|
|
fn grab_action(indexer: &MockServer, transmission: &MockServer) -> GrabAction {
|
|
GrabAction::new(
|
|
ProwlarrClient::new(indexer.uri(), "key").unwrap(),
|
|
TransmissionClient::new(&transmission.uri()).unwrap(),
|
|
PathBuf::from("/mnt/media/transmission/complete"),
|
|
SeedingRules::new(
|
|
SeedingLimits {
|
|
ratio: 1.5,
|
|
idle_minutes: 60,
|
|
},
|
|
HashMap::new(),
|
|
),
|
|
)
|
|
}
|
|
|
|
/// Issue #107: a `Search` command must reset the backoff, not merely be
|
|
/// accepted onto the channel. `search_attempts` starts at 5 (a 7-day
|
|
/// backoff, nowhere near elapsed) so a normal tick would skip this movie;
|
|
/// ending at 1 rather than 6 proves the reset happened, not just a
|
|
/// one-off bypass of the gate.
|
|
#[tokio::test]
|
|
async fn a_search_command_resets_backoff_and_searches_now() {
|
|
let (_dir, database) = database_with_wanted_movie().await;
|
|
sqlx::query(
|
|
"UPDATE movies SET search_attempts = 5,
|
|
last_searched_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '-6 days')",
|
|
)
|
|
.execute(database.pool())
|
|
.await
|
|
.unwrap();
|
|
|
|
let indexer = MockServer::start().await;
|
|
Mock::given(method("GET"))
|
|
.and(path("/api/v1/indexer"))
|
|
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([
|
|
{"id": 7, "name": "tracker", "enable": true}
|
|
])))
|
|
.mount(&indexer)
|
|
.await;
|
|
Mock::given(method("GET"))
|
|
.and(path("/7/api"))
|
|
.and(query_param("t", "caps"))
|
|
.respond_with(ResponseTemplate::new(200).set_body_string(
|
|
r#"<caps><searching><search available="yes" supportedParams="q"/></searching></caps>"#,
|
|
))
|
|
.mount(&indexer)
|
|
.await;
|
|
Mock::given(method("GET"))
|
|
.and(path("/7/api"))
|
|
.and(query_param("t", "search"))
|
|
.respond_with(
|
|
ResponseTemplate::new(200).set_body_string("<rss><channel></channel></rss>"),
|
|
)
|
|
.mount(&indexer)
|
|
.await;
|
|
let transmission = transmission().await;
|
|
let grab = grab_action(&indexer, &transmission);
|
|
|
|
handle_movie(&grab, &database, MovieCommand::Search { movie_id: 1 })
|
|
.await
|
|
.unwrap();
|
|
|
|
let searched = indexer
|
|
.received_requests()
|
|
.await
|
|
.unwrap()
|
|
.into_iter()
|
|
.any(|request| {
|
|
request.url.path() == "/7/api"
|
|
&& request
|
|
.url
|
|
.query_pairs()
|
|
.any(|(name, value)| name == "t" && value == "search")
|
|
});
|
|
assert!(searched, "the reset backoff must let the search run now");
|
|
let (attempts, last_searched_at): (i64, Option<String>) =
|
|
sqlx::query_as("SELECT search_attempts, last_searched_at FROM movies WHERE id = 1")
|
|
.fetch_one(database.pool())
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(attempts, 1, "reset to 0, then one real attempt — not 6");
|
|
assert!(last_searched_at.is_some());
|
|
}
|
|
|
|
/// Issue #107: a `Grab` command sends the already-chosen release straight
|
|
/// to Transmission — no indexer search at all.
|
|
#[tokio::test]
|
|
async fn a_grab_command_sends_the_chosen_release_without_searching() {
|
|
let (_dir, database) = database_with_wanted_movie().await;
|
|
let indexer = MockServer::start().await;
|
|
Mock::given(method("GET"))
|
|
.and(path_regex("^/dl/"))
|
|
.respond_with(ResponseTemplate::new(302).insert_header(
|
|
"location",
|
|
"magnet:?xt=urn:btih:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
|
|
))
|
|
.mount(&indexer)
|
|
.await;
|
|
let transmission = transmission().await;
|
|
let grab = grab_action(&indexer, &transmission);
|
|
|
|
let release_id = sqlx::query(
|
|
"INSERT INTO releases (indexer_id, guid, name, size, download_url, parsed, score, verdict)
|
|
VALUES (7, 'chosen', 'Dune.Part.Two.2024.2160p.WEB-DL', 1, ?, '{}', 900, 'eligible')",
|
|
)
|
|
.bind(format!("{}/dl/chosen", indexer.uri()))
|
|
.execute(database.pool())
|
|
.await
|
|
.unwrap()
|
|
.last_insert_rowid();
|
|
sqlx::query("INSERT INTO movie_releases (movie_id, release_id) VALUES (1, ?)")
|
|
.bind(release_id)
|
|
.execute(database.pool())
|
|
.await
|
|
.unwrap();
|
|
|
|
handle_movie(
|
|
&grab,
|
|
&database,
|
|
MovieCommand::Grab {
|
|
movie_id: 1,
|
|
release_id,
|
|
},
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
let grabbed: i64 = sqlx::query_scalar("SELECT count(*) FROM grabs WHERE target_id = 1")
|
|
.fetch_one(database.pool())
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(grabbed, 1);
|
|
let searched = indexer
|
|
.received_requests()
|
|
.await
|
|
.unwrap()
|
|
.into_iter()
|
|
.any(|request| {
|
|
request
|
|
.url
|
|
.query_pairs()
|
|
.any(|(name, value)| name == "t" && value == "search")
|
|
});
|
|
assert!(!searched, "a chosen release must not trigger a search");
|
|
}
|
|
}
|