feat: multiprocessing for RabbitMQ worker

Helps with very long tasks
This commit is contained in:
2021-01-04 18:22:35 +03:00
parent 4b035964d6
commit af06978cc9

View File

@@ -2,6 +2,8 @@ from tempfile import TemporaryDirectory
import json
import time
import logging
import multiprocessing
import queue
from src.git import clone_repo
from src.optipng import apply_optipng_recursive
@@ -9,6 +11,17 @@ from src.messages import AnalyzeTask, AnalyzeResponse, create_analyze_response
from src.rabbit import get_channel, send_reply
from src.settings import RABBIT_TASK_QUEUE, init_logging, WORKER_NAME
TASK_DELAY = 0.25
def delay_generator():
for _ in range(10):
yield 0.05
for _ in range(10):
yield 0.1
while True:
yield TASK_DELAY
def process_task(task: AnalyzeTask) -> AnalyzeResponse:
with TemporaryDirectory() as repo_path:
@@ -27,9 +40,7 @@ def process_task(task: AnalyzeTask) -> AnalyzeResponse:
)
def rabbit_callback(ch, method, properties, body):
task = AnalyzeTask.from_data(json.loads(body))
logging.info(f"New task url: {task.url}, level: {task.level}")
def _process_main(results_queue: queue.Queue, task: AnalyzeTask):
try:
res = process_task(task).data
res["status"] = "success"
@@ -40,6 +51,23 @@ def rabbit_callback(ch, method, properties, body):
"message": str(err),
"task_id": task.task_id
}
results_queue.put(res)
def rabbit_callback(ch, method, properties, body):
task = AnalyzeTask.from_data(json.loads(body))
logging.info(f"New task url: {task.url}, level: {task.level}")
res_queue = multiprocessing.Queue()
method_process = multiprocessing.Process(
target=_process_main,
args=(res_queue, task)
)
method_process.start()
for delay in delay_generator():
ch.connection.sleep(delay)
if not method_process.is_alive():
break
res = res_queue.get()
send_reply(ch, json.dumps(res).encode("utf-8"))
ch.basic_ack(delivery_tag=method.delivery_tag)