def _maybe_pubsub_notify_via_tq(result_summary, request):
"""Examines result_summary and enqueues a task to send PubSub message.
Must be called within a transaction.
Raises CommitError on errors (to abort the transaction).
"""
assert ndb.in_transaction()
assert isinstance(result_summary,
task_result.TaskResultSummary), result_summary
assert isinstance(request, task_request.TaskRequest), request
if (result_summary.state in task_result.State.STATES_NOT_RUNNING
and request.pubsub_topic):
task_id = task_pack.pack_result_summary_key(result_summary.key)
ok = utils.enqueue_task(url='/internal/taskqueue/pubsub/%s' % task_id,
queue_name='pubsub',
transactional=True,
payload=utils.encode_to_json({
'task_id':
task_id,
'topic':
request.pubsub_topic,
'auth_token':
request.pubsub_auth_token,
'userdata':
request.pubsub_userdata,
}))
if not ok:
raise datastore_utils.CommitError(
'Failed to enqueue task queue task')
def notify_requests(es_cfg, requests, use_tq, is_callback, batch_mode=False):
"""Calls external scheduler to notify it of a task state.
Arguments:
- es_cfg: pools_config.ExternalSchedulerConfig for external scheduler to
notify.
- requests:
A list of (task_request.TaskRequest,
task_result.TaskResultSummary or task_result.TaskRunResult)
tuples.
- use_tq: If true, make this call on a task queue (within the current
datastore transaction).
- is_callback: If true, indicates that this notification was in response
to a external-scheduler-requested callback. This is for
- batch_mode: If true, the notifications will be sent in a batched mode
along with others, to reduce traffic to external scheduler.
Only valid when use_tq and global config's
enable_batch_es_notifications are true.
Returns: Nothing.
"""
logging.debug(
'notify_requests(es_cfg=(%s,%s), requests=%s, use_tq=%s, '
'is_callback=%s, batch_mode=%s)', es_cfg.address, es_cfg.id,
[r.task_id for r, _ in requests], use_tq, is_callback, batch_mode)
req = plugin_pb2.NotifyTasksRequest()
req.is_callback = is_callback
for request, result_summary in requests:
item = req.notifications.add()
# TODO(akeshet): This time should possibly come from the read time from
# datastore, rather than the local server clock.
item.time.FromDatetime(utils.utcnow())
item.task.id = request.task_id
item.task.tags.extend(request.tags)
item.task.enqueued_time.FromDatetime(request.created_ts)
for i in range(request.num_task_slices):
s = request.task_slice(i)
flat_dimensions = task_queues.bot_dimensions_to_flat(
s.properties.dimensions)
s_pb = item.task.slices.add()
s_pb.dimensions.extend(flat_dimensions)
res = swarming_pb2.TaskResult()
result_summary.to_proto(res)
item.task.state = res.state
if result_summary.bot_id:
# TODO(akeshet): We should only actually set this is state is running.
item.task.bot_id = result_summary.bot_id
req.scheduler_id = es_cfg.id
if not use_tq:
# Ignore return value, the response proto is empty.
notify_request_now(es_cfg.address, req)
return
request_json = json_format.MessageToJson(req)
# If enable_batch_es_notifications is true, the notifications will be sent in
# a batched mode along with others, to reduce traffic to external scheduler.
if batch_mode and config.settings().enable_batch_es_notifications:
payload = {'es_host': es_cfg.address, 'request_json': request_json}
req = taskqueue.Task(payload=json.dumps(payload), method='PULL')
if not req.add(queue_name='es-notify-tasks-batch',
transactional=ndb.in_transaction()):
raise datastore_utils.CommitError('Failed to enqueue task')
stats = taskqueue.QueueStatistics.fetch('es-notify-kick')
# Add a kicker task if there are fewer than 10 minutes worth.
if stats.tasks < 600:
job_enqueued = utils.enqueue_task(
'/internal/taskqueue/important/external_scheduler/notify-kick',
'es-notify-kick',
transactional=ndb.in_transaction())
if not job_enqueued:
logging.info('Failed to add a notify-kick for request.')
return
enqueued = utils.enqueue_task(
'/internal/taskqueue/important/external_scheduler/notify-tasks',
'es-notify-tasks',
params={
'es_host': es_cfg.address,
'request_json': request_json
},
transactional=ndb.in_transaction())
if not enqueued:
raise datastore_utils.CommitError('Failed to enqueue task')
def r(*_):
raise datastore_utils.CommitError('Sorry!')
