def register_worker(
self,
company_id: str,
user_id: str,
worker: str,
ip: str = "",
queues: Sequence[str] = None,
timeout: int = 0,
tags: Sequence[str] = None,
) -> WorkerEntry:
"""
Register a worker
:param company_id: worker's company ID
:param user_id: user ID under which this worker is running
:param worker: worker ID
:param ip: the real ip of the worker
:param queues: queues reported as being monitored by the worker
:param timeout: registration expiration timeout in seconds
:param tags: a list of tags for this worker
:raise bad_request.InvalidUserId: in case the calling user or company does not exist
:return: worker entry instance
"""
key = WorkerBLL._get_worker_key(company_id, user_id, worker)
timeout = timeout or DEFAULT_TIMEOUT
queues = queues or []
with translate_errors_context():
query = dict(id=user_id, company=company_id)
user = User.objects(**query).only("id", "name").first()
if not user:
raise bad_request.InvalidUserId(**query)
company = Company.objects(id=company_id).only("id", "name").first()
if not company:
raise server_error.InternalError("invalid company", company=company_id)
queue_objs = Queue.objects(company=company_id, id__in=queues).only("id")
if len(queue_objs) < len(queues):
invalid = set(queues).difference(q.id for q in queue_objs)
raise bad_request.InvalidQueueId(ids=invalid)
now = datetime.utcnow()
entry = WorkerEntry(
key=key,
id=worker,
user=user.to_proper_dict(),
company=company.to_proper_dict(),
ip=ip,
queues=queues,
register_time=now,
register_timeout=timeout,
last_activity_time=now,
tags=tags,
)
self.redis.setex(key, timedelta(seconds=timeout), entry.to_json())
return entry
def _get(
self, company: str, user: str = "*", worker_id: str = "*"
) -> Sequence[WorkerEntry]:
"""Get worker entries matching the company and user, worker patterns"""
match = self._get_worker_key(company, user, worker_id)
with TimingContext("redis", "workers_get_all"):
res = self.redis.scan_iter(match)
return [WorkerEntry.from_json(self.redis.get(r)) for r in res]
def _save_worker(self, entry: WorkerEntry) -> None:
"""Save worker entry in Redis"""
try:
self.redis.setex(
entry.key, timedelta(seconds=entry.register_timeout), entry.to_json()
)
except Exception:
msg = "Failed saving worker entry"
log.exception(msg)
def from_worker_entry(cls, worker: WorkerEntry):
data = worker.to_struct()
queue = data.pop("queue", None) or None
queue_ids = set(data.pop("queues", []))
queues = [QueueEntry(id=id) for id in queue_ids]
if queue:
queue = next((q for q in queues if q.id == queue), None)
return cls(
worker=WorkerResponseEntry(queues=queues, queue=queue, **data),
task_id=worker.task.id if worker.task else None,
queue_ids=queue_ids,
)
def _get_worker(self, company_id: str, user_id: str, worker: str) -> WorkerEntry:
"""
Get a worker entry for the provided worker ID. The entry is loaded from Redis
if it exists (i.e. worker has already been registered), otherwise the worker
is registered and its entry stored into Redis).
:param company_id: worker's company ID
:param user_id: user ID under which this worker is running
:param worker: worker ID
:raise bad_request.InvalidWorkerId: in case the worker id was not found
:return: worker entry instance
"""
key = self._get_worker_key(company_id, user_id, worker)
with TimingContext("redis", "get_worker"):
data = self.redis.get(key)
if data:
try:
entry = WorkerEntry.from_json(data)
if not entry.key:
entry.key = key
self._save_worker(entry)
return entry
except Exception as e:
msg = "Failed parsing worker entry"
log.exception(msg)
raise server_error.DataError(msg, err=e.args[0])
# Failed loading worker from Redis
if config.get("apiserver.workers.auto_register", False):
try:
return self.register_worker(company_id, user_id, worker)
except Exception:
log.error(
"Failed auto registration of {} for company {}".format(
worker, company_id
)
)
raise bad_request.InvalidWorkerId(worker=worker)