Initial commit: AuroraMusic monorepo
This commit is contained in:
Kris
@@ -0,0 +1,23 @@
|
||||
# ruff: noqa: F401
|
||||
from .connections import Connection, get_current_connection, pop_connection, push_connection
|
||||
from .job import Callback, Retry, cancel_job, get_current_job, requeue_job
|
||||
from .queue import Queue
|
||||
from .version import VERSION
|
||||
from .worker import SimpleWorker, Worker
|
||||
|
||||
__all__ = [
|
||||
"Connection",
|
||||
"get_current_connection",
|
||||
"pop_connection",
|
||||
"push_connection",
|
||||
"Callback",
|
||||
"Retry",
|
||||
"cancel_job",
|
||||
"get_current_job",
|
||||
"requeue_job",
|
||||
"Queue",
|
||||
"SimpleWorker",
|
||||
"Worker",
|
||||
]
|
||||
|
||||
__version__ = VERSION
|
||||
BIN
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,6 @@
|
||||
# ruff: noqa: F401 I001
|
||||
from .cli import main
|
||||
|
||||
# TODO: the following imports can be removed when we drop the `rqinfo` and
|
||||
# `rqworkers` commands in favor of just shipping the `rq` command.
|
||||
from .cli import info, worker
|
||||
@@ -0,0 +1,6 @@
|
||||
import sys
|
||||
|
||||
from . import main
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(main())
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,521 @@
|
||||
"""
|
||||
RQ command line tool
|
||||
"""
|
||||
|
||||
import logging
|
||||
import logging.config
|
||||
import os
|
||||
import sys
|
||||
import warnings
|
||||
from typing import List, Type
|
||||
|
||||
import click
|
||||
from redis.exceptions import ConnectionError
|
||||
|
||||
from rq import Connection, Retry
|
||||
from rq import __version__ as version
|
||||
from rq.cli.helpers import (
|
||||
parse_function_args,
|
||||
parse_schedule,
|
||||
pass_cli_config,
|
||||
read_config_file,
|
||||
refresh,
|
||||
setup_loghandlers_from_args,
|
||||
show_both,
|
||||
show_queues,
|
||||
show_workers,
|
||||
)
|
||||
|
||||
# from rq.cli.pool import pool
|
||||
from rq.contrib.legacy import cleanup_ghosts
|
||||
from rq.defaults import (
|
||||
DEFAULT_JOB_MONITORING_INTERVAL,
|
||||
DEFAULT_LOGGING_DATE_FORMAT,
|
||||
DEFAULT_LOGGING_FORMAT,
|
||||
DEFAULT_MAINTENANCE_TASK_INTERVAL,
|
||||
DEFAULT_RESULT_TTL,
|
||||
DEFAULT_WORKER_TTL,
|
||||
)
|
||||
from rq.exceptions import InvalidJobOperationError
|
||||
from rq.job import Job, JobStatus
|
||||
from rq.logutils import blue
|
||||
from rq.registry import FailedJobRegistry, clean_registries
|
||||
from rq.serializers import DefaultSerializer
|
||||
from rq.suspension import is_suspended
|
||||
from rq.suspension import resume as connection_resume
|
||||
from rq.suspension import suspend as connection_suspend
|
||||
from rq.utils import get_call_string, import_attribute
|
||||
from rq.worker import Worker
|
||||
from rq.worker_pool import WorkerPool
|
||||
from rq.worker_registration import clean_worker_registry
|
||||
|
||||
|
||||
@click.group()
|
||||
@click.version_option(version)
|
||||
def main():
|
||||
"""RQ command line tool."""
|
||||
pass
|
||||
|
||||
|
||||
@main.command()
|
||||
@click.option('--all', '-a', is_flag=True, help='Empty all queues')
|
||||
@click.argument('queues', nargs=-1)
|
||||
@pass_cli_config
|
||||
def empty(cli_config, all, queues, serializer, **options):
|
||||
"""Empty given queues."""
|
||||
|
||||
if all:
|
||||
queues = cli_config.queue_class.all(
|
||||
connection=cli_config.connection,
|
||||
job_class=cli_config.job_class,
|
||||
death_penalty_class=cli_config.death_penalty_class,
|
||||
serializer=serializer,
|
||||
)
|
||||
else:
|
||||
queues = [
|
||||
cli_config.queue_class(
|
||||
queue, connection=cli_config.connection, job_class=cli_config.job_class, serializer=serializer
|
||||
)
|
||||
for queue in queues
|
||||
]
|
||||
|
||||
if not queues:
|
||||
click.echo('Nothing to do')
|
||||
sys.exit(0)
|
||||
|
||||
for queue in queues:
|
||||
num_jobs = queue.empty()
|
||||
click.echo('{0} jobs removed from {1} queue'.format(num_jobs, queue.name))
|
||||
|
||||
|
||||
@main.command()
|
||||
@click.option('--all', '-a', is_flag=True, help='Requeue all failed jobs')
|
||||
@click.option('--queue', required=True, type=str)
|
||||
@click.argument('job_ids', nargs=-1)
|
||||
@pass_cli_config
|
||||
def requeue(cli_config, queue, all, job_class, serializer, job_ids, **options):
|
||||
"""Requeue failed jobs."""
|
||||
|
||||
failed_job_registry = FailedJobRegistry(
|
||||
queue, connection=cli_config.connection, job_class=job_class, serializer=serializer
|
||||
)
|
||||
if all:
|
||||
job_ids = failed_job_registry.get_job_ids()
|
||||
|
||||
if not job_ids:
|
||||
click.echo('Nothing to do')
|
||||
sys.exit(0)
|
||||
|
||||
click.echo('Requeueing {0} jobs from failed queue'.format(len(job_ids)))
|
||||
fail_count = 0
|
||||
with click.progressbar(job_ids) as job_ids:
|
||||
for job_id in job_ids:
|
||||
try:
|
||||
failed_job_registry.requeue(job_id)
|
||||
except InvalidJobOperationError:
|
||||
fail_count += 1
|
||||
|
||||
if fail_count > 0:
|
||||
click.secho('Unable to requeue {0} jobs from failed job registry'.format(fail_count), fg='red')
|
||||
|
||||
|
||||
@main.command()
|
||||
@click.option('--interval', '-i', type=float, help='Updates stats every N seconds (default: don\'t poll)')
|
||||
@click.option('--raw', '-r', is_flag=True, help='Print only the raw numbers, no bar charts')
|
||||
@click.option('--only-queues', '-Q', is_flag=True, help='Show only queue info')
|
||||
@click.option('--only-workers', '-W', is_flag=True, help='Show only worker info')
|
||||
@click.option('--by-queue', '-R', is_flag=True, help='Shows workers by queue')
|
||||
@click.argument('queues', nargs=-1)
|
||||
@pass_cli_config
|
||||
def info(cli_config, interval, raw, only_queues, only_workers, by_queue, queues, **options):
|
||||
"""RQ command-line monitor."""
|
||||
|
||||
if only_queues:
|
||||
func = show_queues
|
||||
elif only_workers:
|
||||
func = show_workers
|
||||
else:
|
||||
func = show_both
|
||||
|
||||
try:
|
||||
with Connection(cli_config.connection):
|
||||
if queues:
|
||||
qs = list(map(cli_config.queue_class, queues))
|
||||
else:
|
||||
qs = cli_config.queue_class.all()
|
||||
|
||||
for queue in qs:
|
||||
clean_registries(queue)
|
||||
clean_worker_registry(queue)
|
||||
|
||||
refresh(interval, func, qs, raw, by_queue, cli_config.queue_class, cli_config.worker_class)
|
||||
except ConnectionError as e:
|
||||
click.echo(e)
|
||||
sys.exit(1)
|
||||
except KeyboardInterrupt:
|
||||
click.echo()
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
@main.command()
|
||||
@click.option('--burst', '-b', is_flag=True, help='Run in burst mode (quit after all work is done)')
|
||||
@click.option('--logging_level', type=str, default="INFO", help='Set logging level')
|
||||
@click.option('--log-format', type=str, default=DEFAULT_LOGGING_FORMAT, help='Set the format of the logs')
|
||||
@click.option('--date-format', type=str, default=DEFAULT_LOGGING_DATE_FORMAT, help='Set the date format of the logs')
|
||||
@click.option('--name', '-n', help='Specify a different name')
|
||||
@click.option('--results-ttl', type=int, default=DEFAULT_RESULT_TTL, help='Default results timeout to be used')
|
||||
@click.option('--worker-ttl', type=int, default=DEFAULT_WORKER_TTL, help='Worker timeout to be used')
|
||||
@click.option(
|
||||
'--maintenance-interval',
|
||||
type=int,
|
||||
default=DEFAULT_MAINTENANCE_TASK_INTERVAL,
|
||||
help='Maintenance task interval (in seconds) to be used',
|
||||
)
|
||||
@click.option(
|
||||
'--job-monitoring-interval',
|
||||
type=int,
|
||||
default=DEFAULT_JOB_MONITORING_INTERVAL,
|
||||
help='Default job monitoring interval to be used',
|
||||
)
|
||||
@click.option('--disable-job-desc-logging', is_flag=True, help='Turn off description logging.')
|
||||
@click.option('--verbose', '-v', is_flag=True, help='Show more output')
|
||||
@click.option('--quiet', '-q', is_flag=True, help='Show less output')
|
||||
@click.option('--sentry-ca-certs', envvar='RQ_SENTRY_CA_CERTS', help='Path to CRT file for Sentry DSN')
|
||||
@click.option('--sentry-debug', envvar='RQ_SENTRY_DEBUG', help='Enable debug')
|
||||
@click.option('--sentry-dsn', envvar='RQ_SENTRY_DSN', help='Report exceptions to this Sentry DSN')
|
||||
@click.option('--exception-handler', help='Exception handler(s) to use', multiple=True)
|
||||
@click.option('--pid', help='Write the process ID number to a file at the specified path')
|
||||
@click.option('--disable-default-exception-handler', '-d', is_flag=True, help='Disable RQ\'s default exception handler')
|
||||
@click.option('--max-jobs', type=int, default=None, help='Maximum number of jobs to execute')
|
||||
@click.option('--max-idle-time', type=int, default=None, help='Maximum seconds to stay alive without jobs to execute')
|
||||
@click.option('--with-scheduler', '-s', is_flag=True, help='Run worker with scheduler')
|
||||
@click.option('--serializer', '-S', default=None, help='Run worker with custom serializer')
|
||||
@click.option(
|
||||
'--dequeue-strategy', '-ds', default='default', help='Sets a custom stratey to dequeue from multiple queues'
|
||||
)
|
||||
@click.argument('queues', nargs=-1)
|
||||
@pass_cli_config
|
||||
def worker(
|
||||
cli_config,
|
||||
burst,
|
||||
logging_level,
|
||||
name,
|
||||
results_ttl,
|
||||
worker_ttl,
|
||||
maintenance_interval,
|
||||
job_monitoring_interval,
|
||||
disable_job_desc_logging,
|
||||
verbose,
|
||||
quiet,
|
||||
sentry_ca_certs,
|
||||
sentry_debug,
|
||||
sentry_dsn,
|
||||
exception_handler,
|
||||
pid,
|
||||
disable_default_exception_handler,
|
||||
max_jobs,
|
||||
max_idle_time,
|
||||
with_scheduler,
|
||||
queues,
|
||||
log_format,
|
||||
date_format,
|
||||
serializer,
|
||||
dequeue_strategy,
|
||||
**options,
|
||||
):
|
||||
"""Starts an RQ worker."""
|
||||
settings = read_config_file(cli_config.config) if cli_config.config else {}
|
||||
# Worker specific default arguments
|
||||
queues = queues or settings.get('QUEUES', ['default'])
|
||||
sentry_ca_certs = sentry_ca_certs or settings.get('SENTRY_CA_CERTS')
|
||||
sentry_debug = sentry_debug or settings.get('SENTRY_DEBUG')
|
||||
sentry_dsn = sentry_dsn or settings.get('SENTRY_DSN')
|
||||
name = name or settings.get('NAME')
|
||||
dict_config = settings.get('DICT_CONFIG')
|
||||
|
||||
if dict_config:
|
||||
logging.config.dictConfig(dict_config)
|
||||
|
||||
if pid:
|
||||
with open(os.path.expanduser(pid), "w") as fp:
|
||||
fp.write(str(os.getpid()))
|
||||
|
||||
worker_name = cli_config.worker_class.__qualname__
|
||||
if worker_name in ["RoundRobinWorker", "RandomWorker"]:
|
||||
strategy_alternative = "random" if worker_name == "RandomWorker" else "round_robin"
|
||||
msg = f"WARNING: {worker_name} is deprecated. Use `--dequeue-strategy {strategy_alternative}` instead."
|
||||
warnings.warn(msg, DeprecationWarning)
|
||||
click.secho(msg, fg='yellow')
|
||||
|
||||
if dequeue_strategy not in ("default", "random", "round_robin"):
|
||||
click.secho(
|
||||
"ERROR: Dequeue Strategy can only be one of `default`, `random` or `round_robin`.", err=True, fg='red'
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
setup_loghandlers_from_args(verbose, quiet, date_format, log_format)
|
||||
|
||||
try:
|
||||
cleanup_ghosts(cli_config.connection, worker_class=cli_config.worker_class)
|
||||
exception_handlers = []
|
||||
for h in exception_handler:
|
||||
exception_handlers.append(import_attribute(h))
|
||||
|
||||
if is_suspended(cli_config.connection):
|
||||
click.secho('RQ is currently suspended, to resume job execution run "rq resume"', fg='red')
|
||||
sys.exit(1)
|
||||
|
||||
queues = [
|
||||
cli_config.queue_class(
|
||||
queue, connection=cli_config.connection, job_class=cli_config.job_class, serializer=serializer
|
||||
)
|
||||
for queue in queues
|
||||
]
|
||||
worker = cli_config.worker_class(
|
||||
queues,
|
||||
name=name,
|
||||
connection=cli_config.connection,
|
||||
default_worker_ttl=worker_ttl,
|
||||
default_result_ttl=results_ttl,
|
||||
maintenance_interval=maintenance_interval,
|
||||
job_monitoring_interval=job_monitoring_interval,
|
||||
job_class=cli_config.job_class,
|
||||
queue_class=cli_config.queue_class,
|
||||
exception_handlers=exception_handlers or None,
|
||||
disable_default_exception_handler=disable_default_exception_handler,
|
||||
log_job_description=not disable_job_desc_logging,
|
||||
serializer=serializer,
|
||||
)
|
||||
|
||||
# Should we configure Sentry?
|
||||
if sentry_dsn:
|
||||
sentry_opts = {"ca_certs": sentry_ca_certs, "debug": sentry_debug}
|
||||
from rq.contrib.sentry import register_sentry
|
||||
|
||||
register_sentry(sentry_dsn, **sentry_opts)
|
||||
|
||||
# if --verbose or --quiet, override --logging_level
|
||||
if verbose or quiet:
|
||||
logging_level = None
|
||||
|
||||
worker.work(
|
||||
burst=burst,
|
||||
logging_level=logging_level,
|
||||
date_format=date_format,
|
||||
log_format=log_format,
|
||||
max_jobs=max_jobs,
|
||||
max_idle_time=max_idle_time,
|
||||
with_scheduler=with_scheduler,
|
||||
dequeue_strategy=dequeue_strategy,
|
||||
)
|
||||
except ConnectionError as e:
|
||||
logging.error(e)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
@main.command()
|
||||
@click.option('--duration', help='Seconds you want the workers to be suspended. Default is forever.', type=int)
|
||||
@pass_cli_config
|
||||
def suspend(cli_config, duration, **options):
|
||||
"""Suspends all workers, to resume run `rq resume`"""
|
||||
|
||||
if duration is not None and duration < 1:
|
||||
click.echo("Duration must be an integer greater than 1")
|
||||
sys.exit(1)
|
||||
|
||||
connection_suspend(cli_config.connection, duration)
|
||||
|
||||
if duration:
|
||||
msg = """Suspending workers for {0} seconds. No new jobs will be started during that time, but then will
|
||||
automatically resume""".format(
|
||||
duration
|
||||
)
|
||||
click.echo(msg)
|
||||
else:
|
||||
click.echo("Suspending workers. No new jobs will be started. But current jobs will be completed")
|
||||
|
||||
|
||||
@main.command()
|
||||
@pass_cli_config
|
||||
def resume(cli_config, **options):
|
||||
"""Resumes processing of queues, that were suspended with `rq suspend`"""
|
||||
connection_resume(cli_config.connection)
|
||||
click.echo("Resuming workers.")
|
||||
|
||||
|
||||
@main.command()
|
||||
@click.option('--queue', '-q', help='The name of the queue.', default='default')
|
||||
@click.option(
|
||||
'--timeout', help='Specifies the maximum runtime of the job before it is interrupted and marked as failed.'
|
||||
)
|
||||
@click.option('--result-ttl', help='Specifies how long successful jobs and their results are kept.')
|
||||
@click.option('--ttl', help='Specifies the maximum queued time of the job before it is discarded.')
|
||||
@click.option('--failure-ttl', help='Specifies how long failed jobs are kept.')
|
||||
@click.option('--description', help='Additional description of the job')
|
||||
@click.option(
|
||||
'--depends-on', help='Specifies another job id that must complete before this job will be queued.', multiple=True
|
||||
)
|
||||
@click.option('--job-id', help='The id of this job')
|
||||
@click.option('--at-front', is_flag=True, help='Will place the job at the front of the queue, instead of the end')
|
||||
@click.option('--retry-max', help='Maximum amount of retries', default=0, type=int)
|
||||
@click.option('--retry-interval', help='Interval between retries in seconds', multiple=True, type=int, default=[0])
|
||||
@click.option('--schedule-in', help='Delay until the function is enqueued (e.g. 10s, 5m, 2d).')
|
||||
@click.option(
|
||||
'--schedule-at',
|
||||
help='Schedule job to be enqueued at a certain time formatted in ISO 8601 without '
|
||||
'timezone (e.g. 2021-05-27T21:45:00).',
|
||||
)
|
||||
@click.option('--quiet', is_flag=True, help='Only logs errors.')
|
||||
@click.argument('function')
|
||||
@click.argument('arguments', nargs=-1)
|
||||
@pass_cli_config
|
||||
def enqueue(
|
||||
cli_config,
|
||||
queue,
|
||||
timeout,
|
||||
result_ttl,
|
||||
ttl,
|
||||
failure_ttl,
|
||||
description,
|
||||
depends_on,
|
||||
job_id,
|
||||
at_front,
|
||||
retry_max,
|
||||
retry_interval,
|
||||
schedule_in,
|
||||
schedule_at,
|
||||
quiet,
|
||||
serializer,
|
||||
function,
|
||||
arguments,
|
||||
**options,
|
||||
):
|
||||
"""Enqueues a job from the command line"""
|
||||
args, kwargs = parse_function_args(arguments)
|
||||
function_string = get_call_string(function, args, kwargs)
|
||||
description = description or function_string
|
||||
|
||||
retry = None
|
||||
if retry_max > 0:
|
||||
retry = Retry(retry_max, retry_interval)
|
||||
|
||||
schedule = parse_schedule(schedule_in, schedule_at)
|
||||
|
||||
with Connection(cli_config.connection):
|
||||
queue = cli_config.queue_class(queue, serializer=serializer)
|
||||
|
||||
if schedule is None:
|
||||
job = queue.enqueue_call(
|
||||
function,
|
||||
args,
|
||||
kwargs,
|
||||
timeout,
|
||||
result_ttl,
|
||||
ttl,
|
||||
failure_ttl,
|
||||
description,
|
||||
depends_on,
|
||||
job_id,
|
||||
at_front,
|
||||
None,
|
||||
retry,
|
||||
)
|
||||
else:
|
||||
job = queue.create_job(
|
||||
function,
|
||||
args,
|
||||
kwargs,
|
||||
timeout,
|
||||
result_ttl,
|
||||
ttl,
|
||||
failure_ttl,
|
||||
description,
|
||||
depends_on,
|
||||
job_id,
|
||||
None,
|
||||
JobStatus.SCHEDULED,
|
||||
retry,
|
||||
)
|
||||
queue.schedule_job(job, schedule)
|
||||
|
||||
if not quiet:
|
||||
click.echo('Enqueued %s with job-id \'%s\'.' % (blue(function_string), job.id))
|
||||
|
||||
|
||||
@main.command()
|
||||
@click.option('--burst', '-b', is_flag=True, help='Run in burst mode (quit after all work is done)')
|
||||
@click.option('--logging-level', '-l', type=str, default="INFO", help='Set logging level')
|
||||
@click.option('--sentry-ca-certs', envvar='RQ_SENTRY_CA_CERTS', help='Path to CRT file for Sentry DSN')
|
||||
@click.option('--sentry-debug', envvar='RQ_SENTRY_DEBUG', help='Enable debug')
|
||||
@click.option('--sentry-dsn', envvar='RQ_SENTRY_DSN', help='Report exceptions to this Sentry DSN')
|
||||
@click.option('--verbose', '-v', is_flag=True, help='Show more output')
|
||||
@click.option('--quiet', '-q', is_flag=True, help='Show less output')
|
||||
@click.option('--log-format', type=str, default=DEFAULT_LOGGING_FORMAT, help='Set the format of the logs')
|
||||
@click.option('--date-format', type=str, default=DEFAULT_LOGGING_DATE_FORMAT, help='Set the date format of the logs')
|
||||
@click.option('--job-class', type=str, default=None, help='Dotted path to a Job class')
|
||||
@click.argument('queues', nargs=-1)
|
||||
@click.option('--num-workers', '-n', type=int, default=1, help='Number of workers to start')
|
||||
@pass_cli_config
|
||||
def worker_pool(
|
||||
cli_config,
|
||||
burst: bool,
|
||||
logging_level,
|
||||
queues,
|
||||
serializer,
|
||||
sentry_ca_certs,
|
||||
sentry_debug,
|
||||
sentry_dsn,
|
||||
verbose,
|
||||
quiet,
|
||||
log_format,
|
||||
date_format,
|
||||
worker_class,
|
||||
job_class,
|
||||
num_workers,
|
||||
**options,
|
||||
):
|
||||
"""Starts a RQ worker pool"""
|
||||
settings = read_config_file(cli_config.config) if cli_config.config else {}
|
||||
# Worker specific default arguments
|
||||
queue_names: List[str] = queues or settings.get('QUEUES', ['default'])
|
||||
sentry_ca_certs = sentry_ca_certs or settings.get('SENTRY_CA_CERTS')
|
||||
sentry_debug = sentry_debug or settings.get('SENTRY_DEBUG')
|
||||
sentry_dsn = sentry_dsn or settings.get('SENTRY_DSN')
|
||||
|
||||
setup_loghandlers_from_args(verbose, quiet, date_format, log_format)
|
||||
|
||||
if serializer:
|
||||
serializer_class: Type[DefaultSerializer] = import_attribute(serializer)
|
||||
else:
|
||||
serializer_class = DefaultSerializer
|
||||
|
||||
if worker_class:
|
||||
worker_class = import_attribute(worker_class)
|
||||
else:
|
||||
worker_class = Worker
|
||||
|
||||
if job_class:
|
||||
job_class = import_attribute(job_class)
|
||||
else:
|
||||
job_class = Job
|
||||
|
||||
pool = WorkerPool(
|
||||
queue_names,
|
||||
connection=cli_config.connection,
|
||||
num_workers=num_workers,
|
||||
serializer=serializer_class,
|
||||
worker_class=worker_class,
|
||||
job_class=job_class,
|
||||
)
|
||||
pool.start(burst=burst, logging_level=logging_level)
|
||||
|
||||
# Should we configure Sentry?
|
||||
if sentry_dsn:
|
||||
sentry_opts = {"ca_certs": sentry_ca_certs, "debug": sentry_debug}
|
||||
from rq.contrib.sentry import register_sentry
|
||||
|
||||
register_sentry(sentry_dsn, **sentry_opts)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,424 @@
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from ast import literal_eval
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from enum import Enum
|
||||
from functools import partial, update_wrapper
|
||||
from json import JSONDecodeError, loads
|
||||
from shutil import get_terminal_size
|
||||
|
||||
import click
|
||||
from redis import Redis
|
||||
from redis.sentinel import Sentinel
|
||||
|
||||
from rq.defaults import (
|
||||
DEFAULT_CONNECTION_CLASS,
|
||||
DEFAULT_DEATH_PENALTY_CLASS,
|
||||
DEFAULT_JOB_CLASS,
|
||||
DEFAULT_QUEUE_CLASS,
|
||||
DEFAULT_SERIALIZER_CLASS,
|
||||
DEFAULT_WORKER_CLASS,
|
||||
)
|
||||
from rq.logutils import setup_loghandlers
|
||||
from rq.utils import import_attribute, parse_timeout
|
||||
from rq.worker import WorkerStatus
|
||||
|
||||
red = partial(click.style, fg='red')
|
||||
green = partial(click.style, fg='green')
|
||||
yellow = partial(click.style, fg='yellow')
|
||||
|
||||
|
||||
def read_config_file(module):
|
||||
"""Reads all UPPERCASE variables defined in the given module file."""
|
||||
settings = importlib.import_module(module)
|
||||
return dict([(k, v) for k, v in settings.__dict__.items() if k.upper() == k])
|
||||
|
||||
|
||||
def get_redis_from_config(settings, connection_class=Redis):
|
||||
"""Returns a StrictRedis instance from a dictionary of settings.
|
||||
To use redis sentinel, you must specify a dictionary in the configuration file.
|
||||
Example of a dictionary with keys without values:
|
||||
SENTINEL = {'INSTANCES':, 'SOCKET_TIMEOUT':, 'USERNAME':, 'PASSWORD':, 'DB':, 'MASTER_NAME':, 'SENTINEL_KWARGS':}
|
||||
"""
|
||||
if settings.get('REDIS_URL') is not None:
|
||||
return connection_class.from_url(settings['REDIS_URL'])
|
||||
|
||||
elif settings.get('SENTINEL') is not None:
|
||||
instances = settings['SENTINEL'].get('INSTANCES', [('localhost', 26379)])
|
||||
master_name = settings['SENTINEL'].get('MASTER_NAME', 'mymaster')
|
||||
|
||||
connection_kwargs = {
|
||||
'db': settings['SENTINEL'].get('DB', 0),
|
||||
'username': settings['SENTINEL'].get('USERNAME', None),
|
||||
'password': settings['SENTINEL'].get('PASSWORD', None),
|
||||
'socket_timeout': settings['SENTINEL'].get('SOCKET_TIMEOUT', None),
|
||||
'ssl': settings['SENTINEL'].get('SSL', False),
|
||||
}
|
||||
connection_kwargs.update(settings['SENTINEL'].get('CONNECTION_KWARGS', {}))
|
||||
sentinel_kwargs = settings['SENTINEL'].get('SENTINEL_KWARGS', {})
|
||||
|
||||
sn = Sentinel(instances, sentinel_kwargs=sentinel_kwargs, **connection_kwargs)
|
||||
return sn.master_for(master_name)
|
||||
|
||||
ssl = settings.get('REDIS_SSL', False)
|
||||
if isinstance(ssl, str):
|
||||
if ssl.lower() in ['y', 'yes', 't', 'true']:
|
||||
ssl = True
|
||||
elif ssl.lower() in ['n', 'no', 'f', 'false', '']:
|
||||
ssl = False
|
||||
else:
|
||||
raise ValueError('REDIS_SSL is a boolean and must be "True" or "False".')
|
||||
|
||||
kwargs = {
|
||||
'host': settings.get('REDIS_HOST', 'localhost'),
|
||||
'port': settings.get('REDIS_PORT', 6379),
|
||||
'db': settings.get('REDIS_DB', 0),
|
||||
'password': settings.get('REDIS_PASSWORD', None),
|
||||
'ssl': ssl,
|
||||
'ssl_ca_certs': settings.get('REDIS_SSL_CA_CERTS', None),
|
||||
'ssl_cert_reqs': settings.get('REDIS_SSL_CERT_REQS', 'required'),
|
||||
}
|
||||
|
||||
return connection_class(**kwargs)
|
||||
|
||||
|
||||
def pad(s, pad_to_length):
|
||||
"""Pads the given string to the given length."""
|
||||
return ('%-' + '%ds' % pad_to_length) % (s,)
|
||||
|
||||
|
||||
def get_scale(x):
|
||||
"""Finds the lowest scale where x <= scale."""
|
||||
scales = [20, 50, 100, 200, 400, 600, 800, 1000]
|
||||
for scale in scales:
|
||||
if x <= scale:
|
||||
return scale
|
||||
return x
|
||||
|
||||
|
||||
def state_symbol(state):
|
||||
symbols = {
|
||||
WorkerStatus.BUSY: red('busy'),
|
||||
WorkerStatus.IDLE: green('idle'),
|
||||
WorkerStatus.SUSPENDED: yellow('suspended'),
|
||||
}
|
||||
try:
|
||||
return symbols[state]
|
||||
except KeyError:
|
||||
return state
|
||||
|
||||
|
||||
def show_queues(queues, raw, by_queue, queue_class, worker_class):
|
||||
num_jobs = 0
|
||||
termwidth = get_terminal_size().columns
|
||||
chartwidth = min(20, termwidth - 20)
|
||||
|
||||
max_count = 0
|
||||
counts = dict()
|
||||
for q in queues:
|
||||
count = q.count
|
||||
counts[q] = count
|
||||
max_count = max(max_count, count)
|
||||
scale = get_scale(max_count)
|
||||
ratio = chartwidth * 1.0 / scale
|
||||
|
||||
for q in queues:
|
||||
count = counts[q]
|
||||
if not raw:
|
||||
chart = green('|' + '█' * int(ratio * count))
|
||||
line = '%-12s %s %d, %d executing, %d finished, %d failed' % (
|
||||
q.name,
|
||||
chart,
|
||||
count,
|
||||
q.started_job_registry.count,
|
||||
q.finished_job_registry.count,
|
||||
q.failed_job_registry.count,
|
||||
)
|
||||
else:
|
||||
line = 'queue %s %d, %d executing, %d finished, %d failed' % (
|
||||
q.name,
|
||||
count,
|
||||
q.started_job_registry.count,
|
||||
q.finished_job_registry.count,
|
||||
q.failed_job_registry.count,
|
||||
)
|
||||
click.echo(line)
|
||||
|
||||
num_jobs += count
|
||||
|
||||
# print summary when not in raw mode
|
||||
if not raw:
|
||||
click.echo('%d queues, %d jobs total' % (len(queues), num_jobs))
|
||||
|
||||
|
||||
def show_workers(queues, raw, by_queue, queue_class, worker_class):
|
||||
workers = set()
|
||||
if queues:
|
||||
for queue in queues:
|
||||
for worker in worker_class.all(queue=queue):
|
||||
workers.add(worker)
|
||||
else:
|
||||
for worker in worker_class.all():
|
||||
workers.add(worker)
|
||||
|
||||
if not by_queue:
|
||||
for worker in workers:
|
||||
queue_names = ', '.join(worker.queue_names())
|
||||
name = '%s (%s %s %s)' % (worker.name, worker.hostname, worker.ip_address, worker.pid)
|
||||
if not raw:
|
||||
line = '%s: %s %s. jobs: %d finished, %d failed' % (
|
||||
name,
|
||||
state_symbol(worker.get_state()),
|
||||
queue_names,
|
||||
worker.successful_job_count,
|
||||
worker.failed_job_count,
|
||||
)
|
||||
click.echo(line)
|
||||
else:
|
||||
line = 'worker %s %s %s. jobs: %d finished, %d failed' % (
|
||||
name,
|
||||
worker.get_state(),
|
||||
queue_names,
|
||||
worker.successful_job_count,
|
||||
worker.failed_job_count,
|
||||
)
|
||||
click.echo(line)
|
||||
|
||||
else:
|
||||
# Display workers by queue
|
||||
queue_dict = {}
|
||||
for queue in queues:
|
||||
queue_dict[queue] = worker_class.all(queue=queue)
|
||||
|
||||
if queue_dict:
|
||||
max_length = max(len(q.name) for q, in queue_dict.keys())
|
||||
else:
|
||||
max_length = 0
|
||||
|
||||
for queue in queue_dict:
|
||||
if queue_dict[queue]:
|
||||
queues_str = ", ".join(
|
||||
sorted(map(lambda w: '%s (%s)' % (w.name, state_symbol(w.get_state())), queue_dict[queue]))
|
||||
)
|
||||
else:
|
||||
queues_str = '–'
|
||||
click.echo('%s %s' % (pad(queue.name + ':', max_length + 1), queues_str))
|
||||
|
||||
if not raw:
|
||||
click.echo('%d workers, %d queues' % (len(workers), len(queues)))
|
||||
|
||||
|
||||
def show_both(queues, raw, by_queue, queue_class, worker_class):
|
||||
show_queues(queues, raw, by_queue, queue_class, worker_class)
|
||||
if not raw:
|
||||
click.echo('')
|
||||
show_workers(queues, raw, by_queue, queue_class, worker_class)
|
||||
if not raw:
|
||||
click.echo('')
|
||||
import datetime
|
||||
|
||||
click.echo('Updated: %s' % datetime.datetime.now())
|
||||
|
||||
|
||||
def refresh(interval, func, *args):
|
||||
while True:
|
||||
if interval:
|
||||
click.clear()
|
||||
func(*args)
|
||||
if interval:
|
||||
time.sleep(interval)
|
||||
else:
|
||||
break
|
||||
|
||||
|
||||
def setup_loghandlers_from_args(verbose, quiet, date_format, log_format):
|
||||
if verbose and quiet:
|
||||
raise RuntimeError("Flags --verbose and --quiet are mutually exclusive.")
|
||||
|
||||
if verbose:
|
||||
level = 'DEBUG'
|
||||
elif quiet:
|
||||
level = 'WARNING'
|
||||
else:
|
||||
level = 'INFO'
|
||||
setup_loghandlers(level, date_format=date_format, log_format=log_format)
|
||||
|
||||
|
||||
def parse_function_arg(argument, arg_pos):
|
||||
class ParsingMode(Enum):
|
||||
PLAIN_TEXT = 0
|
||||
JSON = 1
|
||||
LITERAL_EVAL = 2
|
||||
|
||||
keyword = None
|
||||
if argument.startswith(':'): # no keyword, json
|
||||
mode = ParsingMode.JSON
|
||||
value = argument[1:]
|
||||
elif argument.startswith('%'): # no keyword, literal_eval
|
||||
mode = ParsingMode.LITERAL_EVAL
|
||||
value = argument[1:]
|
||||
else:
|
||||
index = argument.find('=')
|
||||
if index > 0:
|
||||
if ':' in argument and argument.index(':') + 1 == index: # keyword, json
|
||||
mode = ParsingMode.JSON
|
||||
keyword = argument[: index - 1]
|
||||
elif '%' in argument and argument.index('%') + 1 == index: # keyword, literal_eval
|
||||
mode = ParsingMode.LITERAL_EVAL
|
||||
keyword = argument[: index - 1]
|
||||
else: # keyword, text
|
||||
mode = ParsingMode.PLAIN_TEXT
|
||||
keyword = argument[:index]
|
||||
value = argument[index + 1 :]
|
||||
else: # no keyword, text
|
||||
mode = ParsingMode.PLAIN_TEXT
|
||||
value = argument
|
||||
|
||||
if value.startswith('@'):
|
||||
try:
|
||||
with open(value[1:], 'r') as file:
|
||||
value = file.read()
|
||||
except FileNotFoundError:
|
||||
raise click.FileError(value[1:], 'Not found')
|
||||
|
||||
if mode == ParsingMode.JSON: # json
|
||||
try:
|
||||
value = loads(value)
|
||||
except JSONDecodeError:
|
||||
raise click.BadParameter('Unable to parse %s as JSON.' % (keyword or '%s. non keyword argument' % arg_pos))
|
||||
elif mode == ParsingMode.LITERAL_EVAL: # literal_eval
|
||||
try:
|
||||
value = literal_eval(value)
|
||||
except Exception:
|
||||
raise click.BadParameter(
|
||||
'Unable to eval %s as Python object. See '
|
||||
'https://docs.python.org/3/library/ast.html#ast.literal_eval'
|
||||
% (keyword or '%s. non keyword argument' % arg_pos)
|
||||
)
|
||||
|
||||
return keyword, value
|
||||
|
||||
|
||||
def parse_function_args(arguments):
|
||||
args = []
|
||||
kwargs = {}
|
||||
|
||||
for argument in arguments:
|
||||
keyword, value = parse_function_arg(argument, len(args) + 1)
|
||||
if keyword is not None:
|
||||
if keyword in kwargs:
|
||||
raise click.BadParameter('You can\'t specify multiple values for the same keyword.')
|
||||
kwargs[keyword] = value
|
||||
else:
|
||||
args.append(value)
|
||||
return args, kwargs
|
||||
|
||||
|
||||
def parse_schedule(schedule_in, schedule_at):
|
||||
if schedule_in is not None:
|
||||
if schedule_at is not None:
|
||||
raise click.BadArgumentUsage('You can\'t specify both --schedule-in and --schedule-at')
|
||||
return datetime.now(timezone.utc) + timedelta(seconds=parse_timeout(schedule_in))
|
||||
elif schedule_at is not None:
|
||||
return datetime.strptime(schedule_at, '%Y-%m-%dT%H:%M:%S')
|
||||
|
||||
|
||||
class CliConfig:
|
||||
"""A helper class to be used with click commands, to handle shared options"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url=None,
|
||||
config=None,
|
||||
worker_class=DEFAULT_WORKER_CLASS,
|
||||
job_class=DEFAULT_JOB_CLASS,
|
||||
death_penalty_class=DEFAULT_DEATH_PENALTY_CLASS,
|
||||
queue_class=DEFAULT_QUEUE_CLASS,
|
||||
connection_class=DEFAULT_CONNECTION_CLASS,
|
||||
path=None,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
self._connection = None
|
||||
self.url = url
|
||||
self.config = config
|
||||
|
||||
if path:
|
||||
for pth in path:
|
||||
sys.path.append(pth)
|
||||
|
||||
try:
|
||||
self.worker_class = import_attribute(worker_class)
|
||||
except (ImportError, AttributeError) as exc:
|
||||
raise click.BadParameter(str(exc), param_hint='--worker-class')
|
||||
try:
|
||||
self.job_class = import_attribute(job_class)
|
||||
except (ImportError, AttributeError) as exc:
|
||||
raise click.BadParameter(str(exc), param_hint='--job-class')
|
||||
|
||||
try:
|
||||
self.death_penalty_class = import_attribute(death_penalty_class)
|
||||
except (ImportError, AttributeError) as exc:
|
||||
raise click.BadParameter(str(exc), param_hint='--death-penalty-class')
|
||||
|
||||
try:
|
||||
self.queue_class = import_attribute(queue_class)
|
||||
except (ImportError, AttributeError) as exc:
|
||||
raise click.BadParameter(str(exc), param_hint='--queue-class')
|
||||
|
||||
try:
|
||||
self.connection_class = import_attribute(connection_class)
|
||||
except (ImportError, AttributeError) as exc:
|
||||
raise click.BadParameter(str(exc), param_hint='--connection-class')
|
||||
|
||||
@property
|
||||
def connection(self):
|
||||
if self._connection is None:
|
||||
if self.url:
|
||||
self._connection = self.connection_class.from_url(self.url)
|
||||
elif self.config:
|
||||
settings = read_config_file(self.config) if self.config else {}
|
||||
self._connection = get_redis_from_config(settings, self.connection_class)
|
||||
else:
|
||||
self._connection = get_redis_from_config(os.environ, self.connection_class)
|
||||
return self._connection
|
||||
|
||||
|
||||
shared_options = [
|
||||
click.option('--url', '-u', envvar='RQ_REDIS_URL', help='URL describing Redis connection details.'),
|
||||
click.option('--config', '-c', envvar='RQ_CONFIG', help='Module containing RQ settings.'),
|
||||
click.option(
|
||||
'--worker-class', '-w', envvar='RQ_WORKER_CLASS', default=DEFAULT_WORKER_CLASS, help='RQ Worker class to use'
|
||||
),
|
||||
click.option('--job-class', '-j', envvar='RQ_JOB_CLASS', default=DEFAULT_JOB_CLASS, help='RQ Job class to use'),
|
||||
click.option('--queue-class', envvar='RQ_QUEUE_CLASS', default=DEFAULT_QUEUE_CLASS, help='RQ Queue class to use'),
|
||||
click.option(
|
||||
'--connection-class',
|
||||
envvar='RQ_CONNECTION_CLASS',
|
||||
default=DEFAULT_CONNECTION_CLASS,
|
||||
help='Redis client class to use',
|
||||
),
|
||||
click.option('--path', '-P', default=['.'], help='Specify the import path.', multiple=True),
|
||||
click.option(
|
||||
'--serializer',
|
||||
'-S',
|
||||
default=DEFAULT_SERIALIZER_CLASS,
|
||||
help='Path to serializer, defaults to rq.serializers.DefaultSerializer',
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def pass_cli_config(func):
|
||||
# add all the shared options to the command
|
||||
for option in shared_options:
|
||||
func = option(func)
|
||||
|
||||
# pass the cli config object into the command
|
||||
def wrapper(*args, **kwargs):
|
||||
ctx = click.get_current_context()
|
||||
cli_config = CliConfig(**kwargs)
|
||||
return ctx.invoke(func, cli_config, *args[1:], **kwargs)
|
||||
|
||||
return update_wrapper(wrapper, func)
|
||||
@@ -0,0 +1,143 @@
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
from typing import TYPE_CHECKING, Any, Dict
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from redis import Redis
|
||||
|
||||
from .worker import Worker
|
||||
|
||||
from rq.exceptions import InvalidJobOperation
|
||||
from rq.job import Job
|
||||
|
||||
PUBSUB_CHANNEL_TEMPLATE = 'rq:pubsub:%s'
|
||||
|
||||
|
||||
def send_command(connection: 'Redis', worker_name: str, command: str, **kwargs):
|
||||
"""
|
||||
Sends a command to a worker.
|
||||
A command is just a string, availble commands are:
|
||||
- `shutdown`: Shuts down a worker
|
||||
- `kill-horse`: Command for the worker to kill the current working horse
|
||||
- `stop-job`: A command for the worker to stop the currently running job
|
||||
|
||||
The command string will be parsed into a dictionary and send to a PubSub Topic.
|
||||
Workers listen to the PubSub, and `handle` the specific command.
|
||||
|
||||
Args:
|
||||
connection (Redis): A Redis Connection
|
||||
worker_name (str): The Job ID
|
||||
"""
|
||||
payload = {'command': command}
|
||||
if kwargs:
|
||||
payload.update(kwargs)
|
||||
connection.publish(PUBSUB_CHANNEL_TEMPLATE % worker_name, json.dumps(payload))
|
||||
|
||||
|
||||
def parse_payload(payload: Dict[Any, Any]) -> Dict[Any, Any]:
|
||||
"""
|
||||
Returns a dict of command data
|
||||
|
||||
Args:
|
||||
payload (dict): Parses the payload dict.
|
||||
"""
|
||||
return json.loads(payload.get('data').decode())
|
||||
|
||||
|
||||
def send_shutdown_command(connection: 'Redis', worker_name: str):
|
||||
"""
|
||||
Sends a command to shutdown a worker.
|
||||
|
||||
Args:
|
||||
connection (Redis): A Redis Connection
|
||||
worker_name (str): The Job ID
|
||||
"""
|
||||
send_command(connection, worker_name, 'shutdown')
|
||||
|
||||
|
||||
def send_kill_horse_command(connection: 'Redis', worker_name: str):
|
||||
"""
|
||||
Tell worker to kill it's horse
|
||||
|
||||
Args:
|
||||
connection (Redis): A Redis Connection
|
||||
worker_name (str): The Job ID
|
||||
"""
|
||||
send_command(connection, worker_name, 'kill-horse')
|
||||
|
||||
|
||||
def send_stop_job_command(connection: 'Redis', job_id: str, serializer=None):
|
||||
"""
|
||||
Instruct a worker to stop a job
|
||||
|
||||
Args:
|
||||
connection (Redis): A Redis Connection
|
||||
job_id (str): The Job ID
|
||||
serializer (): The serializer
|
||||
"""
|
||||
job = Job.fetch(job_id, connection=connection, serializer=serializer)
|
||||
if not job.worker_name:
|
||||
raise InvalidJobOperation('Job is not currently executing')
|
||||
send_command(connection, job.worker_name, 'stop-job', job_id=job_id)
|
||||
|
||||
|
||||
def handle_command(worker: 'Worker', payload: Dict[Any, Any]):
|
||||
"""Parses payload and routes commands to the worker.
|
||||
|
||||
Args:
|
||||
worker (Worker): The worker to use
|
||||
payload (Dict[Any, Any]): The Payload
|
||||
"""
|
||||
if payload['command'] == 'stop-job':
|
||||
handle_stop_job_command(worker, payload)
|
||||
elif payload['command'] == 'shutdown':
|
||||
handle_shutdown_command(worker)
|
||||
elif payload['command'] == 'kill-horse':
|
||||
handle_kill_worker_command(worker, payload)
|
||||
|
||||
|
||||
def handle_shutdown_command(worker: 'Worker'):
|
||||
"""Perform shutdown command.
|
||||
|
||||
Args:
|
||||
worker (Worker): The worker to use.
|
||||
"""
|
||||
worker.log.info('Received shutdown command, sending SIGINT signal.')
|
||||
pid = os.getpid()
|
||||
os.kill(pid, signal.SIGINT)
|
||||
|
||||
|
||||
def handle_kill_worker_command(worker: 'Worker', payload: Dict[Any, Any]):
|
||||
"""
|
||||
Stops work horse
|
||||
|
||||
Args:
|
||||
worker (Worker): The worker to stop
|
||||
payload (Dict[Any, Any]): The payload.
|
||||
"""
|
||||
|
||||
worker.log.info('Received kill horse command.')
|
||||
if worker.horse_pid:
|
||||
worker.log.info('Kiling horse...')
|
||||
worker.kill_horse()
|
||||
else:
|
||||
worker.log.info('Worker is not working, kill horse command ignored')
|
||||
|
||||
|
||||
def handle_stop_job_command(worker: 'Worker', payload: Dict[Any, Any]):
|
||||
"""Handles stop job command.
|
||||
|
||||
Args:
|
||||
worker (Worker): The worker to use
|
||||
payload (Dict[Any, Any]): The payload.
|
||||
"""
|
||||
job_id = payload.get('job_id')
|
||||
worker.log.debug('Received command to stop job %s', job_id)
|
||||
if job_id and worker.get_current_job_id() == job_id:
|
||||
# Sets the '_stopped_job_id' so that the job failure handler knows it
|
||||
# was intentional.
|
||||
worker._stopped_job_id = job_id
|
||||
worker.kill_horse()
|
||||
else:
|
||||
worker.log.info('Not working on job %s, command ignored.', job_id)
|
||||
@@ -0,0 +1,131 @@
|
||||
import warnings
|
||||
from contextlib import contextmanager
|
||||
from typing import Optional, Tuple, Type
|
||||
|
||||
from redis import Connection as RedisConnection
|
||||
from redis import Redis
|
||||
|
||||
from .local import LocalStack
|
||||
|
||||
|
||||
class NoRedisConnectionException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@contextmanager
|
||||
def Connection(connection: Optional['Redis'] = None): # noqa
|
||||
"""The context manager for handling connections in a clean way.
|
||||
It will push the connection to the LocalStack, and pop the connection
|
||||
when leaving the context
|
||||
|
||||
Example:
|
||||
|
||||
..codeblock:python::
|
||||
|
||||
with Connection():
|
||||
w = Worker()
|
||||
w.work()
|
||||
|
||||
This method is deprecated on version 1.12.0 and will be removed in the future.
|
||||
Pass the connection to the worker explicitly to handle Redis Connections.
|
||||
|
||||
Args:
|
||||
connection (Optional[Redis], optional): A Redis Connection instance. Defaults to None.
|
||||
"""
|
||||
warnings.warn(
|
||||
"The Connection context manager is deprecated. Use the `connection` parameter instead.",
|
||||
DeprecationWarning,
|
||||
)
|
||||
if connection is None:
|
||||
connection = Redis()
|
||||
push_connection(connection)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
popped = pop_connection()
|
||||
assert (
|
||||
popped == connection
|
||||
), 'Unexpected Redis connection was popped off the stack. Check your Redis connection setup.'
|
||||
|
||||
|
||||
def push_connection(redis: 'Redis'):
|
||||
"""
|
||||
Pushes the given connection to the stack.
|
||||
|
||||
Args:
|
||||
redis (Redis): A Redis connection
|
||||
"""
|
||||
warnings.warn(
|
||||
"The `push_connection` function is deprecated. Pass the `connection` explicitly instead.",
|
||||
DeprecationWarning,
|
||||
)
|
||||
_connection_stack.push(redis)
|
||||
|
||||
|
||||
def pop_connection() -> 'Redis':
|
||||
"""
|
||||
Pops the topmost connection from the stack.
|
||||
|
||||
Returns:
|
||||
redis (Redis): A Redis connection
|
||||
"""
|
||||
warnings.warn(
|
||||
"The `pop_connection` function is deprecated. Pass the `connection` explicitly instead.",
|
||||
DeprecationWarning,
|
||||
)
|
||||
return _connection_stack.pop()
|
||||
|
||||
|
||||
def get_current_connection() -> 'Redis':
|
||||
"""
|
||||
Returns the current Redis connection (i.e. the topmost on the
|
||||
connection stack).
|
||||
|
||||
Returns:
|
||||
Redis: A Redis Connection
|
||||
"""
|
||||
warnings.warn(
|
||||
"The `get_current_connection` function is deprecated. Pass the `connection` explicitly instead.",
|
||||
DeprecationWarning,
|
||||
)
|
||||
return _connection_stack.top
|
||||
|
||||
|
||||
def resolve_connection(connection: Optional['Redis'] = None) -> 'Redis':
|
||||
"""
|
||||
Convenience function to resolve the given or the current connection.
|
||||
Raises an exception if it cannot resolve a connection now.
|
||||
|
||||
Args:
|
||||
connection (Optional[Redis], optional): A Redis connection. Defaults to None.
|
||||
|
||||
Raises:
|
||||
NoRedisConnectionException: If connection couldn't be resolved.
|
||||
|
||||
Returns:
|
||||
Redis: A Redis Connection
|
||||
"""
|
||||
warnings.warn(
|
||||
"The `resolve_connection` function is deprecated. Pass the `connection` explicitly instead.",
|
||||
DeprecationWarning,
|
||||
)
|
||||
if connection is not None:
|
||||
return connection
|
||||
|
||||
connection = get_current_connection()
|
||||
if connection is None:
|
||||
raise NoRedisConnectionException('Could not resolve a Redis connection')
|
||||
return connection
|
||||
|
||||
|
||||
def parse_connection(connection: Redis) -> Tuple[Type[Redis], Type[RedisConnection], dict]:
|
||||
connection_pool_kwargs = connection.connection_pool.connection_kwargs.copy()
|
||||
connection_pool_class = connection.connection_pool.connection_class
|
||||
|
||||
return connection.__class__, connection_pool_class, connection_pool_kwargs
|
||||
|
||||
|
||||
_connection_stack = LocalStack()
|
||||
|
||||
|
||||
__all__ = ['Connection', 'get_current_connection', 'push_connection', 'pop_connection']
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,23 @@
|
||||
import logging
|
||||
|
||||
from rq import Worker, get_current_connection
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def cleanup_ghosts(conn=None, worker_class=Worker):
|
||||
"""
|
||||
RQ versions < 0.3.6 suffered from a race condition where workers, when
|
||||
abruptly terminated, did not have a chance to clean up their worker
|
||||
registration, leading to reports of ghosted workers in `rqinfo`. Since
|
||||
0.3.6, new worker registrations automatically expire, and the worker will
|
||||
make sure to refresh the registrations as long as it's alive.
|
||||
|
||||
This function will clean up any of such legacy ghosted workers.
|
||||
"""
|
||||
conn = conn if conn else get_current_connection()
|
||||
for worker in worker_class.all(connection=conn):
|
||||
if conn.ttl(worker.key) == -1:
|
||||
ttl = worker.worker_ttl
|
||||
conn.expire(worker.key, ttl)
|
||||
logger.info('Marked ghosted worker {0} to expire in {1} seconds.'.format(worker.name, ttl))
|
||||
@@ -0,0 +1,8 @@
|
||||
def register_sentry(sentry_dsn, **opts):
|
||||
"""Given a Raven client and an RQ worker, registers exception handlers
|
||||
with the worker so exceptions are logged to Sentry.
|
||||
"""
|
||||
import sentry_sdk
|
||||
from sentry_sdk.integrations.rq import RqIntegration
|
||||
|
||||
sentry_sdk.init(sentry_dsn, integrations=[RqIntegration()], **opts)
|
||||
@@ -0,0 +1,124 @@
|
||||
from functools import wraps
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from redis import Redis
|
||||
|
||||
from .job import Retry
|
||||
|
||||
from .defaults import DEFAULT_RESULT_TTL
|
||||
from .job import Callback
|
||||
from .queue import Queue
|
||||
from .utils import backend_class
|
||||
|
||||
|
||||
class job: # noqa
|
||||
queue_class = Queue
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
queue: Union['Queue', str],
|
||||
connection: Optional['Redis'] = None,
|
||||
timeout: Optional[int] = None,
|
||||
result_ttl: int = DEFAULT_RESULT_TTL,
|
||||
ttl: Optional[int] = None,
|
||||
queue_class: Optional[Type['Queue']] = None,
|
||||
depends_on: Optional[List[Any]] = None,
|
||||
at_front: bool = False,
|
||||
meta: Optional[Dict[Any, Any]] = None,
|
||||
description: Optional[str] = None,
|
||||
failure_ttl: Optional[int] = None,
|
||||
retry: Optional['Retry'] = None,
|
||||
on_failure: Optional[Union[Callback, Callable[..., Any]]] = None,
|
||||
on_success: Optional[Union[Callback, Callable[..., Any]]] = None,
|
||||
on_stopped: Optional[Union[Callback, Callable[..., Any]]] = None,
|
||||
):
|
||||
"""A decorator that adds a ``delay`` method to the decorated function,
|
||||
which in turn creates a RQ job when called. Accepts a required
|
||||
``queue`` argument that can be either a ``Queue`` instance or a string
|
||||
denoting the queue name. For example::
|
||||
|
||||
..codeblock:python::
|
||||
|
||||
>>> @job(queue='default')
|
||||
>>> def simple_add(x, y):
|
||||
>>> return x + y
|
||||
>>> ...
|
||||
>>> # Puts `simple_add` function into queue
|
||||
>>> simple_add.delay(1, 2)
|
||||
|
||||
Args:
|
||||
queue (Union['Queue', str]): The queue to use, can be the Queue class itself, or the queue name (str)
|
||||
connection (Optional[Redis], optional): Redis Connection. Defaults to None.
|
||||
timeout (Optional[int], optional): Job timeout. Defaults to None.
|
||||
result_ttl (int, optional): Result time to live. Defaults to DEFAULT_RESULT_TTL.
|
||||
ttl (Optional[int], optional): Time to live. Defaults to None.
|
||||
queue_class (Optional[Queue], optional): A custom class that inherits from `Queue`. Defaults to None.
|
||||
depends_on (Optional[List[Any]], optional): A list of dependents jobs. Defaults to None.
|
||||
at_front (Optional[bool], optional): Whether to enqueue the job at front of the queue. Defaults to None.
|
||||
meta (Optional[Dict[Any, Any]], optional): Arbitraty metadata about the job. Defaults to None.
|
||||
description (Optional[str], optional): Job description. Defaults to None.
|
||||
failure_ttl (Optional[int], optional): Failture time to live. Defaults to None.
|
||||
retry (Optional[Retry], optional): A Retry object. Defaults to None.
|
||||
on_failure (Optional[Union[Callback, Callable[..., Any]]], optional): Callable to run on failure. Defaults
|
||||
to None.
|
||||
on_success (Optional[Union[Callback, Callable[..., Any]]], optional): Callable to run on success. Defaults
|
||||
to None.
|
||||
on_stopped (Optional[Union[Callback, Callable[..., Any]]], optional): Callable to run when stopped. Defaults
|
||||
to None.
|
||||
"""
|
||||
self.queue = queue
|
||||
self.queue_class = backend_class(self, 'queue_class', override=queue_class)
|
||||
self.connection = connection
|
||||
self.timeout = timeout
|
||||
self.result_ttl = result_ttl
|
||||
self.ttl = ttl
|
||||
self.meta = meta
|
||||
self.depends_on = depends_on
|
||||
self.at_front = at_front
|
||||
self.description = description
|
||||
self.failure_ttl = failure_ttl
|
||||
self.retry = retry
|
||||
self.on_success = on_success
|
||||
self.on_failure = on_failure
|
||||
self.on_stopped = on_stopped
|
||||
|
||||
def __call__(self, f):
|
||||
@wraps(f)
|
||||
def delay(*args, **kwargs):
|
||||
if isinstance(self.queue, str):
|
||||
queue = self.queue_class(name=self.queue, connection=self.connection)
|
||||
else:
|
||||
queue = self.queue
|
||||
|
||||
depends_on = kwargs.pop('depends_on', None)
|
||||
job_id = kwargs.pop('job_id', None)
|
||||
at_front = kwargs.pop('at_front', False)
|
||||
|
||||
if not depends_on:
|
||||
depends_on = self.depends_on
|
||||
|
||||
if not at_front:
|
||||
at_front = self.at_front
|
||||
|
||||
return queue.enqueue_call(
|
||||
f,
|
||||
args=args,
|
||||
kwargs=kwargs,
|
||||
timeout=self.timeout,
|
||||
result_ttl=self.result_ttl,
|
||||
ttl=self.ttl,
|
||||
depends_on=depends_on,
|
||||
job_id=job_id,
|
||||
at_front=at_front,
|
||||
meta=self.meta,
|
||||
description=self.description,
|
||||
failure_ttl=self.failure_ttl,
|
||||
retry=self.retry,
|
||||
on_failure=self.on_failure,
|
||||
on_success=self.on_success,
|
||||
on_stopped=self.on_stopped,
|
||||
)
|
||||
|
||||
f.delay = delay
|
||||
return f
|
||||
@@ -0,0 +1,102 @@
|
||||
DEFAULT_JOB_CLASS = 'rq.job.Job'
|
||||
""" The path for the default Job class to use.
|
||||
Defaults to the main `Job` class within the `rq.job` module
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_QUEUE_CLASS = 'rq.Queue'
|
||||
""" The path for the default Queue class to use.
|
||||
Defaults to the main `Queue` class within the `rq.queue` module
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_WORKER_CLASS = 'rq.Worker'
|
||||
""" The path for the default Worker class to use.
|
||||
Defaults to the main `Worker` class within the `rq.worker` module
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_SERIALIZER_CLASS = 'rq.serializers.DefaultSerializer'
|
||||
""" The path for the default Serializer class to use.
|
||||
Defaults to the main `DefaultSerializer` class within the `rq.serializers` module
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_CONNECTION_CLASS = 'redis.Redis'
|
||||
""" The path for the default Redis client class to use.
|
||||
Defaults to the main `Redis` class within the `redis` module
|
||||
As imported like `from redis import Redis`
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_WORKER_TTL = 420
|
||||
""" The default Time To Live (TTL) for the Worker in seconds
|
||||
Defines the effective timeout period for a worker
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_JOB_MONITORING_INTERVAL = 30
|
||||
""" The interval in seconds for Job monitoring
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_RESULT_TTL = 500
|
||||
""" The Time To Live (TTL) in seconds to keep job results
|
||||
Means that the results will be expired from Redis
|
||||
after `DEFAULT_RESULT_TTL` seconds
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_FAILURE_TTL = 31536000
|
||||
""" The Time To Live (TTL) in seconds to keep job failure information
|
||||
Means that the failure information will be expired from Redis
|
||||
after `DEFAULT_FAILURE_TTL` seconds.
|
||||
Defaults to 1 YEAR in seconds
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_SCHEDULER_FALLBACK_PERIOD = 120
|
||||
""" The amount in seconds it will take for a new scheduler
|
||||
to pickup tasks after a scheduler has died.
|
||||
This is used as a safety net to avoid race conditions and duplicates
|
||||
when using multiple schedulers
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_MAINTENANCE_TASK_INTERVAL = 10 * 60
|
||||
""" The interval to run maintenance tasks
|
||||
in seconds. Defaults to 10 minutes.
|
||||
"""
|
||||
|
||||
|
||||
CALLBACK_TIMEOUT = 60
|
||||
""" The timeout period in seconds for Callback functions
|
||||
Means that Functions used in `success_callback`, `stopped_callback`,
|
||||
and `failure_callback` will timeout after N seconds
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_LOGGING_DATE_FORMAT = '%H:%M:%S'
|
||||
""" The Date Format to use for RQ logging.
|
||||
Defaults to Hour:Minute:Seconds on 24hour format
|
||||
eg.: `15:45:23`
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_LOGGING_FORMAT = '%(asctime)s %(message)s'
|
||||
""" The default Logging Format to use
|
||||
Uses Python's default attributes as defined
|
||||
https://docs.python.org/3/library/logging.html#logrecord-attributes
|
||||
"""
|
||||
|
||||
|
||||
DEFAULT_DEATH_PENALTY_CLASS = 'rq.timeouts.UnixSignalDeathPenalty'
|
||||
""" The path for the default Death Penalty class to use.
|
||||
Defaults to the `UnixSignalDeathPenalty` class within the `rq.timeouts` module
|
||||
"""
|
||||
|
||||
|
||||
UNSERIALIZABLE_RETURN_VALUE_PAYLOAD = 'Unserializable return value'
|
||||
""" The value that we store in the job's _result property or in the Result's return_value
|
||||
in case the return value of the actual job is not serializable
|
||||
"""
|
||||
@@ -0,0 +1,27 @@
|
||||
from typing import List
|
||||
|
||||
from redis.client import Pipeline
|
||||
from redis.exceptions import WatchError
|
||||
|
||||
from .job import Job
|
||||
|
||||
|
||||
class Dependency:
|
||||
@classmethod
|
||||
def get_jobs_with_met_dependencies(cls, jobs: List['Job'], pipeline: Pipeline):
|
||||
jobs_with_met_dependencies = []
|
||||
jobs_with_unmet_dependencies = []
|
||||
for job in jobs:
|
||||
while True:
|
||||
try:
|
||||
pipeline.watch(*[Job.key_for(dependency_id) for dependency_id in job._dependency_ids])
|
||||
job.register_dependency(pipeline=pipeline)
|
||||
if job.dependencies_are_met(pipeline=pipeline):
|
||||
jobs_with_met_dependencies.append(job)
|
||||
else:
|
||||
jobs_with_unmet_dependencies.append(job)
|
||||
pipeline.execute()
|
||||
except WatchError:
|
||||
continue
|
||||
break
|
||||
return jobs_with_met_dependencies, jobs_with_unmet_dependencies
|
||||
@@ -0,0 +1,36 @@
|
||||
class NoSuchJobError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class DeserializationError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class InvalidJobDependency(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class InvalidJobOperationError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class InvalidJobOperation(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class DequeueTimeout(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class ShutDownImminentException(Exception):
|
||||
def __init__(self, msg, extra_info):
|
||||
self.extra_info = extra_info
|
||||
super().__init__(msg)
|
||||
|
||||
|
||||
class TimeoutFormatError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class AbandonedJobError(Exception):
|
||||
pass
|
||||
@@ -0,0 +1,119 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
from redis import Redis
|
||||
|
||||
from rq.utils import now
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .queue import Queue
|
||||
from .worker import BaseWorker
|
||||
|
||||
|
||||
class IntermediateQueue(object):
|
||||
|
||||
def __init__(self, queue_key: str, connection: Redis):
|
||||
self.queue_key = queue_key
|
||||
self.key = self.get_intermediate_queue_key(queue_key)
|
||||
self.connection = connection
|
||||
|
||||
@classmethod
|
||||
def get_intermediate_queue_key(cls, queue_key: str) -> str:
|
||||
"""Returns the intermediate queue key for a given queue key.
|
||||
|
||||
Args:
|
||||
key (str): The queue key
|
||||
|
||||
Returns:
|
||||
str: The intermediate queue key
|
||||
"""
|
||||
return f'{queue_key}:intermediate'
|
||||
|
||||
def get_first_seen_key(self, job_id: str) -> str:
|
||||
"""Returns the first seen key for a given job ID.
|
||||
|
||||
Args:
|
||||
job_id (str): The job ID
|
||||
|
||||
Returns:
|
||||
str: The first seen key
|
||||
"""
|
||||
return f'{self.key}:first_seen:{job_id}'
|
||||
|
||||
def set_first_seen(self, job_id: str) -> bool:
|
||||
"""Sets the first seen timestamp for a job.
|
||||
|
||||
Args:
|
||||
job_id (str): The job ID
|
||||
timestamp (float): The timestamp
|
||||
"""
|
||||
# TODO: job_id should be changed to execution ID in 2.0
|
||||
return bool(self.connection.set(self.get_first_seen_key(job_id), now().timestamp(), nx=True, ex=3600 * 24))
|
||||
|
||||
def get_first_seen(self, job_id: str) -> Optional[datetime]:
|
||||
"""Returns the first seen timestamp for a job.
|
||||
|
||||
Args:
|
||||
job_id (str): The job ID
|
||||
|
||||
Returns:
|
||||
Optional[datetime]: The timestamp
|
||||
"""
|
||||
timestamp = self.connection.get(self.get_first_seen_key(job_id))
|
||||
if timestamp:
|
||||
return datetime.fromtimestamp(float(timestamp), tz=timezone.utc)
|
||||
return None
|
||||
|
||||
def should_be_cleaned_up(self, job_id: str) -> bool:
|
||||
"""Returns whether a job should be cleaned up.
|
||||
A job in intermediate queue should be cleaned up if it has been there for more than 1 minute.
|
||||
|
||||
Args:
|
||||
job_id (str): The job ID
|
||||
|
||||
Returns:
|
||||
bool: Whether the job should be cleaned up
|
||||
"""
|
||||
# TODO: should be changed to execution ID in 2.0
|
||||
first_seen = self.get_first_seen(job_id)
|
||||
if not first_seen:
|
||||
return False
|
||||
return now() - first_seen > timedelta(minutes=1)
|
||||
|
||||
def get_job_ids(self) -> List[str]:
|
||||
"""Returns the job IDs in the intermediate queue.
|
||||
|
||||
Returns:
|
||||
List[str]: The job IDs
|
||||
"""
|
||||
return [job_id.decode() for job_id in self.connection.lrange(self.key, 0, -1)]
|
||||
|
||||
def remove(self, job_id: str) -> None:
|
||||
"""Removes a job from the intermediate queue.
|
||||
|
||||
Args:
|
||||
job_id (str): The job ID
|
||||
"""
|
||||
self.connection.lrem(self.key, 1, job_id)
|
||||
|
||||
def cleanup(self, worker: 'BaseWorker', queue: 'Queue') -> None:
|
||||
job_ids = self.get_job_ids()
|
||||
|
||||
for job_id in job_ids:
|
||||
job = queue.fetch_job(job_id)
|
||||
|
||||
if job_id not in queue.started_job_registry:
|
||||
|
||||
if not job:
|
||||
# If the job doesn't exist in the queue, we can safely remove it from the intermediate queue.
|
||||
self.remove(job_id)
|
||||
continue
|
||||
|
||||
# If this is the first time we've seen this job, do nothing.
|
||||
# `set_first_seen` will return `True` if the key was set, `False` if it already existed.
|
||||
if self.set_first_seen(job_id):
|
||||
continue
|
||||
|
||||
if self.should_be_cleaned_up(job_id):
|
||||
worker.handle_job_failure(job, queue, exc_string='Job was stuck in intermediate queue.')
|
||||
self.remove(job_id)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,356 @@
|
||||
# ruff: noqa: E731
|
||||
"""
|
||||
werkzeug.local
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
This module implements context-local objects.
|
||||
|
||||
:copyright: (c) 2011 by the Werkzeug Team, see AUTHORS for more details.
|
||||
:license: BSD, see LICENSE for more details.
|
||||
"""
|
||||
# Since each thread has its own greenlet we can just use those as identifiers
|
||||
# for the context. If greenlets are not available we fall back to the
|
||||
# current thread ident.
|
||||
try:
|
||||
from greenlet import getcurrent as get_ident
|
||||
except ImportError:
|
||||
try:
|
||||
from threading import get_ident
|
||||
except ImportError:
|
||||
try:
|
||||
from _thread import get_ident
|
||||
except ImportError:
|
||||
from dummy_thread import get_ident
|
||||
|
||||
|
||||
def release_local(local):
|
||||
"""Releases the contents of the local for the current context.
|
||||
This makes it possible to use locals without a manager.
|
||||
|
||||
Example::
|
||||
|
||||
>>> loc = Local()
|
||||
>>> loc.foo = 42
|
||||
>>> release_local(loc)
|
||||
>>> hasattr(loc, 'foo')
|
||||
False
|
||||
|
||||
With this function one can release :class:`Local` objects as well
|
||||
as :class:`StackLocal` objects. However it is not possible to
|
||||
release data held by proxies that way, one always has to retain
|
||||
a reference to the underlying local object in order to be able
|
||||
to release it.
|
||||
|
||||
.. versionadded:: 0.6.1
|
||||
"""
|
||||
local.__release_local__()
|
||||
|
||||
|
||||
class Local:
|
||||
__slots__ = ('__storage__', '__ident_func__')
|
||||
|
||||
def __init__(self):
|
||||
object.__setattr__(self, '__storage__', {})
|
||||
object.__setattr__(self, '__ident_func__', get_ident)
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self.__storage__.items())
|
||||
|
||||
def __call__(self, proxy):
|
||||
"""Create a proxy for a name."""
|
||||
return LocalProxy(self, proxy)
|
||||
|
||||
def __release_local__(self):
|
||||
self.__storage__.pop(self.__ident_func__(), None)
|
||||
|
||||
def __getattr__(self, name):
|
||||
try:
|
||||
return self.__storage__[self.__ident_func__()][name]
|
||||
except KeyError:
|
||||
raise AttributeError(name)
|
||||
|
||||
def __setattr__(self, name, value):
|
||||
ident = self.__ident_func__()
|
||||
storage = self.__storage__
|
||||
try:
|
||||
storage[ident][name] = value
|
||||
except KeyError:
|
||||
storage[ident] = {name: value}
|
||||
|
||||
def __delattr__(self, name):
|
||||
try:
|
||||
del self.__storage__[self.__ident_func__()][name]
|
||||
except KeyError:
|
||||
raise AttributeError(name)
|
||||
|
||||
|
||||
class LocalStack:
|
||||
"""This class works similar to a :class:`Local` but keeps a stack
|
||||
of objects instead. This is best explained with an example::
|
||||
|
||||
>>> ls = LocalStack()
|
||||
>>> ls.push(42)
|
||||
>>> ls.top
|
||||
42
|
||||
>>> ls.push(23)
|
||||
>>> ls.top
|
||||
23
|
||||
>>> ls.pop()
|
||||
23
|
||||
>>> ls.top
|
||||
42
|
||||
|
||||
They can be force released by using a :class:`LocalManager` or with
|
||||
the :func:`release_local` function but the correct way is to pop the
|
||||
item from the stack after using. When the stack is empty it will
|
||||
no longer be bound to the current context (and as such released).
|
||||
|
||||
By calling the stack without arguments it returns a proxy that resolves to
|
||||
the topmost item on the stack.
|
||||
|
||||
.. versionadded:: 0.6.1
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._local = Local()
|
||||
|
||||
def __release_local__(self):
|
||||
self._local.__release_local__()
|
||||
|
||||
def _get__ident_func__(self):
|
||||
return self._local.__ident_func__
|
||||
|
||||
def _set__ident_func__(self, value):
|
||||
object.__setattr__(self._local, '__ident_func__', value)
|
||||
|
||||
__ident_func__ = property(_get__ident_func__, _set__ident_func__)
|
||||
del _get__ident_func__, _set__ident_func__
|
||||
|
||||
def __call__(self):
|
||||
def _lookup():
|
||||
rv = self.top
|
||||
if rv is None:
|
||||
raise RuntimeError('object unbound')
|
||||
return rv
|
||||
|
||||
return LocalProxy(_lookup)
|
||||
|
||||
def push(self, obj):
|
||||
"""Pushes a new item to the stack"""
|
||||
rv = getattr(self._local, 'stack', None)
|
||||
if rv is None:
|
||||
self._local.stack = rv = []
|
||||
rv.append(obj)
|
||||
return rv
|
||||
|
||||
def pop(self):
|
||||
"""Removes the topmost item from the stack, will return the
|
||||
old value or `None` if the stack was already empty.
|
||||
"""
|
||||
stack = getattr(self._local, 'stack', None)
|
||||
if stack is None:
|
||||
return None
|
||||
elif len(stack) == 1:
|
||||
release_local(self._local)
|
||||
return stack[-1]
|
||||
else:
|
||||
return stack.pop()
|
||||
|
||||
@property
|
||||
def top(self):
|
||||
"""The topmost item on the stack. If the stack is empty,
|
||||
`None` is returned.
|
||||
"""
|
||||
try:
|
||||
return self._local.stack[-1]
|
||||
except (AttributeError, IndexError):
|
||||
return None
|
||||
|
||||
def __len__(self):
|
||||
stack = getattr(self._local, 'stack', None)
|
||||
if stack is None:
|
||||
return 0
|
||||
return len(stack)
|
||||
|
||||
|
||||
class LocalManager:
|
||||
"""Local objects cannot manage themselves. For that you need a local
|
||||
manager. You can pass a local manager multiple locals or add them later
|
||||
by appending them to `manager.locals`. Everytime the manager cleans up
|
||||
it, will clean up all the data left in the locals for this context.
|
||||
|
||||
The `ident_func` parameter can be added to override the default ident
|
||||
function for the wrapped locals.
|
||||
|
||||
.. versionchanged:: 0.6.1
|
||||
Instead of a manager the :func:`release_local` function can be used
|
||||
as well.
|
||||
|
||||
.. versionchanged:: 0.7
|
||||
`ident_func` was added.
|
||||
"""
|
||||
|
||||
def __init__(self, locals=None, ident_func=None):
|
||||
if locals is None:
|
||||
self.locals = []
|
||||
elif isinstance(locals, Local):
|
||||
self.locals = [locals]
|
||||
else:
|
||||
self.locals = list(locals)
|
||||
if ident_func is not None:
|
||||
self.ident_func = ident_func
|
||||
for local in self.locals:
|
||||
object.__setattr__(local, '__ident_func__', ident_func)
|
||||
else:
|
||||
self.ident_func = get_ident
|
||||
|
||||
def get_ident(self):
|
||||
"""Return the context identifier the local objects use internally for
|
||||
this context. You cannot override this method to change the behavior
|
||||
but use it to link other context local objects (such as SQLAlchemy's
|
||||
scoped sessions) to the Werkzeug locals.
|
||||
|
||||
.. versionchanged:: 0.7
|
||||
You can pass a different ident function to the local manager that
|
||||
will then be propagated to all the locals passed to the
|
||||
constructor.
|
||||
"""
|
||||
return self.ident_func()
|
||||
|
||||
def cleanup(self):
|
||||
"""Manually clean up the data in the locals for this context. Call
|
||||
this at the end of the request or use `make_middleware()`.
|
||||
"""
|
||||
for local in self.locals:
|
||||
release_local(local)
|
||||
|
||||
def __repr__(self):
|
||||
return '<%s storages: %d>' % (self.__class__.__name__, len(self.locals))
|
||||
|
||||
|
||||
class LocalProxy:
|
||||
"""Acts as a proxy for a werkzeug local. Forwards all operations to
|
||||
a proxied object. The only operations not supported for forwarding
|
||||
are right handed operands and any kind of assignment.
|
||||
|
||||
Example usage::
|
||||
|
||||
from werkzeug.local import Local
|
||||
l = Local()
|
||||
|
||||
# these are proxies
|
||||
request = l('request')
|
||||
user = l('user')
|
||||
|
||||
|
||||
from werkzeug.local import LocalStack
|
||||
_response_local = LocalStack()
|
||||
|
||||
# this is a proxy
|
||||
response = _response_local()
|
||||
|
||||
Whenever something is bound to l.user / l.request the proxy objects
|
||||
will forward all operations. If no object is bound a :exc:`RuntimeError`
|
||||
will be raised.
|
||||
|
||||
To create proxies to :class:`Local` or :class:`LocalStack` objects,
|
||||
call the object as shown above. If you want to have a proxy to an
|
||||
object looked up by a function, you can (as of Werkzeug 0.6.1) pass
|
||||
a function to the :class:`LocalProxy` constructor::
|
||||
|
||||
session = LocalProxy(lambda: get_current_request().session)
|
||||
|
||||
.. versionchanged:: 0.6.1
|
||||
The class can be instanciated with a callable as well now.
|
||||
"""
|
||||
|
||||
__slots__ = ('__local', '__dict__', '__name__')
|
||||
|
||||
def __init__(self, local, name=None):
|
||||
object.__setattr__(self, '_LocalProxy__local', local)
|
||||
object.__setattr__(self, '__name__', name)
|
||||
|
||||
def _get_current_object(self):
|
||||
"""Return the current object. This is useful if you want the real
|
||||
object behind the proxy at a time for performance reasons or because
|
||||
you want to pass the object into a different context.
|
||||
"""
|
||||
if not hasattr(self.__local, '__release_local__'):
|
||||
return self.__local()
|
||||
try:
|
||||
return getattr(self.__local, self.__name__)
|
||||
except AttributeError:
|
||||
raise RuntimeError('no object bound to %s' % self.__name__)
|
||||
|
||||
@property
|
||||
def __dict__(self):
|
||||
try:
|
||||
return self._get_current_object().__dict__
|
||||
except RuntimeError:
|
||||
raise AttributeError('__dict__')
|
||||
|
||||
def __repr__(self):
|
||||
try:
|
||||
obj = self._get_current_object()
|
||||
except RuntimeError:
|
||||
return '<%s unbound>' % self.__class__.__name__
|
||||
return repr(obj)
|
||||
|
||||
def __dir__(self):
|
||||
try:
|
||||
return dir(self._get_current_object())
|
||||
except RuntimeError:
|
||||
return []
|
||||
|
||||
def __getattr__(self, name):
|
||||
if name == '__members__':
|
||||
return dir(self._get_current_object())
|
||||
return getattr(self._get_current_object(), name)
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
self._get_current_object()[key] = value
|
||||
|
||||
def __delitem__(self, key):
|
||||
del self._get_current_object()[key]
|
||||
|
||||
__setattr__ = lambda x, n, v: setattr(x._get_current_object(), n, v)
|
||||
__delattr__ = lambda x, n: delattr(x._get_current_object(), n)
|
||||
__str__ = lambda x: str(x._get_current_object())
|
||||
__lt__ = lambda x, o: x._get_current_object() < o
|
||||
__le__ = lambda x, o: x._get_current_object() <= o
|
||||
__eq__ = lambda x, o: x._get_current_object() == o
|
||||
__ne__ = lambda x, o: x._get_current_object() != o
|
||||
__gt__ = lambda x, o: x._get_current_object() > o
|
||||
__ge__ = lambda x, o: x._get_current_object() >= o
|
||||
__hash__ = lambda x: hash(x._get_current_object())
|
||||
__call__ = lambda x, *a, **kw: x._get_current_object()(*a, **kw)
|
||||
__len__ = lambda x: len(x._get_current_object())
|
||||
__getitem__ = lambda x, i: x._get_current_object()[i]
|
||||
__iter__ = lambda x: iter(x._get_current_object())
|
||||
__contains__ = lambda x, i: i in x._get_current_object()
|
||||
__add__ = lambda x, o: x._get_current_object() + o
|
||||
__sub__ = lambda x, o: x._get_current_object() - o
|
||||
__mul__ = lambda x, o: x._get_current_object() * o
|
||||
__floordiv__ = lambda x, o: x._get_current_object() // o
|
||||
__mod__ = lambda x, o: x._get_current_object() % o
|
||||
__divmod__ = lambda x, o: x._get_current_object().__divmod__(o)
|
||||
__pow__ = lambda x, o: x._get_current_object() ** o
|
||||
__lshift__ = lambda x, o: x._get_current_object() << o
|
||||
__rshift__ = lambda x, o: x._get_current_object() >> o
|
||||
__and__ = lambda x, o: x._get_current_object() & o
|
||||
__xor__ = lambda x, o: x._get_current_object() ^ o
|
||||
__or__ = lambda x, o: x._get_current_object() | o
|
||||
__div__ = lambda x, o: x._get_current_object().__div__(o)
|
||||
__truediv__ = lambda x, o: x._get_current_object().__truediv__(o)
|
||||
__neg__ = lambda x: -(x._get_current_object())
|
||||
__pos__ = lambda x: +(x._get_current_object())
|
||||
__abs__ = lambda x: abs(x._get_current_object())
|
||||
__invert__ = lambda x: ~(x._get_current_object())
|
||||
__complex__ = lambda x: complex(x._get_current_object())
|
||||
__int__ = lambda x: int(x._get_current_object())
|
||||
__float__ = lambda x: float(x._get_current_object())
|
||||
__oct__ = lambda x: oct(x._get_current_object())
|
||||
__hex__ = lambda x: hex(x._get_current_object())
|
||||
__index__ = lambda x: x._get_current_object().__index__()
|
||||
__enter__ = lambda x: x._get_current_object().__enter__()
|
||||
__exit__ = lambda x, *a, **kw: x._get_current_object().__exit__(*a, **kw)
|
||||
@@ -0,0 +1,158 @@
|
||||
import logging
|
||||
import sys
|
||||
from typing import Union
|
||||
|
||||
from rq.defaults import DEFAULT_LOGGING_DATE_FORMAT, DEFAULT_LOGGING_FORMAT
|
||||
|
||||
|
||||
class _Colorizer:
|
||||
def __init__(self):
|
||||
esc = "\x1b["
|
||||
|
||||
self.codes = {}
|
||||
self.codes[""] = ""
|
||||
self.codes["reset"] = esc + "39;49;00m"
|
||||
|
||||
self.codes["bold"] = esc + "01m"
|
||||
self.codes["faint"] = esc + "02m"
|
||||
self.codes["standout"] = esc + "03m"
|
||||
self.codes["underline"] = esc + "04m"
|
||||
self.codes["blink"] = esc + "05m"
|
||||
self.codes["overline"] = esc + "06m"
|
||||
|
||||
dark_colors = ["black", "darkred", "darkgreen", "brown", "darkblue", "purple", "teal", "lightgray"]
|
||||
light_colors = ["darkgray", "red", "green", "yellow", "blue", "fuchsia", "turquoise", "white"]
|
||||
|
||||
x = 30
|
||||
for dark, light in zip(dark_colors, light_colors):
|
||||
self.codes[dark] = esc + "%im" % x
|
||||
self.codes[light] = esc + "%i;01m" % x
|
||||
x += 1
|
||||
|
||||
del dark, light, x
|
||||
|
||||
self.codes["darkteal"] = self.codes["turquoise"]
|
||||
self.codes["darkyellow"] = self.codes["brown"]
|
||||
self.codes["fuscia"] = self.codes["fuchsia"]
|
||||
self.codes["white"] = self.codes["bold"]
|
||||
|
||||
if hasattr(sys.stdout, "isatty"):
|
||||
self.notty = not sys.stdout.isatty()
|
||||
else:
|
||||
self.notty = True
|
||||
|
||||
def colorize(self, color_key, text):
|
||||
if self.notty:
|
||||
return text
|
||||
else:
|
||||
return self.codes[color_key] + text + self.codes["reset"]
|
||||
|
||||
|
||||
colorizer = _Colorizer()
|
||||
|
||||
|
||||
def make_colorizer(color: str):
|
||||
"""Creates a function that colorizes text with the given color.
|
||||
|
||||
For example::
|
||||
|
||||
..codeblock::python
|
||||
|
||||
>>> green = make_colorizer('darkgreen')
|
||||
>>> red = make_colorizer('red')
|
||||
>>>
|
||||
>>> # You can then use:
|
||||
>>> print("It's either " + green('OK') + ' or ' + red('Oops'))
|
||||
"""
|
||||
|
||||
def inner(text):
|
||||
return colorizer.colorize(color, text)
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
green = make_colorizer('darkgreen')
|
||||
yellow = make_colorizer('darkyellow')
|
||||
blue = make_colorizer('darkblue')
|
||||
red = make_colorizer('darkred')
|
||||
|
||||
|
||||
class ColorizingStreamHandler(logging.StreamHandler):
|
||||
levels = {
|
||||
logging.WARNING: yellow,
|
||||
logging.ERROR: red,
|
||||
logging.CRITICAL: red,
|
||||
}
|
||||
|
||||
def __init__(self, exclude=None, *args, **kwargs):
|
||||
self.exclude = exclude
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
@property
|
||||
def is_tty(self):
|
||||
isatty = getattr(self.stream, 'isatty', None)
|
||||
return isatty and isatty()
|
||||
|
||||
def format(self, record):
|
||||
message = logging.StreamHandler.format(self, record)
|
||||
if self.is_tty:
|
||||
# Don't colorize any traceback
|
||||
parts = message.split('\n', 1)
|
||||
parts[0] = " ".join([parts[0].split(" ", 1)[0], parts[0].split(" ", 1)[1]])
|
||||
|
||||
message = '\n'.join(parts)
|
||||
|
||||
return message
|
||||
|
||||
|
||||
def setup_loghandlers(
|
||||
level: Union[int, str, None] = None,
|
||||
date_format: str = DEFAULT_LOGGING_DATE_FORMAT,
|
||||
log_format: str = DEFAULT_LOGGING_FORMAT,
|
||||
name: str = 'rq.worker',
|
||||
):
|
||||
"""Sets up a log handler.
|
||||
|
||||
Args:
|
||||
level (Union[int, str, None], optional): The log level.
|
||||
Access an integer level (10-50) or a string level ("info", "debug" etc). Defaults to None.
|
||||
date_format (str, optional): The date format to use. Defaults to DEFAULT_LOGGING_DATE_FORMAT ('%H:%M:%S').
|
||||
log_format (str, optional): The log format to use.
|
||||
Defaults to DEFAULT_LOGGING_FORMAT ('%(asctime)s %(message)s').
|
||||
name (str, optional): The looger name. Defaults to 'rq.worker'.
|
||||
"""
|
||||
logger = logging.getLogger(name)
|
||||
|
||||
if not _has_effective_handler(logger):
|
||||
formatter = logging.Formatter(fmt=log_format, datefmt=date_format)
|
||||
handler = ColorizingStreamHandler(stream=sys.stdout)
|
||||
handler.setFormatter(formatter)
|
||||
handler.addFilter(lambda record: record.levelno < logging.ERROR)
|
||||
error_handler = ColorizingStreamHandler(stream=sys.stderr)
|
||||
error_handler.setFormatter(formatter)
|
||||
error_handler.addFilter(lambda record: record.levelno >= logging.ERROR)
|
||||
logger.addHandler(handler)
|
||||
logger.addHandler(error_handler)
|
||||
|
||||
if level is not None:
|
||||
# The level may be a numeric value (e.g. when using the logging module constants)
|
||||
# Or a string representation of the logging level
|
||||
logger.setLevel(level if isinstance(level, int) else level.upper())
|
||||
|
||||
|
||||
def _has_effective_handler(logger) -> bool:
|
||||
"""
|
||||
Checks if a logger has a handler that will catch its messages in its logger hierarchy.
|
||||
|
||||
Args:
|
||||
logger (logging.Logger): The logger to be checked.
|
||||
|
||||
Returns:
|
||||
is_configured (bool): True if a handler is found for the logger, False otherwise.
|
||||
"""
|
||||
while True:
|
||||
if logger.handlers:
|
||||
return True
|
||||
if not logger.parent:
|
||||
return False
|
||||
logger = logger.parent
|
||||
@@ -0,0 +1,27 @@
|
||||
import warnings
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .intermediate_queue import IntermediateQueue
|
||||
from .queue import Queue
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .worker import BaseWorker
|
||||
|
||||
|
||||
def clean_intermediate_queue(worker: 'BaseWorker', queue: Queue) -> None:
|
||||
"""
|
||||
Check whether there are any jobs stuck in the intermediate queue.
|
||||
|
||||
A job may be stuck in the intermediate queue if a worker has successfully dequeued a job
|
||||
but was not able to push it to the StartedJobRegistry. This may happen in rare cases
|
||||
of hardware or network failure.
|
||||
|
||||
We consider a job to be stuck in the intermediate queue if it doesn't exist in the StartedJobRegistry.
|
||||
"""
|
||||
warnings.warn(
|
||||
"clean_intermediate_queue is deprecated. Use IntermediateQueue.cleanup instead.",
|
||||
DeprecationWarning,
|
||||
)
|
||||
intermediate_queue = IntermediateQueue(queue.key, connection=queue.connection)
|
||||
intermediate_queue.cleanup(worker, queue)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,459 @@
|
||||
import calendar
|
||||
import logging
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any, List, Optional, Type, Union
|
||||
|
||||
from rq.serializers import resolve_serializer
|
||||
|
||||
from .timeouts import BaseDeathPenalty, UnixSignalDeathPenalty
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from redis import Redis
|
||||
from redis.client import Pipeline
|
||||
|
||||
from .connections import resolve_connection
|
||||
from .defaults import DEFAULT_FAILURE_TTL
|
||||
from .exceptions import AbandonedJobError, InvalidJobOperation, NoSuchJobError
|
||||
from .job import Job, JobStatus
|
||||
from .queue import Queue
|
||||
from .utils import as_text, backend_class, current_timestamp
|
||||
|
||||
logger = logging.getLogger("rq.registry")
|
||||
|
||||
|
||||
class BaseRegistry:
|
||||
"""
|
||||
Base implementation of a job registry, implemented in Redis sorted set.
|
||||
Each job is stored as a key in the registry, scored by expiration time
|
||||
(unix timestamp).
|
||||
"""
|
||||
|
||||
job_class = Job
|
||||
death_penalty_class = UnixSignalDeathPenalty
|
||||
key_template = 'rq:registry:{0}'
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str = 'default',
|
||||
connection: Optional['Redis'] = None,
|
||||
job_class: Optional[Type['Job']] = None,
|
||||
queue: Optional['Queue'] = None,
|
||||
serializer: Any = None,
|
||||
death_penalty_class: Optional[Type[BaseDeathPenalty]] = None,
|
||||
):
|
||||
if queue:
|
||||
self.name = queue.name
|
||||
self.connection = queue.connection or resolve_connection()
|
||||
self.serializer = queue.serializer
|
||||
else:
|
||||
self.name = name
|
||||
self.connection = connection or resolve_connection()
|
||||
self.serializer = resolve_serializer(serializer)
|
||||
|
||||
self.key = self.key_template.format(self.name)
|
||||
self.job_class = backend_class(self, 'job_class', override=job_class)
|
||||
self.death_penalty_class = backend_class(self, 'death_penalty_class', override=death_penalty_class)
|
||||
|
||||
def __len__(self):
|
||||
"""Returns the number of jobs in this registry"""
|
||||
return self.count
|
||||
|
||||
def __eq__(self, other):
|
||||
return (
|
||||
self.name == other.name
|
||||
and self.connection.connection_pool.connection_kwargs == other.connection.connection_pool.connection_kwargs
|
||||
)
|
||||
|
||||
def __contains__(self, item: Union[str, 'Job']) -> bool:
|
||||
"""
|
||||
Returns a boolean indicating registry contains the given
|
||||
job instance or job id.
|
||||
|
||||
Args:
|
||||
item (Union[str, Job]): A Job ID or a Job.
|
||||
"""
|
||||
job_id = item
|
||||
if isinstance(item, self.job_class):
|
||||
job_id = item.id
|
||||
return self.connection.zscore(self.key, job_id) is not None
|
||||
|
||||
@property
|
||||
def count(self) -> int:
|
||||
"""Returns the number of jobs in this registry
|
||||
|
||||
Returns:
|
||||
int: _description_
|
||||
"""
|
||||
self.cleanup()
|
||||
return self.connection.zcard(self.key)
|
||||
|
||||
def add(self, job: 'Job', ttl=0, pipeline: Optional['Pipeline'] = None, xx: bool = False) -> int:
|
||||
"""Adds a job to a registry with expiry time of now + ttl, unless it's -1 which is set to +inf
|
||||
|
||||
Args:
|
||||
job (Job): The Job to add
|
||||
ttl (int, optional): The time to live. Defaults to 0.
|
||||
pipeline (Optional[Pipeline], optional): The Redis Pipeline. Defaults to None.
|
||||
xx (bool, optional): .... Defaults to False.
|
||||
|
||||
Returns:
|
||||
result (int): The ZADD command result
|
||||
"""
|
||||
score = ttl if ttl < 0 else current_timestamp() + ttl
|
||||
if score == -1:
|
||||
score = '+inf'
|
||||
if pipeline is not None:
|
||||
return pipeline.zadd(self.key, {job.id: score}, xx=xx)
|
||||
|
||||
return self.connection.zadd(self.key, {job.id: score}, xx=xx)
|
||||
|
||||
def remove(self, job: 'Job', pipeline: Optional['Pipeline'] = None, delete_job: bool = False):
|
||||
"""Removes job from registry and deletes it if `delete_job == True`
|
||||
|
||||
Args:
|
||||
job (Job): The Job to remove from the registry
|
||||
pipeline (Optional[Pipeline], optional): The Redis Pipeline. Defaults to None.
|
||||
delete_job (bool, optional): If should delete the job.. Defaults to False.
|
||||
"""
|
||||
connection = pipeline if pipeline is not None else self.connection
|
||||
job_id = job.id if isinstance(job, self.job_class) else job
|
||||
result = connection.zrem(self.key, job_id)
|
||||
if delete_job:
|
||||
if isinstance(job, self.job_class):
|
||||
job_instance = job
|
||||
else:
|
||||
job_instance = Job.fetch(job_id, connection=connection, serializer=self.serializer)
|
||||
job_instance.delete()
|
||||
return result
|
||||
|
||||
def get_expired_job_ids(self, timestamp: Optional[float] = None):
|
||||
"""Returns job ids whose score are less than current timestamp.
|
||||
|
||||
Returns ids for jobs with an expiry time earlier than timestamp,
|
||||
specified as seconds since the Unix epoch. timestamp defaults to call
|
||||
time if unspecified.
|
||||
"""
|
||||
score = timestamp if timestamp is not None else current_timestamp()
|
||||
expired_jobs = self.connection.zrangebyscore(self.key, 0, score)
|
||||
return [as_text(job_id) for job_id in expired_jobs]
|
||||
|
||||
def get_job_ids(self, start: int = 0, end: int = -1):
|
||||
"""Returns list of all job ids.
|
||||
|
||||
Args:
|
||||
start (int, optional): _description_. Defaults to 0.
|
||||
end (int, optional): _description_. Defaults to -1.
|
||||
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
self.cleanup()
|
||||
return [as_text(job_id) for job_id in self.connection.zrange(self.key, start, end)]
|
||||
|
||||
def get_queue(self):
|
||||
"""Returns Queue object associated with this registry."""
|
||||
return Queue(self.name, connection=self.connection, serializer=self.serializer)
|
||||
|
||||
def get_expiration_time(self, job: 'Job') -> datetime:
|
||||
"""Returns job's expiration time.
|
||||
|
||||
Args:
|
||||
job (Job): The Job to get the expiration
|
||||
"""
|
||||
score = self.connection.zscore(self.key, job.id)
|
||||
return datetime.utcfromtimestamp(score)
|
||||
|
||||
def requeue(self, job_or_id: Union['Job', str], at_front: bool = False) -> 'Job':
|
||||
"""Requeues the job with the given job ID.
|
||||
|
||||
Args:
|
||||
job_or_id (Union['Job', str]): The Job or the Job ID
|
||||
at_front (bool, optional): If the Job should be put at the front of the queue. Defaults to False.
|
||||
|
||||
Raises:
|
||||
InvalidJobOperation: If nothing is returned from the `ZREM` operation.
|
||||
|
||||
Returns:
|
||||
Job: The Requeued Job.
|
||||
"""
|
||||
if isinstance(job_or_id, self.job_class):
|
||||
job = job_or_id
|
||||
serializer = job.serializer
|
||||
else:
|
||||
serializer = self.serializer
|
||||
job = self.job_class.fetch(job_or_id, connection=self.connection, serializer=serializer)
|
||||
|
||||
result = self.connection.zrem(self.key, job.id)
|
||||
if not result:
|
||||
raise InvalidJobOperation
|
||||
|
||||
with self.connection.pipeline() as pipeline:
|
||||
queue = Queue(job.origin, connection=self.connection, job_class=self.job_class, serializer=serializer)
|
||||
job.started_at = None
|
||||
job.ended_at = None
|
||||
job._exc_info = ''
|
||||
job.save()
|
||||
job = queue._enqueue_job(job, pipeline=pipeline, at_front=at_front)
|
||||
pipeline.execute()
|
||||
return job
|
||||
|
||||
|
||||
class StartedJobRegistry(BaseRegistry):
|
||||
"""
|
||||
Registry of currently executing jobs. Each queue maintains a
|
||||
StartedJobRegistry. Jobs in this registry are ones that are currently
|
||||
being executed.
|
||||
|
||||
Jobs are added to registry right before they are executed and removed
|
||||
right after completion (success or failure).
|
||||
"""
|
||||
|
||||
key_template = 'rq:wip:{0}'
|
||||
|
||||
def cleanup(self, timestamp: Optional[float] = None):
|
||||
"""Remove abandoned jobs from registry and add them to FailedJobRegistry.
|
||||
|
||||
Removes jobs with an expiry time earlier than timestamp, specified as
|
||||
seconds since the Unix epoch. timestamp defaults to call time if
|
||||
unspecified. Removed jobs are added to the global failed job queue.
|
||||
|
||||
Args:
|
||||
timestamp (datetime): The datetime to use as the limit.
|
||||
"""
|
||||
score = timestamp if timestamp is not None else current_timestamp()
|
||||
job_ids = self.get_expired_job_ids(score)
|
||||
|
||||
if job_ids:
|
||||
failed_job_registry = FailedJobRegistry(self.name, self.connection, serializer=self.serializer)
|
||||
|
||||
with self.connection.pipeline() as pipeline:
|
||||
for job_id in job_ids:
|
||||
try:
|
||||
job = self.job_class.fetch(job_id, connection=self.connection, serializer=self.serializer)
|
||||
except NoSuchJobError:
|
||||
continue
|
||||
|
||||
job.execute_failure_callback(
|
||||
self.death_penalty_class, AbandonedJobError, AbandonedJobError(), traceback.extract_stack()
|
||||
)
|
||||
|
||||
retry = job.retries_left and job.retries_left > 0
|
||||
|
||||
if retry:
|
||||
queue = self.get_queue()
|
||||
job.retry(queue, pipeline)
|
||||
|
||||
else:
|
||||
exc_string = f"due to {AbandonedJobError.__name__}"
|
||||
logger.warning(
|
||||
f'{self.__class__.__name__} cleanup: Moving job to {FailedJobRegistry.__name__} '
|
||||
f'({exc_string})'
|
||||
)
|
||||
job.set_status(JobStatus.FAILED)
|
||||
job._exc_info = f"Moved to {FailedJobRegistry.__name__}, {exc_string}, at {datetime.now()}"
|
||||
job.save(pipeline=pipeline, include_meta=False)
|
||||
job.cleanup(ttl=-1, pipeline=pipeline)
|
||||
failed_job_registry.add(job, job.failure_ttl)
|
||||
|
||||
pipeline.zremrangebyscore(self.key, 0, score)
|
||||
pipeline.execute()
|
||||
|
||||
return job_ids
|
||||
|
||||
|
||||
class FinishedJobRegistry(BaseRegistry):
|
||||
"""
|
||||
Registry of jobs that have been completed. Jobs are added to this
|
||||
registry after they have successfully completed for monitoring purposes.
|
||||
"""
|
||||
|
||||
key_template = 'rq:finished:{0}'
|
||||
|
||||
def cleanup(self, timestamp: Optional[float] = None):
|
||||
"""Remove expired jobs from registry.
|
||||
|
||||
Removes jobs with an expiry time earlier than timestamp, specified as
|
||||
seconds since the Unix epoch. timestamp defaults to call time if
|
||||
unspecified.
|
||||
"""
|
||||
score = timestamp if timestamp is not None else current_timestamp()
|
||||
self.connection.zremrangebyscore(self.key, 0, score)
|
||||
|
||||
|
||||
class FailedJobRegistry(BaseRegistry):
|
||||
"""
|
||||
Registry of containing failed jobs.
|
||||
"""
|
||||
|
||||
key_template = 'rq:failed:{0}'
|
||||
|
||||
def cleanup(self, timestamp: Optional[float] = None):
|
||||
"""Remove expired jobs from registry.
|
||||
|
||||
Removes jobs with an expiry time earlier than timestamp, specified as
|
||||
seconds since the Unix epoch. timestamp defaults to call time if
|
||||
unspecified.
|
||||
"""
|
||||
score = timestamp if timestamp is not None else current_timestamp()
|
||||
self.connection.zremrangebyscore(self.key, 0, score)
|
||||
|
||||
def add(
|
||||
self,
|
||||
job: 'Job',
|
||||
ttl=None,
|
||||
exc_string: str = '',
|
||||
pipeline: Optional['Pipeline'] = None,
|
||||
_save_exc_to_job: bool = False,
|
||||
):
|
||||
"""
|
||||
Adds a job to a registry with expiry time of now + ttl.
|
||||
`ttl` defaults to DEFAULT_FAILURE_TTL if not specified.
|
||||
"""
|
||||
if ttl is None:
|
||||
ttl = DEFAULT_FAILURE_TTL
|
||||
score = ttl if ttl < 0 else current_timestamp() + ttl
|
||||
|
||||
if pipeline:
|
||||
p = pipeline
|
||||
else:
|
||||
p = self.connection.pipeline()
|
||||
|
||||
job._exc_info = exc_string
|
||||
job.save(pipeline=p, include_meta=False, include_result=_save_exc_to_job)
|
||||
job.cleanup(ttl=ttl, pipeline=p)
|
||||
p.zadd(self.key, {job.id: score})
|
||||
|
||||
if not pipeline:
|
||||
p.execute()
|
||||
|
||||
|
||||
class DeferredJobRegistry(BaseRegistry):
|
||||
"""
|
||||
Registry of deferred jobs (waiting for another job to finish).
|
||||
"""
|
||||
|
||||
key_template = 'rq:deferred:{0}'
|
||||
|
||||
def cleanup(self):
|
||||
"""This method is only here to prevent errors because this method is
|
||||
automatically called by `count()` and `get_job_ids()` methods
|
||||
implemented in BaseRegistry."""
|
||||
pass
|
||||
|
||||
|
||||
class ScheduledJobRegistry(BaseRegistry):
|
||||
"""
|
||||
Registry of scheduled jobs.
|
||||
"""
|
||||
|
||||
key_template = 'rq:scheduled:{0}'
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
# The underlying implementation of get_jobs_to_enqueue() is
|
||||
# the same as get_expired_job_ids, but get_expired_job_ids() doesn't
|
||||
# make sense in this context
|
||||
self.get_jobs_to_enqueue = self.get_expired_job_ids
|
||||
|
||||
def schedule(self, job: 'Job', scheduled_datetime, pipeline: Optional['Pipeline'] = None):
|
||||
"""
|
||||
Adds job to registry, scored by its execution time (in UTC).
|
||||
If datetime has no tzinfo, it will assume localtimezone.
|
||||
"""
|
||||
# If datetime has no timezone, assume server's local timezone
|
||||
if not scheduled_datetime.tzinfo:
|
||||
tz = timezone(timedelta(seconds=-(time.timezone if time.daylight == 0 else time.altzone)))
|
||||
scheduled_datetime = scheduled_datetime.replace(tzinfo=tz)
|
||||
|
||||
timestamp = calendar.timegm(scheduled_datetime.utctimetuple())
|
||||
return self.connection.zadd(self.key, {job.id: timestamp})
|
||||
|
||||
def cleanup(self):
|
||||
"""This method is only here to prevent errors because this method is
|
||||
automatically called by `count()` and `get_job_ids()` methods
|
||||
implemented in BaseRegistry."""
|
||||
pass
|
||||
|
||||
def remove_jobs(self, timestamp: Optional[datetime] = None, pipeline: Optional['Pipeline'] = None):
|
||||
"""Remove jobs whose timestamp is in the past from registry.
|
||||
|
||||
Args:
|
||||
timestamp (Optional[datetime], optional): The timestamp. Defaults to None.
|
||||
pipeline (Optional[Pipeline], optional): The Redis pipeline. Defaults to None.
|
||||
"""
|
||||
connection = pipeline if pipeline is not None else self.connection
|
||||
score = timestamp if timestamp is not None else current_timestamp()
|
||||
return connection.zremrangebyscore(self.key, 0, score)
|
||||
|
||||
def get_jobs_to_schedule(self, timestamp: Optional[datetime] = None, chunk_size: int = 1000) -> List[str]:
|
||||
"""Get's a list of job IDs that should be scheduled.
|
||||
|
||||
Args:
|
||||
timestamp (Optional[datetime], optional): _description_. Defaults to None.
|
||||
chunk_size (int, optional): _description_. Defaults to 1000.
|
||||
|
||||
Returns:
|
||||
jobs (List[str]): A list of Job ids
|
||||
"""
|
||||
score = timestamp if timestamp is not None else current_timestamp()
|
||||
jobs_to_schedule = self.connection.zrangebyscore(self.key, 0, score, start=0, num=chunk_size)
|
||||
return [as_text(job_id) for job_id in jobs_to_schedule]
|
||||
|
||||
def get_scheduled_time(self, job_or_id: Union['Job', str]) -> datetime:
|
||||
"""Returns datetime (UTC) at which job is scheduled to be enqueued
|
||||
|
||||
Args:
|
||||
job_or_id (Union[Job, str]): The Job instance or Job ID
|
||||
|
||||
Raises:
|
||||
NoSuchJobError: If the job was not found
|
||||
|
||||
Returns:
|
||||
datetime (datetime): The scheduled time as datetime object
|
||||
"""
|
||||
if isinstance(job_or_id, self.job_class):
|
||||
job_id = job_or_id.id
|
||||
else:
|
||||
job_id = job_or_id
|
||||
|
||||
score = self.connection.zscore(self.key, job_id)
|
||||
if not score:
|
||||
raise NoSuchJobError
|
||||
|
||||
return datetime.fromtimestamp(score, tz=timezone.utc)
|
||||
|
||||
|
||||
class CanceledJobRegistry(BaseRegistry):
|
||||
key_template = 'rq:canceled:{0}'
|
||||
|
||||
def get_expired_job_ids(self, timestamp: Optional[datetime] = None):
|
||||
raise NotImplementedError
|
||||
|
||||
def cleanup(self):
|
||||
"""This method is only here to prevent errors because this method is
|
||||
automatically called by `count()` and `get_job_ids()` methods
|
||||
implemented in BaseRegistry."""
|
||||
pass
|
||||
|
||||
|
||||
def clean_registries(queue: 'Queue'):
|
||||
"""Cleans StartedJobRegistry, FinishedJobRegistry and FailedJobRegistry of a queue.
|
||||
|
||||
Args:
|
||||
queue (Queue): The queue to clean
|
||||
"""
|
||||
registry = FinishedJobRegistry(
|
||||
name=queue.name, connection=queue.connection, job_class=queue.job_class, serializer=queue.serializer
|
||||
)
|
||||
registry.cleanup()
|
||||
registry = StartedJobRegistry(
|
||||
name=queue.name, connection=queue.connection, job_class=queue.job_class, serializer=queue.serializer
|
||||
)
|
||||
registry.cleanup()
|
||||
|
||||
registry = FailedJobRegistry(
|
||||
name=queue.name, connection=queue.connection, job_class=queue.job_class, serializer=queue.serializer
|
||||
)
|
||||
registry.cleanup()
|
||||
@@ -0,0 +1,208 @@
|
||||
import zlib
|
||||
from base64 import b64decode, b64encode
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
from typing import Any, Optional
|
||||
|
||||
from redis import Redis
|
||||
|
||||
from .defaults import UNSERIALIZABLE_RETURN_VALUE_PAYLOAD
|
||||
from .job import Job
|
||||
from .serializers import resolve_serializer
|
||||
from .utils import decode_redis_hash, now
|
||||
|
||||
|
||||
def get_key(job_id):
|
||||
return 'rq:results:%s' % job_id
|
||||
|
||||
|
||||
class Result:
|
||||
class Type(Enum):
|
||||
SUCCESSFUL = 1
|
||||
FAILED = 2
|
||||
STOPPED = 3
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
job_id: str,
|
||||
type: Type,
|
||||
connection: Redis,
|
||||
id: Optional[str] = None,
|
||||
created_at: Optional[datetime] = None,
|
||||
return_value: Optional[Any] = None,
|
||||
exc_string: Optional[str] = None,
|
||||
serializer=None,
|
||||
):
|
||||
self.return_value = return_value
|
||||
self.exc_string = exc_string
|
||||
self.type = type
|
||||
self.created_at = created_at if created_at else now()
|
||||
self.serializer = resolve_serializer(serializer)
|
||||
self.connection = connection
|
||||
self.job_id = job_id
|
||||
self.id = id
|
||||
|
||||
def __repr__(self):
|
||||
return f'Result(id={self.id}, type={self.Type(self.type).name})'
|
||||
|
||||
def __eq__(self, other):
|
||||
try:
|
||||
return self.id == other.id
|
||||
except AttributeError:
|
||||
return False
|
||||
|
||||
def __bool__(self):
|
||||
return bool(self.id)
|
||||
|
||||
@classmethod
|
||||
def create(cls, job, type, ttl, return_value=None, exc_string=None, pipeline=None):
|
||||
result = cls(
|
||||
job_id=job.id,
|
||||
type=type,
|
||||
connection=job.connection,
|
||||
return_value=return_value,
|
||||
exc_string=exc_string,
|
||||
serializer=job.serializer,
|
||||
)
|
||||
result.save(ttl=ttl, pipeline=pipeline)
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def create_failure(cls, job, ttl, exc_string, pipeline=None):
|
||||
result = cls(
|
||||
job_id=job.id,
|
||||
type=cls.Type.FAILED,
|
||||
connection=job.connection,
|
||||
exc_string=exc_string,
|
||||
serializer=job.serializer,
|
||||
)
|
||||
result.save(ttl=ttl, pipeline=pipeline)
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def all(cls, job: Job, serializer=None):
|
||||
"""Returns all results for job"""
|
||||
# response = job.connection.zrange(cls.get_key(job.id), 0, 10, desc=True, withscores=True)
|
||||
response = job.connection.xrevrange(cls.get_key(job.id), '+', '-')
|
||||
results = []
|
||||
for result_id, payload in response:
|
||||
results.append(
|
||||
cls.restore(job.id, result_id.decode(), payload, connection=job.connection, serializer=serializer)
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
@classmethod
|
||||
def count(cls, job: Job) -> int:
|
||||
"""Returns the number of job results"""
|
||||
return job.connection.xlen(cls.get_key(job.id))
|
||||
|
||||
@classmethod
|
||||
def delete_all(cls, job: Job) -> None:
|
||||
"""Delete all job results"""
|
||||
job.connection.delete(cls.get_key(job.id))
|
||||
|
||||
@classmethod
|
||||
def restore(cls, job_id: str, result_id: str, payload: dict, connection: Redis, serializer=None) -> 'Result':
|
||||
"""Create a Result object from given Redis payload"""
|
||||
created_at = datetime.fromtimestamp(int(result_id.split('-')[0]) / 1000, tz=timezone.utc)
|
||||
payload = decode_redis_hash(payload)
|
||||
# data, timestamp = payload
|
||||
# result_data = json.loads(data)
|
||||
# created_at = datetime.fromtimestamp(timestamp, tz=timezone.utc)
|
||||
|
||||
serializer = resolve_serializer(serializer)
|
||||
return_value = payload.get('return_value')
|
||||
if return_value is not None:
|
||||
return_value = serializer.loads(b64decode(return_value.decode()))
|
||||
|
||||
exc_string = payload.get('exc_string')
|
||||
if exc_string:
|
||||
exc_string = zlib.decompress(b64decode(exc_string)).decode()
|
||||
|
||||
return Result(
|
||||
job_id,
|
||||
Result.Type(int(payload['type'])),
|
||||
connection=connection,
|
||||
id=result_id,
|
||||
created_at=created_at,
|
||||
return_value=return_value,
|
||||
exc_string=exc_string,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def fetch(cls, job: Job, serializer=None) -> Optional['Result']:
|
||||
"""Fetch a result that matches a given job ID. The current sorted set
|
||||
based implementation does not allow us to fetch a given key by ID
|
||||
so we need to iterate through results, deserialize the payload and
|
||||
look for a matching ID.
|
||||
|
||||
Future Redis streams based implementation may make this more efficient
|
||||
and scalable.
|
||||
"""
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def fetch_latest(cls, job: Job, serializer=None, timeout: int = 0) -> Optional['Result']:
|
||||
"""Returns the latest result for given job instance or ID.
|
||||
|
||||
If a non-zero timeout is provided, block for a result until timeout is reached.
|
||||
"""
|
||||
if timeout:
|
||||
# Unlike blpop, xread timeout is in miliseconds. "0-0" is the special value for the
|
||||
# first item in the stream, like '-' for xrevrange.
|
||||
timeout_ms = timeout * 1000
|
||||
response = job.connection.xread({cls.get_key(job.id): "0-0"}, block=timeout_ms)
|
||||
if not response:
|
||||
return None
|
||||
response = response[0] # Querying single stream only.
|
||||
response = response[1] # Xread also returns Result.id, which we don't need.
|
||||
result_id, payload = response[-1] # Take most recent result.
|
||||
|
||||
else:
|
||||
# If not blocking, use xrevrange to load a single result (as xread will load them all).
|
||||
response = job.connection.xrevrange(cls.get_key(job.id), '+', '-', count=1)
|
||||
if not response:
|
||||
return None
|
||||
result_id, payload = response[0]
|
||||
|
||||
res = cls.restore(job.id, result_id.decode(), payload, connection=job.connection, serializer=serializer)
|
||||
return res
|
||||
|
||||
@classmethod
|
||||
def get_key(cls, job_id):
|
||||
return 'rq:results:%s' % job_id
|
||||
|
||||
def save(self, ttl, pipeline=None):
|
||||
"""Save result data to Redis"""
|
||||
key = self.get_key(self.job_id)
|
||||
|
||||
connection = pipeline if pipeline is not None else self.connection
|
||||
# result = connection.zadd(key, {self.serialize(): self.created_at.timestamp()})
|
||||
result = connection.xadd(key, self.serialize(), maxlen=10)
|
||||
# If xadd() is called in a pipeline, it returns a pipeline object instead of stream ID
|
||||
if pipeline is None:
|
||||
self.id = result.decode()
|
||||
if ttl is not None:
|
||||
if ttl == -1:
|
||||
connection.persist(key)
|
||||
else:
|
||||
connection.expire(key, ttl)
|
||||
return self.id
|
||||
|
||||
def serialize(self):
|
||||
data = {'type': self.type.value}
|
||||
|
||||
if self.exc_string is not None:
|
||||
data['exc_string'] = b64encode(zlib.compress(self.exc_string.encode())).decode()
|
||||
|
||||
try:
|
||||
serialized = self.serializer.dumps(self.return_value)
|
||||
except: # noqa
|
||||
serialized = self.serializer.dumps(UNSERIALIZABLE_RETURN_VALUE_PAYLOAD)
|
||||
|
||||
if self.return_value is not None:
|
||||
data['return_value'] = b64encode(serialized).decode()
|
||||
|
||||
# return json.dumps(data)
|
||||
return data
|
||||
@@ -0,0 +1,227 @@
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from multiprocessing import Process
|
||||
from typing import List, Set
|
||||
|
||||
from redis import ConnectionPool, Redis
|
||||
|
||||
from .connections import parse_connection
|
||||
from .defaults import DEFAULT_LOGGING_DATE_FORMAT, DEFAULT_LOGGING_FORMAT, DEFAULT_SCHEDULER_FALLBACK_PERIOD
|
||||
from .job import Job
|
||||
from .logutils import setup_loghandlers
|
||||
from .queue import Queue
|
||||
from .registry import ScheduledJobRegistry
|
||||
from .serializers import resolve_serializer
|
||||
from .utils import current_timestamp, parse_names
|
||||
|
||||
SCHEDULER_KEY_TEMPLATE = 'rq:scheduler:%s'
|
||||
SCHEDULER_LOCKING_KEY_TEMPLATE = 'rq:scheduler-lock:%s'
|
||||
|
||||
|
||||
class SchedulerStatus(str, Enum):
|
||||
STARTED = 'started'
|
||||
WORKING = 'working'
|
||||
STOPPED = 'stopped'
|
||||
|
||||
|
||||
class RQScheduler:
|
||||
# STARTED: scheduler has been started but sleeping
|
||||
# WORKING: scheduler is in the midst of scheduling jobs
|
||||
# STOPPED: scheduler is in stopped condition
|
||||
|
||||
Status = SchedulerStatus
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
queues,
|
||||
connection: Redis,
|
||||
interval=1,
|
||||
logging_level=logging.INFO,
|
||||
date_format=DEFAULT_LOGGING_DATE_FORMAT,
|
||||
log_format=DEFAULT_LOGGING_FORMAT,
|
||||
serializer=None,
|
||||
):
|
||||
self._queue_names = set(parse_names(queues))
|
||||
self._acquired_locks: Set[str] = set()
|
||||
self._scheduled_job_registries: List[ScheduledJobRegistry] = []
|
||||
self.lock_acquisition_time = None
|
||||
self._connection_class, self._pool_class, self._pool_kwargs = parse_connection(connection)
|
||||
self.serializer = resolve_serializer(serializer)
|
||||
|
||||
self._connection = None
|
||||
self.interval = interval
|
||||
self._stop_requested = False
|
||||
self._status = self.Status.STOPPED
|
||||
self._process = None
|
||||
self.log = logging.getLogger(__name__)
|
||||
setup_loghandlers(
|
||||
level=logging_level,
|
||||
name=__name__,
|
||||
log_format=log_format,
|
||||
date_format=date_format,
|
||||
)
|
||||
|
||||
@property
|
||||
def connection(self):
|
||||
if self._connection:
|
||||
return self._connection
|
||||
self._connection = self._connection_class(
|
||||
connection_pool=ConnectionPool(connection_class=self._pool_class, **self._pool_kwargs)
|
||||
)
|
||||
return self._connection
|
||||
|
||||
@property
|
||||
def acquired_locks(self):
|
||||
return self._acquired_locks
|
||||
|
||||
@property
|
||||
def status(self):
|
||||
return self._status
|
||||
|
||||
@property
|
||||
def should_reacquire_locks(self):
|
||||
"""Returns True if lock_acquisition_time is longer than 10 minutes ago"""
|
||||
if self._queue_names == self.acquired_locks:
|
||||
return False
|
||||
if not self.lock_acquisition_time:
|
||||
return True
|
||||
return (datetime.now() - self.lock_acquisition_time).total_seconds() > DEFAULT_SCHEDULER_FALLBACK_PERIOD
|
||||
|
||||
def acquire_locks(self, auto_start=False):
|
||||
"""Returns names of queue it successfully acquires lock on"""
|
||||
successful_locks = set()
|
||||
pid = os.getpid()
|
||||
self.log.debug('Trying to acquire locks for %s', ', '.join(self._queue_names))
|
||||
for name in self._queue_names:
|
||||
if self.connection.set(self.get_locking_key(name), pid, nx=True, ex=self.interval + 60):
|
||||
successful_locks.add(name)
|
||||
|
||||
# Always reset _scheduled_job_registries when acquiring locks
|
||||
self._scheduled_job_registries = []
|
||||
self._acquired_locks = self._acquired_locks.union(successful_locks)
|
||||
self.lock_acquisition_time = datetime.now()
|
||||
|
||||
# If auto_start is requested and scheduler is not started,
|
||||
# run self.start()
|
||||
if self._acquired_locks and auto_start:
|
||||
if not self._process or not self._process.is_alive():
|
||||
self.start()
|
||||
|
||||
return successful_locks
|
||||
|
||||
def prepare_registries(self, queue_names: str = None):
|
||||
"""Prepare scheduled job registries for use"""
|
||||
self._scheduled_job_registries = []
|
||||
if not queue_names:
|
||||
queue_names = self._acquired_locks
|
||||
for name in queue_names:
|
||||
self._scheduled_job_registries.append(
|
||||
ScheduledJobRegistry(name, connection=self.connection, serializer=self.serializer)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_locking_key(cls, name: str):
|
||||
"""Returns scheduler key for a given queue name"""
|
||||
return SCHEDULER_LOCKING_KEY_TEMPLATE % name
|
||||
|
||||
def enqueue_scheduled_jobs(self):
|
||||
"""Enqueue jobs whose timestamp is in the past"""
|
||||
self._status = self.Status.WORKING
|
||||
|
||||
if not self._scheduled_job_registries and self._acquired_locks:
|
||||
self.prepare_registries()
|
||||
|
||||
for registry in self._scheduled_job_registries:
|
||||
timestamp = current_timestamp()
|
||||
|
||||
# TODO: try to use Lua script to make get_jobs_to_schedule()
|
||||
# and remove_jobs() atomic
|
||||
job_ids = registry.get_jobs_to_schedule(timestamp)
|
||||
|
||||
if not job_ids:
|
||||
continue
|
||||
|
||||
queue = Queue(registry.name, connection=self.connection, serializer=self.serializer)
|
||||
|
||||
with self.connection.pipeline() as pipeline:
|
||||
jobs = Job.fetch_many(job_ids, connection=self.connection, serializer=self.serializer)
|
||||
for job in jobs:
|
||||
if job is not None:
|
||||
queue._enqueue_job(job, pipeline=pipeline, at_front=bool(job.enqueue_at_front))
|
||||
registry.remove(job, pipeline=pipeline)
|
||||
pipeline.execute()
|
||||
self._status = self.Status.STARTED
|
||||
|
||||
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._stop_requested = True
|
||||
|
||||
def heartbeat(self):
|
||||
"""Updates the TTL on scheduler keys and the locks"""
|
||||
self.log.debug('Scheduler sending heartbeat to %s', ', '.join(self.acquired_locks))
|
||||
if len(self._acquired_locks) > 1:
|
||||
with self.connection.pipeline() as pipeline:
|
||||
for name in self._acquired_locks:
|
||||
key = self.get_locking_key(name)
|
||||
pipeline.expire(key, self.interval + 60)
|
||||
pipeline.execute()
|
||||
elif self._acquired_locks:
|
||||
key = self.get_locking_key(next(iter(self._acquired_locks)))
|
||||
self.connection.expire(key, self.interval + 60)
|
||||
|
||||
def stop(self):
|
||||
self.log.info('Scheduler stopping, releasing locks for %s...', ', '.join(self._acquired_locks))
|
||||
self.release_locks()
|
||||
self._status = self.Status.STOPPED
|
||||
|
||||
def release_locks(self):
|
||||
"""Release acquired locks"""
|
||||
keys = [self.get_locking_key(name) for name in self._acquired_locks]
|
||||
self.connection.delete(*keys)
|
||||
self._acquired_locks = set()
|
||||
|
||||
def start(self):
|
||||
self._status = self.Status.STARTED
|
||||
# Redis instance can't be pickled across processes so we need to
|
||||
# clean this up before forking
|
||||
self._connection = None
|
||||
self._process = Process(target=run, args=(self,), name='Scheduler')
|
||||
self._process.start()
|
||||
return self._process
|
||||
|
||||
def work(self):
|
||||
self._install_signal_handlers()
|
||||
|
||||
while True:
|
||||
if self._stop_requested:
|
||||
self.stop()
|
||||
break
|
||||
|
||||
if self.should_reacquire_locks:
|
||||
self.acquire_locks()
|
||||
|
||||
self.enqueue_scheduled_jobs()
|
||||
self.heartbeat()
|
||||
time.sleep(self.interval)
|
||||
|
||||
|
||||
def run(scheduler):
|
||||
scheduler.log.info('Scheduler for %s started with PID %s', ', '.join(scheduler._queue_names), os.getpid())
|
||||
try:
|
||||
scheduler.work()
|
||||
except: # noqa
|
||||
scheduler.log.error('Scheduler [PID %s] raised an exception.\n%s', os.getpid(), traceback.format_exc())
|
||||
raise
|
||||
scheduler.log.info('Scheduler with PID %d has stopped', os.getpid())
|
||||
@@ -0,0 +1,48 @@
|
||||
import json
|
||||
import pickle
|
||||
from functools import partial
|
||||
from typing import Optional, Type, Union
|
||||
|
||||
from .utils import import_attribute
|
||||
|
||||
|
||||
class DefaultSerializer:
|
||||
dumps = partial(pickle.dumps, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
loads = pickle.loads
|
||||
|
||||
|
||||
class JSONSerializer:
|
||||
@staticmethod
|
||||
def dumps(*args, **kwargs):
|
||||
return json.dumps(*args, **kwargs).encode('utf-8')
|
||||
|
||||
@staticmethod
|
||||
def loads(s, *args, **kwargs):
|
||||
return json.loads(s.decode('utf-8'), *args, **kwargs)
|
||||
|
||||
|
||||
def resolve_serializer(serializer: Optional[Union[Type[DefaultSerializer], str]] = None) -> Type[DefaultSerializer]:
|
||||
"""This function checks the user defined serializer for ('dumps', 'loads') methods
|
||||
It returns a default pickle serializer if not found else it returns a MySerializer
|
||||
The returned serializer objects implement ('dumps', 'loads') methods
|
||||
Also accepts a string path to serializer that will be loaded as the serializer.
|
||||
|
||||
Args:
|
||||
serializer (Callable): The serializer to resolve.
|
||||
|
||||
Returns:
|
||||
serializer (Callable): An object that implements the SerializerProtocol
|
||||
"""
|
||||
if not serializer:
|
||||
return DefaultSerializer
|
||||
|
||||
if isinstance(serializer, str):
|
||||
serializer = import_attribute(serializer)
|
||||
|
||||
default_serializer_methods = ('dumps', 'loads')
|
||||
|
||||
for instance_method in default_serializer_methods:
|
||||
if not hasattr(serializer, instance_method):
|
||||
raise NotImplementedError('Serializer should have (dumps, loads) methods.')
|
||||
|
||||
return serializer
|
||||
@@ -0,0 +1,49 @@
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from redis import Redis
|
||||
|
||||
from rq.worker import Worker
|
||||
|
||||
|
||||
WORKERS_SUSPENDED = 'rq:suspended'
|
||||
|
||||
|
||||
def is_suspended(connection: 'Redis', worker: Optional['Worker'] = None):
|
||||
"""Checks whether a Worker is suspendeed on a given connection
|
||||
PS: pipeline returns a list of responses
|
||||
Ref: https://github.com/andymccurdy/redis-py#pipelines
|
||||
|
||||
Args:
|
||||
connection (Redis): The Redis Connection
|
||||
worker (Optional[Worker], optional): The Worker. Defaults to None.
|
||||
"""
|
||||
with connection.pipeline() as pipeline:
|
||||
if worker is not None:
|
||||
worker.heartbeat(pipeline=pipeline)
|
||||
pipeline.exists(WORKERS_SUSPENDED)
|
||||
return pipeline.execute()[-1]
|
||||
|
||||
|
||||
def suspend(connection: 'Redis', ttl: Optional[int] = None):
|
||||
"""
|
||||
Suspends.
|
||||
TTL of 0 will invalidate right away.
|
||||
|
||||
Args:
|
||||
connection (Redis): The Redis connection to use..
|
||||
ttl (Optional[int], optional): time to live in seconds. Defaults to `None`
|
||||
"""
|
||||
connection.set(WORKERS_SUSPENDED, 1)
|
||||
if ttl is not None:
|
||||
connection.expire(WORKERS_SUSPENDED, ttl)
|
||||
|
||||
|
||||
def resume(connection: 'Redis'):
|
||||
"""
|
||||
Resumes.
|
||||
|
||||
Args:
|
||||
connection (Redis): The Redis connection to use..
|
||||
"""
|
||||
return connection.delete(WORKERS_SUSPENDED)
|
||||
@@ -0,0 +1,123 @@
|
||||
import ctypes
|
||||
import signal
|
||||
import threading
|
||||
|
||||
|
||||
class BaseTimeoutException(Exception):
|
||||
"""Base exception for timeouts."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class JobTimeoutException(BaseTimeoutException):
|
||||
"""Raised when a job takes longer to complete than the allowed maximum
|
||||
timeout value.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class HorseMonitorTimeoutException(BaseTimeoutException):
|
||||
"""Raised when waiting for a horse exiting takes longer than the maximum
|
||||
timeout value.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class BaseDeathPenalty:
|
||||
"""Base class to setup job timeouts."""
|
||||
|
||||
def __init__(self, timeout, exception=BaseTimeoutException, **kwargs):
|
||||
self._timeout = timeout
|
||||
self._exception = exception
|
||||
|
||||
def __enter__(self):
|
||||
self.setup_death_penalty()
|
||||
|
||||
def __exit__(self, type, value, traceback):
|
||||
# Always cancel immediately, since we're done
|
||||
try:
|
||||
self.cancel_death_penalty()
|
||||
except BaseTimeoutException:
|
||||
# Weird case: we're done with the with body, but now the alarm is
|
||||
# fired. We may safely ignore this situation and consider the
|
||||
# body done.
|
||||
pass
|
||||
|
||||
# __exit__ may return True to supress further exception handling. We
|
||||
# don't want to suppress any exceptions here, since all errors should
|
||||
# just pass through, BaseTimeoutException being handled normally to the
|
||||
# invoking context.
|
||||
return False
|
||||
|
||||
def setup_death_penalty(self):
|
||||
raise NotImplementedError()
|
||||
|
||||
def cancel_death_penalty(self):
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class UnixSignalDeathPenalty(BaseDeathPenalty):
|
||||
def handle_death_penalty(self, signum, frame):
|
||||
raise self._exception('Task exceeded maximum timeout value ({0} seconds)'.format(self._timeout))
|
||||
|
||||
def setup_death_penalty(self):
|
||||
"""Sets up an alarm signal and a signal handler that raises
|
||||
an exception after the timeout amount (expressed in seconds).
|
||||
"""
|
||||
signal.signal(signal.SIGALRM, self.handle_death_penalty)
|
||||
signal.alarm(self._timeout)
|
||||
|
||||
def cancel_death_penalty(self):
|
||||
"""Removes the death penalty alarm and puts back the system into
|
||||
default signal handling.
|
||||
"""
|
||||
signal.alarm(0)
|
||||
signal.signal(signal.SIGALRM, signal.SIG_DFL)
|
||||
|
||||
|
||||
class TimerDeathPenalty(BaseDeathPenalty):
|
||||
def __init__(self, timeout, exception=JobTimeoutException, **kwargs):
|
||||
super().__init__(timeout, exception, **kwargs)
|
||||
self._target_thread_id = threading.current_thread().ident
|
||||
self._timer = None
|
||||
|
||||
# Monkey-patch exception with the message ahead of time
|
||||
# since PyThreadState_SetAsyncExc can only take a class
|
||||
def init_with_message(self, *args, **kwargs): # noqa
|
||||
super(exception, self).__init__("Task exceeded maximum timeout value ({0} seconds)".format(timeout))
|
||||
|
||||
self._exception.__init__ = init_with_message
|
||||
|
||||
def new_timer(self):
|
||||
"""Returns a new timer since timers can only be used once."""
|
||||
return threading.Timer(self._timeout, self.handle_death_penalty)
|
||||
|
||||
def handle_death_penalty(self):
|
||||
"""Raises an asynchronous exception in another thread.
|
||||
|
||||
Reference http://docs.python.org/c-api/init.html#PyThreadState_SetAsyncExc for more info.
|
||||
"""
|
||||
ret = ctypes.pythonapi.PyThreadState_SetAsyncExc(
|
||||
ctypes.c_long(self._target_thread_id), ctypes.py_object(self._exception)
|
||||
)
|
||||
if ret == 0:
|
||||
raise ValueError("Invalid thread ID {}".format(self._target_thread_id))
|
||||
elif ret > 1:
|
||||
ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(self._target_thread_id), 0)
|
||||
raise SystemError("PyThreadState_SetAsyncExc failed")
|
||||
|
||||
def setup_death_penalty(self):
|
||||
"""Starts the timer."""
|
||||
if self._timeout <= 0:
|
||||
return
|
||||
self._timer = self.new_timer()
|
||||
self._timer.start()
|
||||
|
||||
def cancel_death_penalty(self):
|
||||
"""Cancels the timer."""
|
||||
if self._timeout <= 0:
|
||||
return
|
||||
self._timer.cancel()
|
||||
self._timer = None
|
||||
@@ -0,0 +1,17 @@
|
||||
from typing import TYPE_CHECKING, Any, Callable, List, TypeVar, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .job import Dependency, Job
|
||||
|
||||
|
||||
FunctionReferenceType = TypeVar('FunctionReferenceType', str, Callable[..., Any])
|
||||
"""Custom type definition for what a `func` is in the context of a job.
|
||||
A `func` can be a string with the function import path (eg.: `myfile.mymodule.myfunc`)
|
||||
or a direct callable (function/method).
|
||||
"""
|
||||
|
||||
|
||||
JobDependencyType = TypeVar('JobDependencyType', 'Dependency', 'Job', str, List[Union['Dependency', 'Job']])
|
||||
"""Custom type definition for a job dependencies.
|
||||
A simple helper definition for the `depends_on` parameter when creating a job.
|
||||
"""
|
||||
@@ -0,0 +1,381 @@
|
||||
"""
|
||||
Miscellaneous helper functions.
|
||||
|
||||
The formatter for ANSI colored console output is heavily based on Pygments
|
||||
terminal colorizing code, originally by Georg Brandl.
|
||||
"""
|
||||
|
||||
import calendar
|
||||
import datetime
|
||||
import datetime as dt
|
||||
import importlib
|
||||
import logging
|
||||
import numbers
|
||||
from collections.abc import Iterable
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from redis import Redis
|
||||
|
||||
from .queue import Queue
|
||||
|
||||
from redis.exceptions import ResponseError
|
||||
|
||||
from .exceptions import TimeoutFormatError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def compact(lst: List[Any]) -> List[Any]:
|
||||
"""Excludes `None` values from a list-like object.
|
||||
|
||||
Args:
|
||||
lst (list): A list (or list-like) oject
|
||||
|
||||
Returns:
|
||||
object (list): The list without None values
|
||||
"""
|
||||
return [item for item in lst if item is not None]
|
||||
|
||||
|
||||
def as_text(v: Union[bytes, str]) -> str:
|
||||
"""Converts a bytes value to a string using `utf-8`.
|
||||
|
||||
Args:
|
||||
v (Union[bytes, str]): The value (bytes or string)
|
||||
|
||||
Raises:
|
||||
ValueError: If the value is not bytes or string
|
||||
|
||||
Returns:
|
||||
value (Optional[str]): Either the decoded string or None
|
||||
"""
|
||||
if isinstance(v, bytes):
|
||||
return v.decode('utf-8')
|
||||
elif isinstance(v, str):
|
||||
return v
|
||||
else:
|
||||
raise ValueError('Unknown type %r' % type(v))
|
||||
|
||||
|
||||
def decode_redis_hash(h) -> Dict[str, Any]:
|
||||
"""Decodes the Redis hash, ensuring that keys are strings
|
||||
Most importantly, decodes bytes strings, ensuring the dict has str keys.
|
||||
|
||||
Args:
|
||||
h (Dict[Any, Any]): The Redis hash
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: The decoded Redis data (Dictionary)
|
||||
"""
|
||||
return dict((as_text(k), h[k]) for k in h)
|
||||
|
||||
|
||||
def import_attribute(name: str) -> Callable[..., Any]:
|
||||
"""Returns an attribute from a dotted path name. Example: `path.to.func`.
|
||||
|
||||
When the attribute we look for is a staticmethod, module name in its
|
||||
dotted path is not the last-before-end word
|
||||
|
||||
E.g.: package_a.package_b.module_a.ClassA.my_static_method
|
||||
|
||||
Thus we remove the bits from the end of the name until we can import it
|
||||
|
||||
Args:
|
||||
name (str): The name (reference) to the path.
|
||||
|
||||
Raises:
|
||||
ValueError: If no module is found or invalid attribute name.
|
||||
|
||||
Returns:
|
||||
Any: An attribute (normally a Callable)
|
||||
"""
|
||||
name_bits = name.split('.')
|
||||
module_name_bits, attribute_bits = name_bits[:-1], [name_bits[-1]]
|
||||
module = None
|
||||
while len(module_name_bits):
|
||||
try:
|
||||
module_name = '.'.join(module_name_bits)
|
||||
module = importlib.import_module(module_name)
|
||||
break
|
||||
except ImportError:
|
||||
attribute_bits.insert(0, module_name_bits.pop())
|
||||
|
||||
if module is None:
|
||||
# maybe it's a builtin
|
||||
try:
|
||||
return __builtins__[name]
|
||||
except KeyError:
|
||||
raise ValueError('Invalid attribute name: %s' % name)
|
||||
|
||||
attribute_name = '.'.join(attribute_bits)
|
||||
if hasattr(module, attribute_name):
|
||||
return getattr(module, attribute_name)
|
||||
# staticmethods
|
||||
attribute_name = attribute_bits.pop()
|
||||
attribute_owner_name = '.'.join(attribute_bits)
|
||||
try:
|
||||
attribute_owner = getattr(module, attribute_owner_name)
|
||||
except: # noqa
|
||||
raise ValueError('Invalid attribute name: %s' % attribute_name)
|
||||
|
||||
if not hasattr(attribute_owner, attribute_name):
|
||||
raise ValueError('Invalid attribute name: %s' % name)
|
||||
return getattr(attribute_owner, attribute_name)
|
||||
|
||||
|
||||
def utcnow():
|
||||
return datetime.datetime.utcnow()
|
||||
|
||||
|
||||
def now():
|
||||
"""Return now in UTC"""
|
||||
return datetime.datetime.now(datetime.timezone.utc)
|
||||
|
||||
|
||||
_TIMESTAMP_FORMAT = '%Y-%m-%dT%H:%M:%S.%fZ'
|
||||
|
||||
|
||||
def utcformat(dt: dt.datetime) -> str:
|
||||
return dt.strftime(as_text(_TIMESTAMP_FORMAT))
|
||||
|
||||
|
||||
def utcparse(string: str) -> dt.datetime:
|
||||
try:
|
||||
return datetime.datetime.strptime(string, _TIMESTAMP_FORMAT)
|
||||
except ValueError:
|
||||
# This catches any jobs remain with old datetime format
|
||||
return datetime.datetime.strptime(string, '%Y-%m-%dT%H:%M:%SZ')
|
||||
|
||||
|
||||
def first(iterable: Iterable, default=None, key=None):
|
||||
"""Return first element of `iterable` that evaluates true, else return None
|
||||
(or an optional default value).
|
||||
|
||||
>>> first([0, False, None, [], (), 42])
|
||||
42
|
||||
|
||||
>>> first([0, False, None, [], ()]) is None
|
||||
True
|
||||
|
||||
>>> first([0, False, None, [], ()], default='ohai')
|
||||
'ohai'
|
||||
|
||||
>>> import re
|
||||
>>> m = first(re.match(regex, 'abc') for regex in ['b.*', 'a(.*)'])
|
||||
>>> m.group(1)
|
||||
'bc'
|
||||
|
||||
The optional `key` argument specifies a one-argument predicate function
|
||||
like that used for `filter()`. The `key` argument, if supplied, must be
|
||||
in keyword form. For example:
|
||||
|
||||
>>> first([1, 1, 3, 4, 5], key=lambda x: x % 2 == 0)
|
||||
4
|
||||
|
||||
Args:
|
||||
iterable (t.Iterable): _description_
|
||||
default (_type_, optional): _description_. Defaults to None.
|
||||
key (_type_, optional): _description_. Defaults to None.
|
||||
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
if key is None:
|
||||
for el in iterable:
|
||||
if el:
|
||||
return el
|
||||
else:
|
||||
for el in iterable:
|
||||
if key(el):
|
||||
return el
|
||||
|
||||
return default
|
||||
|
||||
|
||||
def is_nonstring_iterable(obj: Any) -> bool:
|
||||
"""Returns whether the obj is an iterable, but not a string
|
||||
|
||||
Args:
|
||||
obj (Any): _description_
|
||||
|
||||
Returns:
|
||||
bool: _description_
|
||||
"""
|
||||
return isinstance(obj, Iterable) and not isinstance(obj, str)
|
||||
|
||||
|
||||
def ensure_list(obj: Any) -> List:
|
||||
"""When passed an iterable of objects, does nothing, otherwise, it returns
|
||||
a list with just that object in it.
|
||||
|
||||
Args:
|
||||
obj (Any): _description_
|
||||
|
||||
Returns:
|
||||
List: _description_
|
||||
"""
|
||||
return obj if is_nonstring_iterable(obj) else [obj]
|
||||
|
||||
|
||||
def current_timestamp() -> int:
|
||||
"""Returns current UTC timestamp
|
||||
|
||||
Returns:
|
||||
int: _description_
|
||||
"""
|
||||
return calendar.timegm(datetime.datetime.utcnow().utctimetuple())
|
||||
|
||||
|
||||
def backend_class(holder, default_name, override=None):
|
||||
"""Get a backend class using its default attribute name or an override
|
||||
|
||||
Args:
|
||||
holder (_type_): _description_
|
||||
default_name (_type_): _description_
|
||||
override (_type_, optional): _description_. Defaults to None.
|
||||
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
if override is None:
|
||||
return getattr(holder, default_name)
|
||||
elif isinstance(override, str):
|
||||
return import_attribute(override)
|
||||
else:
|
||||
return override
|
||||
|
||||
|
||||
def str_to_date(date_str: Optional[str]) -> Union[dt.datetime, Any]:
|
||||
if not date_str:
|
||||
return
|
||||
else:
|
||||
return utcparse(date_str.decode())
|
||||
|
||||
|
||||
def parse_timeout(timeout: Union[int, float, str]) -> int:
|
||||
"""Transfer all kinds of timeout format to an integer representing seconds"""
|
||||
if not isinstance(timeout, numbers.Integral) and timeout is not None:
|
||||
try:
|
||||
timeout = int(timeout)
|
||||
except ValueError:
|
||||
digit, unit = timeout[:-1], (timeout[-1:]).lower()
|
||||
unit_second = {'d': 86400, 'h': 3600, 'm': 60, 's': 1}
|
||||
try:
|
||||
timeout = int(digit) * unit_second[unit]
|
||||
except (ValueError, KeyError):
|
||||
raise TimeoutFormatError(
|
||||
'Timeout must be an integer or a string representing an integer, or '
|
||||
'a string with format: digits + unit, unit can be "d", "h", "m", "s", '
|
||||
'such as "1h", "23m".'
|
||||
)
|
||||
|
||||
return timeout
|
||||
|
||||
|
||||
def get_version(connection: 'Redis') -> Tuple[int, int, int]:
|
||||
"""
|
||||
Returns tuple of Redis server version.
|
||||
This function also correctly handles 4 digit redis server versions.
|
||||
|
||||
Args:
|
||||
connection (Redis): The Redis connection.
|
||||
|
||||
Returns:
|
||||
version (Tuple[int, int, int]): A tuple representing the semantic versioning format (eg. (5, 0, 9))
|
||||
"""
|
||||
try:
|
||||
# Getting the connection info for each job tanks performance, we can cache it on the connection object
|
||||
if not getattr(connection, "__rq_redis_server_version", None):
|
||||
setattr(
|
||||
connection,
|
||||
"__rq_redis_server_version",
|
||||
tuple(int(i) for i in str(connection.info("server")["redis_version"]).split('.')[:3]),
|
||||
)
|
||||
return getattr(connection, "__rq_redis_server_version")
|
||||
except ResponseError: # fakeredis doesn't implement Redis' INFO command
|
||||
return (5, 0, 9)
|
||||
|
||||
|
||||
def ceildiv(a, b):
|
||||
"""Ceiling division. Returns the ceiling of the quotient of a division operation
|
||||
|
||||
Args:
|
||||
a (_type_): _description_
|
||||
b (_type_): _description_
|
||||
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
return -(-a // b)
|
||||
|
||||
|
||||
def split_list(a_list: List[Any], segment_size: int):
|
||||
"""Splits a list into multiple smaller lists having size `segment_size`
|
||||
|
||||
Args:
|
||||
a_list (List[Any]): A list to split
|
||||
segment_size (int): The segment size to split into
|
||||
|
||||
Yields:
|
||||
list: The splitted listed
|
||||
"""
|
||||
for i in range(0, len(a_list), segment_size):
|
||||
yield a_list[i : i + segment_size]
|
||||
|
||||
|
||||
def truncate_long_string(data: str, max_length: Optional[int] = None) -> str:
|
||||
"""Truncate arguments with representation longer than max_length
|
||||
|
||||
Args:
|
||||
data (str): The data to truncate
|
||||
max_length (Optional[int], optional): The max length. Defaults to None.
|
||||
|
||||
Returns:
|
||||
truncated (str): The truncated string
|
||||
"""
|
||||
if max_length is None:
|
||||
return data
|
||||
return (data[:max_length] + '...') if len(data) > max_length else data
|
||||
|
||||
|
||||
def get_call_string(
|
||||
func_name: Optional[str], args: Any, kwargs: Dict[Any, Any], max_length: Optional[int] = None
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Returns a string representation of the call, formatted as a regular
|
||||
Python function invocation statement. If max_length is not None, truncate
|
||||
arguments with representation longer than max_length.
|
||||
|
||||
Args:
|
||||
func_name (str): The funtion name
|
||||
args (Any): The function arguments
|
||||
kwargs (Dict[Any, Any]): The function kwargs
|
||||
max_length (int, optional): The max length. Defaults to None.
|
||||
|
||||
Returns:
|
||||
str: A String representation of the function call.
|
||||
"""
|
||||
if func_name is None:
|
||||
return None
|
||||
|
||||
arg_list = [as_text(truncate_long_string(repr(arg), max_length)) for arg in args]
|
||||
|
||||
list_kwargs = ['{0}={1}'.format(k, as_text(truncate_long_string(repr(v), max_length))) for k, v in kwargs.items()]
|
||||
arg_list += sorted(list_kwargs)
|
||||
args = ', '.join(arg_list)
|
||||
|
||||
return '{0}({1})'.format(func_name, args)
|
||||
|
||||
|
||||
def parse_names(queues_or_names: List[Union[str, 'Queue']]) -> List[str]:
|
||||
"""Given a list of strings or queues, returns queue names"""
|
||||
from .queue import Queue
|
||||
|
||||
names = []
|
||||
for queue_or_name in queues_or_names:
|
||||
if isinstance(queue_or_name, Queue):
|
||||
names.append(queue_or_name.name)
|
||||
else:
|
||||
names.append(str(queue_or_name))
|
||||
return names
|
||||
@@ -0,0 +1 @@
|
||||
VERSION = '1.16.2'
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,260 @@
|
||||
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, Type, Union
|
||||
from uuid import uuid4
|
||||
|
||||
from redis import ConnectionPool, Redis
|
||||
|
||||
from rq.serializers import DefaultSerializer
|
||||
|
||||
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._pool_class, self._pool_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 get_worker_process(
|
||||
self,
|
||||
name: str,
|
||||
burst: bool,
|
||||
_sleep: float = 0,
|
||||
logging_level: str = "INFO",
|
||||
) -> Process:
|
||||
"""Returns the worker process"""
|
||||
return Process(
|
||||
target=run_worker,
|
||||
args=(name, self._queue_names, self._connection_class, self._pool_class, self._pool_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})',
|
||||
)
|
||||
|
||||
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 = self.get_worker_process(name, burst=burst, _sleep=_sleep, logging_level=logging_level)
|
||||
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_pool_class,
|
||||
connection_pool_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_pool=ConnectionPool(connection_class=connection_pool_class, **connection_pool_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, with_scheduler=True, logging_level=logging_level)
|
||||
@@ -0,0 +1,104 @@
|
||||
from typing import TYPE_CHECKING, Optional, Set
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from redis import Redis
|
||||
from redis.client import Pipeline
|
||||
|
||||
from .queue import Queue
|
||||
from .worker import Worker
|
||||
|
||||
from rq.utils import split_list
|
||||
|
||||
from .utils import as_text
|
||||
|
||||
WORKERS_BY_QUEUE_KEY = 'rq:workers:%s'
|
||||
REDIS_WORKER_KEYS = 'rq:workers'
|
||||
MAX_KEYS = 1000
|
||||
|
||||
|
||||
def register(worker: 'Worker', pipeline: Optional['Pipeline'] = None):
|
||||
"""
|
||||
Store worker key in Redis so we can easily discover active workers.
|
||||
|
||||
Args:
|
||||
worker (Worker): The Worker
|
||||
pipeline (Optional[Pipeline], optional): The Redis Pipeline. Defaults to None.
|
||||
"""
|
||||
connection = pipeline if pipeline is not None else worker.connection
|
||||
connection.sadd(worker.redis_workers_keys, worker.key)
|
||||
for name in worker.queue_names():
|
||||
redis_key = WORKERS_BY_QUEUE_KEY % name
|
||||
connection.sadd(redis_key, worker.key)
|
||||
|
||||
|
||||
def unregister(worker: 'Worker', pipeline: Optional['Pipeline'] = None):
|
||||
"""Remove Worker key from Redis
|
||||
|
||||
Args:
|
||||
worker (Worker): The Worker
|
||||
pipeline (Optional[Pipeline], optional): Redis Pipeline. Defaults to None.
|
||||
"""
|
||||
if pipeline is None:
|
||||
connection = worker.connection.pipeline()
|
||||
else:
|
||||
connection = pipeline
|
||||
|
||||
connection.srem(worker.redis_workers_keys, worker.key)
|
||||
for name in worker.queue_names():
|
||||
redis_key = WORKERS_BY_QUEUE_KEY % name
|
||||
connection.srem(redis_key, worker.key)
|
||||
|
||||
if pipeline is None:
|
||||
connection.execute()
|
||||
|
||||
|
||||
def get_keys(queue: Optional['Queue'] = None, connection: Optional['Redis'] = None) -> Set[str]:
|
||||
"""Returns a list of worker keys for a given queue.
|
||||
|
||||
Args:
|
||||
queue (Optional['Queue'], optional): The Queue. Defaults to None.
|
||||
connection (Optional['Redis'], optional): The Redis Connection. Defaults to None.
|
||||
|
||||
Raises:
|
||||
ValueError: If no Queue or Connection is provided.
|
||||
|
||||
Returns:
|
||||
set: A Set with keys.
|
||||
"""
|
||||
if queue is None and connection is None:
|
||||
raise ValueError('"Queue" or "connection" argument is required')
|
||||
|
||||
if queue:
|
||||
redis = queue.connection
|
||||
redis_key = WORKERS_BY_QUEUE_KEY % queue.name
|
||||
else:
|
||||
redis = connection # type: ignore
|
||||
redis_key = REDIS_WORKER_KEYS
|
||||
|
||||
return {as_text(key) for key in redis.smembers(redis_key)}
|
||||
|
||||
|
||||
def clean_worker_registry(queue: 'Queue'):
|
||||
"""Delete invalid worker keys in registry.
|
||||
|
||||
Args:
|
||||
queue (Queue): The Queue
|
||||
"""
|
||||
keys = list(get_keys(queue))
|
||||
|
||||
with queue.connection.pipeline() as pipeline:
|
||||
for key in keys:
|
||||
pipeline.exists(key)
|
||||
results = pipeline.execute()
|
||||
|
||||
invalid_keys = []
|
||||
|
||||
for i, key_exists in enumerate(results):
|
||||
if not key_exists:
|
||||
invalid_keys.append(keys[i])
|
||||
|
||||
if invalid_keys:
|
||||
for invalid_subset in split_list(invalid_keys, MAX_KEYS):
|
||||
pipeline.srem(WORKERS_BY_QUEUE_KEY % queue.name, *invalid_subset)
|
||||
pipeline.srem(REDIS_WORKER_KEYS, *invalid_subset)
|
||||
pipeline.execute()
|
||||
Reference in New Issue
Block a user