Aria2c: Improve download progress and error handling

This commit is contained in:
rlaphoenix 2024-02-15 19:19:37 +00:00
parent e8b07bf03a
commit 4e12b867f1
1 changed files with 74 additions and 31 deletions

View File

@ -17,17 +17,19 @@ from rich.text import Text
from devine.core.config import config from devine.core.config import config
from devine.core.console import console from devine.core.console import console
from devine.core.constants import DOWNLOAD_CANCELLED
from devine.core.utilities import get_binary_path, get_free_port from devine.core.utilities import get_binary_path, get_free_port
def rpc(caller: Callable, secret: str, method: str, *params: Any) -> dict[str, Any]: def rpc(caller: Callable, secret: str, method: str, params: Optional[list[Any]] = None) -> Any:
"""Make a call to Aria2's JSON-RPC API.""" """Make a call to Aria2's JSON-RPC API."""
try:
rpc_res = caller( rpc_res = caller(
json={ json={
"jsonrpc": "2.0", "jsonrpc": "2.0",
"id": get_random_bytes(16).hex(), "id": get_random_bytes(16).hex(),
"method": method, "method": method,
"params": [f"token:{secret}", *params] "params": [f"token:{secret}", *(params or [])]
} }
).json() ).json()
if rpc_res.get("code"): if rpc_res.get("code"):
@ -39,6 +41,9 @@ def rpc(caller: Callable, secret: str, method: str, *params: Any) -> dict[str, A
)) ))
console.log(Text.from_ansi("\n[Aria2c]: " + error_pretty)) console.log(Text.from_ansi("\n[Aria2c]: " + error_pretty))
return rpc_res["result"] return rpc_res["result"]
except requests.exceptions.ConnectionError:
# absorb, process likely ended as it was calling RPC
return
def download( def download(
@ -176,6 +181,8 @@ def download(
continue continue
arguments.extend(["--header", f"{header}: {value}"]) arguments.extend(["--header", f"{header}: {value}"])
yield dict(total=len(urls))
try: try:
p = subprocess.Popen( p = subprocess.Popen(
[ [
@ -190,28 +197,50 @@ def download(
p.stdin.close() p.stdin.close()
while p.poll() is None: while p.poll() is None:
global_stats = rpc( global_stats: dict[str, Any] = rpc(
caller=partial(rpc_session.post, url=rpc_uri), caller=partial(rpc_session.post, url=rpc_uri),
secret=rpc_secret, secret=rpc_secret,
method="aria2.getGlobalStat" method="aria2.getGlobalStat"
) )
if global_stats: if global_stats:
active = int(global_stats["numActive"])
waiting = int(global_stats["numWaiting"])
stopped = int(global_stats["numStopped"])
total = active + waiting + stopped
yield dict( yield dict(
total=total,
completed=stopped,
downloaded=f"{filesize.decimal(int(global_stats['downloadSpeed']))}/s" downloaded=f"{filesize.decimal(int(global_stats['downloadSpeed']))}/s"
) )
if total == stopped:
stopped_downloads: list[dict[str, Any]] = rpc(
caller=partial(rpc_session.post, url=rpc_uri),
secret=rpc_secret,
method="aria2.tellStopped",
params=[0, 999999]
)
for dl in stopped_downloads or []:
if dl["status"] == "complete":
yield dict(advance=1)
elif dl["status"] == "error":
used_uri = next(
uri["uri"]
for file in dl["files"]
if file["selected"] == "true"
for uri in file["uris"]
if uri["status"] == "used"
)
error = f"Download Error (#{dl['gid']}): {dl['errorMessage']} ({dl['errorCode']}), {used_uri}"
error_pretty = "\n ".join(textwrap.wrap(
error,
width=console.width - 20,
initial_indent=""
))
console.log(Text.from_ansi("\n[Aria2c]: " + error_pretty))
raise ValueError(error)
if len(stopped_downloads) == len(urls):
rpc( rpc(
caller=partial(rpc_session.post, url=rpc_uri), caller=partial(rpc_session.post, url=rpc_uri),
secret=rpc_secret, secret=rpc_secret,
method="aria2.shutdown" method="aria2.shutdown"
) )
break break
time.sleep(1) time.sleep(1)
p.wait() p.wait()
@ -227,6 +256,20 @@ def download(
# 0xC000013A is when it never got the chance to # 0xC000013A is when it never got the chance to
raise KeyboardInterrupt() raise KeyboardInterrupt()
raise raise
except KeyboardInterrupt:
DOWNLOAD_CANCELLED.set() # skip pending track downloads
yield dict(downloaded="[yellow]CANCELLED")
raise
except Exception:
DOWNLOAD_CANCELLED.set() # skip pending track downloads
yield dict(downloaded="[red]FAILED")
raise
finally:
rpc(
caller=partial(rpc_session.post, url=rpc_uri),
secret=rpc_secret,
method="aria2.shutdown"
)
def aria2c( def aria2c(