Add pause/resume, enchance tora-cli
This commit is contained in:
+29
-1
@@ -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,
|
||||
|
||||
@@ -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
@@ -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 {}))
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user