1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
|
from .compat import as_text
WORKERS_BY_QUEUE_KEY = 'rq:workers:%s'
REDIS_WORKER_KEYS = 'rq:workers'
def register(worker, pipeline=None):
"""Store worker key in Redis so we can easily discover active workers."""
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, pipeline=None):
"""Remove worker key from Redis."""
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=None, connection=None):
"""Returnes a list of worker keys for a queue"""
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
redis_key = REDIS_WORKER_KEYS
return {as_text(key) for key in redis.smembers(redis_key)}
def clean_worker_registry(queue):
"""Delete invalid worker keys in registry"""
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:
pipeline.srem(WORKERS_BY_QUEUE_KEY % queue.name, *invalid_keys)
pipeline.srem(REDIS_WORKER_KEYS, *invalid_keys)
pipeline.execute()
|
