x360-web-queue/app/bt.py

"""Supervise anacrolix-based BitTorrent subprocesses (no aria2 for torrents).

Torrent payload data starts only after file selection. All subprocess state is
written as atomic JSON so a restart can detect interrupted work and resume.
"""
import asyncio
import json
import os
import time
from pathlib import Path

from . import store

BT_HELPER = '/usr/local/bin/x360-bt'
_INSPECTIONS = {}
_DOWNLOADS = {}


class BitTorrentError(RuntimeError):
    pass


def exe():
    candidate = os.getenv('BT_HELPER', BT_HELPER)
    if not Path(candidate).is_file():
        raise BitTorrentError(f'Brak silnika BitTorrent: {candidate} (przebuduj obraz Docker)')
    return candidate


def _read(path):
    try:
        value = json.loads(Path(path).read_text(encoding='utf-8'))
        return value if isinstance(value, dict) else {}
    except (FileNotFoundError, ValueError, OSError):
        return {}


def inspect_state(preview_id):
    state = _read(store.DATA / 'metadata' / preview_id / 'state.json')
    process = _INSPECTIONS.get(preview_id)
    if not state and (not process or process.returncode is not None):
        return {'status': 'error', 'error': 'Klient BitTorrent przerwał odczyt metadanych. Sprawdź engine.log podglądu lub ponów.'}
    return state


def job_state(job_id):
    return _read(store.DATA / 'bt' / (job_id + '.json'))


async def _spawn(args, log_path):
    log_path.parent.mkdir(parents=True, exist_ok=True)
    with log_path.open('ab') as log:
        try:
            return await asyncio.create_subprocess_exec(*args, stdin=asyncio.subprocess.DEVNULL,
                                                        stdout=log, stderr=asyncio.subprocess.STDOUT)
        except (OSError, ValueError) as exc:
            raise BitTorrentError('Nie można uruchomić silnika BT: ' + str(exc)) from exc


async def start_inspection(preview_id, magnet):
    from . import torrent_meta
    folder = torrent_meta.inspect_dir(preview_id)
    args = [exe(), 'inspect', '-magnet', magnet, '-out', str(folder / 'metadata.torrent'),
            '-status', str(folder / 'state.json'), '-timeout', '5m', '-port', '0']
    process = await _spawn(args, folder / 'engine.log')
    _INSPECTIONS[preview_id] = process
    return process


async def _stop(process):
    if process and process.returncode is None:
        process.terminate()
        try:
            await asyncio.wait_for(process.wait(), 6)
        except asyncio.TimeoutError:
            process.kill()
            await process.wait()


async def stop_inspection(preview_id):
    await _stop(_INSPECTIONS.pop(preview_id, None))


async def start_download(job):
    jid = job['id']
    old = _DOWNLOADS.get(jid)
    if old and old.returncode is None:
        return
    path = store.DATA / 'torrents' / (jid + '.torrent')
    if not path.is_file():
        raise BitTorrentError('Brak zapisanego pliku .torrent dla zadania')
    try:
        spec = json.loads(job['source'])
        selection = spec['selection']
    except (ValueError, KeyError, TypeError) as exc:
        raise BitTorrentError('Uszkodzony zapis wyboru plików torrenta') from exc
    output = store.DATA / 'raw' / jid
    output.mkdir(parents=True, exist_ok=True)
    status_path = store.DATA / 'bt' / (jid + '.json')
    status_path.parent.mkdir(parents=True, exist_ok=True)
    status_path.unlink(missing_ok=True)
    args = [exe(), 'fetch', '-torrent', str(path), '-select', selection, '-out', str(output),
            '-status', str(status_path), '-port', os.getenv('BT_PORT', '6881')]
    _DOWNLOADS[jid] = await _spawn(args, store.DATA / 'logs' / (jid + '.bt.log'))
    store.update(jid, gid='', state='downloading', note='BitTorrent: uruchomiono klienta')


async def poll_download(job):
    jid = job['id']
    process = _DOWNLOADS.get(jid)
    if process is None:
        # A container restart loses the subprocess; persistently downloaded
        # pieces in data/raw are verified and resumed by the helper.
        await start_download(job)
        return
    state = job_state(jid)
    done, total = int(state.get('downloaded', 0) or 0), int(state.get('total', 0) or 0)
    speed = int(state.get('speed', 0) or 0)
    peers = int(state.get('connections', 0) or 0)
    known = int(state.get('known_peers', 0) or 0)
    if state.get('status') == 'done' and process.returncode == 0:
        await process.wait()
        _DOWNLOADS.pop(jid, None)
        store.update(jid, state='preparing', progress=100, speed=0,
                     note='BitTorrent ukończony, przygotowanie plików')
        return
    if state.get('status') == 'error' or process.returncode not in (None, 0):
        await process.wait()
        _DOWNLOADS.pop(jid, None)
        raise BitTorrentError(state.get('error') or 'Silnik BitTorrent zakończył się niepowodzeniem; sprawdź log BT')
    if process.returncode == 0 and state.get('status') != 'done':
        _DOWNLOADS.pop(jid, None)
        raise BitTorrentError('Silnik BitTorrent zakończył pracę bez potwierdzenia ukończenia danych')
    percent = round(done / total * 100, 2) if total else 0
    store.update(jid, progress=percent, speed=speed, total=total, completed=done,
                 note=f'BitTorrent: {peers} połączeń, {known} znanych peerów')


async def stop_download(job_id):
    await _stop(_DOWNLOADS.pop(job_id, None))


async def shutdown():
    jobs = list(_INSPECTIONS.values()) + list(_DOWNLOADS.values())
    _INSPECTIONS.clear()
    _DOWNLOADS.clear()
    await asyncio.gather(*(_stop(p) for p in jobs))