config.go raw

   1  package retry
   2  
   3  import (
   4  	"strconv"
   5  	"time"
   6  
   7  	"google.golang.org/grpc/codes"
   8  )
   9  
  10  type RetryConfig struct {
  11  	mc              map[nameConfig]*methodConfig
  12  	retryThrottling *retryThrottling
  13  }
  14  
  15  type RetryOption func(c *RetryConfig)
  16  
  17  type grpcRetryPolicy struct {
  18  	MethodConfig    []*methodConfig `json:"methodConfig"`
  19  	RetryThrottling retryThrottling `json:"retryThrottling"`
  20  }
  21  
  22  type methodConfig struct {
  23  	NameConfig   []nameConfig `json:"name"`
  24  	RetryPolicy  retryPolicy  `json:"retryPolicy"`
  25  	WaitForReady bool         `json:"waitForReady"`
  26  }
  27  
  28  type nameConfig struct {
  29  	Service string `json:"service,omitempty"`
  30  	Method  string `json:"method,omitempty"`
  31  }
  32  
  33  type retryPolicy struct {
  34  	MaxAttempts          int      `json:"maxAttempts"`
  35  	InitialBackoff       Duration `json:"initialBackoff"`
  36  	MaxBackoff           Duration `json:"maxBackoff"`
  37  	BackoffMultiplier    float64  `json:"backoffMultiplier"`
  38  	RetryableStatusCodes []string `json:"retryableStatusCodes"`
  39  }
  40  
  41  type retryThrottling struct {
  42  	MaxTokens  int     `json:"maxTokens"`
  43  	TokenRatio float64 `json:"tokenRatio"`
  44  }
  45  
  46  type GRPCKeepAliveConfig struct {
  47  	Time                time.Duration `yaml:"time" validate:"required"`
  48  	Timeout             time.Duration `yaml:"timeout" validate:"required"`
  49  	PermitWithoutStream bool          `yaml:"permit_without_stream"`
  50  }
  51  
  52  // ThrottlingMode provides the mode SDK will retry request.
  53  type ThrottlingMode string
  54  
  55  const (
  56  	// ThrottlingModePersistent model provides retry attempts and budget for persistent environments.
  57  	// This mode is suitable when you use SDK in your server application, or any long-lived applications.
  58  	ThrottlingModePersistent ThrottlingMode = "persistent"
  59  
  60  	// ThrottlingModeTemporary model provides retry attempts and budget for temporary environments.
  61  	// This mode is suitable when you use SDK in some CI scripts, or any short-lived applications.
  62  	ThrottlingModeTemporary ThrottlingMode = "temporary"
  63  )
  64  
  65  func parseThrottlingMode(v string) ThrottlingMode {
  66  	switch v {
  67  	case "persistent":
  68  		return ThrottlingModePersistent
  69  	case "temporary":
  70  		return ThrottlingModeTemporary
  71  	default:
  72  		return ThrottlingModeTemporary
  73  	}
  74  }
  75  
  76  func DefaultRetryConfig() *RetryConfig {
  77  	return &RetryConfig{}
  78  }
  79  
  80  func DefaultNameConfig() nameConfig {
  81  	return nameConfig{}
  82  }
  83  
  84  func NewNameConfig(service, method string) nameConfig {
  85  	return nameConfig{
  86  		Service: service,
  87  		Method:  method,
  88  	}
  89  }
  90  
  91  func WithDefaultRetryConfig() RetryOption {
  92  	return func(c *RetryConfig) { // nolint:staticcheck,ineffassign
  93  		c = defaultRetryConfig() // nolint:ineffassign,staticcheck
  94  	}
  95  }
  96  
  97  func WithRetries(nm nameConfig, n int) RetryOption {
  98  	return func(c *RetryConfig) {
  99  		if c == nil {
 100  			c = defaultRetryConfig()
 101  		}
 102  
 103  		if _, ok := c.mc[nm]; !ok {
 104  			c.mc[nm] = defaultMethodConfig()
 105  			c.mc[nm].NameConfig = []nameConfig{nm}
 106  		}
 107  
 108  		c.mc[nm].RetryPolicy.MaxAttempts = n
 109  	}
 110  }
 111  
 112  func WithRetryableStatusCodes(nm nameConfig, codes ...codes.Code) RetryOption {
 113  	return func(c *RetryConfig) {
 114  		if c == nil {
 115  			c = defaultRetryConfig()
 116  		}
 117  
 118  		names := make([]string, len(codes))
 119  		for i, code := range codes {
 120  			names[i] = canonicalString(code)
 121  		}
 122  
 123  		if _, ok := c.mc[nm]; !ok {
 124  			c.mc[nm] = defaultMethodConfig()
 125  			c.mc[nm].NameConfig = []nameConfig{nm}
 126  		}
 127  
 128  		c.mc[nm].RetryPolicy.RetryableStatusCodes = names
 129  	}
 130  }
 131  
 132  func WithThrottlingMode(mode ThrottlingMode) RetryOption {
 133  	return func(c *RetryConfig) {
 134  		if c == nil {
 135  			c = defaultRetryConfig()
 136  		}
 137  
 138  		tm := parseThrottlingMode(string(mode))
 139  
 140  		c.retryThrottling = defaultRetryThrottling(tm)
 141  	}
 142  }
 143  
 144  func canonicalString(c codes.Code) string {
 145  	switch c {
 146  	case codes.OK:
 147  		return "OK"
 148  	case codes.Canceled:
 149  		return "CANCELLED"
 150  	case codes.Unknown:
 151  		return "UNKNOWN"
 152  	case codes.InvalidArgument:
 153  		return "INVALID_ARGUMENT"
 154  	case codes.DeadlineExceeded:
 155  		return "DEADLINE_EXCEEDED"
 156  	case codes.NotFound:
 157  		return "NOT_FOUND"
 158  	case codes.AlreadyExists:
 159  		return "ALREADY_EXISTS"
 160  	case codes.PermissionDenied:
 161  		return "PERMISSION_DENIED"
 162  	case codes.ResourceExhausted:
 163  		return "RESOURCE_EXHAUSTED"
 164  	case codes.FailedPrecondition:
 165  		return "FAILED_PRECONDITION"
 166  	case codes.Aborted:
 167  		return "ABORTED"
 168  	case codes.OutOfRange:
 169  		return "OUT_OF_RANGE"
 170  	case codes.Unimplemented:
 171  		return "UNIMPLEMENTED"
 172  	case codes.Internal:
 173  		return "INTERNAL"
 174  	case codes.Unavailable:
 175  		return "UNAVAILABLE"
 176  	case codes.DataLoss:
 177  		return "DATA_LOSS"
 178  	case codes.Unauthenticated:
 179  		return "UNAUTHENTICATED"
 180  	default:
 181  		return "CODE(" + strconv.FormatInt(int64(c), 10) + ")"
 182  	}
 183  }
 184