feat: support min timeout with breaker

core/stores/redis/breakerhook.go

@@ -2,12 +2,17 @@ package redis
 
 import (
  "context"
+ "errors"
+ "time"
 
  red "github.com/redis/go-redis/v9"
  "github.com/zeromicro/go-zero/core/breaker"
  "github.com/zeromicro/go-zero/core/lang"
+ "github.com/zeromicro/go-zero/core/timex"
 )
 
+const minTimeout = time.Millisecond * 100
+
 var ignoreCmds = map[string]lang.PlaceholderType{
  "blpop": {},
 }
@@ -26,16 +31,32 @@ func (h breakerHook) ProcessHook(next red.ProcessHook) red.ProcessHook {
  return next(ctx, cmd)
  }
 
+ start := timex.Now()
  return h.brk.DoWithAcceptable(func() error {
  return next(ctx, cmd)
- }, acceptable)
+ }, protectedAcceptable(start))
  }
 }
 
 func (h breakerHook) ProcessPipelineHook(next red.ProcessPipelineHook) red.ProcessPipelineHook {
  return func(ctx context.Context, cmds []red.Cmder) error {
+ start := timex.Now()
  return h.brk.DoWithAcceptable(func() error {
  return next(ctx, cmds)
- }, acceptable)
+ }, protectedAcceptable(start))
+ }
+}
+
+func acceptable(err error) bool {
+ return err == nil || errors.Is(err, red.Nil) || errors.Is(err, context.Canceled)
+}
+
+func protectedAcceptable(start time.Duration) breaker.Acceptable {
+ return func(err error) bool {
+ if acceptable(err) {
+ return true
+ }
+
+ return errors.Is(err, context.DeadlineExceeded) && timex.Since(start) < minTimeout
  }
 }

core/stores/redis/redis.go

@@ -2371,10 +2371,6 @@ func withHook(hook red.Hook) Option {
  }
 }
 
-func acceptable(err error) bool {
- return err == nil || errors.Is(err, red.Nil) || errors.Is(err, context.Canceled)
-}
-
 func getRedis(r *Redis) (RedisNode, error) {
  switch r.Type {
  case ClusterType:

core/stores/sqlx/sqlconn.go

@@ -120,7 +120,7 @@ func (db *commonSqlConn) ExecCtx(ctx context.Context, q string, args ...any) (
  endSpan(span, err)
  }()
 
- err = db.brk.DoWithAcceptable(func() error {
+ err = db.run(func() error {
  var conn *sql.DB
  conn, err = db.connProv()
  if err != nil {
@@ -148,7 +148,7 @@ func (db *commonSqlConn) PrepareCtx(ctx context.Context, query string) (stmt Stm
  endSpan(span, err)
  }()
 
- err = db.brk.DoWithAcceptable(func() error {
+ err = db.run(func() error {
  var conn *sql.DB
  conn, err = db.connProv()
  if err != nil {
@@ -256,7 +256,7 @@ func (db *commonSqlConn) TransactCtx(ctx context.Context, fn func(context.Contex
  endSpan(span, err)
  }()
 
- err = db.brk.DoWithAcceptable(func() error {
+ err = db.run(func() error {
  return transact(ctx, db, db.beginTx, fn)
  }, db.acceptable)
  if errors.Is(err, breaker.ErrServiceUnavailable) {
@@ -287,7 +287,7 @@ func (db *commonSqlConn) acceptable(err error) bool {
 func (db *commonSqlConn) queryRows(ctx context.Context, scanner func(*sql.Rows) error,
  q string, args ...any) (err error) {
  var scanFailed bool
- err = db.brk.DoWithAcceptable(func() error {
+ err = db.run(func() error {
  conn, err := db.connProv()
  if err != nil {
  db.onError(ctx, err)
@@ -311,6 +311,10 @@ func (db *commonSqlConn) queryRows(ctx context.Context, scanner func(*sql.Rows)
  return
 }
 
+func (db *commonSqlConn) run(fn func() error, acceptable breaker.Acceptable) error {
+ return runWithBreaker(db.brk, fn, acceptable)
+}
+
 // WithAcceptable returns a SqlOption that setting the acceptable function.
 // acceptable is the func to check if the error can be accepted.
 func WithAcceptable(acceptable func(err error) bool) SqlOption {

core/stores/sqlx/stmt.go

@@ -65,7 +65,7 @@ func (s statement) ExecCtx(ctx context.Context, args ...any) (result sql.Result,
  endSpan(span, err)
  }()
 
- err = s.brk.DoWithAcceptable(func() error {
+ err = s.run(func() error {
  result, err = execStmt(ctx, s.stmt, s.query, args...)
  return err
  }, func(err error) bool {
@@ -141,7 +141,7 @@ func (s statement) QueryRowsPartialCtx(ctx context.Context, v any, args ...any)
 func (s statement) queryRows(ctx context.Context, scanFn func(any, rowsScanner) error,
  v any, args ...any) error {
  var scanFailed bool
- err := s.brk.DoWithAcceptable(func() error {
+ err := s.run(func() error {
  return queryStmt(ctx, s.stmt, func(rows *sql.Rows) error {
  err := scanFn(v, rows)
  if isScanFailed(err) {
@@ -159,6 +159,10 @@ func (s statement) queryRows(ctx context.Context, scanFn func(any, rowsScanner)
  return err
 }
 
+func (s statement) run(fn func() error, acceptable breaker.Acceptable) error {
+ return runWithBreaker(s.brk, fn, acceptable)
+}
+
 // DisableLog disables logging of sql statements, includes info and slow logs.
 func DisableLog() {
  logSql.Set(false)

core/stores/sqlx/utils.go

@@ -8,10 +8,14 @@ import (
  "strings"
  "time"
 
+ "github.com/zeromicro/go-zero/core/breaker"
  "github.com/zeromicro/go-zero/core/logx"
  "github.com/zeromicro/go-zero/core/mapping"
+ "github.com/zeromicro/go-zero/core/timex"
 )
 
+const minTimeout = time.Millisecond * 100
+
 var errUnbalancedEscape = errors.New("no char after escape char")
 
 func desensitize(datasource string) string {
@@ -148,6 +152,17 @@ func logSqlError(ctx context.Context, stmt string, err error) {
  }
 }
 
+func runWithBreaker(brk breaker.Breaker, fn func() error, acceptable breaker.Acceptable) error {
+ start := timex.Now()
+ return brk.DoWithAcceptable(fn, func(err error) bool {
+ if acceptable(err) {
+ return true
+ }
+
+ return errors.Is(err, context.DeadlineExceeded) && timex.Since(start) < minTimeout
+ })
+}
+
 func writeValue(buf *strings.Builder, arg any) {
  switch v := arg.(type) {
  case bool:

zrpc/internal/clientinterceptors/breakerinterceptor.go

@@ -2,18 +2,30 @@
 
 import (
  "context"
+ "errors"
  "path"
+ "time"
 
  "github.com/zeromicro/go-zero/core/breaker"
+ "github.com/zeromicro/go-zero/core/timex"
  "github.com/zeromicro/go-zero/zrpc/internal/codes"
  "google.golang.org/grpc"
 )
 
+const minTimeout = time.Millisecond * 100
+
 // BreakerInterceptor is an interceptor that acts as a circuit breaker.
 func BreakerInterceptor(ctx context.Context, method string, req, reply any,
  cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
  breakerName := path.Join(cc.Target(), method)
+ start := timex.Now()
  return breaker.DoWithAcceptable(breakerName, func() error {
  return invoker(ctx, method, req, reply, cc, opts...)
- }, codes.Acceptable)
+ }, func(err error) bool {
+ if codes.Acceptable(err) {
+ return true
+ }
+
+ return errors.Is(err, context.DeadlineExceeded) && timex.Since(start) < minTimeout
+ })
 }

zrpc/internal/clientinterceptors/breakerinterceptor_test.go

@@ -4,6 +4,7 @@ import (
  "context"
  "errors"
  "testing"
+ "time"
 
  "github.com/stretchr/testify/assert"
  "github.com/zeromicro/go-zero/core/breaker"
@@ -79,3 +80,19 @@ func TestBreakerInterceptor(t *testing.T) {
  })
  }
 }
+
+func TestBreakerTimeout(t *testing.T) {
+ t.Run("not timeout", func(t *testing.T) {
+ cc := new(grpc.ClientConn)
+ for i := 0; i < 1000; i++ {
+ ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
+ err := BreakerInterceptor(ctx, "/foo", nil, nil, cc,
+ func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn,
+ opts ...grpc.CallOption) error {
+ return context.DeadlineExceeded
+ })
+ assert.ErrorIs(t, err, context.DeadlineExceeded)
+ cancel()
+ }
+ })
+}