task_queues: Add more scaffolding

appengine/swarming/handlers_backend.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.
 
 """Main entry point for Swarming backend handlers."""
 
 import datetime
 import json
 import logging
 
 import webapp2
 from google.appengine.api import datastore_errors
 from google.appengine.ext import ndb
 
 from components import utils
 
 import mapreduce_jobs
 from components import decorators
 from components import datastore_utils
 from components import machine_provider
 from server import bot_management
 from server import config
 from server import lease_management
 from server import stats
 from server import task_pack
+from server import task_queues
 from server import task_result
 from server import task_scheduler
 import ts_mon_metrics
 
 
 class CronBotDiedHandler(webapp2.RequestHandler):
   @decorators.require_cronjob
   def get(self):
     try:
       task_scheduler.cron_handle_bot_died(self.request.host_url)
     except datastore_errors.NeedIndexError as e:
       # When a fresh new instance is deployed, it takes a few minutes for the
       # composite indexes to be created even if they are empty. Ignore the case
       # where the index is defined but still being created by AppEngine.
       if not str(e).startswith(
           'NeedIndexError: The index for this query is not ready to serve.'):
         raise
     self.response.headers['Content-Type'] = 'text/plain; charset=utf-8'
     self.response.out.write('Success.')
 
 
 class CronAbortExpiredShardToRunHandler(webapp2.RequestHandler):
   @decorators.require_cronjob
   def get(self):
     task_scheduler.cron_abort_expired_task_to_run(self.request.host_url)
     self.response.headers['Content-Type'] = 'text/plain; charset=utf-8'
     self.response.out.write('Success.')
 
 
+class CronTaskQueues(webapp2.RequestHandler):
+  @decorators.require_cronjob
+  def get(self):
+    task_queues.tidy_stale()
+    self.response.headers['Content-Type'] = 'text/plain; charset=utf-8'
+    self.response.out.write('Success.')
+
+
 class CronMachineProviderBotsUtilizationHandler(webapp2.RequestHandler):
   """Determines Machine Provider bot utilization."""
 
   @decorators.require_cronjob
   def get(self):
     if not config.settings().mp.enabled:
       logging.info('MP support is disabled')
       return
 
     lease_management.compute_utilization()
@@ skipping 111 matching lines @@
         continue
       request_obj = request_key.get()
       if not request_obj:
         logging.error('Request for %s was not found.', request_key.id())
         continue
       ok, was_running = task_scheduler.cancel_task(request_obj, result_key)
       logging.info('task %s canceled: %s was running: %s',
                    task_id, ok, was_running)
 
 
+class TaskDimensionsHandler(webapp2.RequestHandler):
+  @decorators.require_taskqueue('task-dimensions')
+  def post(self):
+    self.tidy_stale(self.request.body)
+    self.response.headers['Content-Type'] = 'text/plain; charset=utf-8'
+    self.response.out.write('Success.')
+
+  @staticmethod
+  def tidy_stale(body):
+    payload = body.split('\n')
+    dimensions_hash = int(payload[0])
+    dimensions_flat = payload[1:]
+    task_queues.rebuild_task_cache(dimensions_hash, dimensions_flat)
+
+
 class TaskSendPubSubMessage(webapp2.RequestHandler):
   """Sends PubSub notification about task completion."""
 
   # Add task_id to the URL for better visibility in request logs.
   @decorators.require_taskqueue('pubsub')
   def post(self, task_id):  # pylint: disable=unused-argument
     task_scheduler.task_handle_pubsub_task(json.loads(self.request.body))
     self.response.headers['Content-Type'] = 'text/plain; charset=utf-8'
     self.response.out.write('Success.')
 
@@ skipping 32 matching lines @@
     mapreduce_jobs.launch_job(job_id)
 
 
 ###
 
 
 def get_routes():
   """Returns internal urls that should only be accessible via the backend."""
   routes = [
     # Cron jobs.
-    # TODO(maruel): Rename cron.yaml job url. Doing so is a bit annoying since
-    # the app version has to be running an already compatible version already.
     ('/internal/cron/abort_bot_died', CronBotDiedHandler),
     ('/internal/cron/handle_bot_died', CronBotDiedHandler),
     ('/internal/cron/abort_expired_task_to_run',
         CronAbortExpiredShardToRunHandler),
+    ('/internal/cron/task_queues_tidy', CronTaskQueues),
 
     ('/internal/cron/stats/update', stats.InternalStatsUpdateHandler),
     ('/internal/cron/aggregate_bots_dimensions',
         CronBotsDimensionAggregationHandler),
     ('/internal/cron/aggregate_tasks_tags',
         CronTasksTagsAggregationHandler),
 
     # Machine Provider.
     ('/internal/cron/machine_provider_bot_usage',
         CronMachineProviderBotsUtilizationHandler),
     ('/internal/cron/machine_provider_config',
         CronMachineProviderConfigHandler),
     ('/internal/cron/machine_provider_manage',
         CronMachineProviderManagementHandler),
 
     # Task queues.
     ('/internal/taskqueue/cancel-tasks', CancelTasksHandler),
+    ('/internal/taskqueue/task-dimensions', TaskDimensionsHandler),
     (r'/internal/taskqueue/pubsub/<task_id:[0-9a-f]+>', TaskSendPubSubMessage),
     ('/internal/taskqueue/machine-provider-manage',
         TaskMachineProviderManagementHandler),
     (r'/internal/taskqueue/tsmon/<kind:[0-9A-Za-z_]+>', TaskGlobalMetrics),
 
     # Mapreduce related urls.
     (r'/internal/taskqueue/mapreduce/launch/<job_id:[^\/]+>',
       InternalLaunchMapReduceJobWorkerHandler),
   ]
   return [webapp2.Route(*a) for a in routes]
 
 
 def create_application(debug):
   return webapp2.WSGIApplication(get_routes(), debug=debug)