swarming: add transaction_id to tasks.new request

appengine/swarming/server/task_scheduler.py

 # Copyright 2014 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.
 
 """High level tasks execution scheduling API.
 
 This is the interface closest to the HTTP handlers.
 """
 
 import datetime
 import logging
 import math
 import random
 
+from google.appengine.api import memcache
 from google.appengine.ext import ndb
 
 from components import auth
 from components import datastore_utils
 from components import pubsub
 from components import utils
 
 import event_mon_metrics
 import ts_mon_metrics
 
@@ skipping 442 matching lines @@
 
   # If the user provided a max then use it, otherwise
   # use default
   max_wait = config.settings().max_bot_sleep_time
   if not max_wait:
     # Enforces more frequent polls on staging.
     max_wait = 3. if utils.is_dev() else 60.
   return min(max_wait, math.pow(1.5, min(attempt_num, 10) + 1))
 
 
-def schedule_request(request, secret_bytes, check_acls=True):
+def schedule_request(
+    request, secret_bytes, check_acls=True,
+    memcache_key_for_task_id=None, memcache_ns_for_task_id=None):
   """Creates and stores all the entities to schedule a new task request.
 
   Checks ACLs first. Raises auth.AuthorizationError if caller is not authorized
   to post this request.
 
   The number of entities created is 3: TaskRequest, TaskToRun and
-  TaskResultSummary.
+  TaskResultSummary. If memcache_key_for_task_id is not None, sets it to the
+  task_id right after creation.
 
   All 4 entities in the same entity group (TaskReqest, TaskToRun,
   TaskResultSummary, SecretBytes) are saved as a DB transaction.
 
   Arguments:
   - request: TaskRequest entity to be saved in the DB. It's key must not be set
              and the entity must not be saved in the DB yet.
   - secret_bytes: SecretBytes entity to be saved in the DB. It's key will be set
              and the entity will be stored by this function. None is allowed if
              there are no SecretBytes for this task.
@@ skipping 64 matching lines @@
           'New request %s reusing %s', result_summary.task_id,
           dupe_summary.task_id)
       deduped = True
 
   if not deduped:
     # Storing these entities makes this task live. It is important at this point
     # that the HTTP handler returns as fast as possible, otherwise the task will
     # be run but the client will not know about it.
     datastore_utils.insert(request, get_new_keys,
         extra=filter(bool, [task, result_summary, secret_bytes]))
+    if memcache_key_for_task_id:
+      memcache.set(
+          memcache_key_for_task_id, result_summary.task_id,
+          namespace=memcache_ns_for_task_id)
     logging.debug('New request %s', result_summary.task_id)
 
   # Get parent task details if applicable.
   if request.parent_task_id:
     parent_run_key = task_pack.unpack_run_result_key(request.parent_task_id)
     parent_task_keys = [
       parent_run_key,
       task_pack.run_result_key_to_result_summary_key(parent_run_key),
     ]
 
@@ skipping 429 matching lines @@
 ## Task queue tasks.
 
 
 def task_handle_pubsub_task(payload):
   """Handles task enqueued by _maybe_pubsub_notify_via_tq."""
   # Do not catch errors to trigger task queue task retry. Errors should not
   # happen in normal case.
   _pubsub_notify(
       payload['task_id'], payload['topic'],
       payload['auth_token'], payload['userdata'])