blob: 0bd0281d9b57f009850aff56af71416e5c634b64 [file] [log] [blame]
// Copyright 2022 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 quota
import (
pb ""
// ModuleName is the globally-unique name for this module.
// Useful for registering this module as a dependency of other modules.
var ModuleName = module.RegisterName("")
// Ensure quotaModule implements server.Module at compile-time.
var _ module.Module = &quotaModule{}
// quotaModule implements module.Module.
type quotaModule struct {
opts *ModuleOptions
// Dependencies returns required and optional dependencies for this module.
// Implements module.Module.
func (*quotaModule) Dependencies() []module.Dependency {
return []module.Dependency{
// Initialize initializes this module by ensuring a quotaconfig.Interface
// implementation is available in the serving context, and installing cron
// routes if configured (see ModuleOptions).
// Implements module.Module.
func (m *quotaModule) Initialize(ctx context.Context, host module.Host, opts module.HostOptions) (context.Context, error) {
// server.Main calls Initialize before calling the user-provided callback where
// users of the quota library are likely to provide a quotaconfig.Interface.
// Use a warmup callback, which is called after the user-provided callback has
// been executed, to run initialization logic requiring quotaconfig.Interface.
host.RegisterWarmup(func(ctx context.Context) {
// getInterface panics if a quotaconfig.Interface isn't available (see Use).
// Intentionally force a panic if one isn't available.
cfg := getInterface(ctx)
// Register the cron handler if specified.
if m.opts.ConfigRefreshCronHandlerID != "" {
cron.RegisterHandler(m.opts.ConfigRefreshCronHandlerID, func(ctx context.Context) error {
// Fetch the quotaconfig.Interface from the context each time
// in case the implementation being used changes later on.
return getInterface(ctx).Refresh(ctx)
// Attempt to call Refresh so policy configs are available before serving.
// This is best-effort, users should ensure they're calling Refresh regularly
// (either manually, or by setting m.opts.ConfigRefreshCronHandlerID).
_ = cfg.Refresh(ctx)
if m.opts.AdminServiceReaders != "" && m.opts.AdminServiceWriters != "" {
pb.RegisterQuotaAdminServer(host.ServiceRegistrar(), NewQuotaAdminServer(m.opts.AdminServiceReaders, m.opts.AdminServiceWriters))
return ctx, nil
// Name returns the module.Name for this module.
// Implements module.Module.
func (*quotaModule) Name() module.Name {
return ModuleName
// ModuleOptions is a set of configuration options for the quota module.
type ModuleOptions struct {
// AdminServerReaders is a Chrome Infra Auth group authorized to use read-only
// methods of the quota admin pRPC service. If unspecified, the service will
// not be exposed.
AdminServiceReaders string
// AdminServerWriters is a Chrome Infra Auth group authorized to use all
// methods of the quota admin pRPC service. If unspecified, the service will
// not be exposed.
AdminServiceWriters string
// ConfigRefreshCronHandlerID is the ID for this module's config refresh
// handler. If specified, the module ensures Refresh is called on the
// quotaconfig.Interface in the server context. The handler will be installed
// at <serving-prefix>/<ID>. The serving prefix is controlled by server/cron.
// If unspecified, module users should refresh policy configs periodically by
// manually calling Refresh.
ConfigRefreshCronHandlerID string
// Register adds command line flags for these module options to the given
// *flag.FlagSet. Mutates module options by initializing defaults.
func (o *ModuleOptions) Register(f *flag.FlagSet) {
f.StringVar(&o.AdminServiceReaders, "quota-admin-service-readers", "", "Chrome Infra Auth group authorized to use read-only admin methods.")
f.StringVar(&o.AdminServiceWriters, "quota-admin-service-writers", "", "Chrome Infra Auth group authorized to use all admin methods.")
f.StringVar(&o.ConfigRefreshCronHandlerID, "quota-cron-handler-id", "", "Config refresh handler ID.")
// NewModule returns a module.Module for the quota library initialized from the
// given *ModuleOptions.
func NewModule(opts *ModuleOptions) module.Module {
return &quotaModule{
opts: opts,
// NewModuleFromFlags returns a module.Module for the quota library which can be
// initialized from command line flags.
func NewModuleFromFlags() module.Module {
opts := &ModuleOptions{}
return NewModule(opts)