mirror of https://github.com/peter4431/rq.git
Merge branch 'master' of github.com:rq/rq
Selwin Ong
2 years ago
commit
5f883c6e1c
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,250 @@
|
|||||||
|
import contextlib
|
||||||
|
import errno
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
import signal
|
||||||
|
import time
|
||||||
|
|
||||||
|
from enum import Enum
|
||||||
|
from multiprocessing import Process
|
||||||
|
from typing import Dict, List, NamedTuple, Optional, Set, Type, Union
|
||||||
|
from uuid import uuid4
|
||||||
|
|
||||||
|
from redis import Redis
|
||||||
|
from redis import SSLConnection, UnixDomainSocketConnection
|
||||||
|
from rq.serializers import DefaultSerializer
|
||||||
|
|
||||||
|
from rq.timeouts import HorseMonitorTimeoutException, UnixSignalDeathPenalty
|
||||||
|
|
||||||
|
from .connections import parse_connection
|
||||||
|
from .defaults import DEFAULT_LOGGING_DATE_FORMAT, DEFAULT_LOGGING_FORMAT
|
||||||
|
from .job import Job
|
||||||
|
from .logutils import setup_loghandlers
|
||||||
|
from .queue import Queue
|
||||||
|
from .utils import parse_names
|
||||||
|
from .worker import BaseWorker, Worker
|
||||||
|
|
||||||
|
|
||||||
|
class WorkerData(NamedTuple):
|
||||||
|
name: str
|
||||||
|
pid: int
|
||||||
|
process: Process
|
||||||
|
|
||||||
|
|
||||||
|
class WorkerPool:
|
||||||
|
class Status(Enum):
|
||||||
|
IDLE = 1
|
||||||
|
STARTED = 2
|
||||||
|
STOPPED = 3
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
queues: List[Union[str, Queue]],
|
||||||
|
connection: Redis,
|
||||||
|
num_workers: int = 1,
|
||||||
|
worker_class: Type[BaseWorker] = Worker,
|
||||||
|
serializer: Type[DefaultSerializer] = DefaultSerializer,
|
||||||
|
job_class: Type[Job] = Job,
|
||||||
|
*args,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
self.num_workers: int = num_workers
|
||||||
|
self._workers: List[Worker] = []
|
||||||
|
setup_loghandlers('INFO', DEFAULT_LOGGING_DATE_FORMAT, DEFAULT_LOGGING_FORMAT, name=__name__)
|
||||||
|
self.log: logging.Logger = logging.getLogger(__name__)
|
||||||
|
# self.log: logging.Logger = logger
|
||||||
|
self._queue_names: List[str] = parse_names(queues)
|
||||||
|
self.connection = connection
|
||||||
|
self.name: str = uuid4().hex
|
||||||
|
self._burst: bool = True
|
||||||
|
self._sleep: int = 0
|
||||||
|
self.status: self.Status = self.Status.IDLE # type: ignore
|
||||||
|
self.worker_class: Type[BaseWorker] = worker_class
|
||||||
|
self.serializer: Type[DefaultSerializer] = serializer
|
||||||
|
self.job_class: Type[Job] = job_class
|
||||||
|
|
||||||
|
# A dictionary of WorkerData keyed by worker name
|
||||||
|
self.worker_dict: Dict[str, WorkerData] = {}
|
||||||
|
self._connection_class, _, self._connection_kwargs = parse_connection(connection)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def queues(self) -> List[Queue]:
|
||||||
|
"""Returns a list of Queue objects"""
|
||||||
|
return [Queue(name, connection=self.connection) for name in self._queue_names]
|
||||||
|
|
||||||
|
@property
|
||||||
|
def number_of_active_workers(self) -> int:
|
||||||
|
"""Returns a list of Queue objects"""
|
||||||
|
return len(self.worker_dict)
|
||||||
|
|
||||||
|
def _install_signal_handlers(self):
|
||||||
|
"""Installs signal handlers for handling SIGINT and SIGTERM
|
||||||
|
gracefully.
|
||||||
|
"""
|
||||||
|
signal.signal(signal.SIGINT, self.request_stop)
|
||||||
|
signal.signal(signal.SIGTERM, self.request_stop)
|
||||||
|
|
||||||
|
def request_stop(self, signum=None, frame=None):
|
||||||
|
"""Toggle self._stop_requested that's checked on every loop"""
|
||||||
|
self.log.info('Received SIGINT/SIGTERM, shutting down...')
|
||||||
|
self.status = self.Status.STOPPED
|
||||||
|
self.stop_workers()
|
||||||
|
|
||||||
|
def all_workers_have_stopped(self) -> bool:
|
||||||
|
"""Returns True if all workers have stopped."""
|
||||||
|
self.reap_workers()
|
||||||
|
# `bool(self.worker_dict)` sometimes returns True even if the dict is empty
|
||||||
|
return self.number_of_active_workers == 0
|
||||||
|
|
||||||
|
def reap_workers(self):
|
||||||
|
"""Removes dead workers from worker_dict"""
|
||||||
|
self.log.debug('Reaping dead workers')
|
||||||
|
worker_datas = list(self.worker_dict.values())
|
||||||
|
|
||||||
|
for data in worker_datas:
|
||||||
|
data.process.join(0.1)
|
||||||
|
if data.process.is_alive():
|
||||||
|
self.log.debug('Worker %s with pid %d is alive', data.name, data.pid)
|
||||||
|
else:
|
||||||
|
self.handle_dead_worker(data)
|
||||||
|
continue
|
||||||
|
|
||||||
|
# I'm still not sure why this is sometimes needed, temporarily commenting
|
||||||
|
# this out until I can figure it out.
|
||||||
|
# with contextlib.suppress(HorseMonitorTimeoutException):
|
||||||
|
# with UnixSignalDeathPenalty(1, HorseMonitorTimeoutException):
|
||||||
|
# try:
|
||||||
|
# # If wait4 returns, the process is dead
|
||||||
|
# os.wait4(data.process.pid, 0) # type: ignore
|
||||||
|
# self.handle_dead_worker(data)
|
||||||
|
# except ChildProcessError:
|
||||||
|
# # Process is dead
|
||||||
|
# self.handle_dead_worker(data)
|
||||||
|
# continue
|
||||||
|
|
||||||
|
def handle_dead_worker(self, worker_data: WorkerData):
|
||||||
|
"""
|
||||||
|
Handle a dead worker
|
||||||
|
"""
|
||||||
|
self.log.info('Worker %s with pid %d is dead', worker_data.name, worker_data.pid)
|
||||||
|
with contextlib.suppress(KeyError):
|
||||||
|
self.worker_dict.pop(worker_data.name)
|
||||||
|
|
||||||
|
def check_workers(self, respawn: bool = True) -> None:
|
||||||
|
"""
|
||||||
|
Check whether workers are still alive
|
||||||
|
"""
|
||||||
|
self.log.debug('Checking worker processes')
|
||||||
|
self.reap_workers()
|
||||||
|
# If we have less number of workers than num_workers,
|
||||||
|
# respawn the difference
|
||||||
|
if respawn and self.status != self.Status.STOPPED:
|
||||||
|
delta = self.num_workers - len(self.worker_dict)
|
||||||
|
if delta:
|
||||||
|
for i in range(delta):
|
||||||
|
self.start_worker(burst=self._burst, _sleep=self._sleep)
|
||||||
|
|
||||||
|
def start_worker(
|
||||||
|
self,
|
||||||
|
count: Optional[int] = None,
|
||||||
|
burst: bool = True,
|
||||||
|
_sleep: float = 0,
|
||||||
|
logging_level: str = "INFO",
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Starts a worker and adds the data to worker_datas.
|
||||||
|
* sleep: waits for X seconds before creating worker, for testing purposes
|
||||||
|
"""
|
||||||
|
name = uuid4().hex
|
||||||
|
process = Process(
|
||||||
|
target=run_worker,
|
||||||
|
args=(name, self._queue_names, self._connection_class, self._connection_kwargs),
|
||||||
|
kwargs={
|
||||||
|
'_sleep': _sleep,
|
||||||
|
'burst': burst,
|
||||||
|
'logging_level': logging_level,
|
||||||
|
'worker_class': self.worker_class,
|
||||||
|
'job_class': self.job_class,
|
||||||
|
'serializer': self.serializer,
|
||||||
|
},
|
||||||
|
name=f'Worker {name} (WorkerPool {self.name})',
|
||||||
|
)
|
||||||
|
process.start()
|
||||||
|
worker_data = WorkerData(name=name, pid=process.pid, process=process) # type: ignore
|
||||||
|
self.worker_dict[name] = worker_data
|
||||||
|
self.log.debug('Spawned worker: %s with PID %d', name, process.pid)
|
||||||
|
|
||||||
|
def start_workers(self, burst: bool = True, _sleep: float = 0, logging_level: str = "INFO"):
|
||||||
|
"""
|
||||||
|
Run the workers
|
||||||
|
* sleep: waits for X seconds before creating worker, only for testing purposes
|
||||||
|
"""
|
||||||
|
self.log.debug(f'Spawning {self.num_workers} workers')
|
||||||
|
for i in range(self.num_workers):
|
||||||
|
self.start_worker(i + 1, burst=burst, _sleep=_sleep, logging_level=logging_level)
|
||||||
|
|
||||||
|
def stop_worker(self, worker_data: WorkerData, sig=signal.SIGINT):
|
||||||
|
"""
|
||||||
|
Send stop signal to worker and catch "No such process" error if the worker is already dead.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
os.kill(worker_data.pid, sig)
|
||||||
|
self.log.info('Sent shutdown command to worker with %s', worker_data.pid)
|
||||||
|
except OSError as e:
|
||||||
|
if e.errno == errno.ESRCH:
|
||||||
|
# "No such process" is fine with us
|
||||||
|
self.log.debug('Horse already dead')
|
||||||
|
else:
|
||||||
|
raise
|
||||||
|
|
||||||
|
def stop_workers(self):
|
||||||
|
"""Send SIGINT to all workers"""
|
||||||
|
self.log.info('Sending stop signal to %s workers', len(self.worker_dict))
|
||||||
|
worker_datas = list(self.worker_dict.values())
|
||||||
|
for worker_data in worker_datas:
|
||||||
|
self.stop_worker(worker_data)
|
||||||
|
|
||||||
|
def start(self, burst: bool = False, logging_level: str = "INFO"):
|
||||||
|
self._burst = burst
|
||||||
|
respawn = not burst # Don't respawn workers if burst mode is on
|
||||||
|
setup_loghandlers(logging_level, DEFAULT_LOGGING_DATE_FORMAT, DEFAULT_LOGGING_FORMAT, name=__name__)
|
||||||
|
self.log.info(f'Starting worker pool {self.name} with pid %d...', os.getpid())
|
||||||
|
self.status = self.Status.IDLE
|
||||||
|
self.start_workers(burst=self._burst, logging_level=logging_level)
|
||||||
|
self._install_signal_handlers()
|
||||||
|
while True:
|
||||||
|
if self.status == self.Status.STOPPED:
|
||||||
|
if self.all_workers_have_stopped():
|
||||||
|
self.log.info('All workers stopped, exiting...')
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
self.log.info('Waiting for workers to shutdown...')
|
||||||
|
time.sleep(1)
|
||||||
|
continue
|
||||||
|
else:
|
||||||
|
self.check_workers(respawn=respawn)
|
||||||
|
if burst and self.number_of_active_workers == 0:
|
||||||
|
self.log.info('All workers stopped, exiting...')
|
||||||
|
break
|
||||||
|
|
||||||
|
time.sleep(1)
|
||||||
|
|
||||||
|
|
||||||
|
def run_worker(
|
||||||
|
worker_name: str,
|
||||||
|
queue_names: List[str],
|
||||||
|
connection_class,
|
||||||
|
connection_kwargs: dict,
|
||||||
|
worker_class: Type[BaseWorker] = Worker,
|
||||||
|
serializer: Type[DefaultSerializer] = DefaultSerializer,
|
||||||
|
job_class: Type[Job] = Job,
|
||||||
|
burst: bool = True,
|
||||||
|
logging_level: str = "INFO",
|
||||||
|
_sleep: int = 0,
|
||||||
|
):
|
||||||
|
connection = connection_class(**connection_kwargs)
|
||||||
|
queues = [Queue(name, connection=connection) for name in queue_names]
|
||||||
|
worker = worker_class(queues, name=worker_name, connection=connection, serializer=serializer, job_class=job_class)
|
||||||
|
worker.log.info("Starting worker started with PID %s", os.getpid())
|
||||||
|
time.sleep(_sleep)
|
||||||
|
worker.work(burst=burst, logging_level=logging_level)
|
||||||
@ -0,0 +1,138 @@
|
|||||||
|
import os
|
||||||
|
import signal
|
||||||
|
|
||||||
|
from multiprocessing import Process
|
||||||
|
from time import sleep
|
||||||
|
from rq.job import JobStatus
|
||||||
|
|
||||||
|
from tests import TestCase
|
||||||
|
from tests.fixtures import CustomJob, _send_shutdown_command, long_running_job, say_hello
|
||||||
|
|
||||||
|
from rq.queue import Queue
|
||||||
|
from rq.serializers import JSONSerializer
|
||||||
|
from rq.worker import SimpleWorker
|
||||||
|
from rq.worker_pool import run_worker, WorkerPool
|
||||||
|
|
||||||
|
|
||||||
|
def wait_and_send_shutdown_signal(pid, time_to_wait=0.0):
|
||||||
|
sleep(time_to_wait)
|
||||||
|
os.kill(pid, signal.SIGTERM)
|
||||||
|
|
||||||
|
|
||||||
|
class TestWorkerPool(TestCase):
|
||||||
|
def test_queues(self):
|
||||||
|
"""Test queue parsing"""
|
||||||
|
pool = WorkerPool(['default', 'foo'], connection=self.connection)
|
||||||
|
self.assertEqual(
|
||||||
|
set(pool.queues), {Queue('default', connection=self.connection), Queue('foo', connection=self.connection)}
|
||||||
|
)
|
||||||
|
|
||||||
|
# def test_spawn_workers(self):
|
||||||
|
# """Test spawning workers"""
|
||||||
|
# pool = WorkerPool(['default', 'foo'], connection=self.connection, num_workers=2)
|
||||||
|
# pool.start_workers(burst=False)
|
||||||
|
# self.assertEqual(len(pool.worker_dict.keys()), 2)
|
||||||
|
# pool.stop_workers()
|
||||||
|
|
||||||
|
def test_check_workers(self):
|
||||||
|
"""Test check_workers()"""
|
||||||
|
pool = WorkerPool(['default'], connection=self.connection, num_workers=2)
|
||||||
|
pool.start_workers(burst=False)
|
||||||
|
|
||||||
|
# There should be two workers
|
||||||
|
pool.check_workers()
|
||||||
|
self.assertEqual(len(pool.worker_dict.keys()), 2)
|
||||||
|
|
||||||
|
worker_data = list(pool.worker_dict.values())[0]
|
||||||
|
_send_shutdown_command(worker_data.name, self.connection.connection_pool.connection_kwargs.copy(), delay=0)
|
||||||
|
# 1 worker should be dead since we sent a shutdown command
|
||||||
|
sleep(0.2)
|
||||||
|
pool.check_workers(respawn=False)
|
||||||
|
self.assertEqual(len(pool.worker_dict.keys()), 1)
|
||||||
|
|
||||||
|
# If we call `check_workers` with `respawn=True`, the worker should be respawned
|
||||||
|
pool.check_workers(respawn=True)
|
||||||
|
self.assertEqual(len(pool.worker_dict.keys()), 2)
|
||||||
|
|
||||||
|
pool.stop_workers()
|
||||||
|
|
||||||
|
def test_reap_workers(self):
|
||||||
|
"""Dead workers are removed from worker_dict"""
|
||||||
|
pool = WorkerPool(['default'], connection=self.connection, num_workers=2)
|
||||||
|
pool.start_workers(burst=False)
|
||||||
|
|
||||||
|
# There should be two workers
|
||||||
|
pool.reap_workers()
|
||||||
|
self.assertEqual(len(pool.worker_dict.keys()), 2)
|
||||||
|
|
||||||
|
worker_data = list(pool.worker_dict.values())[0]
|
||||||
|
_send_shutdown_command(worker_data.name, self.connection.connection_pool.connection_kwargs.copy(), delay=0)
|
||||||
|
# 1 worker should be dead since we sent a shutdown command
|
||||||
|
sleep(0.2)
|
||||||
|
pool.reap_workers()
|
||||||
|
self.assertEqual(len(pool.worker_dict.keys()), 1)
|
||||||
|
pool.stop_workers()
|
||||||
|
|
||||||
|
def test_start(self):
|
||||||
|
"""Test start()"""
|
||||||
|
pool = WorkerPool(['default'], connection=self.connection, num_workers=2)
|
||||||
|
|
||||||
|
p = Process(target=wait_and_send_shutdown_signal, args=(os.getpid(), 0.5))
|
||||||
|
p.start()
|
||||||
|
pool.start()
|
||||||
|
self.assertEqual(pool.status, pool.Status.STOPPED)
|
||||||
|
self.assertTrue(pool.all_workers_have_stopped())
|
||||||
|
# We need this line so the test doesn't hang
|
||||||
|
pool.stop_workers()
|
||||||
|
|
||||||
|
def test_pool_ignores_consecutive_shutdown_signals(self):
|
||||||
|
"""If two shutdown signals are sent within one second, only the first one is processed"""
|
||||||
|
# Send two shutdown signals within one second while the worker is
|
||||||
|
# working on a long running job. The job should still complete (not killed)
|
||||||
|
pool = WorkerPool(['foo'], connection=self.connection, num_workers=2)
|
||||||
|
|
||||||
|
process_1 = Process(target=wait_and_send_shutdown_signal, args=(os.getpid(), 0.5))
|
||||||
|
process_1.start()
|
||||||
|
process_2 = Process(target=wait_and_send_shutdown_signal, args=(os.getpid(), 0.5))
|
||||||
|
process_2.start()
|
||||||
|
|
||||||
|
queue = Queue('foo', connection=self.connection)
|
||||||
|
job = queue.enqueue(long_running_job, 1)
|
||||||
|
pool.start(burst=True)
|
||||||
|
|
||||||
|
self.assertEqual(job.get_status(refresh=True), JobStatus.FINISHED)
|
||||||
|
# We need this line so the test doesn't hang
|
||||||
|
pool.stop_workers()
|
||||||
|
|
||||||
|
def test_run_worker(self):
|
||||||
|
"""Ensure run_worker() properly spawns a Worker"""
|
||||||
|
queue = Queue('foo', connection=self.connection)
|
||||||
|
queue.enqueue(say_hello)
|
||||||
|
run_worker(
|
||||||
|
'test-worker', ['foo'], self.connection.__class__, self.connection.connection_pool.connection_kwargs.copy()
|
||||||
|
)
|
||||||
|
# Worker should have processed the job
|
||||||
|
self.assertEqual(len(queue), 0)
|
||||||
|
|
||||||
|
def test_worker_pool_arguments(self):
|
||||||
|
"""Ensure arguments are properly used to create the right workers"""
|
||||||
|
queue = Queue('foo', connection=self.connection)
|
||||||
|
job = queue.enqueue(say_hello)
|
||||||
|
pool = WorkerPool([queue], connection=self.connection, num_workers=2, worker_class=SimpleWorker)
|
||||||
|
pool.start(burst=True)
|
||||||
|
# Worker should have processed the job
|
||||||
|
self.assertEqual(job.get_status(refresh=True), JobStatus.FINISHED)
|
||||||
|
|
||||||
|
queue = Queue('json', connection=self.connection, serializer=JSONSerializer)
|
||||||
|
job = queue.enqueue(say_hello, 'Hello')
|
||||||
|
pool = WorkerPool(
|
||||||
|
[queue], connection=self.connection, num_workers=2, worker_class=SimpleWorker, serializer=JSONSerializer
|
||||||
|
)
|
||||||
|
pool.start(burst=True)
|
||||||
|
# Worker should have processed the job
|
||||||
|
self.assertEqual(job.get_status(refresh=True), JobStatus.FINISHED)
|
||||||
|
|
||||||
|
pool = WorkerPool([queue], connection=self.connection, num_workers=2, job_class=CustomJob)
|
||||||
|
pool.start(burst=True)
|
||||||
|
# Worker should have processed the job
|
||||||
|
self.assertEqual(job.get_status(refresh=True), JobStatus.FINISHED)
|
||||||
Loading…
Reference in New Issue