blob: ee4d82629310c5195391e3ed85ec1955f7212e64 [file] [log] [blame]
// Copyright 2018 The LUCI Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.
// Package backend includes cron and task queue handlers.
package backend
import (
// dspKey is the key to a *tq.Dispatcher in the context.
var dspKey = "dsp"
// withDispatcher returns a new context with the given *tq.Dispatcher installed.
func withDispatcher(c context.Context, dsp *tq.Dispatcher) context.Context {
return context.WithValue(c, &dspKey, dsp)
// getDispatcher returns the *tq.Dispatcher installed in the current context.
func getDispatcher(c context.Context) *tq.Dispatcher {
return c.Value(&dspKey).(*tq.Dispatcher)
// registerTasks registers task handlers with the given *tq.Dispatcher.
func registerTasks(dsp *tq.Dispatcher) {
dsp.RegisterTask(&tasks.CountVMs{}, countVMs, countVMsQueue, nil)
dsp.RegisterTask(&tasks.CreateInstance{}, createInstance, createInstanceQueue, nil)
dsp.RegisterTask(&tasks.CreateVM{}, createVM, createVMQueue, nil)
dsp.RegisterTask(&tasks.DeleteBot{}, deleteBot, deleteBotQueue, nil)
dsp.RegisterTask(&tasks.DestroyInstance{}, destroyInstance, destroyInstanceQueue, nil)
dsp.RegisterTask(&tasks.ExpandConfig{}, expandConfig, expandConfigQueue, nil)
dsp.RegisterTask(&tasks.ManageBot{}, manageBot, manageBotQueue, nil)
dsp.RegisterTask(&tasks.ReportQuota{}, reportQuota, reportQuotaQueue, nil)
dsp.RegisterTask(&tasks.TerminateBot{}, terminateBot, terminateBotQueue, nil)
// gceKey is the key to a *compute.Service in the context.
var gceKey = "gce"
// withCompute returns a new context with the given *compute.Service installed.
func withCompute(c context.Context, gce *compute.Service) context.Context {
return context.WithValue(c, &gceKey, gce)
// getCompute returns the *compute.Service installed in the current context.
func getCompute(c context.Context) *compute.Service {
return c.Value(&gceKey).(*compute.Service)
// newCompute returns a new *compute.Service. Panics on error.
func newCompute(c context.Context) *compute.Service {
t, err := auth.GetRPCTransport(c, auth.AsSelf, auth.WithScopes(compute.ComputeScope))
if err != nil {
gce, err := compute.New(&http.Client{Transport: t})
if err != nil {
return gce
// swrKey is the key to a *swarming.Service in the context.
var swrKey = "swr"
// withSwarming returns a new context with the given *swarming.Service installed.
func withSwarming(c context.Context, swr *swarming.Service) context.Context {
return context.WithValue(c, &swrKey, swr)
// getSwarming returns the *swarming.Service installed in the current context.
func getSwarming(c context.Context, url string) *swarming.Service {
swr := c.Value(&swrKey).(*swarming.Service)
swr.BasePath = url + "/_ah/api/swarming/v1/"
return swr
// newSwarming returns a new *swarming.Service. Panics on error.
func newSwarming(c context.Context) *swarming.Service {
t, err := auth.GetRPCTransport(c, auth.AsSelf)
if err != nil {
swr, err := swarming.New(&http.Client{Transport: t})
if err != nil {
return swr
// InstallHandlers installs HTTP request handlers into the given router.
func InstallHandlers(r *router.Router, mw router.MiddlewareChain) {
dsp := &tq.Dispatcher{}
mw = mw.Extend(func(c *router.Context, next router.Handler) {
var cancel context.CancelFunc
c.Context, cancel = context.WithTimeout(c.Context, 30*time.Second)
defer cancel()
c.Context = withDispatcher(c.Context, dsp)
c.Context = withCompute(c.Context, newCompute(c.Context))
c.Context = withSwarming(c.Context, newSwarming(c.Context))
dsp.InstallRoutes(r, mw)
r.GET("/internal/cron/count-tasks", mw, newHTTPHandler(countTasks))
r.GET("/internal/cron/count-vms", mw, newHTTPHandler(countVMsAsync))
r.GET("/internal/cron/create-instances", mw, newHTTPHandler(createInstancesAsync))
r.GET("/internal/cron/expand-configs", mw, newHTTPHandler(expandConfigsAsync))
r.GET("/internal/cron/manage-bots", mw, newHTTPHandler(manageBotsAsync))
r.GET("/internal/cron/report-quota", mw, newHTTPHandler(reportQuotasAsync))