Add pause/resume, enchance tora-cli

This commit is contained in:
Alexander
2026-07-03 19:10:26 +02:00
parent 80ebf1cb63
commit 0036b19612
19 changed files with 2796 additions and 241 deletions
+29 -1
View File
@@ -38,7 +38,6 @@ pub struct TorrentRow {
pub error_message: Option<String>,
}
/// Insert a new pending torrent, or return the existing row if the info_hash is already tracked.
pub async fn insert_pending(
pool: &PgPool,
info_hash: &str,
@@ -80,6 +79,17 @@ pub async fn list(pool: &PgPool) -> Result<Vec<TorrentRow>> {
Ok(rows)
}
pub async fn delete(pool: &PgPool, id: Uuid) -> Result<Option<TorrentRow>> {
let row = sqlx::query_as::<_, TorrentRow>(
"DELETE FROM torrents WHERE id = $1
RETURNING id, info_hash, name, source, output_path, total_bytes, downloaded_bytes, state, error_message",
)
.bind(id)
.fetch_optional(pool)
.await?;
Ok(row)
}
pub async fn list_pending(pool: &PgPool) -> Result<Vec<TorrentRow>> {
let rows = sqlx::query_as::<_, TorrentRow>(
"SELECT id, info_hash, name, source, output_path, total_bytes, downloaded_bytes, state, error_message
@@ -90,6 +100,24 @@ pub async fn list_pending(pool: &PgPool) -> Result<Vec<TorrentRow>> {
Ok(rows)
}
pub async fn find_ids_by_prefix(pool: &PgPool, prefix: &str) -> Result<Vec<Uuid>> {
let rows =
sqlx::query_scalar::<_, Uuid>("SELECT id FROM torrents WHERE id::text LIKE $1 || '%'")
.bind(prefix)
.fetch_all(pool)
.await?;
Ok(rows)
}
pub async fn set_state(pool: &PgPool, id: Uuid, state: TorrentState) -> Result<()> {
sqlx::query("UPDATE torrents SET state = $2, updated_at = now() WHERE id = $1")
.bind(id)
.bind(state)
.execute(pool)
.await?;
Ok(())
}
pub struct Progress<'a> {
pub name: Option<&'a str>,
pub total_bytes: i64,
+13 -7
View File
@@ -1,11 +1,17 @@
use anyhow::{Context, Result};
use librqbit::Magnet;
/// Extract the BTIH info hash from a magnet link, normalized to lowercase hex.
pub fn info_hash(magnet: &str) -> Result<String> {
let parsed = Magnet::parse(magnet).context("failed to parse magnet link")?;
let id20 = parsed
.as_id20()
.context("magnet link has no v1 (BTIH) info hash")?;
Ok(id20.as_string())
pub fn info_hash(source: &str) -> Result<String> {
if source.starts_with("magnet:") {
let parsed = Magnet::parse(source).context("failed to parse magnet link")?;
let id20 = parsed
.as_id20()
.context("magnet link has no v1 (BTIH) info hash")?;
return Ok(id20.as_string());
}
let bytes =
std::fs::read(source).with_context(|| format!("failed to read torrent file: {source}"))?;
let parsed = librqbit::torrent_from_bytes_ext::<Vec<u8>>(&bytes)
.context("failed to parse .torrent file")?;
Ok(parsed.meta.info_hash.as_string())
}
+250 -21
View File
@@ -4,12 +4,15 @@ use std::sync::Arc;
use std::time::Duration;
use anyhow::{Context, Result};
use librqbit::api::TorrentIdOrHash;
use librqbit::{AddTorrent, AddTorrentOptions, ManagedTorrent, Session, TorrentStatsState};
use librqbit_core::Id20;
use sqlx::PgPool;
use tokio::sync::Mutex;
use tonic::{Request, Response, Status};
use tora_proto::{
AddRequest, AddResponse, ListRequest, ListResponse, State as ProtoState, StatusRequest,
AddRequest, AddResponse, ListRequest, ListResponse, PauseRequest, PauseResponse, RemoveRequest,
RemoveResponse, ResumeRequest, ResumeResponse, State as ProtoState, StatusRequest,
TorrentStatus, Torrents,
};
use tracing::{error, info, warn};
@@ -19,6 +22,13 @@ use crate::db::{self, TorrentRow, TorrentState};
use crate::magnet;
const POLL_INTERVAL: Duration = Duration::from_secs(2);
const BYTES_PER_MIB: f64 = 1_048_576.0;
pub enum ResolveError {
NotFound,
Ambiguous(usize),
Internal(anyhow::Error),
}
pub struct TorrentManager {
pool: PgPool,
@@ -55,12 +65,100 @@ impl TorrentManager {
db::insert_pending(&self.pool, &info_hash, magnet, &output_path).await
}
pub async fn get(&self, id: Uuid) -> Result<Option<TorrentRow>> {
db::get(&self.pool, id).await
pub async fn get_status(&self, id: Uuid) -> Result<Option<TorrentStatus>> {
let Some(row) = db::get(&self.pool, id).await? else {
return Ok(None);
};
let tracked = self.tracked.lock().await;
let live = extract_live_stats(&tracked, &row.id);
Ok(Some(row_to_status(row, live)))
}
pub async fn list(&self) -> Result<Vec<TorrentRow>> {
db::list(&self.pool).await
pub async fn list_statuses(&self) -> Result<Vec<TorrentStatus>> {
let rows = db::list(&self.pool).await?;
let tracked = self.tracked.lock().await;
Ok(rows
.into_iter()
.map(|row| {
let live = extract_live_stats(&tracked, &row.id);
row_to_status(row, live)
})
.collect())
}
pub async fn resolve_id(&self, id_or_prefix: &str) -> Result<Uuid, ResolveError> {
if let Ok(uuid) = Uuid::parse_str(id_or_prefix) {
return Ok(uuid);
}
let matches = db::find_ids_by_prefix(&self.pool, id_or_prefix)
.await
.map_err(ResolveError::Internal)?;
match matches.len() {
0 => Err(ResolveError::NotFound),
1 => Ok(matches[0]),
_ => Err(ResolveError::Ambiguous(matches.len())),
}
}
pub async fn remove(&self, id: Uuid, delete_files: bool) -> Result<Option<TorrentRow>> {
let row = db::get(&self.pool, id).await?;
if let Some(ref row) = row {
let mut tracked = self.tracked.lock().await;
tracked.remove(&id);
drop(tracked);
if let Ok(info_hash) = row.info_hash.parse::<Id20>()
&& let Err(err) = self
.session
.delete(TorrentIdOrHash::Hash(info_hash), delete_files)
.await
{
warn!(id = %id, error = %err, "failed to delete torrent from session");
}
if delete_files {
let _ = std::fs::remove_dir_all(&row.output_path);
}
}
db::delete(&self.pool, id).await
}
pub async fn pause(&self, id: Uuid) -> Result<Option<TorrentRow>> {
let row = db::get(&self.pool, id).await?;
if row.is_none() {
return Ok(None);
}
let tracked = self.tracked.lock().await;
if let Some(handle) = tracked.get(&id) {
self.session
.pause(handle)
.await
.context("failed to pause torrent in librqbit session")?;
} else {
// Torrent not yet tracked by librqbit (still pending in DB).
// Set state to Paused so the poller won't pick it up.
db::set_state(&self.pool, id, TorrentState::Paused).await?;
}
Ok(row)
}
pub async fn resume(&self, id: Uuid) -> Result<Option<TorrentRow>> {
let row = db::get(&self.pool, id).await?;
if row.is_none() {
return Ok(None);
}
let tracked = self.tracked.lock().await;
if let Some(handle) = tracked.get(&id) {
self.session
.unpause(handle)
.await
.context("failed to resume torrent in librqbit session")?;
} else {
// Torrent not tracked by librqbit. Set state to Pending so
// the poller picks it up again.
db::set_state(&self.pool, id, TorrentState::Pending).await?;
}
Ok(row)
}
pub fn spawn_poller(self: Arc<Self>) {
@@ -96,11 +194,14 @@ impl TorrentManager {
overwrite: true,
..Default::default()
};
match self
.session
.add_torrent(AddTorrent::from_url(row.source.clone()), Some(options))
.await
{
let add = match AddTorrent::from_cli_argument(&row.source) {
Ok(add) => add,
Err(err) => {
error!(id = %row.id, source = %row.source, error = %err, "failed to parse torrent source");
continue;
}
};
match self.session.add_torrent(add, Some(options)).await {
Ok(response) => match response.into_handle() {
Some(handle) => {
info!(id = %row.id, info_hash = %row.info_hash, "torrent added to session");
@@ -155,6 +256,59 @@ impl TorrentManager {
}
}
#[derive(Default)]
struct LiveTorrentStats {
download_speed_bps: u64,
eta_seconds: u64,
uploaded_bytes: u64,
upload_speed_bps: u64,
peers: u32,
seeds: u32,
}
fn extract_live_stats(tracked: &HashMap<Uuid, Arc<ManagedTorrent>>, id: &Uuid) -> LiveTorrentStats {
let Some(handle) = tracked.get(id) else {
return LiveTorrentStats::default();
};
let stats = handle.stats();
let uploaded_bytes = stats.uploaded_bytes;
let Some(live) = stats.live else {
return LiveTorrentStats {
uploaded_bytes,
..Default::default()
};
};
let download_bps = if live.download_speed.mbps > 0.0 {
(live.download_speed.mbps * BYTES_PER_MIB) as u64
} else {
0
};
let upload_bps = if live.upload_speed.mbps > 0.0 {
(live.upload_speed.mbps * BYTES_PER_MIB) as u64
} else {
0
};
let total = stats.total_bytes;
let downloaded = stats.progress_bytes;
let remaining = total.saturating_sub(downloaded);
let eta = if download_bps > 0 {
(remaining as f64 / download_bps as f64) as u64
} else {
0
};
LiveTorrentStats {
download_speed_bps: download_bps,
eta_seconds: eta,
uploaded_bytes,
upload_speed_bps: upload_bps,
peers: live.snapshot.peer_stats.live as u32,
seeds: live.snapshot.peer_stats.not_needed as u32,
}
}
fn proto_state(state: TorrentState) -> ProtoState {
match state {
TorrentState::Pending => ProtoState::Pending,
@@ -165,7 +319,7 @@ fn proto_state(state: TorrentState) -> ProtoState {
}
}
fn row_to_status(row: TorrentRow) -> TorrentStatus {
fn row_to_status(row: TorrentRow, live: LiveTorrentStats) -> TorrentStatus {
TorrentStatus {
id: row.id.to_string(),
info_hash: row.info_hash,
@@ -176,6 +330,12 @@ fn row_to_status(row: TorrentRow) -> TorrentStatus {
downloaded_bytes: row.downloaded_bytes as u64,
state: proto_state(row.state) as i32,
error_message: row.error_message.unwrap_or_default(),
download_speed_bps: live.download_speed_bps,
eta_seconds: live.eta_seconds,
uploaded_bytes: live.uploaded_bytes,
upload_speed_bps: live.upload_speed_bps,
peers: live.peers,
seeds: live.seeds,
}
}
@@ -191,6 +351,16 @@ impl GrpcTorrents {
}
}
fn resolve_err(e: ResolveError) -> Status {
match e {
ResolveError::NotFound => Status::not_found("torrent not found"),
ResolveError::Ambiguous(n) => {
Status::invalid_argument(format!("ambiguous prefix, matches {n} torrents"))
}
ResolveError::Internal(err) => Status::internal(err.to_string()),
}
}
#[tonic::async_trait]
impl Torrents for GrpcTorrents {
async fn add(&self, request: Request<AddRequest>) -> Result<Response<AddResponse>, Status> {
@@ -210,25 +380,84 @@ impl Torrents for GrpcTorrents {
&self,
request: Request<StatusRequest>,
) -> Result<Response<TorrentStatus>, Status> {
let id = Uuid::parse_str(&request.into_inner().id)
.map_err(|_| Status::invalid_argument("invalid id"))?;
let row = self
let id = self
.manager
.get(id)
.resolve_id(&request.into_inner().id)
.await
.map_err(resolve_err)?;
let status = self
.manager
.get_status(id)
.await
.map_err(|err| Status::internal(err.to_string()))?
.ok_or_else(|| Status::not_found("torrent not found"))?;
Ok(Response::new(row_to_status(row)))
Ok(Response::new(status))
}
async fn list(&self, _request: Request<ListRequest>) -> Result<Response<ListResponse>, Status> {
let rows = self
let torrents = self
.manager
.list()
.list_statuses()
.await
.map_err(|err| Status::internal(err.to_string()))?;
Ok(Response::new(ListResponse {
torrents: rows.into_iter().map(row_to_status).collect(),
}))
Ok(Response::new(ListResponse { torrents }))
}
async fn remove(
&self,
request: Request<RemoveRequest>,
) -> Result<Response<RemoveResponse>, Status> {
let req = request.into_inner();
let id = self
.manager
.resolve_id(&req.id)
.await
.map_err(resolve_err)?;
let row = self
.manager
.remove(id, req.delete_files)
.await
.map_err(|err| Status::internal(err.to_string()))?
.ok_or_else(|| Status::not_found("torrent not found"))?;
info!(id = %id, name = ?row.name, "torrent removed");
Ok(Response::new(RemoveResponse {}))
}
async fn pause(
&self,
request: Request<PauseRequest>,
) -> Result<Response<PauseResponse>, Status> {
let id = self
.manager
.resolve_id(&request.into_inner().id)
.await
.map_err(resolve_err)?;
let row = self
.manager
.pause(id)
.await
.map_err(|err| Status::internal(err.to_string()))?
.ok_or_else(|| Status::not_found("torrent not found"))?;
info!(id = %id, name = ?row.name, "torrent paused");
Ok(Response::new(PauseResponse {}))
}
async fn resume(
&self,
request: Request<ResumeRequest>,
) -> Result<Response<ResumeResponse>, Status> {
let id = self
.manager
.resolve_id(&request.into_inner().id)
.await
.map_err(resolve_err)?;
let row = self
.manager
.resume(id)
.await
.map_err(|err| Status::internal(err.to_string()))?
.ok_or_else(|| Status::not_found("torrent not found"))?;
info!(id = %id, name = ?row.name, "torrent resumed");
Ok(Response::new(ResumeResponse {}))
}
}