github.com/cloudwego/kitex@v0.9.0/pkg/retry/retryer.go

/*
 * Copyright 2021 CloudWeGo 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 retry implements rpc retry
package retry

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/bytedance/gopkg/cloud/circuitbreaker"

	"github.com/cloudwego/kitex/pkg/circuitbreak"
	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
)

// RPCCallFunc is the definition with wrap rpc call
type RPCCallFunc func(context.Context, Retryer) (rpcinfo rpcinfo.RPCInfo, resp interface{}, err error)

// Retryer is the interface for Retry implements
type Retryer interface {
	// AllowRetry to check if current request satisfy retry condition[eg: circuit, retry times == 0, chain stop, ddl].
	// If not satisfy won't execute Retryer.Do and return the reason message
	// Execute anyway for the first time regardless of able to retry.
	AllowRetry(ctx context.Context) (msg string, ok bool)

	// ShouldRetry to check if retry request can be called, it is checked in retryer.Do.
	// If not satisfy will return the reason message
	ShouldRetry(ctx context.Context, err error, callTimes int, req interface{}, cbKey string) (msg string, ok bool)
	UpdatePolicy(policy Policy) error

	// Retry policy execute func. recycleRI is to decide if the firstRI can be recycled.
	Do(ctx context.Context, rpcCall RPCCallFunc, firstRI rpcinfo.RPCInfo, request interface{}) (lastRI rpcinfo.RPCInfo, recycleRI bool, err error)
	AppendErrMsgIfNeeded(err error, ri rpcinfo.RPCInfo, msg string)

	// Prepare to do something needed before retry call.
	Prepare(ctx context.Context, prevRI, retryRI rpcinfo.RPCInfo)
	Dump() map[string]interface{}
	Type() Type
}

// NewRetryContainerWithCB build Container that doesn't do circuit breaker statistic but get statistic result.
// Which is used in case that circuit breaker is enabled.
// eg:
//
//	   cbs := circuitbreak.NewCBSuite(circuitbreak.RPCInfo2Key)
//	   retryC := retry.NewRetryContainerWithCB(cbs.ServiceControl(), cbs.ServicePanel())
//		  var opts []client.Option
//		  opts = append(opts, client.WithRetryContainer(retryC))
//	   // enable service circuit breaker
//		  opts = append(opts, client.WithMiddleware(cbs.ServiceCBMW()))
func NewRetryContainerWithCB(cc *circuitbreak.Control, cp circuitbreaker.Panel) *Container {
	return NewRetryContainer(WithContainerCBControl(cc), WithContainerCBPanel(cp))
}

func newCBSuite() *circuitbreak.CBSuite {
	return circuitbreak.NewCBSuite(circuitbreak.RPCInfo2Key)
}

// NewRetryContainerWithCBStat build Container that need to do circuit breaker statistic.
// Which is used in case that the service CB key is customized.
// eg:
//
//	cbs := circuitbreak.NewCBSuite(YourGenServiceCBKeyFunc)
//	retry.NewRetryContainerWithCBStat(cbs.ServiceControl(), cbs.ServicePanel())
func NewRetryContainerWithCBStat(cc *circuitbreak.Control, cp circuitbreaker.Panel) *Container {
	return NewRetryContainer(WithContainerCBControl(cc), WithContainerCBPanel(cp), WithContainerCBStat())
}

// NewRetryContainerWithPercentageLimit build a Container to limiting the percentage of retry requests;
// This is the RECOMMENDED initializer if you want to control PRECISELY the percentage of retry requests.
func NewRetryContainerWithPercentageLimit() *Container {
	return NewRetryContainer(WithContainerEnablePercentageLimit())
}

// ContainerOption is used when initializing a Container
type ContainerOption func(rc *Container)

// WithContainerCBSuite specifies the CBSuite used in the retry circuitbreak
// retryer will use its ServiceControl and ServicePanel
// Its priority is lower than WithContainerCBControl and WithContainerCBPanel
func WithContainerCBSuite(cbs *circuitbreak.CBSuite) ContainerOption {
	return func(rc *Container) {
		rc.cbContainer.cbSuite = cbs
	}
}

// WithContainerCBControl is specifies the circuitbreak.Control used in the retry circuitbreaker
// It's user's responsibility to make sure it's paired with panel
func WithContainerCBControl(ctrl *circuitbreak.Control) ContainerOption {
	return func(rc *Container) {
		rc.cbContainer.cbCtl = ctrl
	}
}

// WithContainerCBPanel is specifies the circuitbreaker.Panel used in the retry circuitbreaker
// It's user's responsibility to make sure it's paired with control
func WithContainerCBPanel(panel circuitbreaker.Panel) ContainerOption {
	return func(rc *Container) {
		rc.cbContainer.cbPanel = panel
	}
}

// WithContainerCBStat instructs the circuitbreak.RecordStat is called within the retryer
func WithContainerCBStat() ContainerOption {
	return func(rc *Container) {
		rc.cbContainer.cbStat = true
	}
}

// WithContainerEnablePercentageLimit should be called for limiting the percentage of retry requests
func WithContainerEnablePercentageLimit() ContainerOption {
	return func(rc *Container) {
		rc.cbContainer.enablePercentageLimit = true
	}
}

// NewRetryContainer build Container that need to build circuit breaker and do circuit breaker statistic.
// The caller is responsible for calling Container.Close() to release resources referenced.
func NewRetryContainer(opts ...ContainerOption) *Container {
	rc := &Container{
		cbContainer: &cbContainer{
			cbSuite: nil,
		},
		retryerMap: sync.Map{},
	}
	for _, opt := range opts {
		opt(rc)
	}

	if rc.cbContainer.enablePercentageLimit {
		// ignore cbSuite/cbCtl/cbPanel options
		rc.cbContainer = &cbContainer{
			enablePercentageLimit: true,
			cbSuite:               newCBSuite(),
		}
	}

	container := rc.cbContainer
	if container.cbCtl == nil && container.cbPanel == nil {
		if container.cbSuite == nil {
			container.cbSuite = newCBSuite()
			container.cbStat = true
		}
		container.cbCtl = container.cbSuite.ServiceControl()
		container.cbPanel = container.cbSuite.ServicePanel()
	}
	if !container.IsValid() {
		panic("KITEX: invalid container")
	}
	return rc
}

// Container is a wrapper for Retryer.
type Container struct {
	hasCodeCfg  bool
	retryerMap  sync.Map // <method: retryer>
	cbContainer *cbContainer
	msg         string
	sync.RWMutex

	// shouldResultRetry is only used with FailureRetry
	shouldResultRetry *ShouldResultRetry
}

// Recommended usage: NewRetryContainerWithPercentageLimit()
// For more details, refer to the following comments for each field.
type cbContainer struct {
	// In NewRetryContainer, if cbCtrl & cbPanel are not set, Kitex will use cbSuite.ServiceControl() and
	// cbSuite.ServicePanel(); If cbSuite is nil, Kitex will create one.
	cbSuite *circuitbreak.CBSuite

	// It's more recommended to rely on the cbSuite than specifying cbCtl & cbPanel with corresponding options,
	// since cbCtl & cbPanel should be correctly paired, and with the cbSuite, Kitex will ensure it by using the
	// cbSuite.ServiceControl() and cbSuite.ServicePanel().
	cbCtl   *circuitbreak.Control
	cbPanel circuitbreaker.Panel

	// If cbStat && !enablePercentageLimit, retryer will call `circuitbreak.RecordStat` after rpcCall to record
	// rpc failures/timeouts, for cutting down on the retry requests when the error rate is beyond the threshold.
	cbStat bool

	// If enabled, Kitex will always create a cbSuite and use its cbCtl & cbPanel, and retryer will call
	// recordRetryStat before rpcCall, to precisely control the percentage of retry requests over all requests.
	enablePercentageLimit bool
}

// IsValid returns true when both cbCtl & cbPanel are not nil
// It's the user's responsibility to guarantee that cbCtl & cbPanel are correctly paired.
func (c *cbContainer) IsValid() bool {
	return c.cbCtl != nil && c.cbPanel != nil
}

// InitWithPolicies to init Retryer with methodPolicies
// Notice, InitWithPolicies is export func, the lock should be added inside
func (rc *Container) InitWithPolicies(methodPolicies map[string]Policy) error {
	if methodPolicies == nil {
		return nil
	}
	rc.Lock()
	defer rc.Unlock()
	var inited bool
	for m := range methodPolicies {
		if methodPolicies[m].Enable {
			inited = true
			if _, ok := rc.retryerMap.Load(m); ok {
				// NotifyPolicyChange may happen before
				continue
			}
			if err := rc.initRetryer(m, methodPolicies[m]); err != nil {
				rc.msg = err.Error()
				return err
			}
		}
	}
	rc.hasCodeCfg = inited
	return nil
}

// DeletePolicy to delete the method by method.
func (rc *Container) DeletePolicy(method string) {
	rc.Lock()
	defer rc.Unlock()
	rc.msg = ""
	if rc.hasCodeCfg {
		// the priority of user setup code policy is higher than remote config
		return
	}
	_, ok := rc.retryerMap.Load(method)
	if ok {
		rc.retryerMap.Delete(method)
		rc.msg = fmt.Sprintf("delete retryer[%s] at %s", method, time.Now())
	}
}

// NotifyPolicyChange to receive policy when it changes
func (rc *Container) NotifyPolicyChange(method string, p Policy) {
	rc.Lock()
	defer rc.Unlock()
	rc.msg = ""
	if rc.hasCodeCfg {
		// the priority of user setup code policy is higher than remote config
		return
	}
	r, ok := rc.retryerMap.Load(method)
	if ok && r != nil {
		retryer, ok := r.(Retryer)
		if ok {
			if retryer.Type() == p.Type {
				retryer.UpdatePolicy(p)
				rc.msg = fmt.Sprintf("update retryer[%s-%s] at %s", method, retryer.Type(), time.Now())
				return
			}
			rc.retryerMap.Delete(method)
			rc.msg = fmt.Sprintf("delete retryer[%s-%s] at %s", method, retryer.Type(), time.Now())
		}
	}
	rc.initRetryer(method, p)
}

// Init to build Retryer with code config.
func (rc *Container) Init(mp map[string]Policy, rr *ShouldResultRetry) (err error) {
	// NotifyPolicyChange func may execute before Init func.
	// Because retry Container is built before Client init, NotifyPolicyChange can be triggered first
	rc.updateRetryer(rr)
	if err = rc.InitWithPolicies(mp); err != nil {
		return fmt.Errorf("NewRetryer in Init failed, err=%w", err)
	}
	return nil
}

// PrepareRetryContext adds necessary keys to context for retry
// These keys should be added to `ctx` no matter whether there's a need to retry, to avoid sharing the same
// object objects with another method call, since `ctx` might be reused in user-defined middlewares.
func PrepareRetryContext(ctx context.Context) context.Context {
	// reqOp can be used to avoid multiple writes to the request object.
	// If a blocking write is needed, implement a lock based on it (spin-lock for example).
	reqOp := OpNo
	ctx = context.WithValue(ctx, CtxReqOp, &reqOp)

	// `respOp` is used to avoid concurrent write/read on the response object, especially for backup requests.
	// If `respOp` is modified by one request of this method call, all other requests will skip decoding.
	respOp := OpNo
	ctx = context.WithValue(ctx, CtxRespOp, &respOp)
	return ctx
}

// WithRetryIfNeeded to check if there is a retryer can be used and if current call can retry.
// When the retry condition is satisfied, use retryer to call
func (rc *Container) WithRetryIfNeeded(ctx context.Context, callOptRetry *Policy, rpcCall RPCCallFunc, ri rpcinfo.RPCInfo, request interface{}) (lastRI rpcinfo.RPCInfo, recycleRI bool, err error) {
	var retryer Retryer
	if callOptRetry != nil && callOptRetry.Enable {
		// build retryer for call level if retry policy is set up with callopt
		if retryer, err = NewRetryer(*callOptRetry, nil, rc.cbContainer); err != nil {
			klog.Warnf("KITEX: new callopt retryer[%s] failed, err=%w", callOptRetry.Type, err)
		}
	} else {
		retryer = rc.getRetryer(ri)
	}

	// case 1(default, fast path): no retry policy
	if retryer == nil {
		if _, _, err = rpcCall(ctx, nil); err == nil {
			return ri, true, nil
		}
		return ri, false, err
	}

	// case 2: setup retry policy, but not satisfy retry condition eg: circuit, retry times == 0, chain stop, ddl
	if msg, ok := retryer.AllowRetry(ctx); !ok {
		if _, _, err = rpcCall(ctx, retryer); err == nil {
			return ri, true, err
		}
		if msg != "" {
			retryer.AppendErrMsgIfNeeded(err, ri, msg)
		}
		return ri, false, err
	}

	// case 3: do rpc call with retry policy
	lastRI, recycleRI, err = retryer.Do(ctx, rpcCall, ri, request)
	return
}

// NewRetryer build a retryer with policy
func NewRetryer(p Policy, r *ShouldResultRetry, cbC *cbContainer) (retryer Retryer, err error) {
	// just one retry policy can be enabled at same time
	if p.Type == BackupType {
		retryer, err = newBackupRetryer(p, cbC)
	} else {
		retryer, err = newFailureRetryer(p, r, cbC)
	}
	return
}

func (rc *Container) getRetryer(ri rpcinfo.RPCInfo) Retryer {
	// the priority of specific method is high
	r, ok := rc.retryerMap.Load(ri.To().Method())
	if ok {
		return r.(Retryer)
	}
	r, ok = rc.retryerMap.Load(Wildcard)
	if ok {
		return r.(Retryer)
	}
	return nil
}

// Dump is used to show current retry policy
func (rc *Container) Dump() interface{} {
	rc.RLock()
	dm := make(map[string]interface{})
	dm["has_code_cfg"] = rc.hasCodeCfg
	rc.retryerMap.Range(func(key, value interface{}) bool {
		if r, ok := value.(Retryer); ok {
			dm[key.(string)] = r.Dump()
		}
		return true
	})
	if rc.msg != "" {
		dm["msg"] = rc.msg
	}
	rc.RUnlock()
	return dm
}

func (rc *Container) initRetryer(method string, p Policy) error {
	retryer, err := NewRetryer(p, rc.shouldResultRetry, rc.cbContainer)
	if err != nil {
		errMsg := fmt.Sprintf("new retryer[%s-%s] failed, err=%s, at %s", method, p.Type, err.Error(), time.Now())
		rc.msg = errMsg
		klog.Warnf(errMsg)
		return err
	}

	rc.retryerMap.Store(method, retryer)
	if p.Enable {
		rc.msg = fmt.Sprintf("new retryer[%s-%s] at %s", method, retryer.Type(), time.Now())
	} else {
		rc.msg = fmt.Sprintf("disable retryer[%s-%s](enable=%t) %s", method, p.Type, p.Enable, time.Now())
	}
	return nil
}

func (rc *Container) updateRetryer(rr *ShouldResultRetry) {
	rc.Lock()
	defer rc.Unlock()

	rc.shouldResultRetry = rr
	if rc.shouldResultRetry != nil {
		rc.retryerMap.Range(func(key, value interface{}) bool {
			if fr, ok := value.(*failureRetryer); ok {
				fr.setSpecifiedResultRetryIfNeeded(rc.shouldResultRetry)
			}
			return true
		})
	}
}

// Close releases all possible resources referenced.
func (rc *Container) Close() (err error) {
	if rc.cbContainer != nil && rc.cbContainer.cbSuite != nil {
		err = rc.cbContainer.cbSuite.Close()
	}
	return
}