swarming: add transaction_id to tasks.new request

appengine/swarming/handlers_endpoints.py

 # Copyright 2015 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """This module defines Swarming Server endpoints handlers."""
 
 import logging
 import os
 
 from google.appengine.api import datastore_errors
@@ skipping 32 matching lines @@
 _old_decode_field = protojson.ProtoJson.decode_field
 def _decode_field(self, field, value):
   if (isinstance(field, messages.BooleanField) and
       isinstance(value, basestring)):
     return value.lower() == 'true'
   return _old_decode_field(self, field, value)
 protojson.ProtoJson.decode_field = _decode_field
 
 
 def get_request_and_result(task_id):
-  """Provides the key and TaskRequest corresponding to a task ID.
+  """Provides the TaskRequest and result corresponding to a task ID.
 
   Enforces the ACL for users. Allows bots all access for the moment.
 
   Returns:
     tuple(TaskRequest, result): result can be either for a TaskRunResult or a
-                                TaskResultSummay.
+                                TaskResultSummary.
   """
   try:
     request_key, result_key = task_pack.get_request_and_result_keys(task_id)
     request, result = ndb.get_multi((request_key, result_key))
     if not request or not result:
       raise endpoints.NotFoundException('%s not found.' % task_id)
     if not acl.is_bot() and not request.has_access:
       raise endpoints.ForbiddenException('%s is not accessible.' % task_id)
     return request, result
   except ValueError:
     raise endpoints.BadRequestException('%s is an invalid key.' % task_id)
 
 
+def get_request_and_result_by_transaction_id(transaction_id):
+  """Provides the TaskRequest and result corresponding to a transaction ID.
+
+  Enforces the ACL for users. Allows bots all access for the moment.
+
+  May return (None, None) if the transaction, request or result is not found.
+
+  Returns:
+    tuple(TaskRequest, result): result can be either for a TaskRunResult or a
+                                TaskResultSummary.
+  """
+  task_id = memcache.get(transaction_id_to_memcache_key(transaction_id))

[M-A Ruel, 2017/05/03 01:33:20]

use a namespace, it's simpler. Then no need for
transaction_id_to_memcache_key().

[nodir, 2017/05/05 18:11:38]

Done.

+  if not task_id:
+    return None, None

[M-A Ruel, 2017/05/03 01:33:20]

This makes the contract a bit scary though. :/ I mean, memcache can be flushed
anytime. I guess it's fine for the general case.

For example, using a duplicate transaction id once per day may work when
memcache is small, but fail when used at a slightly higher frequency.

Technically speaking the "right" thing is to do an DB get but it's much more
expensive, since on steady state it is a cold (uncached) miss.

Overall, I'm fine, it just needs to be documented with a lot of caveats.

[nodir, 2017/05/05 18:11:38]

maybe set timeout to 1min?

+  request_key, result_key = task_pack.get_request_and_result_keys(task_id)
+  request, result = ndb.get_multi((request_key, result_key))
+  if not request or not result:
+    return None, None
+  if not acl.is_bot() and not request.has_access:
+    raise endpoints.ForbiddenException(
+        '%s, created for transaction %s, is not accessible.' %
+        (task_id, transaction_id))
+  return request, result
+
+
+def transaction_id_to_memcache_key(transaction_id):
+  assert transaction_id
+  return 'tasks.new/txn/%s' % transaction_id
+
+
 def get_or_raise(key):
   """Returns an entity or raises an endpoints exception if it does not exist."""
   result = key.get()
   if not result:
     raise endpoints.NotFoundException('%s not found.' % key.id())
   return result
 
 
 def apply_property_defaults(properties):
   """Fills ndb task properties with default values read from server settings."""
@@ skipping 169 matching lines @@
         who=obj.who.to_bytes() if obj.who else None,
         when=obj.created_ts,
         version=str(obj.version))
 
 
 TaskId = endpoints.ResourceContainer(
     message_types.VoidMessage,
     task_id=messages.StringField(1, required=True))
 
 
+TaskIdWithTransactionID = endpoints.ResourceContainer(
+    message_types.VoidMessage,
+    task_id=messages.StringField(1),
+    transaction_id=messages.StringField(2))
+
+
 TaskIdWithPerf = endpoints.ResourceContainer(
     message_types.VoidMessage,
     task_id=messages.StringField(1, required=True),
     include_performance_stats=messages.BooleanField(2, default=False))
 
 
 @swarming_api.api_class(resource_name='task', path='task')
 class SwarmingTaskService(remote.Service):
   """Swarming's task-related API."""
   @gae_ts_mon.instrument_endpoint()
@@ skipping 25 matching lines @@
       http_method='GET')
   @auth.require(acl.is_bot_or_user)
   def request(self, request):
     """Returns the task request corresponding to a task ID."""
     logging.info('%s', request)
     request_obj, _ = get_request_and_result(request.task_id)
     return message_conversion.task_request_to_rpc(request_obj)
 
   @gae_ts_mon.instrument_endpoint()
   @auth.endpoints_method(
-      TaskId, swarming_rpcs.CancelResponse,
+      TaskIdWithTransactionID, swarming_rpcs.CancelResponse,
       name='cancel',
       path='{task_id}/cancel')
   @auth.require(acl.is_bot_or_user)
   def cancel(self, request):
     """Cancels a task.
 
     If a bot was running the task, the bot will forcibly cancel the task.
     """
     logging.info('%s', request)
-    request_obj, result = get_request_and_result(request.task_id)
+    if request.task_id and request.transaction_id:
+      raise endpoints.BadRequestException(
+          'only one of task_id or transaction_id must be specified')

[M-A Ruel, 2017/05/03 01:33:20]

That can't work because path is '{task_id}/cancel'. The test works because it
bypasses the URI router.

[nodir, 2017/05/05 18:11:38]

ugh, should I define a new API then? cancel_by_transaction_id or something?
the problem we are trying to solve here is not knowing the task id and passing
dummy task id is not great

[M-A Ruel, 2017/05/05 18:15:33]

Yes I prefer a new API. If we get rid of the feature then it's simpler to delete
a single method. :)

[nodir, 2017/05/05 18:41:37]

Done.

+    if request.task_id:
+      request_obj, result = get_request_and_result(request.task_id)
+    else:
+      request_obj, result = get_request_and_result_by_transaction_id(
+          request.transaction_id)
+      if not request_obj or not result:
+        raise endpoints.NotFoundException(
+            'request for transaction %s not found.' % request.transaction_id)
+
     ok, was_running = task_scheduler.cancel_task(request_obj, result.key)
     return swarming_rpcs.CancelResponse(ok=ok, was_running=was_running)
 
   @gae_ts_mon.instrument_endpoint()
   @auth.endpoints_method(
       TaskId, swarming_rpcs.TaskOutput,
       name='stdout',
       path='{task_id}/stdout',
       http_method='GET')
   @auth.require(acl.is_bot_or_user)
@@ skipping 21 matching lines @@
   def new(self, request):
     """Creates a new task.
 
     The task will be enqueued in the tasks list and will be executed at the
     earliest opportunity by a bot that has at least the dimensions as described
     in the task request.
     """
     sb = (request.properties.secret_bytes
           if request.properties is not None else None)
     if sb is not None:
-      request.properties.secret_bytes = "HIDDEN"
+      request.properties.secret_bytes = 'HIDDEN'
     logging.info('%s', request)
     if sb is not None:
       request.properties.secret_bytes = sb
 
     try:
-      request, secret_bytes = message_conversion.new_task_request_from_rpc(
+      request_obj, secret_bytes = message_conversion.new_task_request_from_rpc(
           request, utils.utcnow())
-      apply_property_defaults(request.properties)
+      apply_property_defaults(request_obj.properties)
       task_request.init_new_request(
-          request, acl.can_schedule_high_priority_tasks(), secret_bytes)
+          request_obj, acl.can_schedule_high_priority_tasks(), secret_bytes)
+    except (TypeError, ValueError) as e:
+      raise endpoints.BadRequestException(e.message)
 
-      result_summary = task_scheduler.schedule_request(request, secret_bytes)
-    except (datastore_errors.BadValueError, TypeError, ValueError) as e:
+    # Check if a request was already created with the same transaction.
+    transaction_memcache_key = None
+    if request.transaction_id:
+      transaction_memcache_key = transaction_id_to_memcache_key(
+          request.transaction_id)
+      txn_request_obj, txn_result_summary = (
+          get_request_and_result_by_transaction_id(request.transaction_id))
+      if txn_request_obj and txn_result_summary:
+        todict = lambda r: r.to_dict(['created_ts', 'expiration_ts'])
+        if todict(txn_request_obj) != todict(request_obj):
+          raise endpoints.BadRequestException(
+              'A request with different parameters was created for '
+              'transaction %s. Send the same tasks.new request '
+              'within the same transaction!' % request.transaction_id)
+        task_id = task_pack.pack_result_summary_key(txn_result_summary.key)
+        logging.info(
+            'reusing existing %s for transaction %s',
+            task_id, request.transaction_id)
+        return swarming_rpcs.TaskRequestMetadata(
+            request=message_conversion.task_request_to_rpc(txn_request_obj),
+            task_id=task_id,
+            task_result=message_conversion.task_result_to_rpc(
+                txn_result_summary, False))
+
+    try:
+      result_summary = task_scheduler.schedule_request(
+          request_obj, secret_bytes,
+          memcache_key_for_task_id=transaction_memcache_key)
+    except datastore_errors.BadValueError as e:
       raise endpoints.BadRequestException(e.message)
 
     previous_result = None
     if result_summary.deduped_from:
       previous_result = message_conversion.task_result_to_rpc(
           result_summary, False)
 
     return swarming_rpcs.TaskRequestMetadata(
-        request=message_conversion.task_request_to_rpc(request),
+        request=message_conversion.task_request_to_rpc(request_obj),
         task_id=task_pack.pack_result_summary_key(result_summary.key),
         task_result=previous_result)
 
   @gae_ts_mon.instrument_endpoint()
   @auth.endpoints_method(
       swarming_rpcs.TasksRequest, swarming_rpcs.TaskList,
       http_method='GET')
   @auth.require(acl.is_privileged_user)
   def list(self, request):
     """Returns tasks results based on the filters.
@@ skipping 442 matching lines @@
 def get_routes():
   return (
     endpoints_webapp2.api_routes(SwarmingServerService) +
     endpoints_webapp2.api_routes(SwarmingTaskService) +
     endpoints_webapp2.api_routes(SwarmingTasksService) +
     endpoints_webapp2.api_routes(SwarmingBotService) +
     endpoints_webapp2.api_routes(SwarmingBotsService) +
     # components.config endpoints for validation and configuring of luci-config
     # service URL.
     endpoints_webapp2.api_routes(config.ConfigApi))