增加概率通过逻辑

internal/utils/utils.go

@@ -21,7 +21,7 @@ func NewUUID() string {
 	return hex.EncodeToString(uux[:])
 }
 
-func MpKB[K comparable](a []K) (mp map[K]bool) {
+func MapKxB[K comparable](a []K) (mp map[K]bool) {
 	mp = make(map[K]bool, len(a))
 	for _, v := range a {
 		mp[v] = true

internal/utils/utils_test.go

@@ -6,8 +6,8 @@ import (
 	"github.com/yyle88/neatjson/neatjsons"
 )
 
-func TestMpKB(t *testing.T) {
-	t.Log(neatjsons.S(MpKB([]string{"a", "b", "c"})))
+func TestMapKxB(t *testing.T) {
+	t.Log(neatjsons.S(MapKxB([]string{"a", "b", "c"})))
 }
 
 func TestSample(t *testing.T) {

passkratosrandom/pass_kratos_random.go

@@ -0,0 +1,87 @@
+package passkratosrandom
+
+import (
+	"context"
+	"math/rand"
+	"net/http"
+
+	"github.com/go-kratos/kratos/v2/errors"
+	"github.com/go-kratos/kratos/v2/log"
+	"github.com/go-kratos/kratos/v2/middleware"
+	"github.com/go-kratos/kratos/v2/middleware/selector"
+	"github.com/orzkratos/authkratos/authkratospath"
+)
+
+type Config struct {
+	rateMap map[authkratospath.Path]float64
+	rate    float64
+	enable  bool
+}
+
+func NewConfig(
+	rateMap map[authkratospath.Path]float64,
+	rate float64,
+) *Config {
+	return &Config{
+		rateMap: rateMap,
+		rate:    rate,
+		enable:  true,
+	}
+}
+
+func (a *Config) SetEnable(v bool) {
+	a.enable = v
+}
+
+func (a *Config) IsEnable() bool {
+	if a != nil {
+		return a.enable
+	}
+	return false
+}
+
+// NewMiddleware 让接口有一定概率失败
+func NewMiddleware(cfg *Config, LOGGER log.Logger) middleware.Middleware {
+	LOG := log.NewHelper(LOGGER)
+	LOG.Infof(
+		"new rate_pass middleware enable=%v operations=%v rate=%v",
+		cfg.IsEnable(),
+		len(cfg.rateMap),
+		cfg.rate,
+	)
+
+	return selector.Server(middlewareFunc()).Match(matchFunc(cfg, LOGGER)).Build()
+}
+
+func matchFunc(cfg *Config, LOGGER log.Logger) selector.MatchFunc {
+	LOG := log.NewHelper(LOGGER)
+
+	return func(ctx context.Context, operation string) bool {
+		if !cfg.enable {
+			return false
+		}
+		if len(cfg.rateMap) > 0 {
+			path := authkratospath.New(operation)
+			if rate, ok := cfg.rateMap[path]; ok {
+				pass := rand.Float64() < rate //比如设置0.6就是有60%的概率通过
+				LOG.Debugf("operation=%s in rate_map rate_pass rate=%v pass=%v", operation, rate, pass)
+				return !pass
+			}
+		}
+		//这里不是else，而是默认的，就是没配置通过率的，就是用这个默认的通过率
+		pass := rand.Float64() < cfg.rate //设置0.6就是有60%的概率通过
+		LOG.Debugf("operation=%s rate_pass rate=%v pass=%v", operation, cfg.rate, pass)
+		return !pass //当不通过时才执行 middlewareFunc
+	}
+}
+
+func middlewareFunc() middleware.Middleware {
+	erk := errors.New(http.StatusServiceUnavailable, "RANDOM_RATE_NOT_PASS", "random rate not pass")
+
+	//当已经命中概率的时候，就直接返回错误
+	return func(handleFunc middleware.Handler) middleware.Handler {
+		return func(ctx context.Context, req interface{}) (interface{}, error) {
+			return nil, erk
+		}
+	}
+}

passkratosrandom/pass_kratos_random_test.go

@@ -0,0 +1,9 @@
+package passkratosrandom
+
+import (
+	"testing"
+)
+
+func TestMain(m *testing.M) {
+	m.Run()
+}

slowkratoshandle/slowkratoshandle.go → slowkratoshandle/slow_kratos_handle.go

@@ -11,26 +11,26 @@ import (
 	"github.com/orzkratos/authkratos/internal/utils"
 )
 
-type SlowFastConfig struct {
+type Config struct {
 	fastTimeoutGap time.Duration //快速超时的时间
 	fastOperations []authkratospath.Path
 	slowOperations []authkratospath.Path
 }
 
-func NewSlowFastConfig(
+func NewConfig(
 	fastTimeoutGap time.Duration,
 	fastOperations authkratospath.Paths,
 	slowOperations authkratospath.Paths,
-) *SlowFastConfig {
-	return &SlowFastConfig{
+) *Config {
+	return &Config{
 		fastTimeoutGap: fastTimeoutGap,
 		fastOperations: fastOperations,
 		slowOperations: slowOperations,
 	}
 }
 
 // NewMiddleware 有时接口分为快速返回和耗时返回两种，我们可以单独设置它们的timeout时间，否则假如把超时都设置为10分钟，则某些小接口卡住时也不行
-func NewMiddleware(cfg *SlowFastConfig, LOGGER log.Logger) middleware.Middleware {
+func NewMiddleware(cfg *Config, LOGGER log.Logger) middleware.Middleware {
 	LOG := log.NewHelper(LOGGER)
 	LOG.Infof(
 		"new slow_fast middleware slow=%v fast=%v fast_timeout=%v",
@@ -42,10 +42,10 @@ func NewMiddleware(cfg *SlowFastConfig, LOGGER log.Logger) middleware.Middleware
 	return selector.Server(middlewareFunc(cfg)).Match(matchFunc(cfg, LOGGER)).Build()
 }
 
-func matchFunc(cfg *SlowFastConfig, LOGGER log.Logger) selector.MatchFunc {
+func matchFunc(cfg *Config, LOGGER log.Logger) selector.MatchFunc {
 	LOG := log.NewHelper(LOGGER)
-	qMap := utils.MpKB(cfg.fastOperations)
-	sMap := utils.MpKB(cfg.slowOperations)
+	qMap := utils.MapKxB(cfg.fastOperations)
+	sMap := utils.MapKxB(cfg.slowOperations)
 	return func(ctx context.Context, operation string) bool {
 		path := authkratospath.New(operation)
 		if qMap[path] {
@@ -61,7 +61,7 @@ func matchFunc(cfg *SlowFastConfig, LOGGER log.Logger) selector.MatchFunc {
 	}
 }
 
-func middlewareFunc(cfg *SlowFastConfig) middleware.Middleware {
+func middlewareFunc(cfg *Config) middleware.Middleware {
 	return func(handleFunc middleware.Handler) middleware.Handler {
 		return func(ctx context.Context, req interface{}) (interface{}, error) {
 			//设置新超时时间，因此需要外面的超时时间更长些，选择部分接口设置快速超时