refactoring for new arch, minor fixes

This commit is contained in:
2023-08-24 16:45:55 +03:00
committed by nikili0n
parent 952589d647
commit a2001a0346
6 changed files with 198 additions and 66 deletions

View File

@ -1,85 +1,77 @@
import asyncio
import json
import concurrent.futures as pool
import multiprocessing
from functools import partial
from multiprocessing import freeze_support
from collections import deque
from fastapi import HTTPException
from aio_pika import connect
from aio_pika.abc import AbstractIncomingMessage
from src.core.send import body
from src.core.rabbitmq import get_messages
from src.core.result import Result, ResultTypeEnum
from src.core.ydl import VideoDownloader
from src.exceptions.download_exceptions import SiteNotImplementedException
async def on_message(message: AbstractIncomingMessage, queue) -> None:
async with message.process():
# print(f" [x] Received message {message!r}")
await queue.put(json.loads(message.body))
print(f" Message body is: {message.body!r}")
async def get_messages(inner_queue) -> None:
# Perform connection
connection = await connect("amqp://guest:guest@localhost/")
async with connection:
# Creating a channel
channel = await connection.channel()
await channel.set_qos(prefetch_count=1)
# Declaring queue
queue = await channel.declare_queue(
"hello",
durable=True,
)
# Start listening the queue with name 'task_queue'
await queue.consume(partial(on_message, queue=inner_queue))
print(" [*] Waiting for messages. To exit press CTRL+C")
await asyncio.Future()
def get_url_for_download_video(data: str):
file_name = VideoDownloader.get_unique_video_filename()
ydl_opts = {
"format": data["format"],
"merge_output_format": data["merge_output_format"],
'outtmpl': data["outtmpl"],
}
downloader = VideoDownloader(link=data["link"], ydl_opts=ydl_opts)
try:
result = downloader.download()
except SiteNotImplementedException as ex:
raise HTTPException(
status_code=400,
detail=ex.message
)
async def create_workers(queue):
link = await queue.get()
get_url_for_download_video(link)
class MasterService:
def __init__(self):
self.loop = asyncio.get_event_loop()
self.MAX_EXECUTOR_WORKERS = 8
self.executor = pool.ProcessPoolExecutor(max_workers=self.MAX_EXECUTOR_WORKERS)
self.executor = pool.ProcessPoolExecutor(max_workers=self.MAX_EXECUTOR_WORKERS,
initializer=executor_initializer)
self.queue = asyncio.Queue()
self.rabbit_consumer = get_messages
self.currently_underway = {} # contains currently in progress videos
async def run(self):
loop = asyncio.get_event_loop()
tasks = [loop.run_in_executor(self.executor, create_workers, self.queue) for i in range(self.MAX_EXECUTOR_WORKERS)]
# TODO add cleanup procedure for existing exector procs
tasks = [self.loop.create_task(self.create_workers()) for i in range(self.MAX_EXECUTOR_WORKERS + 1)]
await asyncio.gather(self.rabbit_consumer(self.queue), *tasks)
async def create_workers(self):
while True:
video_params = await self.queue.get()
self.currently_underway[video_params['link']] = video_params
# TODO save current copy of self.currently_underway to DB\file\etc
download_task = self.loop.run_in_executor(self.executor, partial(
MasterService.video_processing_executor, video_params=video_params
)
)
if download_task.done():
result = download_task.result()
# process result
self.queue.task_done()
@staticmethod
def video_download(video_params: dict):
ydl_opts = {
"format": video_params["format"],
"merge_output_format": video_params["merge_output_format"],
'outtmpl': video_params["outtmpl"],
}
downloader = VideoDownloader(link=video_params["link"], ydl_opts=ydl_opts)
try:
result = downloader.download()
return result
except SiteNotImplementedException as ex:
raise HTTPException(
status_code=400,
detail=ex.message
)
@staticmethod
def video_processing_executor(video_params: dict):
# TODO check if video exists on server etc. Result(type=ResultType.EXISTS, value=False)
try:
MasterService.video_download(video_params=video_params)
return Result(result_type=ResultTypeEnum.DONE)
except Exception as ex:
return Result(result_type=ResultTypeEnum.EXCEPTION, value=ex)
# TODO upload to server
def executor_initializer():
import setproctitle
setproctitle.setproctitle(f'video_downloader_executor_process')
return True
if __name__ == '__main__':
freeze_support()
ms = MasterService()
asyncio.run(ms.run())

32
src/core/rabbitmq.py Normal file
View File

@ -0,0 +1,32 @@
import asyncio
import json
from functools import partial
from aio_pika import connect
from aio_pika.abc import AbstractIncomingMessage
async def on_message(message: AbstractIncomingMessage, queue) -> None:
async with message.process():
await queue.put(json.loads(message.body))
print(f" Message body is: {message.body!r}")
async def get_messages(inner_queue) -> None:
# Perform connection
async with await connect("amqp://guest:guest@localhost/") as connection:
# Creating a channel
channel = await connection.channel()
await channel.set_qos(prefetch_count=1)
# Declaring queue
queue = await channel.declare_queue(
"hello",
durable=True,
)
# Start listening the queue with name 'task_queue'
await queue.consume(partial(on_message, queue=inner_queue))
print(" [*] Waiting for messages. To exit press CTRL+C")
await asyncio.Future()

15
src/core/result.py Normal file
View File

@ -0,0 +1,15 @@
from enum import Enum
class ResultTypeEnum(Enum):
EXCEPTION = "Error"
DONE = "Done"
class Result:
def __init__(self, result_type: ResultTypeEnum, value: Exception = None):
self.result_type = result_type
self.value = value
def __repr__(self):
return f'Result: {self.result_type.value}. Traceback: {self.value if self.value else None}'