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 18 matching lines @@
 from server import acl
 from server import bot_code
 from server import bot_management
 from server import config
 from server import task_pack
 from server import task_request
 from server import task_result
 from server import task_scheduler
 
 
+TRANSACTION_ID_MEMCACHE_NAMESPACE = 'tasks.new/transaction_id'
+
+
 ### Helper Methods
 
 
 # Add support for BooleanField in protorpc in endpoints GET requests.
 _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, namespace=TRANSACTION_ID_MEMCACHE_NAMESPACE)
+  if not task_id:
+    return None, None
+  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 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))
 
 
+TransactionId = endpoints.ResourceContainer(
+    message_types.VoidMessage,
+    transaction_id=messages.StringField(1))
+
+
 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 77 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.
+    if 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=request.transaction_id,
+          memcache_ns_for_task_id=TRANSACTION_ID_MEMCACHE_NAMESPACE)
+    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 104 matching lines @@
     else:
       logging.info('No tasks to cancel.')
 
     return swarming_rpcs.TasksCancelResponse(
         cursor=cursor,
         matched=len(tasks),
         now=now)
 
   @gae_ts_mon.instrument_endpoint()
   @auth.endpoints_method(
+      TransactionId, swarming_rpcs.CancelResponse,
+      name='cancel_by_transaction_id',
+      path='cancel_by_transaction_id/{transaction_id}')
+  @auth.require(acl.is_bot_or_user)
+  def cancel_by_transaction_id(self, request):
+    """Cancels a task by a transaction ID.
+
+    Useful when id of a new task is unknown because the request failed.
+    The transaction ID here is the one supplied by the user when creating the
+    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_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(
+        task_id=request_obj.task_id,
+        ok=ok,
+        was_running=was_running)
+
+  @gae_ts_mon.instrument_endpoint()
+  @auth.endpoints_method(
       swarming_rpcs.TasksCountRequest, swarming_rpcs.TasksCount,
       http_method='GET')
   @auth.require(acl.is_privileged_user)
   def count(self, request):
     """Counts number of tasks in a given state."""
     logging.info('%s', request)
     if not request.start:
       raise endpoints.BadRequestException('start (as epoch) is required')
     now = utils.utcnow()
     mem_key = self._memcache_key(request, now)
@@ skipping 318 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))