blob: 796a8781344e17952907f77fdf5535c523632208 [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 rpcquota
import (
luciflag ""
quota ""
var quotaTrackOnlyKey = ""
var ModuleName = module.RegisterName("")
var _ module.Module = &rpcquotaModule{}
// rpcquotaModule implements module.Module.
type rpcquotaModule struct {
opts *ModuleOptions
const (
// Disabled: Quota is fully off. No quotaconfig loaded, no quota
// amounts deducted in redis, no UnaryServerInterceptor registered.
RPCQuotaModeDisabled = "disabled"
// Track only: deduct quota but don't fail requests. Intended for dark
// launch.
RPCQuotaModeTrackOnly = "track-only"
// RPC quota fully enabled. Requests that exceed quota will be failed.
RPCQuotaModeEnforce = "enforce"
type ModuleOptions struct {
// RPCQuotaMode determines whether quota is off, active (but not
// enforced), or active and enforced.
RPCQuotaMode string
func (o *ModuleOptions) Register(fs *flag.FlagSet) {
if o.RPCQuotaMode == "" {
o.RPCQuotaMode = RPCQuotaModeDisabled
luciflag.NewChoice(&o.RPCQuotaMode, RPCQuotaModeDisabled, RPCQuotaModeTrackOnly, RPCQuotaModeEnforce),
"RPC quota mode. Options are: disabled, track-only, enforce.")
// NewModule returns a module.Module for the rpcquota library initialized from
// the given *ModuleOptions.
func NewModule(opts *ModuleOptions) module.Module {
return &rpcquotaModule{
opts: opts,
// NewModuleFromFlags returns a module.Module for the rpcquota library which
// can be initialized from command line flags.
func NewModuleFromFlags() module.Module {
opts := &ModuleOptions{}
return NewModule(opts)
// Dependencies returns required and optional dependencies for this module.
// Implements module.Module.
func (*rpcquotaModule) Dependencies() []module.Dependency {
return []module.Dependency{
// NOTE: cfgmodule and rpcquota must be initialized before
// quota, so it declaring quota.ModuleName as a dependency
// doesn't work. Instead rpcquota users must pass rpcquota
// directly to MainWithModules to ensure it gets initialized
// first.
// Name returns the module.Name for this module.
// Implements module.Module.
func (*rpcquotaModule) Name() module.Name {
return ModuleName
var ErrInvalidOptions = errors.New("Invalid options")
// Initialize initializes this module by adding a quota.Use implementation to
// the serving context, according to the RPCQuotaMode in the ModuleOptions.
func (m *rpcquotaModule) Initialize(ctx context.Context, host module.Host, opts module.HostOptions) (context.Context, error) {
switch m.opts.RPCQuotaMode {
case RPCQuotaModeDisabled:
// Install a stub config so that luci/server/quota Initialize
// does not panic, but otherwise there's nothing to do.
cfg, err := quotaconfig.NewMemory(ctx, nil)
if err != nil {
return nil, err
return quota.Use(ctx, cfg), nil
case RPCQuotaModeTrackOnly:
// Leave a flag in the context that will be checked by
// UpdateUserQuota before returning ErrInsufficientQuota.
// Note that this is not the common path: we expect that
// 1. nearly all requests will be within quota
// 2. this option is only set during the "dark" launch phase.
ctx = context.WithValue(ctx, &quotaTrackOnlyKey, true)
case RPCQuotaModeEnforce:
return nil, ErrInvalidOptions
if !opts.Prod {
logging.Warningf(ctx, "RPC quota not disabled, but not running in prod.")
ss, err := config.ServiceSet("luci-resultdb")
if err != nil {
return nil, err
cfgsvc := configservice.New(ctx, ss, "rpcquota.cfg")
if err := cfgsvc.Refresh(ctx); err != nil {
return nil, err
ctx = quota.Use(ctx, cfgsvc)
return ctx, nil