chore: optimize lock in discov.etcd (#4275)

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

.github/FUNDING.yml

@@ -1,13 +1,3 @@
 # These are supported funding model platforms
 
-github: # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]
-patreon: # Replace with a single Patreon username
-open_collective: # Replace with a single Open Collective username
-ko_fi: # Replace with a single Ko-fi username
-tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel
-community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry
-liberapay: # Replace with a single Liberapay username
-issuehunt: # Replace with a single IssueHunt username
-otechie: # Replace with a single Otechie username
-custom: # https://gitee.com/kevwan/static/raw/master/images/sponsor.jpg
-ethereum: # 0x5052b7f6B937B02563996D23feb69b38D06Ca150 | kevwan
+github: [zeromicro]

.github/workflows/go.yml

@@ -17,7 +17,7 @@ jobs:
       - name: Set up Go 1.x
         uses: actions/setup-go@v5
         with:
-          go-version: '1.19'
+          go-version: '1.20'
           check-latest: true
           cache: true
         id: go
@@ -52,8 +52,8 @@ jobs:
       - name: Set up Go 1.x
         uses: actions/setup-go@v5
         with:
-          # use 1.19 to guarantee Go 1.19 compatibility
-          go-version: '1.19'
+          # use 1.20 to guarantee Go 1.20 compatibility
+          go-version: '1.20'
           check-latest: true
           cache: true
 

SECURITY.md

@@ -0,0 +1,16 @@
+# Security Policy
+
+## Supported Versions
+
+We publish releases monthly.
+
+| Version  | Supported          |
+| -------  | ------------------ |
+| >= 1.4.4 | :white_check_mark: |
+| < 1.4.4  | :x:                |
+
+## Reporting a Vulnerability
+
+https://github.com/zeromicro/go-zero/security/advisories
+
+Accepted vulnerabilities are expected to be fixed within a month.

core/bloom/bloom.go

@@ -2,6 +2,7 @@ package bloom
 
 import (
 	"context"
+	_ "embed"
 	"errors"
 	"strconv"
 
@@ -17,19 +18,13 @@ var (
 	// ErrTooLargeOffset indicates the offset is too large in bitset.
 	ErrTooLargeOffset = errors.New("too large offset")
 
-	setScript = redis.NewScript(`
-for _, offset in ipairs(ARGV) do
-	redis.call("setbit", KEYS[1], offset, 1)
-end
-`)
-	testScript = redis.NewScript(`
-for _, offset in ipairs(ARGV) do
-	if tonumber(redis.call("getbit", KEYS[1], offset)) == 0 then
-		return false
-	end
-end
-return true
-`)
+	//go:embed setscript.lua
+	setLuaScript string
+	setScript    = redis.NewScript(setLuaScript)
+
+	//go:embed testscript.lua
+	testLuaScript string
+	testScript    = redis.NewScript(testLuaScript)
 )
 
 type (
@@ -110,7 +105,7 @@ func newRedisBitSet(store *redis.Redis, key string, bits uint) *redisBitSet {
 }
 
 func (r *redisBitSet) buildOffsetArgs(offsets []uint) ([]string, error) {
-	var args []string
+	args := make([]string, 0, len(offsets))
 
 	for _, offset := range offsets {
 		if offset >= r.bits {
@@ -130,7 +125,7 @@ func (r *redisBitSet) check(ctx context.Context, offsets []uint) (bool, error) {
 	}
 
 	resp, err := r.store.ScriptRunCtx(ctx, testScript, []string{r.key}, args)
-	if err == redis.Nil {
+	if errors.Is(err, redis.Nil) {
 		return false, nil
 	} else if err != nil {
 		return false, err
@@ -162,7 +157,7 @@ func (r *redisBitSet) set(ctx context.Context, offsets []uint) error {
 	}
 
 	_, err = r.store.ScriptRunCtx(ctx, setScript, []string{r.key}, args)
-	if err == redis.Nil {
+	if errors.Is(err, redis.Nil) {
 		return nil
 	}
 

core/bloom/setscript.lua

@@ -0,0 +1,3 @@
+for _, offset in ipairs(ARGV) do
+    redis.call("setbit", KEYS[1], offset, 1)
+end

core/bloom/testscript.lua

@@ -0,0 +1,6 @@
+for _, offset in ipairs(ARGV) do
+    if tonumber(redis.call("getbit", KEYS[1], offset)) == 0 then
+        return false
+    end
+end
+return true

core/breaker/breaker.go

@@ -1,6 +1,7 @@
 package breaker
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strings"
@@ -31,16 +32,21 @@ type (
 		Name() string
 
 		// Allow checks if the request is allowed.
-		// If allowed, a promise will be returned, the caller needs to call promise.Accept()
-		// on success, or call promise.Reject() on failure.
-		// If not allow, ErrServiceUnavailable will be returned.
+		// If allowed, a promise will be returned,
+		// otherwise ErrServiceUnavailable will be returned as the error.
+		// The caller needs to call promise.Accept() on success,
+		// or call promise.Reject() on failure.
 		Allow() (Promise, error)
+		// AllowCtx checks if the request is allowed when ctx isn't done.
+		AllowCtx(ctx context.Context) (Promise, error)
 
 		// Do runs the given request if the Breaker accepts it.
 		// Do returns an error instantly if the Breaker rejects the request.
 		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
 		Do(req func() error) error
+		// DoCtx runs the given request if the Breaker accepts it when ctx isn't done.
+		DoCtx(ctx context.Context, req func() error) error
 
 		// DoWithAcceptable runs the given request if the Breaker accepts it.
 		// DoWithAcceptable returns an error instantly if the Breaker rejects the request.
@@ -48,12 +54,16 @@ type (
 		// and causes the same panic again.
 		// acceptable checks if it's a successful call, even if the error is not nil.
 		DoWithAcceptable(req func() error, acceptable Acceptable) error
+		// DoWithAcceptableCtx runs the given request if the Breaker accepts it when ctx isn't done.
+		DoWithAcceptableCtx(ctx context.Context, req func() error, acceptable Acceptable) error
 
 		// DoWithFallback runs the given request if the Breaker accepts it.
 		// DoWithFallback runs the fallback if the Breaker rejects the request.
 		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
 		DoWithFallback(req func() error, fallback Fallback) error
+		// DoWithFallbackCtx runs the given request if the Breaker accepts it when ctx isn't done.
+		DoWithFallbackCtx(ctx context.Context, req func() error, fallback Fallback) error
 
 		// DoWithFallbackAcceptable runs the given request if the Breaker accepts it.
 		// DoWithFallbackAcceptable runs the fallback if the Breaker rejects the request.
@@ -61,6 +71,9 @@ type (
 		// and causes the same panic again.
 		// acceptable checks if it's a successful call, even if the error is not nil.
 		DoWithFallbackAcceptable(req func() error, fallback Fallback, acceptable Acceptable) error
+		// DoWithFallbackAcceptableCtx runs the given request if the Breaker accepts it when ctx isn't done.
+		DoWithFallbackAcceptableCtx(ctx context.Context, req func() error, fallback Fallback,
+			acceptable Acceptable) error
 	}
 
 	// Fallback is the func to be called if the request is rejected.
@@ -117,23 +130,71 @@ func (cb *circuitBreaker) Allow() (Promise, error) {
 	return cb.throttle.allow()
 }
 
+func (cb *circuitBreaker) AllowCtx(ctx context.Context) (Promise, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+		return cb.Allow()
+	}
+}
+
 func (cb *circuitBreaker) Do(req func() error) error {
 	return cb.throttle.doReq(req, nil, defaultAcceptable)
 }
 
+func (cb *circuitBreaker) DoCtx(ctx context.Context, req func() error) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return cb.Do(req)
+	}
+}
+
 func (cb *circuitBreaker) DoWithAcceptable(req func() error, acceptable Acceptable) error {
 	return cb.throttle.doReq(req, nil, acceptable)
 }
 
+func (cb *circuitBreaker) DoWithAcceptableCtx(ctx context.Context, req func() error,
+	acceptable Acceptable) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return cb.DoWithAcceptable(req, acceptable)
+	}
+}
+
 func (cb *circuitBreaker) DoWithFallback(req func() error, fallback Fallback) error {
 	return cb.throttle.doReq(req, fallback, defaultAcceptable)
 }
 
+func (cb *circuitBreaker) DoWithFallbackCtx(ctx context.Context, req func() error,
+	fallback Fallback) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return cb.DoWithFallback(req, fallback)
+	}
+}
+
 func (cb *circuitBreaker) DoWithFallbackAcceptable(req func() error, fallback Fallback,
 	acceptable Acceptable) error {
 	return cb.throttle.doReq(req, fallback, acceptable)
 }
 
+func (cb *circuitBreaker) DoWithFallbackAcceptableCtx(ctx context.Context, req func() error,
+	fallback Fallback, acceptable Acceptable) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return cb.DoWithFallbackAcceptable(req, fallback, acceptable)
+	}
+}
+
 func (cb *circuitBreaker) Name() string {
 	return cb.name
 }
@@ -208,7 +269,7 @@ func (ew *errorWindow) add(reason string) {
 }
 
 func (ew *errorWindow) String() string {
-	var reasons []string
+	reasons := make([]string, 0, ew.count)
 
 	ew.lock.Lock()
 	// reverse order

core/breaker/breaker_test.go

@@ -1,11 +1,13 @@
 package breaker
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strconv"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/stat"
@@ -16,10 +18,274 @@ func init() {
 }
 
 func TestCircuitBreaker_Allow(t *testing.T) {
-	b := NewBreaker()
-	assert.True(t, len(b.Name()) > 0)
-	_, err := b.Allow()
-	assert.Nil(t, err)
+	t.Run("allow", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		_, err := b.Allow()
+		assert.Nil(t, err)
+	})
+
+	t.Run("allow with ctx", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		_, err := b.AllowCtx(context.Background())
+		assert.Nil(t, err)
+	})
+
+	t.Run("allow with ctx timeout", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		ctx, cancel := context.WithTimeout(context.Background(), time.Microsecond)
+		defer cancel()
+		time.Sleep(time.Millisecond)
+		_, err := b.AllowCtx(ctx)
+		assert.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+
+	t.Run("allow with ctx cancel", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		for i := 0; i < 100; i++ {
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			cancel()
+			_, err := b.AllowCtx(ctx)
+			assert.ErrorIs(t, err, context.Canceled)
+		}
+		_, err := b.AllowCtx(context.Background())
+		assert.NoError(t, err)
+	})
+}
+
+func TestCircuitBreaker_Do(t *testing.T) {
+	t.Run("do", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		err := b.Do(func() error {
+			return nil
+		})
+		assert.Nil(t, err)
+	})
+
+	t.Run("do with ctx", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		err := b.DoCtx(context.Background(), func() error {
+			return nil
+		})
+		assert.Nil(t, err)
+	})
+
+	t.Run("do with ctx timeout", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		ctx, cancel := context.WithTimeout(context.Background(), time.Microsecond)
+		defer cancel()
+		time.Sleep(time.Millisecond)
+		err := b.DoCtx(ctx, func() error {
+			return nil
+		})
+		assert.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+
+	t.Run("do with ctx cancel", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		for i := 0; i < 100; i++ {
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			cancel()
+			err := b.DoCtx(ctx, func() error {
+				return nil
+			})
+			assert.ErrorIs(t, err, context.Canceled)
+		}
+		assert.NoError(t, b.DoCtx(context.Background(), func() error {
+			return nil
+		}))
+	})
+}
+
+func TestCircuitBreaker_DoWithAcceptable(t *testing.T) {
+	t.Run("doWithAcceptable", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		err := b.DoWithAcceptable(func() error {
+			return nil
+		}, func(err error) bool {
+			return true
+		})
+		assert.Nil(t, err)
+	})
+
+	t.Run("doWithAcceptable with ctx", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		err := b.DoWithAcceptableCtx(context.Background(), func() error {
+			return nil
+		}, func(err error) bool {
+			return true
+		})
+		assert.Nil(t, err)
+	})
+
+	t.Run("doWithAcceptable with ctx timeout", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		ctx, cancel := context.WithTimeout(context.Background(), time.Microsecond)
+		defer cancel()
+		time.Sleep(time.Millisecond)
+		err := b.DoWithAcceptableCtx(ctx, func() error {
+			return nil
+		}, func(err error) bool {
+			return true
+		})
+		assert.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+
+	t.Run("doWithAcceptable with ctx cancel", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		for i := 0; i < 100; i++ {
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			cancel()
+			err := b.DoWithAcceptableCtx(ctx, func() error {
+				return nil
+			}, func(err error) bool {
+				return true
+			})
+			assert.ErrorIs(t, err, context.Canceled)
+		}
+		assert.NoError(t, b.DoWithAcceptableCtx(context.Background(), func() error {
+			return nil
+		}, func(err error) bool {
+			return true
+		}))
+	})
+}
+
+func TestCircuitBreaker_DoWithFallback(t *testing.T) {
+	t.Run("doWithFallback", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		err := b.DoWithFallback(func() error {
+			return nil
+		}, func(err error) error {
+			return err
+		})
+		assert.Nil(t, err)
+	})
+
+	t.Run("doWithFallback with ctx", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		err := b.DoWithFallbackCtx(context.Background(), func() error {
+			return nil
+		}, func(err error) error {
+			return err
+		})
+		assert.Nil(t, err)
+	})
+
+	t.Run("doWithFallback with ctx timeout", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		ctx, cancel := context.WithTimeout(context.Background(), time.Microsecond)
+		defer cancel()
+		time.Sleep(time.Millisecond)
+		err := b.DoWithFallbackCtx(ctx, func() error {
+			return nil
+		}, func(err error) error {
+			return err
+		})
+		assert.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+
+	t.Run("doWithFallback with ctx cancel", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		for i := 0; i < 100; i++ {
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			cancel()
+			err := b.DoWithFallbackCtx(ctx, func() error {
+				return nil
+			}, func(err error) error {
+				return err
+			})
+			assert.ErrorIs(t, err, context.Canceled)
+		}
+		assert.NoError(t, b.DoWithFallbackCtx(context.Background(), func() error {
+			return nil
+		}, func(err error) error {
+			return err
+		}))
+	})
+}
+
+func TestCircuitBreaker_DoWithFallbackAcceptable(t *testing.T) {
+	t.Run("doWithFallbackAcceptable", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		err := b.DoWithFallbackAcceptable(func() error {
+			return nil
+		}, func(err error) error {
+			return err
+		}, func(err error) bool {
+			return true
+		})
+		assert.Nil(t, err)
+	})
+
+	t.Run("doWithFallbackAcceptable with ctx", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		err := b.DoWithFallbackAcceptableCtx(context.Background(), func() error {
+			return nil
+		}, func(err error) error {
+			return err
+		}, func(err error) bool {
+			return true
+		})
+		assert.Nil(t, err)
+	})
+
+	t.Run("doWithFallbackAcceptable with ctx timeout", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		ctx, cancel := context.WithTimeout(context.Background(), time.Microsecond)
+		defer cancel()
+		time.Sleep(time.Millisecond)
+		err := b.DoWithFallbackAcceptableCtx(ctx, func() error {
+			return nil
+		}, func(err error) error {
+			return err
+		}, func(err error) bool {
+			return true
+		})
+		assert.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+
+	t.Run("doWithFallbackAcceptable with ctx cancel", func(t *testing.T) {
+		b := NewBreaker()
+		assert.True(t, len(b.Name()) > 0)
+		for i := 0; i < 100; i++ {
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			cancel()
+			err := b.DoWithFallbackAcceptableCtx(ctx, func() error {
+				return nil
+			}, func(err error) error {
+				return err
+			}, func(err error) bool {
+				return true
+			})
+			assert.ErrorIs(t, err, context.Canceled)
+		}
+		assert.NoError(t, b.DoWithFallbackAcceptableCtx(context.Background(), func() error {
+			return nil
+		}, func(err error) error {
+			return err
+		}, func(err error) bool {
+			return true
+		}))
+	})
 }
 
 func TestLogReason(t *testing.T) {

core/breaker/breakers.go

@@ -1,6 +1,9 @@
 package breaker
 
-import "sync"
+import (
+	"context"
+	"sync"
+)
 
 var (
 	lock     sync.RWMutex
@@ -14,6 +17,13 @@ func Do(name string, req func() error) error {
 	})
 }
 
+// DoCtx calls Breaker.DoCtx on the Breaker with given name.
+func DoCtx(ctx context.Context, name string, req func() error) error {
+	return do(name, func(b Breaker) error {
+		return b.DoCtx(ctx, req)
+	})
+}
+
 // DoWithAcceptable calls Breaker.DoWithAcceptable on the Breaker with given name.
 func DoWithAcceptable(name string, req func() error, acceptable Acceptable) error {
 	return do(name, func(b Breaker) error {
@@ -21,6 +31,14 @@ func DoWithAcceptable(name string, req func() error, acceptable Acceptable) erro
 	})
 }
 
+// DoWithAcceptableCtx calls Breaker.DoWithAcceptableCtx on the Breaker with given name.
+func DoWithAcceptableCtx(ctx context.Context, name string, req func() error,
+	acceptable Acceptable) error {
+	return do(name, func(b Breaker) error {
+		return b.DoWithAcceptableCtx(ctx, req, acceptable)
+	})
+}
+
 // DoWithFallback calls Breaker.DoWithFallback on the Breaker with given name.
 func DoWithFallback(name string, req func() error, fallback Fallback) error {
 	return do(name, func(b Breaker) error {
@@ -28,6 +46,13 @@ func DoWithFallback(name string, req func() error, fallback Fallback) error {
 	})
 }
 
+// DoWithFallbackCtx calls Breaker.DoWithFallbackCtx on the Breaker with given name.
+func DoWithFallbackCtx(ctx context.Context, name string, req func() error, fallback Fallback) error {
+	return do(name, func(b Breaker) error {
+		return b.DoWithFallbackCtx(ctx, req, fallback)
+	})
+}
+
 // DoWithFallbackAcceptable calls Breaker.DoWithFallbackAcceptable on the Breaker with given name.
 func DoWithFallbackAcceptable(name string, req func() error, fallback Fallback,
 	acceptable Acceptable) error {
@@ -36,6 +61,14 @@ func DoWithFallbackAcceptable(name string, req func() error, fallback Fallback,
 	})
 }
 
+// DoWithFallbackAcceptableCtx calls Breaker.DoWithFallbackAcceptableCtx on the Breaker with given name.
+func DoWithFallbackAcceptableCtx(ctx context.Context, name string, req func() error,
+	fallback Fallback, acceptable Acceptable) error {
+	return do(name, func(b Breaker) error {
+		return b.DoWithFallbackAcceptableCtx(ctx, req, fallback, acceptable)
+	})
+}
+
 // GetBreaker returns the Breaker with the given name.
 func GetBreaker(name string) Breaker {
 	lock.RLock()

core/breaker/breakers_test.go

@@ -1,6 +1,7 @@
 package breaker
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"testing"
@@ -22,6 +23,9 @@ func TestBreakersDo(t *testing.T) {
 	assert.Equal(t, errDummy, Do("any", func() error {
 		return errDummy
 	}))
+	assert.Equal(t, errDummy, DoCtx(context.Background(), "any", func() error {
+		return errDummy
+	}))
 }
 
 func TestBreakersDoWithAcceptable(t *testing.T) {
@@ -38,6 +42,13 @@ func TestBreakersDoWithAcceptable(t *testing.T) {
 			return nil
 		}) == nil
 	})
+	verify(t, func() bool {
+		return DoWithAcceptableCtx(context.Background(), "anyone", func() error {
+			return nil
+		}, func(err error) bool {
+			return true
+		}) == nil
+	})
 
 	for i := 0; i < 10000; i++ {
 		err := DoWithAcceptable("another", func() error {
@@ -76,6 +87,12 @@ func TestBreakersFallback(t *testing.T) {
 			return nil
 		})
 		assert.True(t, err == nil || errors.Is(err, errDummy))
+		err = DoWithFallbackCtx(context.Background(), "fallback", func() error {
+			return errDummy
+		}, func(err error) error {
+			return nil
+		})
+		assert.True(t, err == nil || errors.Is(err, errDummy))
 	}
 	verify(t, func() bool {
 		return errors.Is(Do("fallback", func() error {
@@ -86,7 +103,7 @@ func TestBreakersFallback(t *testing.T) {
 
 func TestBreakersAcceptableFallback(t *testing.T) {
 	errDummy := errors.New("any")
-	for i := 0; i < 10000; i++ {
+	for i := 0; i < 5000; i++ {
 		err := DoWithFallbackAcceptable("acceptablefallback", func() error {
 			return errDummy
 		}, func(err error) error {
@@ -95,6 +112,14 @@ func TestBreakersAcceptableFallback(t *testing.T) {
 			return err == nil
 		})
 		assert.True(t, err == nil || errors.Is(err, errDummy))
+		err = DoWithFallbackAcceptableCtx(context.Background(), "acceptablefallback", func() error {
+			return errDummy
+		}, func(err error) error {
+			return nil
+		}, func(err error) bool {
+			return err == nil
+		})
+		assert.True(t, err == nil || errors.Is(err, errDummy))
 	}
 	verify(t, func() bool {
 		return errors.Is(Do("acceptablefallback", func() error {
@@ -110,5 +135,5 @@ func verify(t *testing.T, fn func() bool) {
 			count++
 		}
 	}
-	assert.True(t, count >= 80, fmt.Sprintf("should be greater than 80, actual %d", count))
+	assert.True(t, count >= 75, fmt.Sprintf("should be greater than 75, actual %d", count))
 }

core/breaker/bucket.go

@@ -0,0 +1,48 @@
+package breaker
+
+const (
+	success = iota
+	fail
+	drop
+)
+
+// bucket defines the bucket that holds sum and num of additions.
+type bucket struct {
+	Sum     int64
+	Success int64
+	Failure int64
+	Drop    int64
+}
+
+func (b *bucket) Add(v int64) {
+	switch v {
+	case fail:
+		b.fail()
+	case drop:
+		b.drop()
+	default:
+		b.succeed()
+	}
+}
+
+func (b *bucket) Reset() {
+	b.Sum = 0
+	b.Success = 0
+	b.Failure = 0
+	b.Drop = 0
+}
+
+func (b *bucket) drop() {
+	b.Sum++
+	b.Drop++
+}
+
+func (b *bucket) fail() {
+	b.Sum++
+	b.Failure++
+}
+
+func (b *bucket) succeed() {
+	b.Sum++
+	b.Success++
+}

core/breaker/bucket_test.go

@@ -0,0 +1,43 @@
+package breaker
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBucketAdd(t *testing.T) {
+	b := &bucket{}
+
+	// Test succeed
+	b.Add(0) // Using 0 for success
+	assert.Equal(t, int64(1), b.Sum, "Sum should be incremented")
+	assert.Equal(t, int64(1), b.Success, "Success should be incremented")
+	assert.Equal(t, int64(0), b.Failure, "Failure should not be incremented")
+	assert.Equal(t, int64(0), b.Drop, "Drop should not be incremented")
+
+	// Test failure
+	b.Add(fail)
+	assert.Equal(t, int64(2), b.Sum, "Sum should be incremented")
+	assert.Equal(t, int64(1), b.Failure, "Failure should be incremented")
+	assert.Equal(t, int64(0), b.Drop, "Drop should not be incremented")
+
+	// Test drop
+	b.Add(drop)
+	assert.Equal(t, int64(3), b.Sum, "Sum should be incremented")
+	assert.Equal(t, int64(1), b.Drop, "Drop should be incremented")
+}
+
+func TestBucketReset(t *testing.T) {
+	b := &bucket{
+		Sum:     3,
+		Success: 1,
+		Failure: 1,
+		Drop:    1,
+	}
+	b.Reset()
+	assert.Equal(t, int64(0), b.Sum, "Sum should be reset to 0")
+	assert.Equal(t, int64(0), b.Success, "Success should be reset to 0")
+	assert.Equal(t, int64(0), b.Failure, "Failure should be reset to 0")
+	assert.Equal(t, int64(0), b.Drop, "Drop should be reset to 0")
+}

core/breaker/googlebreaker.go

@@ -5,53 +5,83 @@ import (
 
 	"github.com/zeromicro/go-zero/core/collection"
 	"github.com/zeromicro/go-zero/core/mathx"
+	"github.com/zeromicro/go-zero/core/syncx"
+	"github.com/zeromicro/go-zero/core/timex"
 )
 
 const (
 	// 250ms for bucket duration
-	window     = time.Second * 10
-	buckets    = 40
-	k          = 1.5
-	protection = 5
+	window            = time.Second * 10
+	buckets           = 40
+	forcePassDuration = time.Second
+	k                 = 1.5
+	minK              = 1.1
+	protection        = 5
 )
 
 // googleBreaker is a netflixBreaker pattern from google.
 // see Client-Side Throttling section in https://landing.google.com/sre/sre-book/chapters/handling-overload/
-type googleBreaker struct {
-	k     float64
-	stat  *collection.RollingWindow
-	proba *mathx.Proba
-}
+type (
+	googleBreaker struct {
+		k        float64
+		stat     *collection.RollingWindow[int64, *bucket]
+		proba    *mathx.Proba
+		lastPass *syncx.AtomicDuration
+	}
+
+	windowResult struct {
+		accepts        int64
+		total          int64
+		failingBuckets int64
+		workingBuckets int64
+	}
+)
 
 func newGoogleBreaker() *googleBreaker {
 	bucketDuration := time.Duration(int64(window) / int64(buckets))
-	st := collection.NewRollingWindow(buckets, bucketDuration)
+	st := collection.NewRollingWindow[int64, *bucket](func() *bucket {
+		return new(bucket)
+	}, buckets, bucketDuration)
 	return &googleBreaker{
-		stat:  st,
-		k:     k,
-		proba: mathx.NewProba(),
+		stat:     st,
+		k:        k,
+		proba:    mathx.NewProba(),
+		lastPass: syncx.NewAtomicDuration(),
 	}
 }
 
 func (b *googleBreaker) accept() error {
-	accepts, total := b.history()
-	weightedAccepts := b.k * float64(accepts)
+	var w float64
+	history := b.history()
+	w = b.k - (b.k-minK)*float64(history.failingBuckets)/buckets
+	weightedAccepts := mathx.AtLeast(w, minK) * float64(history.accepts)
 	// https://landing.google.com/sre/sre-book/chapters/handling-overload/#eq2101
-	// for better performance, no need to care about negative ratio
-	dropRatio := (float64(total-protection) - weightedAccepts) / float64(total+1)
+	// for better performance, no need to care about the negative ratio
+	dropRatio := (float64(history.total-protection) - weightedAccepts) / float64(history.total+1)
 	if dropRatio <= 0 {
 		return nil
 	}
 
+	lastPass := b.lastPass.Load()
+	if lastPass > 0 && timex.Since(lastPass) > forcePassDuration {
+		b.lastPass.Set(timex.Now())
+		return nil
+	}
+
+	dropRatio *= float64(buckets-history.workingBuckets) / buckets
+
 	if b.proba.TrueOnProba(dropRatio) {
 		return ErrServiceUnavailable
 	}
 
+	b.lastPass.Set(timex.Now())
+
 	return nil
 }
 
 func (b *googleBreaker) allow() (internalPromise, error) {
 	if err := b.accept(); err != nil {
+		b.markDrop()
 		return nil, err
 	}
 
@@ -62,7 +92,7 @@ func (b *googleBreaker) allow() (internalPromise, error) {
 
 func (b *googleBreaker) doReq(req func() error, fallback Fallback, acceptable Acceptable) error {
 	if err := b.accept(); err != nil {
-		b.markFailure()
+		b.markDrop()
 		if fallback != nil {
 			return fallback(err)
 		}
@@ -70,10 +100,10 @@ func (b *googleBreaker) doReq(req func() error, fallback Fallback, acceptable Ac
 		return err
 	}
 
-	var success bool
+	var succ bool
 	defer func() {
 		// if req() panic, success is false, mark as failure
-		if success {
+		if succ {
 			b.markSuccess()
 		} else {
 			b.markFailure()
@@ -82,27 +112,43 @@ func (b *googleBreaker) doReq(req func() error, fallback Fallback, acceptable Ac
 
 	err := req()
 	if acceptable(err) {
-		success = true
+		succ = true
 	}
 
 	return err
 }
 
-func (b *googleBreaker) markSuccess() {
-	b.stat.Add(1)
+func (b *googleBreaker) markDrop() {
+	b.stat.Add(drop)
 }
 
 func (b *googleBreaker) markFailure() {
-	b.stat.Add(0)
+	b.stat.Add(fail)
 }
 
-func (b *googleBreaker) history() (accepts, total int64) {
-	b.stat.Reduce(func(b *collection.Bucket) {
-		accepts += int64(b.Sum)
-		total += b.Count
+func (b *googleBreaker) markSuccess() {
+	b.stat.Add(success)
+}
+
+func (b *googleBreaker) history() windowResult {
+	var result windowResult
+
+	b.stat.Reduce(func(b *bucket) {
+		result.accepts += b.Success
+		result.total += b.Sum
+		if b.Failure > 0 {
+			result.workingBuckets = 0
+		} else if b.Success > 0 {
+			result.workingBuckets++
+		}
+		if b.Success > 0 {
+			result.failingBuckets = 0
+		} else if b.Failure > 0 {
+			result.failingBuckets++
+		}
 	})
 
-	return
+	return result
 }
 
 type googlePromise struct {

core/breaker/googlebreaker_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/zeromicro/go-zero/core/collection"
 	"github.com/zeromicro/go-zero/core/mathx"
 	"github.com/zeromicro/go-zero/core/stat"
+	"github.com/zeromicro/go-zero/core/syncx"
 )
 
 const (
@@ -22,11 +23,14 @@ func init() {
 }
 
 func getGoogleBreaker() *googleBreaker {
-	st := collection.NewRollingWindow(testBuckets, testInterval)
+	st := collection.NewRollingWindow[int64, *bucket](func() *bucket {
+		return new(bucket)
+	}, testBuckets, testInterval)
 	return &googleBreaker{
-		stat:  st,
-		k:     5,
-		proba: mathx.NewProba(),
+		stat:     st,
+		k:        5,
+		proba:    mathx.NewProba(),
+		lastPass: syncx.NewAtomicDuration(),
 	}
 }
 
@@ -63,6 +67,33 @@ func TestGoogleBreakerOpen(t *testing.T) {
 	})
 }
 
+func TestGoogleBreakerRecover(t *testing.T) {
+	st := collection.NewRollingWindow[int64, *bucket](func() *bucket {
+		return new(bucket)
+	}, testBuckets*2, testInterval)
+	b := &googleBreaker{
+		stat:     st,
+		k:        k,
+		proba:    mathx.NewProba(),
+		lastPass: syncx.NewAtomicDuration(),
+	}
+	for i := 0; i < testBuckets; i++ {
+		for j := 0; j < 100; j++ {
+			b.stat.Add(1)
+		}
+		time.Sleep(testInterval)
+	}
+	for i := 0; i < testBuckets; i++ {
+		for j := 0; j < 100; j++ {
+			b.stat.Add(0)
+		}
+		time.Sleep(testInterval)
+	}
+	verify(t, func() bool {
+		return b.accept() == nil
+	})
+}
+
 func TestGoogleBreakerFallback(t *testing.T) {
 	b := getGoogleBreaker()
 	markSuccess(b, 1)
@@ -89,6 +120,43 @@ func TestGoogleBreakerReject(t *testing.T) {
 	}, nil, defaultAcceptable))
 }
 
+func TestGoogleBreakerMoreFallingBuckets(t *testing.T) {
+	t.Parallel()
+
+	t.Run("more falling buckets", func(t *testing.T) {
+		b := getGoogleBreaker()
+
+		func() {
+			stopChan := time.After(testInterval * 6)
+			for {
+				time.Sleep(time.Millisecond)
+				select {
+				case <-stopChan:
+					return
+				default:
+					assert.Error(t, b.doReq(func() error {
+						return errors.New("foo")
+					}, func(err error) error {
+						return err
+					}, func(err error) bool {
+						return err == nil
+					}))
+				}
+			}
+		}()
+
+		var count int
+		for i := 0; i < 100; i++ {
+			if errors.Is(b.doReq(func() error {
+				return ErrServiceUnavailable
+			}, nil, defaultAcceptable), ErrServiceUnavailable) {
+				count++
+			}
+		}
+		assert.True(t, count > 90)
+	})
+}
+
 func TestGoogleBreakerAcceptable(t *testing.T) {
 	b := getGoogleBreaker()
 	errAcceptable := errors.New("any")
@@ -164,41 +232,38 @@ func TestGoogleBreakerSelfProtection(t *testing.T) {
 }
 
 func TestGoogleBreakerHistory(t *testing.T) {
-	var b *googleBreaker
-	var accepts, total int64
-
 	sleep := testInterval
 	t.Run("accepts == total", func(t *testing.T) {
-		b = getGoogleBreaker()
+		b := getGoogleBreaker()
 		markSuccessWithDuration(b, 10, sleep/2)
-		accepts, total = b.history()
-		assert.Equal(t, int64(10), accepts)
-		assert.Equal(t, int64(10), total)
+		result := b.history()
+		assert.Equal(t, int64(10), result.accepts)
+		assert.Equal(t, int64(10), result.total)
 	})
 
 	t.Run("fail == total", func(t *testing.T) {
-		b = getGoogleBreaker()
+		b := getGoogleBreaker()
 		markFailedWithDuration(b, 10, sleep/2)
-		accepts, total = b.history()
-		assert.Equal(t, int64(0), accepts)
-		assert.Equal(t, int64(10), total)
+		result := b.history()
+		assert.Equal(t, int64(0), result.accepts)
+		assert.Equal(t, int64(10), result.total)
 	})
 
 	t.Run("accepts = 1/2 * total, fail = 1/2 * total", func(t *testing.T) {
-		b = getGoogleBreaker()
+		b := getGoogleBreaker()
 		markFailedWithDuration(b, 5, sleep/2)
 		markSuccessWithDuration(b, 5, sleep/2)
-		accepts, total = b.history()
-		assert.Equal(t, int64(5), accepts)
-		assert.Equal(t, int64(10), total)
+		result := b.history()
+		assert.Equal(t, int64(5), result.accepts)
+		assert.Equal(t, int64(10), result.total)
 	})
 
 	t.Run("auto reset rolling counter", func(t *testing.T) {
-		b = getGoogleBreaker()
+		b := getGoogleBreaker()
 		time.Sleep(testInterval * testBuckets)
-		accepts, total = b.history()
-		assert.Equal(t, int64(0), accepts)
-		assert.Equal(t, int64(0), total)
+		result := b.history()
+		assert.Equal(t, int64(0), result.accepts)
+		assert.Equal(t, int64(0), result.total)
 	})
 }
 

core/breaker/nopbreaker.go

@@ -1,5 +1,7 @@
 package breaker
 
+import "context"
+
 const nopBreakerName = "nopBreaker"
 
 type nopBreaker struct{}
@@ -17,22 +19,43 @@ func (b nopBreaker) Allow() (Promise, error) {
 	return nopPromise{}, nil
 }
 
+func (b nopBreaker) AllowCtx(_ context.Context) (Promise, error) {
+	return nopPromise{}, nil
+}
+
 func (b nopBreaker) Do(req func() error) error {
 	return req()
 }
 
+func (b nopBreaker) DoCtx(_ context.Context, req func() error) error {
+	return req()
+}
+
 func (b nopBreaker) DoWithAcceptable(req func() error, _ Acceptable) error {
 	return req()
 }
 
+func (b nopBreaker) DoWithAcceptableCtx(_ context.Context, req func() error, _ Acceptable) error {
+	return req()
+}
+
 func (b nopBreaker) DoWithFallback(req func() error, _ Fallback) error {
 	return req()
 }
 
+func (b nopBreaker) DoWithFallbackCtx(_ context.Context, req func() error, _ Fallback) error {
+	return req()
+}
+
 func (b nopBreaker) DoWithFallbackAcceptable(req func() error, _ Fallback, _ Acceptable) error {
 	return req()
 }
 
+func (b nopBreaker) DoWithFallbackAcceptableCtx(_ context.Context, req func() error,
+	_ Fallback, _ Acceptable) error {
+	return req()
+}
+
 type nopPromise struct{}
 
 func (p nopPromise) Accept() {

core/breaker/nopbreaker_test.go

@@ -1,6 +1,7 @@
 package breaker
 
 import (
+	"context"
 	"errors"
 	"testing"
 
@@ -12,6 +13,8 @@ func TestNopBreaker(t *testing.T) {
 	assert.Equal(t, nopBreakerName, b.Name())
 	p, err := b.Allow()
 	assert.Nil(t, err)
+	p, err = b.AllowCtx(context.Background())
+	assert.Nil(t, err)
 	p.Accept()
 	for i := 0; i < 1000; i++ {
 		p, err := b.Allow()
@@ -21,18 +24,34 @@ func TestNopBreaker(t *testing.T) {
 	assert.Nil(t, b.Do(func() error {
 		return nil
 	}))
+	assert.Nil(t, b.DoCtx(context.Background(), func() error {
+		return nil
+	}))
 	assert.Nil(t, b.DoWithAcceptable(func() error {
 		return nil
 	}, defaultAcceptable))
+	assert.Nil(t, b.DoWithAcceptableCtx(context.Background(), func() error {
+		return nil
+	}, defaultAcceptable))
 	errDummy := errors.New("any")
 	assert.Equal(t, errDummy, b.DoWithFallback(func() error {
 		return errDummy
 	}, func(err error) error {
 		return nil
 	}))
+	assert.Equal(t, errDummy, b.DoWithFallbackCtx(context.Background(), func() error {
+		return errDummy
+	}, func(err error) error {
+		return nil
+	}))
 	assert.Equal(t, errDummy, b.DoWithFallbackAcceptable(func() error {
 		return errDummy
 	}, func(err error) error {
 		return nil
 	}, defaultAcceptable))
+	assert.Equal(t, errDummy, b.DoWithFallbackAcceptableCtx(context.Background(), func() error {
+		return errDummy
+	}, func(err error) error {
+		return nil
+	}, defaultAcceptable))
 }

core/codec/dh.go

@@ -23,7 +23,7 @@ var (
 	zero = big.NewInt(0)
 )
 
-// DhKey defines the Diffie Hellman key.
+// DhKey defines the Diffie-Hellman key.
 type DhKey struct {
 	PriKey *big.Int
 	PubKey *big.Int
@@ -46,7 +46,7 @@ func ComputeKey(pubKey, priKey *big.Int) (*big.Int, error) {
 	return new(big.Int).Exp(pubKey, priKey, p), nil
 }
 
-// GenerateKey returns a Diffie Hellman key.
+// GenerateKey returns a Diffie-Hellman key.
 func GenerateKey() (*DhKey, error) {
 	var err error
 	var x *big.Int

core/collection/cache.go

@@ -128,8 +128,8 @@ func (c *Cache) Take(key string, fetch func() (any, error)) (any, error) {
 
 	var fresh bool
 	val, err := c.barrier.Do(key, func() (any, error) {
-		// because O(1) on map search in memory, and fetch is an IO query
-		// so we do double check, cache might be taken by another call
+		// because O(1) on map search in memory, and fetch is an IO query,
+		// so we do double-check, cache might be taken by another call
 		if val, ok := c.doGet(key); ok {
 			return val, nil
 		}

core/collection/rollingwindow.go

@@ -4,18 +4,28 @@ import (
 	"sync"
 	"time"
 
+	"github.com/zeromicro/go-zero/core/mathx"
 	"github.com/zeromicro/go-zero/core/timex"
 )
 
 type (
-	// RollingWindowOption let callers customize the RollingWindow.
-	RollingWindowOption func(rollingWindow *RollingWindow)
+	// BucketInterface is the interface that defines the buckets.
+	BucketInterface[T Numerical] interface {
+		Add(v T)
+		Reset()
+	}
 
-	// RollingWindow defines a rolling window to calculate the events in buckets with time interval.
-	RollingWindow struct {
+	// Numerical is the interface that restricts the numerical type.
+	Numerical = mathx.Numerical
+
+	// RollingWindowOption let callers customize the RollingWindow.
+	RollingWindowOption[T Numerical, B BucketInterface[T]] func(rollingWindow *RollingWindow[T, B])
+
+	// RollingWindow defines a rolling window to calculate the events in buckets with the time interval.
+	RollingWindow[T Numerical, B BucketInterface[T]] struct {
 		lock          sync.RWMutex
 		size          int
-		win           *window
+		win           *window[T, B]
 		interval      time.Duration
 		offset        int
 		ignoreCurrent bool
@@ -25,14 +35,15 @@ type (
 
 // NewRollingWindow returns a RollingWindow that with size buckets and time interval,
 // use opts to customize the RollingWindow.
-func NewRollingWindow(size int, interval time.Duration, opts ...RollingWindowOption) *RollingWindow {
+func NewRollingWindow[T Numerical, B BucketInterface[T]](newBucket func() B, size int,
+	interval time.Duration, opts ...RollingWindowOption[T, B]) *RollingWindow[T, B] {
 	if size < 1 {
 		panic("size must be greater than 0")
 	}
 
-	w := &RollingWindow{
+	w := &RollingWindow[T, B]{
 		size:     size,
-		win:      newWindow(size),
+		win:      newWindow[T, B](newBucket, size),
 		interval: interval,
 		lastTime: timex.Now(),
 	}
@@ -43,7 +54,7 @@ func NewRollingWindow(size int, interval time.Duration, opts ...RollingWindowOpt
 }
 
 // Add adds value to current bucket.
-func (rw *RollingWindow) Add(v float64) {
+func (rw *RollingWindow[T, B]) Add(v T) {
 	rw.lock.Lock()
 	defer rw.lock.Unlock()
 	rw.updateOffset()
@@ -51,13 +62,13 @@ func (rw *RollingWindow) Add(v float64) {
 }
 
 // Reduce runs fn on all buckets, ignore current bucket if ignoreCurrent was set.
-func (rw *RollingWindow) Reduce(fn func(b *Bucket)) {
+func (rw *RollingWindow[T, B]) Reduce(fn func(b B)) {
 	rw.lock.RLock()
 	defer rw.lock.RUnlock()
 
 	var diff int
 	span := rw.span()
-	// ignore current bucket, because of partial data
+	// ignore the current bucket, because of partial data
 	if span == 0 && rw.ignoreCurrent {
 		diff = rw.size - 1
 	} else {
@@ -69,7 +80,7 @@ func (rw *RollingWindow) Reduce(fn func(b *Bucket)) {
 	}
 }
 
-func (rw *RollingWindow) span() int {
+func (rw *RollingWindow[T, B]) span() int {
 	offset := int(timex.Since(rw.lastTime) / rw.interval)
 	if 0 <= offset && offset < rw.size {
 		return offset
@@ -78,7 +89,7 @@ func (rw *RollingWindow) span() int {
 	return rw.size
 }
 
-func (rw *RollingWindow) updateOffset() {
+func (rw *RollingWindow[T, B]) updateOffset() {
 	span := rw.span()
 	if span <= 0 {
 		return
@@ -97,54 +108,54 @@ func (rw *RollingWindow) updateOffset() {
 }
 
 // Bucket defines the bucket that holds sum and num of additions.
-type Bucket struct {
-	Sum   float64
+type Bucket[T Numerical] struct {
+	Sum   T
 	Count int64
 }
 
-func (b *Bucket) add(v float64) {
+func (b *Bucket[T]) Add(v T) {
 	b.Sum += v
 	b.Count++
 }
 
-func (b *Bucket) reset() {
+func (b *Bucket[T]) Reset() {
 	b.Sum = 0
 	b.Count = 0
 }
 
-type window struct {
-	buckets []*Bucket
+type window[T Numerical, B BucketInterface[T]] struct {
+	buckets []B
 	size    int
 }
 
-func newWindow(size int) *window {
-	buckets := make([]*Bucket, size)
+func newWindow[T Numerical, B BucketInterface[T]](newBucket func() B, size int) *window[T, B] {
+	buckets := make([]B, size)
 	for i := 0; i < size; i++ {
-		buckets[i] = new(Bucket)
+		buckets[i] = newBucket()
 	}
-	return &window{
+	return &window[T, B]{
 		buckets: buckets,
 		size:    size,
 	}
 }
 
-func (w *window) add(offset int, v float64) {
-	w.buckets[offset%w.size].add(v)
+func (w *window[T, B]) add(offset int, v T) {
+	w.buckets[offset%w.size].Add(v)
 }
 
-func (w *window) reduce(start, count int, fn func(b *Bucket)) {
+func (w *window[T, B]) reduce(start, count int, fn func(b B)) {
 	for i := 0; i < count; i++ {
 		fn(w.buckets[(start+i)%w.size])
 	}
 }
 
-func (w *window) resetBucket(offset int) {
-	w.buckets[offset%w.size].reset()
+func (w *window[T, B]) resetBucket(offset int) {
+	w.buckets[offset%w.size].Reset()
 }
 
 // IgnoreCurrentBucket lets the Reduce call ignore current bucket.
-func IgnoreCurrentBucket() RollingWindowOption {
-	return func(w *RollingWindow) {
+func IgnoreCurrentBucket[T Numerical, B BucketInterface[T]]() RollingWindowOption[T, B] {
+	return func(w *RollingWindow[T, B]) {
 		w.ignoreCurrent = true
 	}
 }

core/collection/rollingwindow_test.go

@@ -12,18 +12,24 @@ import (
 const duration = time.Millisecond * 50
 
 func TestNewRollingWindow(t *testing.T) {
-	assert.NotNil(t, NewRollingWindow(10, time.Second))
+	assert.NotNil(t, NewRollingWindow[int64, *Bucket[int64]](func() *Bucket[int64] {
+		return new(Bucket[int64])
+	}, 10, time.Second))
 	assert.Panics(t, func() {
-		NewRollingWindow(0, time.Second)
+		NewRollingWindow[int64, *Bucket[int64]](func() *Bucket[int64] {
+			return new(Bucket[int64])
+		}, 0, time.Second)
 	})
 }
 
 func TestRollingWindowAdd(t *testing.T) {
 	const size = 3
-	r := NewRollingWindow(size, duration)
+	r := NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+		return new(Bucket[float64])
+	}, size, duration)
 	listBuckets := func() []float64 {
 		var buckets []float64
-		r.Reduce(func(b *Bucket) {
+		r.Reduce(func(b *Bucket[float64]) {
 			buckets = append(buckets, b.Sum)
 		})
 		return buckets
@@ -47,10 +53,12 @@ func TestRollingWindowAdd(t *testing.T) {
 
 func TestRollingWindowReset(t *testing.T) {
 	const size = 3
-	r := NewRollingWindow(size, duration, IgnoreCurrentBucket())
+	r := NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+		return new(Bucket[float64])
+	}, size, duration, IgnoreCurrentBucket[float64, *Bucket[float64]]())
 	listBuckets := func() []float64 {
 		var buckets []float64
-		r.Reduce(func(b *Bucket) {
+		r.Reduce(func(b *Bucket[float64]) {
 			buckets = append(buckets, b.Sum)
 		})
 		return buckets
@@ -72,15 +80,19 @@ func TestRollingWindowReset(t *testing.T) {
 func TestRollingWindowReduce(t *testing.T) {
 	const size = 4
 	tests := []struct {
-		win    *RollingWindow
+		win    *RollingWindow[float64, *Bucket[float64]]
 		expect float64
 	}{
 		{
-			win:    NewRollingWindow(size, duration),
+			win: NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+				return new(Bucket[float64])
+			}, size, duration),
 			expect: 10,
 		},
 		{
-			win:    NewRollingWindow(size, duration, IgnoreCurrentBucket()),
+			win: NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+				return new(Bucket[float64])
+			}, size, duration, IgnoreCurrentBucket[float64, *Bucket[float64]]()),
 			expect: 4,
 		},
 	}
@@ -97,7 +109,7 @@ func TestRollingWindowReduce(t *testing.T) {
 				}
 			}
 			var result float64
-			r.Reduce(func(b *Bucket) {
+			r.Reduce(func(b *Bucket[float64]) {
 				result += b.Sum
 			})
 			assert.Equal(t, test.expect, result)
@@ -108,10 +120,12 @@ func TestRollingWindowReduce(t *testing.T) {
 func TestRollingWindowBucketTimeBoundary(t *testing.T) {
 	const size = 3
 	interval := time.Millisecond * 30
-	r := NewRollingWindow(size, interval)
+	r := NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+		return new(Bucket[float64])
+	}, size, interval)
 	listBuckets := func() []float64 {
 		var buckets []float64
-		r.Reduce(func(b *Bucket) {
+		r.Reduce(func(b *Bucket[float64]) {
 			buckets = append(buckets, b.Sum)
 		})
 		return buckets
@@ -138,7 +152,9 @@ func TestRollingWindowBucketTimeBoundary(t *testing.T) {
 
 func TestRollingWindowDataRace(t *testing.T) {
 	const size = 3
-	r := NewRollingWindow(size, duration)
+	r := NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+		return new(Bucket[float64])
+	}, size, duration)
 	stop := make(chan bool)
 	go func() {
 		for {
@@ -157,7 +173,7 @@ func TestRollingWindowDataRace(t *testing.T) {
 			case <-stop:
 				return
 			default:
-				r.Reduce(func(b *Bucket) {})
+				r.Reduce(func(b *Bucket[float64]) {})
 			}
 		}
 	}()

core/conf/config.go

@@ -133,7 +133,7 @@ func addOrMergeFields(info *fieldInfo, key string, child *fieldInfo, fullName st
 			return newConflictKeyError(fullName)
 		}
 
-		if err := mergeFields(prev, key, child.children, fullName); err != nil {
+		if err := mergeFields(prev, child.children, fullName); err != nil {
 			return err
 		}
 	} else {
@@ -281,7 +281,7 @@ func getTagName(field reflect.StructField) string {
 	return field.Name
 }
 
-func mergeFields(prev *fieldInfo, key string, children map[string]*fieldInfo, fullName string) error {
+func mergeFields(prev *fieldInfo, children map[string]*fieldInfo, fullName string) error {
 	if len(prev.children) == 0 || len(children) == 0 {
 		return newConflictKeyError(fullName)
 	}

core/discov/internal/registry.go

@@ -30,7 +30,7 @@ var (
 // A Registry is a registry that manages the etcd client connections.
 type Registry struct {
 	clusters map[string]*cluster
-	lock     sync.Mutex
+	lock     sync.RWMutex
 }
 
 // GetRegistry returns a global Registry.
@@ -60,12 +60,19 @@ func (r *Registry) Monitor(endpoints []string, key string, l UpdateListener) err
 
 func (r *Registry) getCluster(endpoints []string) (c *cluster, exists bool) {
 	clusterKey := getClusterKey(endpoints)
-	r.lock.Lock()
-	defer r.lock.Unlock()
+	r.lock.RLock()
 	c, exists = r.clusters[clusterKey]
+	r.lock.RUnlock()
+
 	if !exists {
-		c = newCluster(endpoints)
-		r.clusters[clusterKey] = c
+		r.lock.Lock()
+		defer r.lock.Unlock()
+		// double-check locking
+		c, exists = r.clusters[clusterKey]
+		if !exists {
+			c = newCluster(endpoints)
+			r.clusters[clusterKey] = c
+		}
 	}
 
 	return
@@ -78,7 +85,7 @@ type cluster struct {
 	listeners  map[string][]UpdateListener
 	watchGroup *threading.RoutineGroup
 	done       chan lang.PlaceholderType
-	lock       sync.Mutex
+	lock       sync.RWMutex
 }
 
 func newCluster(endpoints []string) *cluster {
@@ -108,8 +115,8 @@ func (c *cluster) getClient() (EtcdClient, error) {
 }
 
 func (c *cluster) getCurrent(key string) []KV {
-	c.lock.Lock()
-	defer c.lock.Unlock()
+	c.lock.RLock()
+	defer c.lock.RUnlock()
 
 	var kvs []KV
 	for k, v := range c.values[key] {
@@ -125,6 +132,7 @@ func (c *cluster) getCurrent(key string) []KV {
 func (c *cluster) handleChanges(key string, kvs []KV) {
 	var add []KV
 	var remove []KV
+
 	c.lock.Lock()
 	listeners := append([]UpdateListener(nil), c.listeners[key]...)
 	vals, ok := c.values[key]
@@ -173,9 +181,9 @@ func (c *cluster) handleChanges(key string, kvs []KV) {
 }
 
 func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
-	c.lock.Lock()
+	c.lock.RLock()
 	listeners := append([]UpdateListener(nil), c.listeners[key]...)
-	c.lock.Unlock()
+	c.lock.RUnlock()
 
 	for _, ev := range events {
 		switch ev.Type {
@@ -222,7 +230,7 @@ func (c *cluster) load(cli EtcdClient, key string) int64 {
 			break
 		}
 
-		logx.Error(err)
+		logx.Errorf("%s, key is %s", err.Error(), key)
 		time.Sleep(coolDownInterval)
 	}
 

core/errorx/check.go

@@ -0,0 +1,14 @@
+package errorx
+
+import "errors"
+
+// In checks if the given err is one of errs.
+func In(err error, errs ...error) bool {
+	for _, each := range errs {
+		if errors.Is(err, each) {
+			return true
+		}
+	}
+
+	return false
+}

core/errorx/check_test.go

@@ -0,0 +1,70 @@
+package errorx
+
+import (
+	"errors"
+	"testing"
+)
+
+func TestIn(t *testing.T) {
+	err1 := errors.New("error 1")
+	err2 := errors.New("error 2")
+	err3 := errors.New("error 3")
+
+	tests := []struct {
+		name string
+		err  error
+		errs []error
+		want bool
+	}{
+		{
+			name: "Error matches one of the errors in the list",
+			err:  err1,
+			errs: []error{err1, err2},
+			want: true,
+		},
+		{
+			name: "Error does not match any errors in the list",
+			err:  err3,
+			errs: []error{err1, err2},
+			want: false,
+		},
+		{
+			name: "Empty error list",
+			err:  err1,
+			errs: []error{},
+			want: false,
+		},
+		{
+			name: "Nil error with non-nil list",
+			err:  nil,
+			errs: []error{err1, err2},
+			want: false,
+		},
+		{
+			name: "Non-nil error with nil in list",
+			err:  err1,
+			errs: []error{nil, err2},
+			want: false,
+		},
+		{
+			name: "Error matches nil error in the list",
+			err:  nil,
+			errs: []error{nil, err2},
+			want: true,
+		},
+		{
+			name: "Nil error with empty list",
+			err:  nil,
+			errs: []error{},
+			want: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := In(tt.err, tt.errs...); got != tt.want {
+				t.Errorf("In() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

core/filex/file.go

@@ -35,6 +35,7 @@ func firstLine(file *os.File) (string, error) {
 	for {
 		buf := make([]byte, bufSize)
 		n, err := file.ReadAt(buf, offset)
+
 		if err != nil && err != io.EOF {
 			return "", err
 		}
@@ -45,6 +46,10 @@ func firstLine(file *os.File) (string, error) {
 			}
 		}
 
+		if err == io.EOF {
+			return string(append(first, buf[:n]...)), nil
+		}
+
 		first = append(first, buf[:n]...)
 		offset += bufSize
 	}
@@ -57,30 +62,42 @@ func lastLine(filename string, file *os.File) (string, error) {
 	}
 
 	var last []byte
+	bufLen := int64(bufSize)
 	offset := info.Size()
-	for {
-		offset -= bufSize
-		if offset < 0 {
+
+	for offset > 0 {
+		if offset < bufLen {
+			bufLen = offset
 			offset = 0
+		} else {
+			offset -= bufLen
 		}
-		buf := make([]byte, bufSize)
+
+		buf := make([]byte, bufLen)
 		n, err := file.ReadAt(buf, offset)
 		if err != nil && err != io.EOF {
 			return "", err
 		}
 
+		if n == 0 {
+			break
+		}
+
 		if buf[n-1] == '\n' {
 			buf = buf[:n-1]
 			n--
 		} else {
 			buf = buf[:n]
 		}
-		for n--; n >= 0; n-- {
-			if buf[n] == '\n' {
-				return string(append(buf[n+1:], last...)), nil
+
+		for i := n - 1; i >= 0; i-- {
+			if buf[i] == '\n' {
+				return string(append(buf[i+1:], last...)), nil
 			}
 		}
 
 		last = append(buf, last...)
 	}
+
+	return string(last), nil
 }

core/filex/file_test.go

@@ -52,6 +52,7 @@ last line`
 second line
 last line
 `
+	emptyContent = ``
 )
 
 func TestFirstLine(t *testing.T) {
@@ -79,6 +80,26 @@ func TestFirstLineError(t *testing.T) {
 	assert.Error(t, err)
 }
 
+func TestFirstLineEmptyFile(t *testing.T) {
+	filename, err := fs.TempFilenameWithText(emptyContent)
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	val, err := FirstLine(filename)
+	assert.Nil(t, err)
+	assert.Equal(t, "", val)
+}
+
+func TestFirstLineWithoutNewline(t *testing.T) {
+	filename, err := fs.TempFilenameWithText(longLine)
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	val, err := FirstLine(filename)
+	assert.Nil(t, err)
+	assert.Equal(t, longLine, val)
+}
+
 func TestLastLine(t *testing.T) {
 	filename, err := fs.TempFilenameWithText(text)
 	assert.Nil(t, err)
@@ -99,6 +120,16 @@ func TestLastLineWithLastNewline(t *testing.T) {
 	assert.Equal(t, longLine, val)
 }
 
+func TestLastLineWithoutLastNewline(t *testing.T) {
+	filename, err := fs.TempFilenameWithText(longLine)
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	val, err := LastLine(filename)
+	assert.Nil(t, err)
+	assert.Equal(t, longLine, val)
+}
+
 func TestLastLineShort(t *testing.T) {
 	filename, err := fs.TempFilenameWithText(shortText)
 	assert.Nil(t, err)
@@ -123,3 +154,67 @@ func TestLastLineError(t *testing.T) {
 	_, err := LastLine("/tmp/does-not-exist")
 	assert.Error(t, err)
 }
+
+func TestLastLineEmptyFile(t *testing.T) {
+	filename, err := fs.TempFilenameWithText(emptyContent)
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	val, err := LastLine(filename)
+	assert.Nil(t, err)
+	assert.Equal(t, "", val)
+}
+
+func TestFirstLineExactlyBufSize(t *testing.T) {
+	content := make([]byte, bufSize)
+	for i := range content {
+		content[i] = 'a'
+	}
+	content[bufSize-1] = '\n' // Ensure there is a newline at the edge
+
+	filename, err := fs.TempFilenameWithText(string(content))
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	val, err := FirstLine(filename)
+	assert.Nil(t, err)
+	assert.Equal(t, string(content[:bufSize-1]), val)
+}
+
+func TestLastLineExactlyBufSize(t *testing.T) {
+	content := make([]byte, bufSize)
+	for i := range content {
+		content[i] = 'a'
+	}
+	content[bufSize-1] = '\n' // Ensure there is a newline at the edge
+
+	filename, err := fs.TempFilenameWithText(string(content))
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	val, err := LastLine(filename)
+	assert.Nil(t, err)
+	assert.Equal(t, string(content[:bufSize-1]), val)
+}
+
+func TestFirstLineLargeFile(t *testing.T) {
+	content := text + text + text + "\n" + "extra"
+	filename, err := fs.TempFilenameWithText(content)
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	val, err := FirstLine(filename)
+	assert.Nil(t, err)
+	assert.Equal(t, "first line", val)
+}
+
+func TestLastLineLargeFile(t *testing.T) {
+	content := text + text + text + "\n" + "extra"
+	filename, err := fs.TempFilenameWithText(content)
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	val, err := LastLine(filename)
+	assert.Nil(t, err)
+	assert.Equal(t, "extra", val)
+}

core/filex/progressscanner.go

@@ -5,7 +5,7 @@ import "gopkg.in/cheggaaa/pb.v1"
 type (
 	// A Scanner is used to read lines.
 	Scanner interface {
-		// Scan checks if has remaining to read.
+		// Scan checks if it has remaining to read.
 		Scan() bool
 		// Text returns next line.
 		Text() string

core/fx/parallel.go

@@ -1,6 +1,9 @@
 package fx
 
-import "github.com/zeromicro/go-zero/core/threading"
+import (
+	"github.com/zeromicro/go-zero/core/errorx"
+	"github.com/zeromicro/go-zero/core/threading"
+)
 
 // Parallel runs fns parallelly and waits for done.
 func Parallel(fns ...func()) {
@@ -10,3 +13,20 @@ func Parallel(fns ...func()) {
 	}
 	group.Wait()
 }
+
+func ParallelErr(fns ...func() error) error {
+	var be errorx.BatchError
+
+	group := threading.NewRoutineGroup()
+	for _, fn := range fns {
+		f := fn
+		group.RunSafe(func() {
+			if err := f(); err != nil {
+				be.Add(err)
+			}
+		})
+	}
+	group.Wait()
+
+	return be.Err()
+}

core/fx/parallel_test.go

@@ -1,6 +1,7 @@
 package fx
 
 import (
+	"errors"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -22,3 +23,54 @@ func TestParallel(t *testing.T) {
 	})
 	assert.Equal(t, int32(6), count)
 }
+
+func TestParallelErr(t *testing.T) {
+	var count int32
+	err := ParallelErr(
+		func() error {
+			time.Sleep(time.Millisecond * 100)
+			atomic.AddInt32(&count, 1)
+			return errors.New("failed to exec #1")
+		},
+		func() error {
+			time.Sleep(time.Millisecond * 100)
+			atomic.AddInt32(&count, 2)
+			return errors.New("failed to exec #2")
+
+		},
+		func() error {
+			time.Sleep(time.Millisecond * 100)
+			atomic.AddInt32(&count, 3)
+			return nil
+		},
+	)
+
+	assert.Equal(t, int32(6), count)
+	assert.Error(t, err)
+	assert.ErrorContains(t, err, "failed to exec #1", "failed to exec #2")
+}
+
+func TestParallelErrErrorNil(t *testing.T) {
+	var count int32
+	err := ParallelErr(
+		func() error {
+			time.Sleep(time.Millisecond * 100)
+			atomic.AddInt32(&count, 1)
+			return nil
+		},
+		func() error {
+			time.Sleep(time.Millisecond * 100)
+			atomic.AddInt32(&count, 2)
+			return nil
+
+		},
+		func() error {
+			time.Sleep(time.Millisecond * 100)
+			atomic.AddInt32(&count, 3)
+			return nil
+		},
+	)
+
+	assert.Equal(t, int32(6), count)
+	assert.NoError(t, err)
+}

core/fx/stream.go

@@ -84,10 +84,10 @@ func Range(source <-chan any) Stream {
 	}
 }
 
-// AllMach returns whether all elements of this stream match the provided predicate.
+// AllMatch returns whether all elements of this stream match the provided predicate.
 // May not evaluate the predicate on all elements if not necessary for determining the result.
 // If the stream is empty then true is returned and the predicate is not evaluated.
-func (s Stream) AllMach(predicate func(item any) bool) bool {
+func (s Stream) AllMatch(predicate func(item any) bool) bool {
 	for item := range s.source {
 		if !predicate(item) {
 			// make sure the former goroutine not block, and current func returns fast.
@@ -99,10 +99,10 @@ func (s Stream) AllMach(predicate func(item any) bool) bool {
 	return true
 }
 
-// AnyMach returns whether any elements of this stream match the provided predicate.
+// AnyMatch returns whether any elements of this stream match the provided predicate.
 // May not evaluate the predicate on all elements if not necessary for determining the result.
 // If the stream is empty then false is returned and the predicate is not evaluated.
-func (s Stream) AnyMach(predicate func(item any) bool) bool {
+func (s Stream) AnyMatch(predicate func(item any) bool) bool {
 	for item := range s.source {
 		if predicate(item) {
 			// make sure the former goroutine not block, and current func returns fast.
@@ -352,7 +352,7 @@ func (s Stream) Parallel(fn ParallelFunc, opts ...Option) {
 	}, opts...).Done()
 }
 
-// Reduce is an utility method to let the caller deal with the underlying channel.
+// Reduce is a utility method to let the caller deal with the underlying channel.
 func (s Stream) Reduce(fn ReduceFunc) (any, error) {
 	return fn(s.source)
 }

core/fx/stream_test.go

@@ -398,16 +398,16 @@ func TestWalk(t *testing.T) {
 
 func TestStream_AnyMach(t *testing.T) {
 	runCheckedTest(t, func(t *testing.T) {
-		assetEqual(t, false, Just(1, 2, 3).AnyMach(func(item any) bool {
+		assetEqual(t, false, Just(1, 2, 3).AnyMatch(func(item any) bool {
 			return item.(int) == 4
 		}))
-		assetEqual(t, false, Just(1, 2, 3).AnyMach(func(item any) bool {
+		assetEqual(t, false, Just(1, 2, 3).AnyMatch(func(item any) bool {
 			return item.(int) == 0
 		}))
-		assetEqual(t, true, Just(1, 2, 3).AnyMach(func(item any) bool {
+		assetEqual(t, true, Just(1, 2, 3).AnyMatch(func(item any) bool {
 			return item.(int) == 2
 		}))
-		assetEqual(t, true, Just(1, 2, 3).AnyMach(func(item any) bool {
+		assetEqual(t, true, Just(1, 2, 3).AnyMatch(func(item any) bool {
 			return item.(int) == 2
 		}))
 	})
@@ -416,17 +416,17 @@ func TestStream_AnyMach(t *testing.T) {
 func TestStream_AllMach(t *testing.T) {
 	runCheckedTest(t, func(t *testing.T) {
 		assetEqual(
-			t, true, Just(1, 2, 3).AllMach(func(item any) bool {
+			t, true, Just(1, 2, 3).AllMatch(func(item any) bool {
 				return true
 			}),
 		)
 		assetEqual(
-			t, false, Just(1, 2, 3).AllMach(func(item any) bool {
+			t, false, Just(1, 2, 3).AllMatch(func(item any) bool {
 				return false
 			}),
 		)
 		assetEqual(
-			t, false, Just(1, 2, 3).AllMach(func(item any) bool {
+			t, false, Just(1, 2, 3).AllMatch(func(item any) bool {
 				return item.(int) == 1
 			}),
 		)

core/iox/pipe.go

@@ -2,7 +2,7 @@ package iox
 
 import "os"
 
-// RedirectInOut redirects stdin to r, stdout to w, and callers need to call restore afterwards.
+// RedirectInOut redirects stdin to r, stdout to w, and callers need to call restore afterward.
 func RedirectInOut() (restore func(), err error) {
 	var r, w *os.File
 	r, w, err = os.Pipe()

core/iox/textfile.go

@@ -9,7 +9,7 @@ import (
 
 const bufSize = 32 * 1024
 
-// CountLines returns the number of lines in file.
+// CountLines returns the number of lines in the file.
 func CountLines(file string) (int, error) {
 	f, err := os.Open(file)
 	if err != nil {

core/iox/textlinescanner.go

@@ -2,11 +2,12 @@ package iox
 
 import (
 	"bufio"
+	"errors"
 	"io"
 	"strings"
 )
 
-// A TextLineScanner is a scanner that can scan lines from given reader.
+// A TextLineScanner is a scanner that can scan lines from the given reader.
 type TextLineScanner struct {
 	reader  *bufio.Reader
 	hasNext bool
@@ -14,7 +15,7 @@ type TextLineScanner struct {
 	err     error
 }
 
-// NewTextLineScanner returns a TextLineScanner with given reader.
+// NewTextLineScanner returns a TextLineScanner with the given reader.
 func NewTextLineScanner(reader io.Reader) *TextLineScanner {
 	return &TextLineScanner{
 		reader:  bufio.NewReader(reader),
@@ -30,7 +31,7 @@ func (scanner *TextLineScanner) Scan() bool {
 
 	line, err := scanner.reader.ReadString('\n')
 	scanner.line = strings.TrimRight(line, "\n")
-	if err == io.EOF {
+	if errors.Is(err, io.EOF) {
 		scanner.hasNext = false
 		return true
 	} else if err != nil {

core/limit/periodlimit.go

@@ -2,6 +2,7 @@ package limit
 
 import (
 	"context"
+	_ "embed"
 	"errors"
 	"strconv"
 	"time"
@@ -28,20 +29,9 @@ var (
 	// ErrUnknownCode is an error that represents unknown status code.
 	ErrUnknownCode = errors.New("unknown status code")
 
-	// to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
-	periodScript = redis.NewScript(`local limit = tonumber(ARGV[1])
-local window = tonumber(ARGV[2])
-local current = redis.call("INCRBY", KEYS[1], 1)
-if current == 1 then
-    redis.call("expire", KEYS[1], window)
-end
-if current < limit then
-    return 1
-elseif current == limit then
-    return 2
-else
-    return 0
-end`)
+	//go:embed periodscript.lua
+	periodLuaScript string
+	periodScript    = redis.NewScript(periodLuaScript)
 )
 
 type (

core/limit/periodscript.lua

@@ -0,0 +1,14 @@
+-- to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
+local limit = tonumber(ARGV[1])
+local window = tonumber(ARGV[2])
+local current = redis.call("INCRBY", KEYS[1], 1)
+if current == 1 then
+    redis.call("expire", KEYS[1], window)
+end
+if current < limit then
+    return 1
+elseif current == limit then
+    return 2
+else
+    return 0
+end

core/limit/tokenlimit.go

@@ -2,6 +2,7 @@ package limit
 
 import (
 	"context"
+	_ "embed"
 	"errors"
 	"fmt"
 	"strconv"
@@ -9,6 +10,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/zeromicro/go-zero/core/errorx"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stores/redis"
 	xrate "golang.org/x/time/rate"
@@ -20,37 +22,11 @@ const (
 	pingInterval    = time.Millisecond * 100
 )
 
-// to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
-// KEYS[1] as tokens_key
-// KEYS[2] as timestamp_key
-var script = redis.NewScript(`local rate = tonumber(ARGV[1])
-local capacity = tonumber(ARGV[2])
-local now = tonumber(ARGV[3])
-local requested = tonumber(ARGV[4])
-local fill_time = capacity/rate
-local ttl = math.floor(fill_time*2)
-local last_tokens = tonumber(redis.call("get", KEYS[1]))
-if last_tokens == nil then
-    last_tokens = capacity
-end
-
-local last_refreshed = tonumber(redis.call("get", KEYS[2]))
-if last_refreshed == nil then
-    last_refreshed = 0
-end
-
-local delta = math.max(0, now-last_refreshed)
-local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
-local allowed = filled_tokens >= requested
-local new_tokens = filled_tokens
-if allowed then
-    new_tokens = filled_tokens - requested
-end
-
-redis.call("setex", KEYS[1], ttl, new_tokens)
-redis.call("setex", KEYS[2], ttl, now)
-
-return allowed`)
+var (
+	//go:embed tokenscript.lua
+	tokenLuaScript string
+	tokenScript    = redis.NewScript(tokenLuaScript)
+)
 
 // A TokenLimiter controls how frequently events are allowed to happen with in one second.
 type TokenLimiter struct {
@@ -112,7 +88,7 @@ func (lim *TokenLimiter) reserveN(ctx context.Context, now time.Time, n int) boo
 	}
 
 	resp, err := lim.store.ScriptRunCtx(ctx,
-		script,
+		tokenScript,
 		[]string{
 			lim.tokenKey,
 			lim.timestampKey,
@@ -125,10 +101,10 @@ func (lim *TokenLimiter) reserveN(ctx context.Context, now time.Time, n int) boo
 		})
 	// redis allowed == false
 	// Lua boolean false -> r Nil bulk reply
-	if err == redis.Nil {
+	if errors.Is(err, redis.Nil) {
 		return false
 	}
-	if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
+	if errorx.In(err, context.DeadlineExceeded, context.Canceled) {
 		logx.Errorf("fail to use rate limiter: %s", err)
 		return false
 	}

core/limit/tokenscript.lua

@@ -0,0 +1,31 @@
+-- to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
+-- KEYS[1] as tokens_key
+-- KEYS[2] as timestamp_key
+local rate = tonumber(ARGV[1])
+local capacity = tonumber(ARGV[2])
+local now = tonumber(ARGV[3])
+local requested = tonumber(ARGV[4])
+local fill_time = capacity/rate
+local ttl = math.floor(fill_time*2)
+local last_tokens = tonumber(redis.call("get", KEYS[1]))
+if last_tokens == nil then
+    last_tokens = capacity
+end
+
+local last_refreshed = tonumber(redis.call("get", KEYS[2]))
+if last_refreshed == nil then
+    last_refreshed = 0
+end
+
+local delta = math.max(0, now-last_refreshed)
+local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
+local allowed = filled_tokens >= requested
+local new_tokens = filled_tokens
+if allowed then
+    new_tokens = filled_tokens - requested
+end
+
+redis.call("setex", KEYS[1], ttl, new_tokens)
+redis.call("setex", KEYS[2], ttl, now)
+
+return allowed

core/load/adaptiveshedder.go

@@ -76,8 +76,8 @@ type (
 		avgFlyingLock   syncx.SpinLock
 		overloadTime    *syncx.AtomicDuration
 		droppedRecently *syncx.AtomicBool
-		passCounter     *collection.RollingWindow
-		rtCounter       *collection.RollingWindow
+		passCounter     *collection.RollingWindow[int64, *collection.Bucket[int64]]
+		rtCounter       *collection.RollingWindow[int64, *collection.Bucket[int64]]
 	}
 )
 
@@ -107,15 +107,16 @@ func NewAdaptiveShedder(opts ...ShedderOption) Shedder {
 		opt(&options)
 	}
 	bucketDuration := options.window / time.Duration(options.buckets)
+	newBucket := func() *collection.Bucket[int64] {
+		return new(collection.Bucket[int64])
+	}
 	return &adaptiveShedder{
 		cpuThreshold:    options.cpuThreshold,
 		windowScale:     float64(time.Second) / float64(bucketDuration) / millisecondsPerSecond,
 		overloadTime:    syncx.NewAtomicDuration(),
 		droppedRecently: syncx.NewAtomicBool(),
-		passCounter: collection.NewRollingWindow(options.buckets, bucketDuration,
-			collection.IgnoreCurrentBucket()),
-		rtCounter: collection.NewRollingWindow(options.buckets, bucketDuration,
-			collection.IgnoreCurrentBucket()),
+		passCounter:     collection.NewRollingWindow[int64, *collection.Bucket[int64]](newBucket, options.buckets, bucketDuration, collection.IgnoreCurrentBucket[int64, *collection.Bucket[int64]]()),
+		rtCounter:       collection.NewRollingWindow[int64, *collection.Bucket[int64]](newBucket, options.buckets, bucketDuration, collection.IgnoreCurrentBucket[int64, *collection.Bucket[int64]]()),
 	}
 }
 
@@ -138,10 +139,10 @@ func (as *adaptiveShedder) Allow() (Promise, error) {
 func (as *adaptiveShedder) addFlying(delta int64) {
 	flying := atomic.AddInt64(&as.flying, delta)
 	// update avgFlying when the request is finished.
-	// this strategy makes avgFlying have a little bit lag against flying, and smoother.
+	// this strategy makes avgFlying have a little bit of lag against flying, and smoother.
 	// when the flying requests increase rapidly, avgFlying increase slower, accept more requests.
 	// when the flying requests drop rapidly, avgFlying drop slower, accept fewer requests.
-	// it makes the service to serve as more requests as possible.
+	// it makes the service to serve as many requests as possible.
 	if delta < 0 {
 		as.avgFlyingLock.Lock()
 		as.avgFlying = as.avgFlying*flyingBeta + float64(flying)*(1-flyingBeta)
@@ -167,15 +168,15 @@ func (as *adaptiveShedder) maxFlight() float64 {
 }
 
 func (as *adaptiveShedder) maxPass() int64 {
-	var result float64 = 1
+	var result int64 = 1
 
-	as.passCounter.Reduce(func(b *collection.Bucket) {
+	as.passCounter.Reduce(func(b *collection.Bucket[int64]) {
 		if b.Sum > result {
 			result = b.Sum
 		}
 	})
 
-	return int64(result)
+	return result
 }
 
 func (as *adaptiveShedder) minRt() float64 {
@@ -183,12 +184,12 @@ func (as *adaptiveShedder) minRt() float64 {
 	// its a reasonable large value to avoid dropping requests.
 	result := defaultMinRt
 
-	as.rtCounter.Reduce(func(b *collection.Bucket) {
+	as.rtCounter.Reduce(func(b *collection.Bucket[int64]) {
 		if b.Count <= 0 {
 			return
 		}
 
-		avg := math.Round(b.Sum / float64(b.Count))
+		avg := math.Round(float64(b.Sum) / float64(b.Count))
 		if avg < result {
 			result = avg
 		}
@@ -200,7 +201,7 @@ func (as *adaptiveShedder) minRt() float64 {
 func (as *adaptiveShedder) overloadFactor() float64 {
 	// as.cpuThreshold must be less than cpuMax
 	factor := (cpuMax - float64(stat.CpuUsage())) / (cpuMax - float64(as.cpuThreshold))
-	// at least accept 10% of acceptable requests even cpu is highly overloaded.
+	// at least accept 10% of acceptable requests, even cpu is highly overloaded.
 	return mathx.Between(factor, overloadFactorLowerBound, 1)
 }
 
@@ -250,14 +251,14 @@ func (as *adaptiveShedder) systemOverloaded() bool {
 	return true
 }
 
-// WithBuckets customizes the Shedder with given number of buckets.
+// WithBuckets customizes the Shedder with the given number of buckets.
 func WithBuckets(buckets int) ShedderOption {
 	return func(opts *shedderOptions) {
 		opts.buckets = buckets
 	}
 }
 
-// WithCpuThreshold customizes the Shedder with given cpu threshold.
+// WithCpuThreshold customizes the Shedder with the given cpu threshold.
 func WithCpuThreshold(threshold int64) ShedderOption {
 	return func(opts *shedderOptions) {
 		opts.cpuThreshold = threshold
@@ -283,6 +284,6 @@ func (p *promise) Fail() {
 func (p *promise) Pass() {
 	rt := float64(timex.Since(p.start)) / float64(time.Millisecond)
 	p.shedder.addFlying(-1)
-	p.shedder.rtCounter.Add(math.Ceil(rt))
+	p.shedder.rtCounter.Add(int64(math.Ceil(rt)))
 	p.shedder.passCounter.Add(1)
 }

core/load/adaptiveshedder_test.go

@@ -58,7 +58,7 @@ func TestAdaptiveShedder(t *testing.T) {
 func TestAdaptiveShedderMaxPass(t *testing.T) {
 	passCounter := newRollingWindow()
 	for i := 1; i <= 10; i++ {
-		passCounter.Add(float64(i * 100))
+		passCounter.Add(int64(i * 100))
 		time.Sleep(bucketDuration)
 	}
 	shedder := &adaptiveShedder{
@@ -83,7 +83,7 @@ func TestAdaptiveShedderMinRt(t *testing.T) {
 			time.Sleep(bucketDuration)
 		}
 		for j := i*10 + 1; j <= i*10+10; j++ {
-			rtCounter.Add(float64(j))
+			rtCounter.Add(int64(j))
 		}
 	}
 	shedder := &adaptiveShedder{
@@ -107,9 +107,9 @@ func TestAdaptiveShedderMaxFlight(t *testing.T) {
 		if i > 0 {
 			time.Sleep(bucketDuration)
 		}
-		passCounter.Add(float64((i + 1) * 100))
+		passCounter.Add(int64((i + 1) * 100))
 		for j := i*10 + 1; j <= i*10+10; j++ {
-			rtCounter.Add(float64(j))
+			rtCounter.Add(int64(j))
 		}
 	}
 	shedder := &adaptiveShedder{
@@ -129,9 +129,9 @@ func TestAdaptiveShedderShouldDrop(t *testing.T) {
 		if i > 0 {
 			time.Sleep(bucketDuration)
 		}
-		passCounter.Add(float64((i + 1) * 100))
+		passCounter.Add(int64((i + 1) * 100))
 		for j := i*10 + 1; j <= i*10+10; j++ {
-			rtCounter.Add(float64(j))
+			rtCounter.Add(int64(j))
 		}
 	}
 	shedder := &adaptiveShedder{
@@ -184,9 +184,9 @@ func TestAdaptiveShedderStillHot(t *testing.T) {
 		if i > 0 {
 			time.Sleep(bucketDuration)
 		}
-		passCounter.Add(float64((i + 1) * 100))
+		passCounter.Add(int64((i + 1) * 100))
 		for j := i*10 + 1; j <= i*10+10; j++ {
-			rtCounter.Add(float64(j))
+			rtCounter.Add(int64(j))
 		}
 	}
 	shedder := &adaptiveShedder{
@@ -248,9 +248,9 @@ func BenchmarkMaxFlight(b *testing.B) {
 		if i > 0 {
 			time.Sleep(bucketDuration)
 		}
-		passCounter.Add(float64((i + 1) * 100))
+		passCounter.Add(int64((i + 1) * 100))
 		for j := i*10 + 1; j <= i*10+10; j++ {
-			rtCounter.Add(float64(j))
+			rtCounter.Add(int64(j))
 		}
 	}
 	shedder := &adaptiveShedder{
@@ -265,6 +265,8 @@ func BenchmarkMaxFlight(b *testing.B) {
 	}
 }
 
-func newRollingWindow() *collection.RollingWindow {
-	return collection.NewRollingWindow(buckets, bucketDuration, collection.IgnoreCurrentBucket())
+func newRollingWindow() *collection.RollingWindow[int64, *collection.Bucket[int64]] {
+	return collection.NewRollingWindow[int64, *collection.Bucket[int64]](func() *collection.Bucket[int64] {
+		return new(collection.Bucket[int64])
+	}, buckets, bucketDuration, collection.IgnoreCurrentBucket[int64, *collection.Bucket[int64]]())
 }

core/load/sheddergroup.go

@@ -6,7 +6,7 @@ import (
 	"github.com/zeromicro/go-zero/core/syncx"
 )
 
-// A ShedderGroup is a manager to manage key based shedders.
+// A ShedderGroup is a manager to manage key-based shedders.
 type ShedderGroup struct {
 	options []ShedderOption
 	manager *syncx.ResourceManager

core/logc/logs.go

@@ -42,7 +42,7 @@ func Debugv(ctx context.Context, v interface{}) {
 	getLogger(ctx).Debugv(v)
 }
 
-// Debugw writes msg along with fields into access log.
+// Debugw writes msg along with fields into the access log.
 func Debugw(ctx context.Context, msg string, fields ...LogField) {
 	getLogger(ctx).Debugw(msg, fields...)
 }
@@ -63,7 +63,7 @@ func Errorv(ctx context.Context, v any) {
 	getLogger(ctx).Errorv(v)
 }
 
-// Errorw writes msg along with fields into error log.
+// Errorw writes msg along with fields into the error log.
 func Errorw(ctx context.Context, msg string, fields ...LogField) {
 	getLogger(ctx).Errorw(msg, fields...)
 }
@@ -88,7 +88,7 @@ func Infov(ctx context.Context, v any) {
 	getLogger(ctx).Infov(v)
 }
 
-// Infow writes msg along with fields into access log.
+// Infow writes msg along with fields into the access log.
 func Infow(ctx context.Context, msg string, fields ...LogField) {
 	getLogger(ctx).Infow(msg, fields...)
 }
@@ -108,10 +108,11 @@ func SetLevel(level uint32) {
 	logx.SetLevel(level)
 }
 
-// SetUp sets up the logx. If already set up, just return nil.
-// we allow SetUp to be called multiple times, because for example
+// SetUp sets up the logx.
+// If already set up, return nil.
+// We allow SetUp to be called multiple times, because, for example,
 // we need to allow different service frameworks to initialize logx respectively.
-// the same logic for SetUp
+// The same logic for SetUp
 func SetUp(c LogConf) error {
 	return logx.SetUp(c)
 }

core/logx/lesslogger.go

@@ -1,6 +1,6 @@
 package logx
 
-// A LessLogger is a logger that control to log once during the given duration.
+// A LessLogger is a logger that controls to log once during the given duration.
 type LessLogger struct {
 	*limitedExecutor
 }

core/logx/logger.go

@@ -7,13 +7,13 @@ import (
 
 // A Logger represents a logger.
 type Logger interface {
-	// Debug logs a message at info level.
+	// Debug logs a message at debug level.
 	Debug(...any)
-	// Debugf logs a message at info level.
+	// Debugf logs a message at debug level.
 	Debugf(string, ...any)
-	// Debugv logs a message at info level.
+	// Debugv logs a message at debug level.
 	Debugv(any)
-	// Debugw logs a message at info level.
+	// Debugw logs a message at debug level.
 	Debugw(string, ...LogField)
 	// Error logs a message at error level.
 	Error(...any)

core/logx/logs.go

@@ -6,6 +6,7 @@ import (
 	"log"
 	"os"
 	"path"
+	"reflect"
 	"runtime/debug"
 	"sync"
 	"sync/atomic"
@@ -51,6 +52,26 @@ type (
 	}
 )
 
+// AddWriter adds a new writer.
+// If there is already a writer, the new writer will be added to the writer chain.
+// For example, to write logs to both file and console, if there is already a file writer,
+// ```go
+// logx.AddWriter(logx.NewWriter(os.Stdout))
+// ```
+func AddWriter(w Writer) {
+	ow := Reset()
+	if ow == nil {
+		SetWriter(w)
+	} else {
+		// no need to check if the existing writer is a comboWriter,
+		// because it is not common to add more than one writer.
+		// even more than one writer, the behavior is the same.
+		SetWriter(comboWriter{
+			writers: []Writer{ow, w},
+		})
+	}
+}
+
 // Alert alerts v in alert level, and the message is written to error log.
 func Alert(v string) {
 	getWriter().Alert(v)
@@ -86,7 +107,7 @@ func Debugv(v any) {
 	}
 }
 
-// Debugw writes msg along with fields into access log.
+// Debugw writes msg along with fields into the access log.
 func Debugw(msg string, fields ...LogField) {
 	if shallLog(DebugLevel) {
 		writeDebug(msg, fields...)
@@ -142,7 +163,7 @@ func Errorv(v any) {
 	}
 }
 
-// Errorw writes msg along with fields into error log.
+// Errorw writes msg along with fields into the error log.
 func Errorw(msg string, fields ...LogField) {
 	if shallLog(ErrorLevel) {
 		writeError(msg, fields...)
@@ -153,11 +174,11 @@ func Errorw(msg string, fields ...LogField) {
 func Field(key string, value any) LogField {
 	switch val := value.(type) {
 	case error:
-		return LogField{Key: key, Value: val.Error()}
+		return LogField{Key: key, Value: encodeError(val)}
 	case []error:
 		var errs []string
 		for _, err := range val {
-			errs = append(errs, err.Error())
+			errs = append(errs, encodeError(err))
 		}
 		return LogField{Key: key, Value: errs}
 	case time.Duration:
@@ -175,11 +196,11 @@ func Field(key string, value any) LogField {
 		}
 		return LogField{Key: key, Value: times}
 	case fmt.Stringer:
-		return LogField{Key: key, Value: val.String()}
+		return LogField{Key: key, Value: encodeStringer(val)}
 	case []fmt.Stringer:
 		var strs []string
 		for _, str := range val {
-			strs = append(strs, str.String())
+			strs = append(strs, encodeStringer(str))
 		}
 		return LogField{Key: key, Value: strs}
 	default:
@@ -208,7 +229,7 @@ func Infov(v any) {
 	}
 }
 
-// Infow writes msg along with fields into access log.
+// Infow writes msg along with fields into the access log.
 func Infow(msg string, fields ...LogField) {
 	if shallLog(InfoLevel) {
 		writeInfo(msg, fields...)
@@ -254,11 +275,12 @@ func SetWriter(w Writer) {
 	}
 }
 
-// SetUp sets up the logx. If already set up, just return nil.
-// we allow SetUp to be called multiple times, because for example
+// SetUp sets up the logx.
+// If already set up, return nil.
+// We allow SetUp to be called multiple times, because, for example,
 // we need to allow different service frameworks to initialize logx respectively.
 func SetUp(c LogConf) (err error) {
-	// Just ignore the subsequent SetUp calls.
+	// Ignore the later SetUp calls.
 	// Because multiple services in one process might call SetUp respectively.
 	// Need to wait for the first caller to complete the execution.
 	setupOnce.Do(func() {
@@ -413,6 +435,32 @@ func createOutput(path string) (io.WriteCloser, error) {
 	return NewLogger(path, rule, options.gzipEnabled)
 }
 
+func encodeError(err error) (ret string) {
+	return encodeWithRecover(err, func() string {
+		return err.Error()
+	})
+}
+
+func encodeStringer(v fmt.Stringer) (ret string) {
+	return encodeWithRecover(v, func() string {
+		return v.String()
+	})
+}
+
+func encodeWithRecover(arg any, fn func() string) (ret string) {
+	defer func() {
+		if err := recover(); err != nil {
+			if v := reflect.ValueOf(arg); v.Kind() == reflect.Ptr && v.IsNil() {
+				ret = nilAngleString
+			} else {
+				ret = fmt.Sprintf("panic: %v", err)
+			}
+		}
+	}()
+
+	return fn()
+}
+
 func getWriter() Writer {
 	w := writer.Load()
 	if w == nil {
@@ -480,7 +528,7 @@ func writeDebug(val any, fields ...LogField) {
 	getWriter().Debug(val, addCaller(fields...)...)
 }
 
-// writeError writes v into error log.
+// writeError writes v into the error log.
 // Not checking shallLog here is for performance consideration.
 // If we check shallLog here, the fmt.Sprint might be called even if the log level is not enabled.
 // The caller should check shallLog before calling this function.
@@ -520,7 +568,7 @@ func writeStack(msg string) {
 	getWriter().Stack(fmt.Sprintf("%s\n%s", msg, string(debug.Stack())))
 }
 
-// writeStat writes v into stat log.
+// writeStat writes v into the stat log.
 // Not checking shallLog here is for performance consideration.
 // If we check shallLog here, the fmt.Sprint might be called even if the log level is not enabled.
 // The caller should check shallLog before calling this function.

core/logx/logs_test.go

@@ -348,6 +348,27 @@ func TestStructedLogInfow(t *testing.T) {
 	})
 }
 
+func TestStructedLogFieldNil(t *testing.T) {
+	w := new(mockWriter)
+	old := writer.Swap(w)
+	defer writer.Store(old)
+
+	assert.NotPanics(t, func() {
+		var s *string
+		Infow("test", Field("bb", s))
+		var d *nilStringer
+		Infow("test", Field("bb", d))
+		var e *nilError
+		Errorw("test", Field("bb", e))
+	})
+	assert.NotPanics(t, func() {
+		var p panicStringer
+		Infow("test", Field("bb", p))
+		var ps innerPanicStringer
+		Infow("test", Field("bb", ps))
+	})
+}
+
 func TestStructedLogInfoConsoleAny(t *testing.T) {
 	w := new(mockWriter)
 	old := writer.Swap(w)
@@ -570,7 +591,7 @@ func TestErrorfWithWrappedError(t *testing.T) {
 	old := writer.Swap(w)
 	defer writer.Store(old)
 
-	Errorf("hello %w", errors.New(message))
+	Errorf("hello %s", errors.New(message))
 	assert.True(t, strings.Contains(w.String(), "hello there"))
 }
 
@@ -658,6 +679,10 @@ func TestSetup(t *testing.T) {
 
 func TestDisable(t *testing.T) {
 	Disable()
+	defer func() {
+		SetLevel(InfoLevel)
+		atomic.StoreUint32(&encoding, jsonEncodingType)
+	}()
 
 	var opt logOptions
 	WithKeepDays(1)(&opt)
@@ -680,6 +705,17 @@ func TestDisableStat(t *testing.T) {
 	assert.Equal(t, 0, w.builder.Len())
 }
 
+func TestAddWriter(t *testing.T) {
+	const message = "hello there"
+	w := new(mockWriter)
+	AddWriter(w)
+	w1 := new(mockWriter)
+	AddWriter(w1)
+	Error(message)
+	assert.Contains(t, w.String(), message)
+	assert.Contains(t, w1.String(), message)
+}
+
 func TestSetWriter(t *testing.T) {
 	atomic.StoreUint32(&logLevel, 0)
 	Reset()
@@ -859,3 +895,36 @@ func validateFields(t *testing.T, content string, fields map[string]any) {
 		}
 	}
 }
+
+type nilError struct {
+	Name string
+}
+
+func (e *nilError) Error() string {
+	return e.Name
+}
+
+type nilStringer struct {
+	Name string
+}
+
+func (s *nilStringer) String() string {
+	return s.Name
+}
+
+type innerPanicStringer struct {
+	Inner *struct {
+		Name string
+	}
+}
+
+func (s innerPanicStringer) String() string {
+	return s.Inner.Name
+}
+
+type panicStringer struct {
+}
+
+func (s panicStringer) String() string {
+	panic("panic")
+}

core/logx/richlogger.go

@@ -141,23 +141,43 @@ func (l *richLogger) WithCallerSkip(skip int) Logger {
 		return l
 	}
 
-	l.callerSkip = skip
-	return l
+	return &richLogger{
+		ctx:        l.ctx,
+		callerSkip: skip,
+		fields:     l.fields,
+	}
 }
 
 func (l *richLogger) WithContext(ctx context.Context) Logger {
-	l.ctx = ctx
-	return l
+	return &richLogger{
+		ctx:        ctx,
+		callerSkip: l.callerSkip,
+		fields:     l.fields,
+	}
 }
 
 func (l *richLogger) WithDuration(duration time.Duration) Logger {
-	l.fields = append(l.fields, Field(durationKey, timex.ReprOfDuration(duration)))
-	return l
+	fields := append(l.fields, Field(durationKey, timex.ReprOfDuration(duration)))
+
+	return &richLogger{
+		ctx:        l.ctx,
+		callerSkip: l.callerSkip,
+		fields:     fields,
+	}
 }
 
 func (l *richLogger) WithFields(fields ...LogField) Logger {
-	l.fields = append(l.fields, fields...)
-	return l
+	if len(fields) == 0 {
+		return l
+	}
+
+	f := append(l.fields, fields...)
+
+	return &richLogger{
+		ctx:        l.ctx,
+		callerSkip: l.callerSkip,
+		fields:     f,
+	}
 }
 
 func (l *richLogger) buildFields(fields ...LogField) []LogField {

core/logx/richlogger_test.go

@@ -287,6 +287,54 @@ func TestLogWithCallerSkip(t *testing.T) {
 	assert.True(t, w.Contains(fmt.Sprintf("%s:%d", file, line+1)))
 }
 
+func TestLogWithCallerSkipCopy(t *testing.T) {
+	log1 := WithCallerSkip(2)
+	log2 := log1.WithCallerSkip(3)
+	log3 := log2.WithCallerSkip(-1)
+	assert.Equal(t, 2, log1.(*richLogger).callerSkip)
+	assert.Equal(t, 3, log2.(*richLogger).callerSkip)
+	assert.Equal(t, 3, log3.(*richLogger).callerSkip)
+}
+
+func TestLogWithContextCopy(t *testing.T) {
+	c1 := context.Background()
+	c2 := context.WithValue(context.Background(), "foo", "bar")
+	log1 := WithContext(c1)
+	log2 := log1.WithContext(c2)
+	assert.Equal(t, c1, log1.(*richLogger).ctx)
+	assert.Equal(t, c2, log2.(*richLogger).ctx)
+}
+
+func TestLogWithDurationCopy(t *testing.T) {
+	log1 := WithContext(context.Background())
+	log2 := log1.WithDuration(time.Second)
+	assert.Empty(t, log1.(*richLogger).fields)
+	assert.Equal(t, 1, len(log2.(*richLogger).fields))
+
+	var w mockWriter
+	old := writer.Swap(&w)
+	defer writer.Store(old)
+	log2.Info("hello")
+	assert.Contains(t, w.String(), `"duration":"1000.0ms"`)
+}
+
+func TestLogWithFieldsCopy(t *testing.T) {
+	log1 := WithContext(context.Background())
+	log2 := log1.WithFields(Field("foo", "bar"))
+	log3 := log1.WithFields()
+	assert.Empty(t, log1.(*richLogger).fields)
+	assert.Equal(t, 1, len(log2.(*richLogger).fields))
+	assert.Equal(t, log1, log3)
+	assert.Empty(t, log3.(*richLogger).fields)
+
+	var w mockWriter
+	old := writer.Swap(&w)
+	defer writer.Store(old)
+
+	log2.Info("hello")
+	assert.Contains(t, w.String(), `"foo":"bar"`)
+}
+
 func TestLoggerWithFields(t *testing.T) {
 	w := new(mockWriter)
 	old := writer.Swap(w)

core/logx/rotatelogger.go

@@ -319,7 +319,7 @@ func (l *RotateLogger) maybeCompressFile(file string) {
 	}()
 
 	if _, err := os.Stat(file); err != nil {
-		// file not exists or other error, ignore compression
+		// file doesn't exist or another error, ignore compression
 		return
 	}
 

core/logx/vars.go

@@ -48,6 +48,7 @@ const (
 	levelDebug  = "debug"
 
 	backupFileDelimiter = "-"
+	nilAngleString      = "<nil>"
 	flags               = 0x0
 )
 

core/logx/writer.go

@@ -13,6 +13,7 @@ import (
 
 	fatihcolor "github.com/fatih/color"
 	"github.com/zeromicro/go-zero/core/color"
+	"github.com/zeromicro/go-zero/core/errorx"
 )
 
 type (
@@ -33,6 +34,10 @@ type (
 		lock   sync.RWMutex
 	}
 
+	comboWriter struct {
+		writers []Writer
+	}
+
 	concreteWriter struct {
 		infoLog   io.WriteCloser
 		errorLog  io.WriteCloser
@@ -88,6 +93,62 @@ func (w *atomicWriter) Swap(v Writer) Writer {
 	return old
 }
 
+func (c comboWriter) Alert(v any) {
+	for _, w := range c.writers {
+		w.Alert(v)
+	}
+}
+
+func (c comboWriter) Close() error {
+	var be errorx.BatchError
+	for _, w := range c.writers {
+		be.Add(w.Close())
+	}
+	return be.Err()
+}
+
+func (c comboWriter) Debug(v any, fields ...LogField) {
+	for _, w := range c.writers {
+		w.Debug(v, fields...)
+	}
+}
+
+func (c comboWriter) Error(v any, fields ...LogField) {
+	for _, w := range c.writers {
+		w.Error(v, fields...)
+	}
+}
+
+func (c comboWriter) Info(v any, fields ...LogField) {
+	for _, w := range c.writers {
+		w.Info(v, fields...)
+	}
+}
+
+func (c comboWriter) Severe(v any) {
+	for _, w := range c.writers {
+		w.Severe(v)
+	}
+}
+
+func (c comboWriter) Slow(v any, fields ...LogField) {
+	for _, w := range c.writers {
+		w.Slow(v, fields...)
+	}
+}
+
+func (c comboWriter) Stack(v any) {
+	for _, w := range c.writers {
+		w.Stack(v)
+	}
+}
+
+func (c comboWriter) Stat(v any, fields ...LogField) {
+	for _, w := range c.writers {
+		w.Stat(v, fields...)
+	}
+}
+
 func newConsoleWriter() Writer {
 	outLog := newLogWriter(log.New(fatihcolor.Output, "", flags))
 	errLog := newLogWriter(log.New(fatihcolor.Error, "", flags))
@@ -254,11 +315,10 @@ func (n nopWriter) Stack(_ any) {
 func (n nopWriter) Stat(_ any, _ ...LogField) {
 }
 
-func buildPlainFields(fields ...LogField) []string {
-	var items []string
-
-	for _, field := range fields {
-		items = append(items, fmt.Sprintf("%s=%+v", field.Key, field.Value))
+func buildPlainFields(fields logEntry) []string {
+	items := make([]string, 0, len(fields))
+	for k, v := range fields {
+		items = append(items, fmt.Sprintf("%s=%+v", k, v))
 	}
 
 	return items
@@ -278,6 +338,20 @@ func combineGlobalFields(fields []LogField) []LogField {
 	return ret
 }
 
+func marshalJson(t interface{}) ([]byte, error) {
+	var buf bytes.Buffer
+	encoder := json.NewEncoder(&buf)
+	encoder.SetEscapeHTML(false)
+	err := encoder.Encode(t)
+	// go 1.5+ will append a newline to the end of the json string
+	// https://github.com/golang/go/issues/13520
+	if l := buf.Len(); l > 0 && buf.Bytes()[l-1] == '\n' {
+		buf.Truncate(l - 1)
+	}
+
+	return buf.Bytes(), err
+}
+
 func output(writer io.Writer, level string, val any, fields ...LogField) {
 	// only truncate string content, don't know how to truncate the values of other types.
 	if v, ok := val.(string); ok {
@@ -289,15 +363,17 @@ func output(writer io.Writer, level string, val any, fields ...LogField) {
 	}
 
 	fields = combineGlobalFields(fields)
+	// +3 for timestamp, level and content
+	entry := make(logEntry, len(fields)+3)
+	for _, field := range fields {
+		entry[field.Key] = field.Value
+	}
 
 	switch atomic.LoadUint32(&encoding) {
 	case plainEncodingType:
-		writePlainAny(writer, level, val, buildPlainFields(fields...)...)
+		plainFields := buildPlainFields(entry)
+		writePlainAny(writer, level, val, plainFields...)
 	default:
-		entry := make(logEntry)
-		for _, field := range fields {
-			entry[field.Key] = field.Value
-		}
 		entry[timestampKey] = getTimestamp()
 		entry[levelKey] = level
 		entry[contentKey] = val
@@ -332,7 +408,7 @@ func wrapLevelWithColor(level string) string {
 }
 
 func writeJson(writer io.Writer, info any) {
-	if content, err := json.Marshal(info); err != nil {
+	if content, err := marshalJson(info); err != nil {
 		log.Printf("err: %s\n\n%s", err.Error(), debug.Stack())
 	} else if writer == nil {
 		log.Println(string(content))

core/logx/writer_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 )
 
 func TestNewWriter(t *testing.T) {
@@ -189,6 +190,41 @@ func TestWritePlainAny(t *testing.T) {
 	assert.Contains(t, buf.String(), "runtime/debug.Stack")
 }
 
+func TestWritePlainDuplicate(t *testing.T) {
+	old := atomic.SwapUint32(&encoding, plainEncodingType)
+	t.Cleanup(func() {
+		atomic.StoreUint32(&encoding, old)
+	})
+
+	var buf bytes.Buffer
+	output(&buf, levelInfo, "foo", LogField{
+		Key:   "first",
+		Value: "a",
+	}, LogField{
+		Key:   "first",
+		Value: "b",
+	})
+	assert.Contains(t, buf.String(), "foo")
+	assert.NotContains(t, buf.String(), "first=a")
+	assert.Contains(t, buf.String(), "first=b")
+
+	buf.Reset()
+	output(&buf, levelInfo, "foo", LogField{
+		Key:   "first",
+		Value: "a",
+	}, LogField{
+		Key:   "first",
+		Value: "b",
+	}, LogField{
+		Key:   "second",
+		Value: "c",
+	})
+	assert.Contains(t, buf.String(), "foo")
+	assert.NotContains(t, buf.String(), "first=a")
+	assert.Contains(t, buf.String(), "first=b")
+	assert.Contains(t, buf.String(), "second=c")
+}
+
 func TestLogWithLimitContentLength(t *testing.T) {
 	maxLen := atomic.LoadUint32(&maxContentLength)
 	atomic.StoreUint32(&maxContentLength, 10)
@@ -219,6 +255,117 @@ func TestLogWithLimitContentLength(t *testing.T) {
 	})
 }
 
+func TestComboWriter(t *testing.T) {
+	var mockWriters []Writer
+	for i := 0; i < 3; i++ {
+		mockWriters = append(mockWriters, new(tracedWriter))
+	}
+
+	cw := comboWriter{
+		writers: mockWriters,
+	}
+
+	t.Run("Alert", func(t *testing.T) {
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).On("Alert", "test alert").Once()
+		}
+		cw.Alert("test alert")
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).AssertCalled(t, "Alert", "test alert")
+		}
+	})
+
+	t.Run("Close", func(t *testing.T) {
+		for i := range cw.writers {
+			if i == 1 {
+				cw.writers[i].(*tracedWriter).On("Close").Return(errors.New("error")).Once()
+			} else {
+				cw.writers[i].(*tracedWriter).On("Close").Return(nil).Once()
+			}
+		}
+		err := cw.Close()
+		assert.Error(t, err)
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).AssertCalled(t, "Close")
+		}
+	})
+
+	t.Run("Debug", func(t *testing.T) {
+		fields := []LogField{{Key: "key", Value: "value"}}
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).On("Debug", "test debug", fields).Once()
+		}
+		cw.Debug("test debug", fields...)
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).AssertCalled(t, "Debug", "test debug", fields)
+		}
+	})
+
+	t.Run("Error", func(t *testing.T) {
+		fields := []LogField{{Key: "key", Value: "value"}}
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).On("Error", "test error", fields).Once()
+		}
+		cw.Error("test error", fields...)
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).AssertCalled(t, "Error", "test error", fields)
+		}
+	})
+
+	t.Run("Info", func(t *testing.T) {
+		fields := []LogField{{Key: "key", Value: "value"}}
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).On("Info", "test info", fields).Once()
+		}
+		cw.Info("test info", fields...)
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).AssertCalled(t, "Info", "test info", fields)
+		}
+	})
+
+	t.Run("Severe", func(t *testing.T) {
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).On("Severe", "test severe").Once()
+		}
+		cw.Severe("test severe")
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).AssertCalled(t, "Severe", "test severe")
+		}
+	})
+
+	t.Run("Slow", func(t *testing.T) {
+		fields := []LogField{{Key: "key", Value: "value"}}
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).On("Slow", "test slow", fields).Once()
+		}
+		cw.Slow("test slow", fields...)
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).AssertCalled(t, "Slow", "test slow", fields)
+		}
+	})
+
+	t.Run("Stack", func(t *testing.T) {
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).On("Stack", "test stack").Once()
+		}
+		cw.Stack("test stack")
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).AssertCalled(t, "Stack", "test stack")
+		}
+	})
+
+	t.Run("Stat", func(t *testing.T) {
+		fields := []LogField{{Key: "key", Value: "value"}}
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).On("Stat", "test stat", fields).Once()
+		}
+		cw.Stat("test stat", fields...)
+		for _, mw := range cw.writers {
+			mw.(*tracedWriter).AssertCalled(t, "Stat", "test stat", fields)
+		}
+	})
+}
+
 type mockedEntry struct {
 	Level     string `json:"level"`
 	Content   string `json:"content"`
@@ -250,3 +397,44 @@ type hardToWriteWriter struct{}
 func (h hardToWriteWriter) Write(_ []byte) (_ int, _ error) {
 	return 0, errors.New("write error")
 }
+
+type tracedWriter struct {
+	mock.Mock
+}
+
+func (w *tracedWriter) Alert(v any) {
+	w.Called(v)
+}
+
+func (w *tracedWriter) Close() error {
+	args := w.Called()
+	return args.Error(0)
+}
+
+func (w *tracedWriter) Debug(v any, fields ...LogField) {
+	w.Called(v, fields)
+}
+
+func (w *tracedWriter) Error(v any, fields ...LogField) {
+	w.Called(v, fields)
+}
+
+func (w *tracedWriter) Info(v any, fields ...LogField) {
+	w.Called(v, fields)
+}
+
+func (w *tracedWriter) Severe(v any) {
+	w.Called(v)
+}
+
+func (w *tracedWriter) Slow(v any, fields ...LogField) {
+	w.Called(v, fields)
+}
+
+func (w *tracedWriter) Stack(v any) {
+	w.Called(v)
+}
+
+func (w *tracedWriter) Stat(v any, fields ...LogField) {
+	w.Called(v, fields)
+}

core/mapping/marshaler.go

@@ -12,7 +12,7 @@ const (
 )
 
 // Marshal marshals the given val and returns the map that contains the fields.
-// optional=another is not implemented, and it's hard to implement and not common used.
+// optional=another is not implemented, and it's hard to implement and not commonly used.
 func Marshal(val any) (map[string]map[string]any, error) {
 	ret := make(map[string]map[string]any)
 	tp := reflect.TypeOf(val)

core/mapping/unmarshaler.go

@@ -39,7 +39,7 @@ var (
 )
 
 type (
-	// Unmarshaler is used to unmarshal with given tag key.
+	// Unmarshaler is used to unmarshal with the given tag key.
 	Unmarshaler struct {
 		key  string
 		opts unmarshalOptions
@@ -69,7 +69,7 @@ func NewUnmarshaler(key string, opts ...UnmarshalOption) *Unmarshaler {
 	return &unmarshaler
 }
 
-// UnmarshalKey unmarshals m into v with tag key.
+// UnmarshalKey unmarshals m into v with the tag key.
 func UnmarshalKey(m map[string]any, v any) error {
 	return keyUnmarshaler.Unmarshal(m, v)
 }
@@ -113,7 +113,8 @@ func (u *Unmarshaler) unmarshalValuer(m Valuer, v any, fullName string) error {
 	return u.unmarshalWithFullName(simpleValuer{current: m}, v, fullName)
 }
 
-func (u *Unmarshaler) fillMap(fieldType reflect.Type, value reflect.Value, mapValue any, fullName string) error {
+func (u *Unmarshaler) fillMap(fieldType reflect.Type, value reflect.Value,
+	mapValue any, fullName string) error {
 	if !value.CanSet() {
 		return errValueNotSettable
 	}
@@ -154,7 +155,8 @@ func (u *Unmarshaler) fillMapFromString(value reflect.Value, mapValue any) error
 	return nil
 }
 
-func (u *Unmarshaler) fillSlice(fieldType reflect.Type, value reflect.Value, mapValue any, fullName string) error {
+func (u *Unmarshaler) fillSlice(fieldType reflect.Type, value reflect.Value,
+	mapValue any, fullName string) error {
 	if !value.CanSet() {
 		return errValueNotSettable
 	}
@@ -307,7 +309,34 @@ func (u *Unmarshaler) fillSliceWithDefault(derefedType reflect.Type, value refle
 	return u.fillSlice(derefedType, value, slice, fullName)
 }
 
-func (u *Unmarshaler) generateMap(keyType, elemType reflect.Type, mapValue any, fullName string) (reflect.Value, error) {
+func (u *Unmarshaler) fillUnmarshalerStruct(fieldType reflect.Type,
+	value reflect.Value, targetValue string) error {
+	if !value.CanSet() {
+		return errValueNotSettable
+	}
+
+	baseType := Deref(fieldType)
+	target := reflect.New(baseType)
+	switch u.key {
+	case jsonTagKey:
+		unmarshaler, ok := target.Interface().(json.Unmarshaler)
+		if !ok {
+			return errUnsupportedType
+		}
+
+		if err := unmarshaler.UnmarshalJSON([]byte(targetValue)); err != nil {
+			return err
+		}
+	default:
+		return errUnsupportedType
+	}
+
+	value.Set(target)
+	return nil
+}
+
+func (u *Unmarshaler) generateMap(keyType, elemType reflect.Type, mapValue any,
+	fullName string) (reflect.Value, error) {
 	mapType := reflect.MapOf(keyType, elemType)
 	valueType := reflect.TypeOf(mapValue)
 	if mapType == valueType {
@@ -399,6 +428,15 @@ func (u *Unmarshaler) generateMap(keyType, elemType reflect.Type, mapValue any,
 	return targetValue, nil
 }
 
+func (u *Unmarshaler) implementsUnmarshaler(t reflect.Type) bool {
+	switch u.key {
+	case jsonTagKey:
+		return t.Implements(reflect.TypeOf((*json.Unmarshaler)(nil)).Elem())
+	default:
+		return false
+	}
+}
+
 func (u *Unmarshaler) parseOptionsWithContext(field reflect.StructField, m Valuer, fullName string) (
 	string, *fieldOptionsWithContext, error) {
 	key, options, err := parseKeyAndOptions(u.key, field)
@@ -576,6 +614,8 @@ func (u *Unmarshaler) processFieldNotFromString(fieldType reflect.Type, value re
 		return u.fillSliceFromString(fieldType, value, mapValue, fullName)
 	case valueKind == reflect.String && derefedFieldType == durationType:
 		return fillDurationValue(fieldType, value, mapValue.(string))
+	case valueKind == reflect.String && typeKind == reflect.Struct && u.implementsUnmarshaler(fieldType):
+		return u.fillUnmarshalerStruct(fieldType, value, mapValue.(string))
 	default:
 		return u.processFieldPrimitive(fieldType, value, mapValue, opts, fullName)
 	}
@@ -629,7 +669,7 @@ func (u *Unmarshaler) processFieldPrimitiveWithJSONNumber(fieldType reflect.Type
 			return err
 		}
 
-		// if value is a pointer, we need to check overflow with the pointer's value.
+		// if the value is a pointer, we need to check overflow with the pointer's value.
 		derefedValue := value
 		for derefedValue.Type().Kind() == reflect.Ptr {
 			derefedValue = derefedValue.Elem()

core/mapping/unmarshaler_test.go

@@ -2,6 +2,7 @@ package mapping
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"reflect"
 	"strconv"
@@ -260,6 +261,7 @@ func TestUnmarshalInt(t *testing.T) {
 		Int64FromStr int64 `key:"int64str,string"`
 		DefaultInt   int64 `key:"defaultint,default=11"`
 		Optional     int   `key:"optional,optional"`
+		IntOptDef    int   `key:"intopt,optional,default=6"`
 	}
 	m := map[string]any{
 		"int":      1,
@@ -288,6 +290,7 @@ func TestUnmarshalInt(t *testing.T) {
 		ast.Equal(int64(9), in.Int64)
 		ast.Equal(int64(10), in.Int64FromStr)
 		ast.Equal(int64(11), in.DefaultInt)
+		ast.Equal(6, in.IntOptDef)
 	}
 }
 
@@ -5760,6 +5763,49 @@ func TestUnmarshalWithIgnoreFields(t *testing.T) {
 	}
 }
 
+func TestUnmarshal_Unmarshaler(t *testing.T) {
+	t.Run("success", func(t *testing.T) {
+		v := struct {
+			Foo *mockUnmarshaler `json:"name"`
+		}{}
+		body := `{"name": "hello"}`
+		assert.NoError(t, UnmarshalJsonBytes([]byte(body), &v))
+		assert.Equal(t, "hello", v.Foo.Name)
+	})
+
+	t.Run("failure", func(t *testing.T) {
+		v := struct {
+			Foo *mockUnmarshalerWithError `json:"name"`
+		}{}
+		body := `{"name": "hello"}`
+		assert.Error(t, UnmarshalJsonBytes([]byte(body), &v))
+	})
+
+	t.Run("not json unmarshaler", func(t *testing.T) {
+		v := struct {
+			Foo *struct {
+				Name string
+			} `key:"name"`
+		}{}
+		u := NewUnmarshaler(defaultKeyName)
+		assert.Error(t, u.Unmarshal(map[string]any{
+			"name": "hello",
+		}, &v))
+	})
+
+	t.Run("not with json key", func(t *testing.T) {
+		v := struct {
+			Foo *mockUnmarshaler `json:"name"`
+		}{}
+		u := NewUnmarshaler(defaultKeyName)
+		// with different key, ignore
+		assert.NoError(t, u.Unmarshal(map[string]any{
+			"name": "hello",
+		}, &v))
+		assert.Nil(t, v.Foo)
+	})
+}
+
 func BenchmarkDefaultValue(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		var a struct {
@@ -5866,10 +5912,27 @@ type mockValuerWithParent struct {
 	ok     bool
 }
 
-func (m mockValuerWithParent) Value(key string) (any, bool) {
+func (m mockValuerWithParent) Value(_ string) (any, bool) {
 	return m.value, m.ok
 }
 
 func (m mockValuerWithParent) Parent() valuerWithParent {
 	return m.parent
 }
+
+type mockUnmarshaler struct {
+	Name string
+}
+
+func (m *mockUnmarshaler) UnmarshalJSON(b []byte) error {
+	m.Name = string(b)
+	return nil
+}
+
+type mockUnmarshalerWithError struct {
+	Name string
+}
+
+func (m *mockUnmarshalerWithError) UnmarshalJSON(b []byte) error {
+	return errors.New("foo")
+}

core/mapping/utils.go

@@ -416,7 +416,7 @@ func parseOption(fieldOpts *fieldOptions, fieldName, option string) error {
 }
 
 // parseOptions parses the given options in tag.
-// for example: `json:"name,options=foo|bar"` or `json:"name,options=[foo,bar]"`
+// for example, `json:"name,options=foo|bar"` or `json:"name,options=[foo,bar]"`
 func parseOptions(val string) []string {
 	if len(val) == 0 {
 		return nil

core/mapping/valuer.go

@@ -26,9 +26,9 @@ type (
 		parent valuerWithParent
 	}
 
-	// mapValuer is a type for map to meet the Valuer interface.
+	// mapValuer is a type for the map to meet the Valuer interface.
 	mapValuer map[string]any
-	// simpleValuer is a type to get value from current node.
+	// simpleValuer is a type to get value from the current node.
 	simpleValuer node
 	// recursiveValuer is a type to get the value recursively from current and parent nodes.
 	recursiveValuer node

core/mathx/range.go

@@ -1,13 +1,13 @@
 package mathx
 
-type numerical interface {
+type Numerical interface {
 	~int | ~int8 | ~int16 | ~int32 | ~int64 |
 		~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 |
 		~float32 | ~float64
 }
 
 // AtLeast returns the greater of x or lower.
-func AtLeast[T numerical](x, lower T) T {
+func AtLeast[T Numerical](x, lower T) T {
 	if x < lower {
 		return lower
 	}
@@ -15,7 +15,7 @@ func AtLeast[T numerical](x, lower T) T {
 }
 
 // AtMost returns the smaller of x or upper.
-func AtMost[T numerical](x, upper T) T {
+func AtMost[T Numerical](x, upper T) T {
 	if x > upper {
 		return upper
 	}
@@ -23,7 +23,7 @@ func AtMost[T numerical](x, upper T) T {
 }
 
 // Between returns the value of x clamped to the range [lower, upper].
-func Between[T numerical](x, lower, upper T) T {
+func Between[T Numerical](x, lower, upper T) T {
 	if x < lower {
 		return lower
 	}

core/mr/mapreduce.go

@@ -363,9 +363,7 @@ func newGuardedWriter[T any](ctx context.Context, channel chan<- T, done <-chan
 func (gw guardedWriter[T]) Write(v T) {
 	select {
 	case <-gw.ctx.Done():
-		return
 	case <-gw.done:
-		return
 	default:
 		gw.channel <- v
 	}

core/proc/goroutines_test.go

@@ -36,6 +36,6 @@ type fakeCreator struct {
 	err  error
 }
 
-func (fc fakeCreator) Create(name string) (file *os.File, err error) {
+func (fc fakeCreator) Create(_ string) (file *os.File, err error) {
 	return fc.file, fc.err
 }

core/queue/queue.go

@@ -76,7 +76,7 @@ func (q *Queue) AddListener(listener Listener) {
 	q.listeners = append(q.listeners, listener)
 }
 
-// Broadcast broadcasts message to all event channels.
+// Broadcast broadcasts the message to all event channels.
 func (q *Queue) Broadcast(message any) {
 	go func() {
 		q.eventLock.Lock()
@@ -202,7 +202,7 @@ func (q *Queue) produce() {
 }
 
 func (q *Queue) produceOne(producer Producer) (string, bool) {
-	// avoid panic quit the producer, just log it and continue
+	// avoid panic quit the producer, log it and continue
 	defer rescue.Recover()
 
 	return producer.Produce()

core/queue/util_test.go

@@ -67,7 +67,7 @@ func (p *mockedPusher) Name() string {
 	return p.name
 }
 
-func (p *mockedPusher) Push(s string) error {
+func (p *mockedPusher) Push(_ string) error {
 	if proba.TrueOnProba(failProba) {
 		return errors.New("dummy")
 	}

core/stat/metrics_test.go

@@ -71,6 +71,6 @@ func (m *mockedWriter) Write(report *StatReport) error {
 
 type badWriter struct{}
 
-func (b *badWriter) Write(report *StatReport) error {
+func (b *badWriter) Write(_ *StatReport) error {
 	return errors.New("bad")
 }

core/stat/remotewriter_test.go

@@ -1,6 +1,7 @@
 package stat
 
 import (
+	"errors"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -28,3 +29,14 @@ func TestRemoteWriterFail(t *testing.T) {
 	})
 	assert.NotNil(t, err)
 }
+
+func TestRemoteWriterError(t *testing.T) {
+	defer gock.Off()
+
+	gock.New("http://foo.com").ReplyError(errors.New("foo"))
+	writer := NewRemoteWriter("http://foo.com")
+	err := writer.Write(&StatReport{
+		Name: "bar",
+	})
+	assert.NotNil(t, err)
+}

core/stat/task.go

@@ -2,7 +2,7 @@ package stat
 
 import "time"
 
-// A Task is a task that is reported to Metrics.
+// A Task is a task reported to Metrics.
 type Task struct {
 	Drop        bool
 	Duration    time.Duration

core/stores/builder/builder.go

@@ -41,7 +41,7 @@ func RawFieldNames(in any, postgreSql ...bool) []string {
 				out = append(out, fmt.Sprintf("`%s`", fi.Name))
 			}
 		default:
-			// get tag name with the tag opton, e.g.:
+			// get tag name with the tag option, e.g.:
 			// `db:"id"`
 			// `db:"id,type=char,length=16"`
 			// `db:",type=char,length=16"`

core/stores/cache/cacheopt.go

@@ -8,7 +8,7 @@ const (
 )
 
 type (
-	// An Options is used to store the cache options.
+	// Options is used to store the cache options.
 	Options struct {
 		Expiry         time.Duration
 		NotFoundExpiry time.Duration

core/stores/mon/clientmanager_test.go

@@ -7,6 +7,10 @@ import (
 	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
 )
 
+func init() {
+	_ = mtest.Setup()
+}
+
 func TestClientManger_getClient(t *testing.T) {
 	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
 	mt.Run("test", func(mt *mtest.T) {

core/stores/mon/collection.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/zeromicro/go-zero/core/breaker"
+	"github.com/zeromicro/go-zero/core/errorx"
 	"github.com/zeromicro/go-zero/core/timex"
 	"go.mongodb.org/mongo-driver/mongo"
 	mopt "go.mongodb.org/mongo-driver/mongo/options"
@@ -15,7 +16,8 @@ import (
 const (
 	defaultSlowThreshold = time.Millisecond * 500
 	// spanName is the span name of the mongo calls.
-	spanName = "mongo"
+	spanName         = "mongo"
+	duplicateKeyCode = 11000
 
 	// mongodb method names
 	aggregate              = "Aggregate"
@@ -141,7 +143,7 @@ func (c *decoratedCollection) Aggregate(ctx context.Context, pipeline any,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		starTime := timex.Now()
 		defer func() {
 			c.logDurationSimple(ctx, aggregate, starTime, err)
@@ -161,7 +163,7 @@ func (c *decoratedCollection) BulkWrite(ctx context.Context, models []mongo.Writ
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDurationSimple(ctx, bulkWrite, startTime, err)
@@ -181,7 +183,7 @@ func (c *decoratedCollection) CountDocuments(ctx context.Context, filter any,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDurationSimple(ctx, countDocuments, startTime, err)
@@ -201,7 +203,7 @@ func (c *decoratedCollection) DeleteMany(ctx context.Context, filter any,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDurationSimple(ctx, deleteMany, startTime, err)
@@ -221,7 +223,7 @@ func (c *decoratedCollection) DeleteOne(ctx context.Context, filter any,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDuration(ctx, deleteOne, startTime, err, filter)
@@ -241,7 +243,7 @@ func (c *decoratedCollection) Distinct(ctx context.Context, fieldName string, fi
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDurationSimple(ctx, distinct, startTime, err)
@@ -261,7 +263,7 @@ func (c *decoratedCollection) EstimatedDocumentCount(ctx context.Context,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDurationSimple(ctx, estimatedDocumentCount, startTime, err)
@@ -281,7 +283,7 @@ func (c *decoratedCollection) Find(ctx context.Context, filter any,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDuration(ctx, find, startTime, err, filter)
@@ -301,7 +303,7 @@ func (c *decoratedCollection) FindOne(ctx context.Context, filter any,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDuration(ctx, findOne, startTime, err, filter)
@@ -322,7 +324,7 @@ func (c *decoratedCollection) FindOneAndDelete(ctx context.Context, filter any,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDuration(ctx, findOneAndDelete, startTime, err, filter)
@@ -344,7 +346,7 @@ func (c *decoratedCollection) FindOneAndReplace(ctx context.Context, filter any,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDuration(ctx, findOneAndReplace, startTime, err, filter, replacement)
@@ -365,7 +367,7 @@ func (c *decoratedCollection) FindOneAndUpdate(ctx context.Context, filter, upda
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDuration(ctx, findOneAndUpdate, startTime, err, filter, update)
@@ -386,7 +388,7 @@ func (c *decoratedCollection) InsertMany(ctx context.Context, documents []any,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDurationSimple(ctx, insertMany, startTime, err)
@@ -406,7 +408,7 @@ func (c *decoratedCollection) InsertOne(ctx context.Context, document any,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDuration(ctx, insertOne, startTime, err, document)
@@ -426,7 +428,7 @@ func (c *decoratedCollection) ReplaceOne(ctx context.Context, filter, replacemen
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDuration(ctx, replaceOne, startTime, err, filter, replacement)
@@ -446,7 +448,7 @@ func (c *decoratedCollection) UpdateByID(ctx context.Context, id, update any,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDuration(ctx, updateByID, startTime, err, id, update)
@@ -466,7 +468,7 @@ func (c *decoratedCollection) UpdateMany(ctx context.Context, filter, update any
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDurationSimple(ctx, updateMany, startTime, err)
@@ -486,7 +488,7 @@ func (c *decoratedCollection) UpdateOne(ctx context.Context, filter, update any,
 		endSpan(span, err)
 	}()
 
-	err = c.brk.DoWithAcceptable(func() error {
+	err = c.brk.DoWithAcceptableCtx(ctx, func() error {
 		startTime := timex.Now()
 		defer func() {
 			c.logDuration(ctx, updateOne, startTime, err, filter, update)
@@ -527,19 +529,20 @@ func (p keepablePromise) keep(err error) error {
 }
 
 func acceptable(err error) bool {
-	return err == nil ||
-		errors.Is(err, mongo.ErrNoDocuments) ||
-		errors.Is(err, mongo.ErrNilValue) ||
-		errors.Is(err, mongo.ErrNilDocument) ||
-		errors.Is(err, mongo.ErrNilCursor) ||
-		errors.Is(err, mongo.ErrEmptySlice) ||
-		// session errors
-		errors.Is(err, session.ErrSessionEnded) ||
-		errors.Is(err, session.ErrNoTransactStarted) ||
-		errors.Is(err, session.ErrTransactInProgress) ||
-		errors.Is(err, session.ErrAbortAfterCommit) ||
-		errors.Is(err, session.ErrAbortTwice) ||
-		errors.Is(err, session.ErrCommitAfterAbort) ||
-		errors.Is(err, session.ErrUnackWCUnsupported) ||
-		errors.Is(err, session.ErrSnapshotTransaction)
+	return err == nil || isDupKeyError(err) ||
+		errorx.In(err, mongo.ErrNoDocuments, mongo.ErrNilValue,
+			mongo.ErrNilDocument, mongo.ErrNilCursor, mongo.ErrEmptySlice,
+			// session errors
+			session.ErrSessionEnded, session.ErrNoTransactStarted, session.ErrTransactInProgress,
+			session.ErrAbortAfterCommit, session.ErrAbortTwice, session.ErrCommitAfterAbort,
+			session.ErrUnackWCUnsupported, session.ErrSnapshotTransaction)
+}
+
+func isDupKeyError(err error) bool {
+	var e mongo.WriteException
+	if !errors.As(err, &e) {
+		return false
+	}
+
+	return e.HasErrorCode(duplicateKeyCode)
 }

core/stores/mon/collection_test.go

@@ -15,6 +15,7 @@ import (
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
 	mopt "go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/session"
 )
 
 var errDummy = errors.New("dummy")
@@ -572,6 +573,56 @@ func TestDecoratedCollection_LogDuration(t *testing.T) {
 	assert.Contains(t, buf.String(), "slowcall")
 }
 
+func TestAcceptable(t *testing.T) {
+	tests := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{"NilError", nil, true},
+		{"NoDocuments", mongo.ErrNoDocuments, true},
+		{"NilValue", mongo.ErrNilValue, true},
+		{"NilDocument", mongo.ErrNilDocument, true},
+		{"NilCursor", mongo.ErrNilCursor, true},
+		{"EmptySlice", mongo.ErrEmptySlice, true},
+		{"SessionEnded", session.ErrSessionEnded, true},
+		{"NoTransactStarted", session.ErrNoTransactStarted, true},
+		{"TransactInProgress", session.ErrTransactInProgress, true},
+		{"AbortAfterCommit", session.ErrAbortAfterCommit, true},
+		{"AbortTwice", session.ErrAbortTwice, true},
+		{"CommitAfterAbort", session.ErrCommitAfterAbort, true},
+		{"UnackWCUnsupported", session.ErrUnackWCUnsupported, true},
+		{"SnapshotTransaction", session.ErrSnapshotTransaction, true},
+		{"DuplicateKeyError", mongo.WriteException{WriteErrors: []mongo.WriteError{{Code: duplicateKeyCode}}}, true},
+		{"OtherError", errors.New("other error"), false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, acceptable(tt.err))
+		})
+	}
+}
+
+func TestIsDupKeyError(t *testing.T) {
+	tests := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{"NilError", nil, false},
+		{"NonDupKeyError", errors.New("some other error"), false},
+		{"DupKeyError", mongo.WriteException{WriteErrors: []mongo.WriteError{{Code: duplicateKeyCode}}}, true},
+		{"OtherMongoError", mongo.WriteException{WriteErrors: []mongo.WriteError{{Code: 12345}}}, false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, isDupKeyError(tt.err))
+		})
+	}
+}
+
 type mockPromise struct {
 	accepted bool
 	reason   string
@@ -595,19 +646,40 @@ func (d *dropBreaker) Allow() (breaker.Promise, error) {
 	return nil, errDummy
 }
 
+func (d *dropBreaker) AllowCtx(_ context.Context) (breaker.Promise, error) {
+	return nil, errDummy
+}
+
 func (d *dropBreaker) Do(_ func() error) error {
 	return nil
 }
 
+func (d *dropBreaker) DoCtx(_ context.Context, _ func() error) error {
+	return nil
+}
+
 func (d *dropBreaker) DoWithAcceptable(_ func() error, _ breaker.Acceptable) error {
 	return errDummy
 }
 
+func (d *dropBreaker) DoWithAcceptableCtx(_ context.Context, _ func() error, _ breaker.Acceptable) error {
+	return errDummy
+}
+
 func (d *dropBreaker) DoWithFallback(_ func() error, _ breaker.Fallback) error {
 	return nil
 }
 
+func (d *dropBreaker) DoWithFallbackCtx(_ context.Context, _ func() error, _ breaker.Fallback) error {
+	return nil
+}
+
 func (d *dropBreaker) DoWithFallbackAcceptable(_ func() error, _ breaker.Fallback,
 	_ breaker.Acceptable) error {
 	return nil
 }
+
+func (d *dropBreaker) DoWithFallbackAcceptableCtx(_ context.Context, _ func() error,
+	_ breaker.Fallback, _ breaker.Acceptable) error {
+	return nil
+}

core/stores/mon/model.go

@@ -69,27 +69,21 @@ func newModel(name string, cli *mongo.Client, coll Collection, brk breaker.Break
 
 // StartSession starts a new session.
 func (m *Model) StartSession(opts ...*mopt.SessionOptions) (sess mongo.Session, err error) {
-	err = m.brk.DoWithAcceptable(func() error {
-		starTime := timex.Now()
-		defer func() {
-			logDuration(context.Background(), m.name, startSession, starTime, err)
-		}()
+	starTime := timex.Now()
+	defer func() {
+		logDuration(context.Background(), m.name, startSession, starTime, err)
+	}()
 
-		session, sessionErr := m.cli.StartSession(opts...)
-		if sessionErr != nil {
-			return sessionErr
-		}
+	session, sessionErr := m.cli.StartSession(opts...)
+	if sessionErr != nil {
+		return nil, sessionErr
+	}
 
-		sess = &wrappedSession{
-			Session: session,
-			name:    m.name,
-			brk:     m.brk,
-		}
-
-		return nil
-	}, acceptable)
-
-	return
+	return &wrappedSession{
+		Session: session,
+		name:    m.name,
+		brk:     m.brk,
+	}, nil
 }
 
 // Aggregate executes an aggregation pipeline.
@@ -184,7 +178,7 @@ func (w *wrappedSession) AbortTransaction(ctx context.Context) (err error) {
 		endSpan(span, err)
 	}()
 
-	return w.brk.DoWithAcceptable(func() error {
+	return w.brk.DoWithAcceptableCtx(ctx, func() error {
 		starTime := timex.Now()
 		defer func() {
 			logDuration(ctx, w.name, abortTransaction, starTime, err)
@@ -201,7 +195,7 @@ func (w *wrappedSession) CommitTransaction(ctx context.Context) (err error) {
 		endSpan(span, err)
 	}()
 
-	return w.brk.DoWithAcceptable(func() error {
+	return w.brk.DoWithAcceptableCtx(ctx, func() error {
 		starTime := timex.Now()
 		defer func() {
 			logDuration(ctx, w.name, commitTransaction, starTime, err)
@@ -222,7 +216,7 @@ func (w *wrappedSession) WithTransaction(
 		endSpan(span, err)
 	}()
 
-	err = w.brk.DoWithAcceptable(func() error {
+	err = w.brk.DoWithAcceptableCtx(ctx, func() error {
 		starTime := timex.Now()
 		defer func() {
 			logDuration(ctx, w.name, withTransaction, starTime, err)
@@ -243,7 +237,7 @@ func (w *wrappedSession) EndSession(ctx context.Context) {
 		endSpan(span, err)
 	}()
 
-	err = w.brk.DoWithAcceptable(func() error {
+	err = w.brk.DoWithAcceptableCtx(ctx, func() error {
 		starTime := timex.Now()
 		defer func() {
 			logDuration(ctx, w.name, endSession, starTime, err)

core/stores/mon/trace.go

@@ -2,8 +2,8 @@ package mon
 
 import (
 	"context"
-	"errors"
 
+	"github.com/zeromicro/go-zero/core/errorx"
 	"github.com/zeromicro/go-zero/core/trace"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.opentelemetry.io/otel/attribute"
@@ -24,8 +24,7 @@ func startSpan(ctx context.Context, cmd string) (context.Context, oteltrace.Span
 func endSpan(span oteltrace.Span, err error) {
 	defer span.End()
 
-	if err == nil || errors.Is(err, mongo.ErrNoDocuments) ||
-		errors.Is(err, mongo.ErrNilValue) || errors.Is(err, mongo.ErrNilDocument) {
+	if err == nil || errorx.In(err, mongo.ErrNoDocuments, mongo.ErrNilValue, mongo.ErrNilDocument) {
 		span.SetStatus(codes.Ok, "")
 		return
 	}

core/stores/redis/breakerhook.go

@@ -0,0 +1,41 @@
+package redis
+
+import (
+	"context"
+
+	red "github.com/redis/go-redis/v9"
+	"github.com/zeromicro/go-zero/core/breaker"
+	"github.com/zeromicro/go-zero/core/lang"
+)
+
+var ignoreCmds = map[string]lang.PlaceholderType{
+	"blpop": {},
+}
+
+type breakerHook struct {
+	brk breaker.Breaker
+}
+
+func (h breakerHook) DialHook(next red.DialHook) red.DialHook {
+	return next
+}
+
+func (h breakerHook) ProcessHook(next red.ProcessHook) red.ProcessHook {
+	return func(ctx context.Context, cmd red.Cmder) error {
+		if _, ok := ignoreCmds[cmd.Name()]; ok {
+			return next(ctx, cmd)
+		}
+
+		return h.brk.DoWithAcceptableCtx(ctx, func() error {
+			return next(ctx, cmd)
+		}, acceptable)
+	}
+}
+
+func (h breakerHook) ProcessPipelineHook(next red.ProcessPipelineHook) red.ProcessPipelineHook {
+	return func(ctx context.Context, cmds []red.Cmder) error {
+		return h.brk.DoWithAcceptableCtx(ctx, func() error {
+			return next(ctx, cmds)
+		}, acceptable)
+	}
+}

core/stores/redis/breakerhook_test.go

@@ -0,0 +1,135 @@
+package redis
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/alicebob/miniredis/v2"
+	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/breaker"
+)
+
+func TestBreakerHook_ProcessHook(t *testing.T) {
+	t.Run("breakerHookOpen", func(t *testing.T) {
+		s := miniredis.RunT(t)
+
+		rds := MustNewRedis(RedisConf{
+			Host: s.Addr(),
+			Type: NodeType,
+		})
+
+		someError := errors.New("ERR some error")
+		s.SetError(someError.Error())
+
+		var err error
+		for i := 0; i < 1000; i++ {
+			_, err = rds.Get("key")
+			if err != nil && err.Error() != someError.Error() {
+				break
+			}
+		}
+		assert.Equal(t, breaker.ErrServiceUnavailable, err)
+	})
+
+	t.Run("breakerHookClose", func(t *testing.T) {
+		s := miniredis.RunT(t)
+
+		rds := MustNewRedis(RedisConf{
+			Host: s.Addr(),
+			Type: NodeType,
+		})
+
+		var err error
+		for i := 0; i < 1000; i++ {
+			_, err = rds.Get("key")
+			if err != nil {
+				break
+			}
+		}
+		assert.NotEqual(t, breaker.ErrServiceUnavailable, err)
+	})
+
+	t.Run("breakerHook_ignoreCmd", func(t *testing.T) {
+		s := miniredis.RunT(t)
+
+		rds := MustNewRedis(RedisConf{
+			Host: s.Addr(),
+			Type: NodeType,
+		})
+
+		someError := errors.New("ERR some error")
+		s.SetError(someError.Error())
+
+		var err error
+
+		node, err := getRedis(rds)
+		assert.NoError(t, err)
+
+		for i := 0; i < 1000; i++ {
+			_, err = rds.Blpop(node, "key")
+			if err != nil && err.Error() != someError.Error() {
+				break
+			}
+		}
+		assert.Equal(t, someError.Error(), err.Error())
+	})
+}
+
+func TestBreakerHook_ProcessPipelineHook(t *testing.T) {
+	t.Run("breakerPipelineHookOpen", func(t *testing.T) {
+		s := miniredis.RunT(t)
+
+		rds := MustNewRedis(RedisConf{
+			Host: s.Addr(),
+			Type: NodeType,
+		})
+
+		someError := errors.New("ERR some error")
+		s.SetError(someError.Error())
+
+		var err error
+		for i := 0; i < 1000; i++ {
+			err = rds.Pipelined(
+				func(pipe Pipeliner) error {
+					pipe.Incr(context.Background(), "pipelined_counter")
+					pipe.Expire(context.Background(), "pipelined_counter", time.Hour)
+					pipe.ZAdd(context.Background(), "zadd", Z{Score: 12, Member: "zadd"})
+					return nil
+				},
+			)
+
+			if err != nil && err.Error() != someError.Error() {
+				break
+			}
+		}
+		assert.Equal(t, breaker.ErrServiceUnavailable, err)
+	})
+
+	t.Run("breakerPipelineHookClose", func(t *testing.T) {
+		s := miniredis.RunT(t)
+
+		rds := MustNewRedis(RedisConf{
+			Host: s.Addr(),
+			Type: NodeType,
+		})
+
+		var err error
+		for i := 0; i < 1000; i++ {
+			err = rds.Pipelined(
+				func(pipe Pipeliner) error {
+					pipe.Incr(context.Background(), "pipelined_counter")
+					pipe.Expire(context.Background(), "pipelined_counter", time.Hour)
+					pipe.ZAdd(context.Background(), "zadd", Z{Score: 12, Member: "zadd"})
+					return nil
+				},
+			)
+
+			if err != nil {
+				break
+			}
+		}
+		assert.NotEqual(t, breaker.ErrServiceUnavailable, err)
+	})
+}

core/stores/redis/conf.go

@@ -47,7 +47,7 @@ func (rc RedisConf) NewRedis() *Redis {
 		opts = append(opts, WithTLS())
 	}
 
-	return New(rc.Host, opts...)
+	return newRedis(rc.Host, opts...)
 }
 
 // Validate validates the RedisConf.

core/stores/redis/delscript.lua

@@ -0,0 +1,5 @@
+if redis.call("GET", KEYS[1]) == ARGV[1] then
+    return redis.call("DEL", KEYS[1])
+else
+    return 0
+end

core/stores/redis/durationhook.go

@@ -23,17 +23,18 @@ import (
 const spanName = "redis"
 
 var (
-	durationHook          = hook{}
+	defaultDurationHook   = durationHook{}
 	redisCmdsAttributeKey = attribute.Key("redis.cmds")
 )
 
-type hook struct{}
+type durationHook struct {
+}
 
-func (h hook) DialHook(next red.DialHook) red.DialHook {
+func (h durationHook) DialHook(next red.DialHook) red.DialHook {
 	return next
 }
 
-func (h hook) ProcessHook(next red.ProcessHook) red.ProcessHook {
+func (h durationHook) ProcessHook(next red.ProcessHook) red.ProcessHook {
 	return func(ctx context.Context, cmd red.Cmder) error {
 		start := timex.Now()
 		ctx, endSpan := h.startSpan(ctx, cmd)
@@ -57,7 +58,7 @@ func (h hook) ProcessHook(next red.ProcessHook) red.ProcessHook {
 	}
 }
 
-func (h hook) ProcessPipelineHook(next red.ProcessPipelineHook) red.ProcessPipelineHook {
+func (h durationHook) ProcessPipelineHook(next red.ProcessPipelineHook) red.ProcessPipelineHook {
 	return func(ctx context.Context, cmds []red.Cmder) error {
 		if len(cmds) == 0 {
 			return next(ctx, cmds)
@@ -83,6 +84,33 @@ func (h hook) ProcessPipelineHook(next red.ProcessPipelineHook) red.ProcessPipel
 	}
 }
 
+func (h durationHook) startSpan(ctx context.Context, cmds ...red.Cmder) (context.Context, func(err error)) {
+	tracer := trace.TracerFromContext(ctx)
+
+	ctx, span := tracer.Start(ctx,
+		spanName,
+		oteltrace.WithSpanKind(oteltrace.SpanKindClient),
+	)
+
+	cmdStrs := make([]string, 0, len(cmds))
+	for _, cmd := range cmds {
+		cmdStrs = append(cmdStrs, cmd.Name())
+	}
+	span.SetAttributes(redisCmdsAttributeKey.StringSlice(cmdStrs))
+
+	return ctx, func(err error) {
+		defer span.End()
+
+		if err == nil || errors.Is(err, red.Nil) {
+			span.SetStatus(codes.Ok, "")
+			return
+		}
+
+		span.SetStatus(codes.Error, err.Error())
+		span.RecordError(err)
+	}
+}
+
 func formatError(err error) string {
 	if err == nil || errors.Is(err, red.Nil) {
 		return ""
@@ -95,7 +123,7 @@ func formatError(err error) string {
 	}
 
 	switch {
-	case err == io.EOF:
+	case errors.Is(err, io.EOF):
 		return "eof"
 	case errors.Is(err, context.DeadlineExceeded):
 		return "context deadline"
@@ -123,30 +151,3 @@ func logDuration(ctx context.Context, cmds []red.Cmder, duration time.Duration)
 	}
 	logx.WithContext(ctx).WithDuration(duration).Slowf("[REDIS] slowcall on executing: %s", buf.String())
 }
-
-func (h hook) startSpan(ctx context.Context, cmds ...red.Cmder) (context.Context, func(err error)) {
-	tracer := trace.TracerFromContext(ctx)
-
-	ctx, span := tracer.Start(ctx,
-		spanName,
-		oteltrace.WithSpanKind(oteltrace.SpanKindClient),
-	)
-
-	cmdStrs := make([]string, 0, len(cmds))
-	for _, cmd := range cmds {
-		cmdStrs = append(cmdStrs, cmd.Name())
-	}
-	span.SetAttributes(redisCmdsAttributeKey.StringSlice(cmdStrs))
-
-	return ctx, func(err error) {
-		defer span.End()
-
-		if err == nil || errors.Is(err, red.Nil) {
-			span.SetStatus(codes.Ok, "")
-			return
-		}
-
-		span.SetStatus(codes.Error, err.Error())
-		span.RecordError(err)
-	}
-}

core/stores/redis/durationhook_test.go

@@ -21,7 +21,7 @@ func TestHookProcessCase1(t *testing.T) {
 	tracetest.NewInMemoryExporter(t)
 	w := logtest.NewCollector(t)
 
-	err := durationHook.ProcessHook(func(ctx context.Context, cmd red.Cmder) error {
+	err := defaultDurationHook.ProcessHook(func(ctx context.Context, cmd red.Cmder) error {
 		assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
 		return nil
 	})(context.Background(), red.NewCmd(context.Background()))
@@ -36,7 +36,7 @@ func TestHookProcessCase2(t *testing.T) {
 	tracetest.NewInMemoryExporter(t)
 	w := logtest.NewCollector(t)
 
-	err := durationHook.ProcessHook(func(ctx context.Context, cmd red.Cmder) error {
+	err := defaultDurationHook.ProcessHook(func(ctx context.Context, cmd red.Cmder) error {
 		assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
 		time.Sleep(slowThreshold.Load() + time.Millisecond)
 		return nil
@@ -54,12 +54,12 @@ func TestHookProcessPipelineCase1(t *testing.T) {
 	tracetest.NewInMemoryExporter(t)
 	w := logtest.NewCollector(t)
 
-	err := durationHook.ProcessPipelineHook(func(ctx context.Context, cmds []red.Cmder) error {
+	err := defaultDurationHook.ProcessPipelineHook(func(ctx context.Context, cmds []red.Cmder) error {
 		return nil
 	})(context.Background(), nil)
 	assert.NoError(t, err)
 
-	err = durationHook.ProcessPipelineHook(func(ctx context.Context, cmds []red.Cmder) error {
+	err = defaultDurationHook.ProcessPipelineHook(func(ctx context.Context, cmds []red.Cmder) error {
 		assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
 		return nil
 	})(context.Background(), []red.Cmder{
@@ -74,7 +74,7 @@ func TestHookProcessPipelineCase2(t *testing.T) {
 	tracetest.NewInMemoryExporter(t)
 	w := logtest.NewCollector(t)
 
-	err := durationHook.ProcessPipelineHook(func(ctx context.Context, cmds []red.Cmder) error {
+	err := defaultDurationHook.ProcessPipelineHook(func(ctx context.Context, cmds []red.Cmder) error {
 		assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
 		time.Sleep(slowThreshold.Load() + time.Millisecond)
 		return nil
@@ -91,7 +91,7 @@ func TestHookProcessPipelineCase2(t *testing.T) {
 func TestHookProcessPipelineCase3(t *testing.T) {
 	te := tracetest.NewInMemoryExporter(t)
 
-	err := durationHook.ProcessPipelineHook(func(ctx context.Context, cmds []red.Cmder) error {
+	err := defaultDurationHook.ProcessPipelineHook(func(ctx context.Context, cmds []red.Cmder) error {
 		assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
 		return assert.AnError
 	})(context.Background(), []red.Cmder{

core/stores/redis/lockscript.lua

@@ -0,0 +1,6 @@
+if redis.call("GET", KEYS[1]) == ARGV[1] then
+    redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2])
+    return "OK"
+else
+    return redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2])
+end

core/stores/redis/metrics_test.go

@@ -19,7 +19,7 @@ func TestRedisMetric(t *testing.T) {
 	cfg := devserver.Config{}
 	_ = conf.FillDefault(&cfg)
 	server := devserver.NewServer(cfg)
-	server.StartAsync()
+	server.StartAsync(cfg)
 	time.Sleep(time.Second)
 
 	metricReqDur.Observe(8, "test-cmd")

core/stores/redis/redis_test.go

@@ -36,7 +36,7 @@ func (m myHook) ProcessHook(next red.ProcessHook) red.ProcessHook {
 		if cmd.Name() == "ping" && !m.includePing {
 			return next(ctx, cmd)
 		}
-		return errors.New("hook error")
+		return errors.New("durationHook error")
 	}
 }
 
@@ -155,12 +155,12 @@ func TestRedis_NonBlock(t *testing.T) {
 
 	t.Run("nonBlock true", func(t *testing.T) {
 		s := miniredis.RunT(t)
-		// use hook to simulate redis ping error
+		// use durationHook to simulate redis ping error
 		_, err := NewRedis(RedisConf{
 			Host:     s.Addr(),
 			NonBlock: true,
 			Type:     NodeType,
-		}, withHook(myHook{includePing: true}))
+		}, WithHook(myHook{includePing: true}))
 		assert.NoError(t, err)
 	})
 
@@ -170,7 +170,7 @@ func TestRedis_NonBlock(t *testing.T) {
 			Host:     s.Addr(),
 			NonBlock: false,
 			Type:     NodeType,
-		}, withHook(myHook{includePing: true}))
+		}, WithHook(myHook{includePing: true}))
 		assert.ErrorContains(t, err, "redis connect error")
 	})
 }

core/stores/redis/redisclientmanager.go

@@ -37,8 +37,11 @@ func getClient(r *Redis) (*red.Client, error) {
 			MinIdleConns: idleConns,
 			TLSConfig:    tlsConfig,
 		})
-		store.AddHook(durationHook)
-		for _, hook := range r.hooks {
+
+		hooks := append([]red.Hook{defaultDurationHook, breakerHook{
+			brk: r.brk,
+		}}, r.hooks...)
+		for _, hook := range hooks {
 			store.AddHook(hook)
 		}
 

core/stores/redis/redisclustermanager.go

@@ -33,8 +33,11 @@ func getCluster(r *Redis) (*red.ClusterClient, error) {
 			MinIdleConns: idleConns,
 			TLSConfig:    tlsConfig,
 		})
-		store.AddHook(durationHook)
-		for _, hook := range r.hooks {
+
+		hooks := append([]red.Hook{defaultDurationHook, breakerHook{
+			brk: r.brk,
+		}}, r.hooks...)
+		for _, hook := range hooks {
 			store.AddHook(hook)
 		}
 

core/stores/redis/redisclustermanager_test.go

@@ -51,7 +51,7 @@ func TestGetCluster(t *testing.T) {
 		Addr:  r.Addr(),
 		Type:  ClusterType,
 		tls:   true,
-		hooks: []red.Hook{durationHook},
+		hooks: []red.Hook{defaultDurationHook},
 	})
 	if assert.NoError(t, err) {
 		assert.NotNil(t, c)

core/stores/redis/redislock.go

@@ -2,6 +2,8 @@ package redis
 
 import (
 	"context"
+	_ "embed"
+	"errors"
 	"math/rand"
 	"strconv"
 	"sync/atomic"
@@ -19,17 +21,13 @@ const (
 )
 
 var (
-	lockScript = NewScript(`if redis.call("GET", KEYS[1]) == ARGV[1] then
-    redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2])
-    return "OK"
-else
-    return redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2])
-end`)
-	delScript = NewScript(`if redis.call("GET", KEYS[1]) == ARGV[1] then
-    return redis.call("DEL", KEYS[1])
-else
-    return 0
-end`)
+	//go:embed lockscript.lua
+	lockLuaScript string
+	lockScript    = NewScript(lockLuaScript)
+
+	//go:embed delscript.lua
+	delLuaScript string
+	delScript    = NewScript(delLuaScript)
 )
 
 // A RedisLock is a redis lock.
@@ -64,7 +62,7 @@ func (rl *RedisLock) AcquireCtx(ctx context.Context) (bool, error) {
 	resp, err := rl.store.ScriptRunCtx(ctx, lockScript, []string{rl.key}, []string{
 		rl.id, strconv.Itoa(int(seconds)*millisPerSecond + tolerance),
 	})
-	if err == red.Nil {
+	if errors.Is(err, red.Nil) {
 		return false, nil
 	} else if err != nil {
 		logx.Errorf("Error on acquiring lock for %s, %s", rl.key, err.Error())

core/stores/sqlc/cachedsql.go

@@ -190,6 +190,17 @@ func (cc CachedConn) QueryRowNoCacheCtx(ctx context.Context, v any, q string,
 	return cc.db.QueryRowCtx(ctx, v, q, args...)
 }
 
+// QueryRowPartialNoCache unmarshals into v with given statement.
+func (cc CachedConn) QueryRowPartialNoCache(v any, q string, args ...any) error {
+	return cc.QueryRowPartialNoCacheCtx(context.Background(), v, q, args...)
+}
+
+// QueryRowPartialNoCacheCtx unmarshals into v with given statement.
+func (cc CachedConn) QueryRowPartialNoCacheCtx(ctx context.Context, v any, q string,
+	args ...any) error {
+	return cc.db.QueryRowPartialCtx(ctx, v, q, args...)
+}
+
 // QueryRowsNoCache unmarshals into v with given statement.
 // It doesn't use cache, because it might cause consistency problem.
 func (cc CachedConn) QueryRowsNoCache(v any, q string, args ...any) error {
@@ -203,6 +214,19 @@ func (cc CachedConn) QueryRowsNoCacheCtx(ctx context.Context, v any, q string,
 	return cc.db.QueryRowsCtx(ctx, v, q, args...)
 }
 
+// QueryRowsPartialNoCache unmarshals into v with given statement.
+// It doesn't use cache, because it might cause consistency problem.
+func (cc CachedConn) QueryRowsPartialNoCache(v any, q string, args ...any) error {
+	return cc.QueryRowsPartialNoCacheCtx(context.Background(), v, q, args...)
+}
+
+// QueryRowsPartialNoCacheCtx unmarshals into v with given statement.
+// It doesn't use cache, because it might cause consistency problem.
+func (cc CachedConn) QueryRowsPartialNoCacheCtx(ctx context.Context, v any, q string,
+	args ...any) error {
+	return cc.db.QueryRowsPartialCtx(ctx, v, q, args...)
+}
+
 // SetCache sets v into cache with given key.
 func (cc CachedConn) SetCache(key string, val any) error {
 	return cc.SetCacheCtx(context.Background(), key, val)

core/stores/sqlc/cachedsql_test.go

@@ -579,6 +579,48 @@ func TestQueryRowNoCache(t *testing.T) {
 	assert.True(t, ran)
 }
 
+func TestQueryRowPartialNoCache(t *testing.T) {
+	r := redistest.CreateRedis(t)
+
+	const (
+		key   = "user"
+		value = "any"
+	)
+	var user string
+	var ran bool
+	conn := dummySqlConn{queryRow: func(v any, q string, args ...any) error {
+		user = value
+		ran = true
+		return nil
+	}}
+	c := NewNodeConn(&conn, r, cache.WithExpiry(time.Second*30))
+	err := c.QueryRowPartialNoCache(&user, key)
+	assert.Nil(t, err)
+	assert.Equal(t, value, user)
+	assert.True(t, ran)
+}
+
+func TestQueryRowsPartialNoCache(t *testing.T) {
+	r := redistest.CreateRedis(t)
+
+	var (
+		key    = "user"
+		values = []string{"any", "any"}
+	)
+	var users []string
+	var ran bool
+	conn := dummySqlConn{queryRows: func(v any, q string, args ...any) error {
+		users = values
+		ran = true
+		return nil
+	}}
+	c := NewNodeConn(&conn, r, cache.WithExpiry(time.Second*30))
+	err := c.QueryRowsPartialNoCache(&users, key)
+	assert.Nil(t, err)
+	assert.Equal(t, values, users)
+	assert.True(t, ran)
+}
+
 func TestNewConnWithCache(t *testing.T) {
 	r := redistest.CreateRedis(t)
 
@@ -716,7 +758,8 @@ func resetStats() {
 }
 
 type dummySqlConn struct {
-	queryRow func(any, string, ...any) error
+	queryRow  func(any, string, ...any) error
+	queryRows func(any, string, ...any) error
 }
 
 func (d dummySqlConn) ExecCtx(_ context.Context, _ string, _ ...any) (sql.Result, error) {
@@ -727,7 +770,11 @@ func (d dummySqlConn) PrepareCtx(_ context.Context, _ string) (sqlx.StmtSession,
 	return nil, nil
 }
 
-func (d dummySqlConn) QueryRowPartialCtx(_ context.Context, _ any, _ string, _ ...any) error {
+func (d dummySqlConn) QueryRowPartialCtx(_ context.Context, v any, query string, args ...any) error {
+	if d.queryRow != nil {
+		return d.queryRow(v, query, args...)
+	}
+
 	return nil
 }
 
@@ -735,7 +782,11 @@ func (d dummySqlConn) QueryRowsCtx(_ context.Context, _ any, _ string, _ ...any)
 	return nil
 }
 
-func (d dummySqlConn) QueryRowsPartialCtx(_ context.Context, _ any, _ string, _ ...any) error {
+func (d dummySqlConn) QueryRowsPartialCtx(_ context.Context, v any, query string, args ...any) error {
+	if d.queryRows != nil {
+		return d.queryRows(v, query, args...)
+	}
+
 	return nil
 }
 

core/stores/sqlx/metrics_test.go

@@ -20,7 +20,7 @@ func TestSqlxMetric(t *testing.T) {
 	_ = conf.FillDefault(&cfg)
 	cfg.Port = 6480
 	server := devserver.NewServer(cfg)
-	server.StartAsync()
+	server.StartAsync(cfg)
 	time.Sleep(time.Second)
 
 	metricReqDur.Observe(8, "test-cmd")

core/stores/sqlx/mysql.go

@@ -13,7 +13,7 @@ const (
 
 // NewMysql returns a mysql connection.
 func NewMysql(datasource string, opts ...SqlOption) SqlConn {
-	opts = append(opts, withMysqlAcceptable())
+	opts = append([]SqlOption{withMysqlAcceptable()}, opts...)
 	return NewSqlConn(mysqlDriverName, datasource, opts...)
 }
 

core/stores/sqlx/mysql_test.go

@@ -2,7 +2,6 @@ package sqlx
 
 import (
 	"errors"
-	"reflect"
 	"testing"
 
 	"github.com/go-sql-driver/mysql"
@@ -38,7 +37,6 @@ func TestBreakerOnNotHandlingDuplicateEntry(t *testing.T) {
 func TestMysqlAcceptable(t *testing.T) {
 	conn := NewMysql("nomysql").(*commonSqlConn)
 	withMysqlAcceptable()(conn)
-	assert.EqualValues(t, reflect.ValueOf(mysqlAcceptable).Pointer(), reflect.ValueOf(conn.accept).Pointer())
 	assert.True(t, mysqlAcceptable(nil))
 	assert.False(t, mysqlAcceptable(errors.New("any")))
 	assert.False(t, mysqlAcceptable(new(mysql.MySQLError)))

core/stores/sqlx/orm.go

@@ -1,6 +1,7 @@
 package sqlx
 
 import (
+	"context"
 	"errors"
 	"reflect"
 	"strings"
@@ -87,6 +88,10 @@ func getValueInterface(value reflect.Value) (any, error) {
 	}
 }
 
+func isScanFailed(err error) bool {
+	return err != nil && !errors.Is(err, context.DeadlineExceeded)
+}
+
 func mapStructFieldsIntoSlice(v reflect.Value, columns []string, strict bool) ([]any, error) {
 	fields := unwrapFields(v)
 	if strict && len(columns) < len(fields) {
@@ -248,7 +253,7 @@ func unmarshalRows(v any, scanner rowsScanner, strict bool) error {
 			return ErrUnsupportedValueType
 		}
 
-		return nil
+		return scanner.Err()
 	default:
 		return ErrUnsupportedValueType
 	}

core/stores/sqlx/sqlconn.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 
 	"github.com/zeromicro/go-zero/core/breaker"
+	"github.com/zeromicro/go-zero/core/errorx"
 	"github.com/zeromicro/go-zero/core/logx"
 )
 
@@ -83,7 +84,7 @@ func NewSqlConn(driverName, datasource string, opts ...SqlOption) SqlConn {
 }
 
 // NewSqlConnFromDB returns a SqlConn with the given sql.DB.
-// Use it with caution, it's provided for other ORM to interact with.
+// Use it with caution; it's provided for other ORM to interact with.
 func NewSqlConnFromDB(db *sql.DB, opts ...SqlOption) SqlConn {
 	conn := &commonSqlConn{
 		connProv: func() (*sql.DB, error) {
@@ -120,7 +121,7 @@ func (db *commonSqlConn) ExecCtx(ctx context.Context, q string, args ...any) (
 		endSpan(span, err)
 	}()
 
-	err = db.brk.DoWithAcceptable(func() error {
+	err = db.brk.DoWithAcceptableCtx(ctx, func() error {
 		var conn *sql.DB
 		conn, err = db.connProv()
 		if err != nil {
@@ -148,7 +149,7 @@ func (db *commonSqlConn) PrepareCtx(ctx context.Context, query string) (stmt Stm
 		endSpan(span, err)
 	}()
 
-	err = db.brk.DoWithAcceptable(func() error {
+	err = db.brk.DoWithAcceptableCtx(ctx, func() error {
 		var conn *sql.DB
 		conn, err = db.connProv()
 		if err != nil {
@@ -256,7 +257,7 @@ func (db *commonSqlConn) TransactCtx(ctx context.Context, fn func(context.Contex
 		endSpan(span, err)
 	}()
 
-	err = db.brk.DoWithAcceptable(func() error {
+	err = db.brk.DoWithAcceptableCtx(ctx, func() error {
 		return transact(ctx, db, db.beginTx, fn)
 	}, db.acceptable)
 	if errors.Is(err, breaker.ErrServiceUnavailable) {
@@ -267,8 +268,7 @@ func (db *commonSqlConn) TransactCtx(ctx context.Context, fn func(context.Contex
 }
 
 func (db *commonSqlConn) acceptable(err error) bool {
-	if err == nil || errors.Is(err, sql.ErrNoRows) || errors.Is(err, sql.ErrTxDone) ||
-		errors.Is(err, context.Canceled) {
+	if err == nil || errorx.In(err, sql.ErrNoRows, sql.ErrTxDone, context.Canceled) {
 		return true
 	}
 
@@ -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.brk.DoWithAcceptableCtx(ctx, func() error {
 		conn, err := db.connProv()
 		if err != nil {
 			db.onError(ctx, err)
@@ -296,7 +296,7 @@ func (db *commonSqlConn) queryRows(ctx context.Context, scanner func(*sql.Rows)
 
 		return query(ctx, conn, func(rows *sql.Rows) error {
 			e := scanner(rows)
-			if e != nil {
+			if isScanFailed(e) {
 				scanFailed = true
 			}
 			return e
@@ -315,6 +315,13 @@ func (db *commonSqlConn) queryRows(ctx context.Context, scanner func(*sql.Rows)
 // acceptable is the func to check if the error can be accepted.
 func WithAcceptable(acceptable func(err error) bool) SqlOption {
 	return func(conn *commonSqlConn) {
-		conn.accept = acceptable
+		if conn.accept == nil {
+			conn.accept = acceptable
+		} else {
+			pre := conn.accept
+			conn.accept = func(err error) bool {
+				return pre(err) || acceptable(err)
+			}
+		}
 	}
 }

core/stores/sqlx/sqlconn_test.go

@@ -264,6 +264,45 @@ func TestBreakerWithScanError(t *testing.T) {
 	})
 }
 
+func TestWithAcceptable(t *testing.T) {
+	var (
+		acceptableErr  = errors.New("acceptable")
+		acceptableErr2 = errors.New("acceptable2")
+		acceptableErr3 = errors.New("acceptable3")
+	)
+	opts := []SqlOption{
+		WithAcceptable(func(err error) bool {
+			if err == nil {
+				return true
+			}
+			return errors.Is(err, acceptableErr)
+		}),
+		WithAcceptable(func(err error) bool {
+			if err == nil {
+				return true
+			}
+			return errors.Is(err, acceptableErr2)
+		}),
+		WithAcceptable(func(err error) bool {
+			if err == nil {
+				return true
+			}
+			return errors.Is(err, acceptableErr3)
+		}),
+	}
+
+	var conn = &commonSqlConn{}
+	for _, opt := range opts {
+		opt(conn)
+	}
+
+	assert.True(t, conn.accept(nil))
+	assert.False(t, conn.accept(assert.AnError))
+	assert.True(t, conn.accept(acceptableErr))
+	assert.True(t, conn.accept(acceptableErr2))
+	assert.True(t, conn.accept(acceptableErr3))
+}
+
 func buildConn() (mock sqlmock.Sqlmock, err error) {
 	_, err = connManager.GetResource(mockedDatasource, func() (io.Closer, error) {
 		var db *sql.DB

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.brk.DoWithAcceptableCtx(ctx, func() error {
 		result, err = execStmt(ctx, s.stmt, s.query, args...)
 		return err
 	}, func(err error) bool {
@@ -141,10 +141,10 @@ 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.brk.DoWithAcceptableCtx(ctx, func() error {
 		return queryStmt(ctx, s.stmt, func(rows *sql.Rows) error {
 			err := scanFn(v, rows)
-			if err != nil {
+			if isScanFailed(err) {
 				scanFailed = true
 			}
 			return err

core/stores/sqlx/stmt_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"errors"
+	"strconv"
 	"testing"
 	"time"
 
@@ -290,6 +291,24 @@ func TestStmtBreaker(t *testing.T) {
 	})
 }
 
+func TestQueryRowsScanTimeout(t *testing.T) {
+	dbtest.RunTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
+		rows := sqlmock.NewRows([]string{"foo"})
+		for i := 0; i < 10000; i++ {
+			rows = rows.AddRow("bar" + strconv.Itoa(i))
+		}
+		mock.ExpectQuery("any").WillReturnRows(rows)
+		var val []struct {
+			Foo string
+		}
+		conn := NewSqlConnFromDB(db)
+		ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*2)
+		err := conn.QueryRowsCtx(ctx, &val, "any")
+		assert.ErrorIs(t, err, context.DeadlineExceeded)
+		cancel()
+	})
+}
+
 type mockedSessionConn struct {
 	lastInsertId int64
 	rowsAffected int64

core/trace/agent.go

@@ -59,7 +59,13 @@ func StartAgent(c Config) {
 
 // StopAgent shuts down the span processors in the order they were registered.
 func StopAgent() {
-	_ = tp.Shutdown(context.Background())
+	lock.Lock()
+	defer lock.Unlock()
+
+	if tp != nil {
+		_ = tp.Shutdown(context.Background())
+		tp = nil
+	}
 }
 
 func createExporter(c Config) (sdktrace.SpanExporter, error) {
@@ -91,9 +97,12 @@ func createExporter(c Config) (sdktrace.SpanExporter, error) {
 	case kindOtlpHttp:
 		// Not support flexible configuration now.
 		opts := []otlptracehttp.Option{
-			otlptracehttp.WithInsecure(),
 			otlptracehttp.WithEndpoint(c.Endpoint),
 		}
+
+		if !c.OtlpHttpSecure {
+			opts = append(opts, otlptracehttp.WithInsecure())
+		}
 		if len(c.OtlpHeaders) > 0 {
 			opts = append(opts, otlptracehttp.WithHeaders(c.OtlpHeaders))
 		}

core/trace/config.go

@@ -17,6 +17,8 @@ type Config struct {
 	// For example
 	// /v1/traces
 	OtlpHttpPath string `json:",optional"`
+	// OtlpHttpSecure represents the scheme to use for OTLP HTTP transport.
+	OtlpHttpSecure bool `json:",optional"`
 	// Disabled indicates whether StartAgent starts the agent.
 	Disabled bool `json:",optional"`
 }

core/trace/message_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/trace"
@@ -13,8 +14,9 @@ import (
 )
 
 func TestMessageType_Event(t *testing.T) {
-	var span mockSpan
-	ctx := trace.ContextWithSpan(context.Background(), &span)
+	ctx, s := otel.Tracer(TraceName).Start(context.Background(), "test")
+	span := mockSpan{Span: s}
+	ctx = trace.ContextWithSpan(ctx, &span)
 	MessageReceived.Event(ctx, 1, "foo")
 	assert.Equal(t, messageEvent, span.name)
 	assert.NotEmpty(t, span.options)
@@ -30,11 +32,12 @@ func TestMessageType_EventProtoMessage(t *testing.T) {
 }
 
 type mockSpan struct {
+	trace.Span
 	name    string
 	options []trace.EventOption
 }
 
-func (m *mockSpan) End(options ...trace.SpanEndOption) {
+func (m *mockSpan) End(_ ...trace.SpanEndOption) {
 }
 
 func (m *mockSpan) AddEvent(name string, options ...trace.EventOption) {
@@ -46,20 +49,20 @@ func (m *mockSpan) IsRecording() bool {
 	return false
 }
 
-func (m *mockSpan) RecordError(err error, options ...trace.EventOption) {
+func (m *mockSpan) RecordError(_ error, _ ...trace.EventOption) {
 }
 
 func (m *mockSpan) SpanContext() trace.SpanContext {
 	panic("implement me")
 }
 
-func (m *mockSpan) SetStatus(code codes.Code, description string) {
+func (m *mockSpan) SetStatus(_ codes.Code, _ string) {
 }
 
-func (m *mockSpan) SetName(name string) {
+func (m *mockSpan) SetName(_ string) {
 }
 
-func (m *mockSpan) SetAttributes(kv ...attribute.KeyValue) {
+func (m *mockSpan) SetAttributes(_ ...attribute.KeyValue) {
 }
 
 func (m *mockSpan) TracerProvider() trace.TracerProvider {

core/trace/vars.go

@@ -4,5 +4,5 @@ import "net/http"
 
 // TraceIdKey is the trace id header.
 // https://www.w3.org/TR/trace-context/#trace-id
-// May change it to trace-id afterwards.
+// May change it to trace-id afterward.
 var TraceIdKey = http.CanonicalHeaderKey("x-trace-id")