This repository has been archived on 2023-11-11. You can view files and clone it, but cannot push or open issues or pull requests.
automated-youtube-dl/server/api/job_tracker.py

123 lines
3.2 KiB
Python

import copy
import datetime
import threading
import uuid
from multiprocessing import Manager
from multiprocessing.managers import BaseManager
from typing import Dict
from server.mysql import query
class Job:
_manager = Manager()
def __init__(self, list_id):
self._list_id = list_id
self._job_id = str(uuid.uuid4())
self._status = 'running'
self._start_time = int(datetime.datetime.now(datetime.timezone.utc).timestamp() * 1e3)
self._end_time = None
self._success = None
self._progresses: dict[str, dict]
self._manager = Manager()
self._progresses = self._manager.dict()
self._completed = self._manager.list()
def success(self, s: bool):
self._success = s
def id(self):
return self._job_id
def list_id(self):
return self._list_id
def status(self):
return self._status
def start_time(self):
return self._start_time
def end_time(self):
return self._end_time
def new_progress_thread(self, video_id: str) -> dict:
self._progresses[video_id] = self._manager.dict({
'downloaded_bytes': -1,
'percent': -1,
'total': -1,
'total_bytes': -1,
'total_bytes_estimate': -1,
'speed': -1,
'size': -1,
'start_time': 0
})
return self._progresses[video_id]
def del_progress_thread(self, video_id: str):
del self._progresses[video_id]
def progresses(self):
return copy.deepcopy(self._progresses)
def add_completed(self, video_id: str):
self._completed.append(video_id)
def finish(self):
self._status = 'finished'
self._end_time = int(datetime.datetime.now(datetime.timezone.utc).timestamp() * 1e3)
query("UPDATE `jobs` SET `result` = 'finished' WHERE `jobs`.`job_id` = %s", (self._job_id,))
def dict(self) -> dict:
return {
'job_id': self._job_id,
'list_id': self._list_id,
'status': self._status,
'start_time': self._start_time,
'end_time': self._end_time,
'progresses': copy.deepcopy({k: v for k, v in self._progresses.items()}),
'completed': list(self._completed),
'success': self._success
}
class JobManager(BaseManager):
pass
class JobTracker:
def __init__(self):
self._lock = threading.Lock()
self._jobs = {}
def del_job(self, job_id):
del self._jobs[job_id]
# def new_job(self, list_id: str):
# job_id = str(uuid.uuid4())
# assert job_id not in self._jobs
# with self.lock:
# self._jobs[job_id] = Job(job_id, list_id)
# return self._jobs[job_id]
def add_job(self, job: Job):
assert job.id() not in self._jobs
with self._lock:
self._jobs[job.id()] = job
def finish_job(self, job_id):
with self._lock:
job = self._jobs.get(job_id)
if job:
job.finish()
def get_job(self, job_id) -> Job:
with self._lock:
return self._jobs.get(job_id)
@property
def jobs(self) -> Dict[str, Job]:
return self._jobs.copy()