go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/filter/featureBreaker/featurebreaker.go

// Copyright 2015 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
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package featureBreaker

import (
	"context"
	"errors"
	"fmt"
	"runtime"
	"strings"
	"sync"
)

// BreakFeatureCallback can be used to break features at the time of the call.
//
// If it return an error, this error will be returned by the corresponding API
// call as is. If returns nil, the call will be performed as usual.
//
// It receives a derivative of a context passed to the original API call which
// can be used to extract any contextual information, if necessary.
//
// The callback will be called often and concurrently. Provide your own
// synchronization if necessary.
type BreakFeatureCallback func(ctx context.Context, feature string) error

// FeatureBreaker is the state-access interface for all Filter* functions in
// this package.  A feature is the Name of some method on the filtered service.
//
// So if you had:
//
//	c, fb := FilterMC(...)
//	mc := gae.GetMC(c)
//
// you could do:
//
//	fb.BreakFeatures(memcache.ErrServerError, "Add", "Set")
//
// and then
//
//	mc.Add(...) and mc.Set(...)
//
// would return the error.
//
// You may also pass nil as the error for BreakFeatures, and the fake will
// provide the DefaultError which you passed to the Filter function.
//
// This interface can only break features which return errors.
type FeatureBreaker interface {
	// BreakFeatures allows you to set an error that should be returned by the
	// corresponding functions.
	//
	// For example
	//   m.BreakFeatures(memcache.ErrServerError, "Add")
	//
	// would make memcache.Add return memcache.ErrServerError. You can reverse
	// this by calling UnbreakFeatures("Add").
	//
	// The only exception to this rule is two "fake" functions that can be used
	// to simulate breaking "RunInTransaction" in a more detailed way:
	//   * Use "BeginTransaction" as a feature name to simulate breaking of a new
	//     transaction attempt. It is called before each individual retry.
	//   * Use "CommitTransaction" as a feature name to simulate breaking the
	//     transaction commit RPC. It is called after the transaction body
	//     completes. Returning datastore.ErrConcurrentTransaction here will cause
	//     a retry.
	//
	// "RunInTransaction" itself is not breakable. Break "BeginTransaction" or
	// "CommitTransaction" instead.
	BreakFeatures(err error, feature ...string)

	// BreakFeaturesWithCallback is like BreakFeatures, except it allows you to
	// decide whether to return an error or not at the time the function call is
	// happening.
	//
	// The callback will be called often and concurrently. Provide your own
	// synchronization if necessary.
	//
	// You may use a callback returned by flaky.Errors(...) to emulate randomly
	// occurring errors.
	//
	// Note that the default error passed to Filter* functions are ignored when
	// using callbacks.
	BreakFeaturesWithCallback(cb BreakFeatureCallback, feature ...string)

	// UnbreakFeatures is the inverse of BreakFeatures/BreakFeaturesWithCallback,
	// and will return the named features back to their original functionality.
	UnbreakFeatures(feature ...string)
}

// errUseDefault is never returned but used as an indicator to use defaultError.
var errUseDefault = errors.New("use default error")

type state struct {
	l      sync.RWMutex
	broken map[string]BreakFeatureCallback

	// defaultError is the default error to return when you call
	// BreakFeatures(nil, ...). If this is unset and the user calls BreakFeatures
	// with nil, BrokenFeatures will return a generic error.
	defaultError error
}

func newState(dflt error) *state {
	return &state{
		broken:       map[string]BreakFeatureCallback{},
		defaultError: dflt,
	}
}

func (s *state) BreakFeatures(err error, feature ...string) {
	if err == nil {
		err = errUseDefault
	}
	s.BreakFeaturesWithCallback(
		func(context.Context, string) error { return err },
		feature...)
}

func (s *state) BreakFeaturesWithCallback(cb BreakFeatureCallback, feature ...string) {
	for _, f := range feature {
		if f == "RunInTransaction" {
			panic("break BeginTransaction or CommitTransaction instead of RunInTransaction")
		}
	}
	s.l.Lock()
	defer s.l.Unlock()
	for _, f := range feature {
		s.broken[f] = cb
	}
}

func (s *state) UnbreakFeatures(feature ...string) {
	s.l.Lock()
	defer s.l.Unlock()
	for _, f := range feature {
		delete(s.broken, f)
	}
}

func (s *state) run(c context.Context, f func() error) error {
	if s.noBrokenFeatures() {
		return f()
	}

	pc, _, _, _ := runtime.Caller(1)
	fullName := runtime.FuncForPC(pc).Name()
	fullNameParts := strings.Split(fullName, ".")
	name := fullNameParts[len(fullNameParts)-1]

	s.l.RLock()
	cb := s.broken[name]
	dflt := s.defaultError
	s.l.RUnlock()

	if cb == nil {
		return f()
	}

	switch err := cb(c, name); {
	case err == nil: // the callback decided not to break the feature
		return f()
	case err == errUseDefault:
		if dflt != nil {
			return dflt
		}
		return fmt.Errorf("feature %q is broken", name)
	default:
		return err
	}
}

func (s *state) noBrokenFeatures() bool {
	s.l.RLock()
	defer s.l.RUnlock()
	return len(s.broken) == 0
}