Update version.go (#4204)

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

.dockerignore

@@ -1 +1,7 @@
 **/.git
+.dockerignore
+Dockerfile
+goctl
+Makefile
+readme.md
+readme-cn.md

.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: kevwan

.github/dependabot.yml

@@ -5,6 +5,14 @@
 
 version: 2
 updates:
+  - package-ecosystem: "docker" # Update image tags in Dockerfile
+    directory: "/"
+    schedule:
+      interval: "weekly"
+  - package-ecosystem: "github-actions" # Update GitHub Actions
+    directory: "/"
+    schedule:
+      interval: "weekly"
   - package-ecosystem: "gomod" # See documentation for possible values
     directory: "/" # Location of package manifests
     schedule:

.github/workflows/codeql-analysis.yml

@@ -39,7 +39,7 @@ jobs:
 
     # Initializes the CodeQL tools for scanning.
     - name: Initialize CodeQL
-      uses: github/codeql-action/init@v2
+      uses: github/codeql-action/init@v3
       with:
         languages: ${{ matrix.language }}
         # If you wish to specify custom queries, you can do so here or in a config file.
@@ -50,7 +50,7 @@ jobs:
     # Autobuild attempts to build any compiled languages  (C/C++, C#, or Java).
     # If this step fails, then you should remove it and run the build manually (see below)
     - name: Autobuild
-      uses: github/codeql-action/autobuild@v2
+      uses: github/codeql-action/autobuild@v3
 
     # ℹ️ Command-line programs to run using the OS shell.
     # 📚 https://git.io/JvXDl
@@ -64,4 +64,4 @@ jobs:
     #   make release
 
     - name: Perform CodeQL Analysis
-      uses: github/codeql-action/analyze@v2
+      uses: github/codeql-action/analyze@v3

.github/workflows/go.yml

@@ -15,9 +15,9 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Set up Go 1.x
-        uses: actions/setup-go@v4
+        uses: actions/setup-go@v5
         with:
-          go-version: 1.19
+          go-version: '1.19'
           check-latest: true
           cache: true
         id: go
@@ -40,7 +40,7 @@ jobs:
         run: go test -race -coverprofile=coverage.txt -covermode=atomic ./...
 
       - name: Codecov
-        uses: codecov/codecov-action@v3
+        uses: codecov/codecov-action@v4
 
   test-win:
     name: Windows
@@ -50,10 +50,10 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Set up Go 1.x
-        uses: actions/setup-go@v4
+        uses: actions/setup-go@v5
         with:
           # use 1.19 to guarantee Go 1.19 compatibility
-          go-version: 1.19
+          go-version: '1.19'
           check-latest: true
           cache: true
 

.github/workflows/issues.yml

@@ -7,7 +7,7 @@ jobs:
   close-issues:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/stale@v8
+      - uses: actions/stale@v9
         with:
           days-before-issue-stale: 365
           days-before-issue-close: 90

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 (
@@ -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,21 +54,31 @@ 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 func(err error) error) error
+		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.
 		// If a panic occurs in the request, the Breaker handles it as an error
 		// 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 func(err error) error, acceptable Acceptable) error
+		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.
+	Fallback func(err error) error
+
 	// Option defines the method to customize a Breaker.
 	Option func(breaker *circuitBreaker)
 
@@ -86,12 +102,12 @@ type (
 
 	internalThrottle interface {
 		allow() (internalPromise, error)
-		doReq(req func() error, fallback func(err error) error, acceptable Acceptable) error
+		doReq(req func() error, fallback Fallback, acceptable Acceptable) error
 	}
 
 	throttle interface {
 		allow() (Promise, error)
-		doReq(req func() error, fallback func(err error) error, acceptable Acceptable) error
+		doReq(req func() error, fallback Fallback, acceptable Acceptable) error
 	}
 )
 
@@ -114,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) DoWithFallback(req func() error, fallback func(err error) error) error {
+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) DoWithFallbackAcceptable(req func() error, fallback func(err error) error,
+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
 }
@@ -168,7 +232,7 @@ func (lt loggedThrottle) allow() (Promise, error) {
 	}, lt.logError(err)
 }
 
-func (lt loggedThrottle) doReq(req func() error, fallback func(err error) error, acceptable Acceptable) error {
+func (lt loggedThrottle) doReq(req func() error, fallback Fallback, acceptable Acceptable) error {
 	return lt.logError(lt.internalThrottle.doReq(req, fallback, func(err error) bool {
 		accept := acceptable(err)
 		if !accept && err != nil {

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,21 +31,44 @@ 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 func(err error) error) error {
+func DoWithFallback(name string, req func() error, fallback Fallback) error {
 	return do(name, func(b Breaker) error {
 		return b.DoWithFallback(req, fallback)
 	})
 }
 
+// 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 func(err error) error,
+func DoWithFallbackAcceptable(name string, req func() error, fallback Fallback,
 	acceptable Acceptable) error {
 	return do(name, func(b Breaker) error {
 		return b.DoWithFallbackAcceptable(req, fallback, acceptable)
 	})
 }
 
+// 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()
@@ -59,7 +92,7 @@ func GetBreaker(name string) Breaker {
 // NoBreakerFor disables the circuit breaker for the given name.
 func NoBreakerFor(name string) {
 	lock.Lock()
-	breakers[name] = newNopBreaker()
+	breakers[name] = NopBreaker()
 	lock.Unlock()
 }
 

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 {

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

@@ -1,57 +1,87 @@
 package breaker
 
 import (
-	"math"
 	"time"
 
 	"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
-	dropRatio := math.Max(0, (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
 	}
 
@@ -60,8 +90,9 @@ func (b *googleBreaker) allow() (internalPromise, error) {
 	}, nil
 }
 
-func (b *googleBreaker) doReq(req func() error, fallback func(err error) error, acceptable Acceptable) error {
+func (b *googleBreaker) doReq(req func() error, fallback Fallback, acceptable Acceptable) error {
 	if err := b.accept(); err != nil {
+		b.markDrop()
 		if fallback != nil {
 			return fallback(err)
 		}
@@ -69,38 +100,55 @@ func (b *googleBreaker) doReq(req func() error, fallback func(err error) error,
 		return err
 	}
 
+	var succ bool
 	defer func() {
-		if e := recover(); e != nil {
+		// if req() panic, success is false, mark as failure
+		if succ {
+			b.markSuccess()
+		} else {
 			b.markFailure()
-			panic(e)
 		}
 	}()
 
 	err := req()
 	if acceptable(err) {
-		b.markSuccess()
-	} else {
-		b.markFailure()
+		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)
 	})
 }
 
@@ -206,7 +271,7 @@ func BenchmarkGoogleBreakerAllow(b *testing.B) {
 	breaker := getGoogleBreaker()
 	b.ResetTimer()
 	for i := 0; i <= b.N; i++ {
-		breaker.accept()
+		_ = breaker.accept()
 		if i%2 == 0 {
 			breaker.markSuccess()
 		} else {
@@ -215,6 +280,16 @@ func BenchmarkGoogleBreakerAllow(b *testing.B) {
 	}
 }
 
+func BenchmarkGoogleBreakerDoReq(b *testing.B) {
+	breaker := getGoogleBreaker()
+	b.ResetTimer()
+	for i := 0; i <= b.N; i++ {
+		_ = breaker.doReq(func() error {
+			return nil
+		}, nil, defaultAcceptable)
+	}
+}
+
 func markSuccess(b *googleBreaker, count int) {
 	for i := 0; i < count; i++ {
 		p, err := b.allow()

core/breaker/nopbreaker.go

@@ -1,10 +1,13 @@
 package breaker
 
+import "context"
+
 const nopBreakerName = "nopBreaker"
 
 type nopBreaker struct{}
 
-func newNopBreaker() Breaker {
+// NopBreaker returns a breaker that never trigger breaker circuit.
+func NopBreaker() Breaker {
 	return nopBreaker{}
 }
 
@@ -16,20 +19,40 @@ 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) DoWithFallback(req func() error, _ func(err error) error) error {
+func (b nopBreaker) DoWithAcceptableCtx(_ context.Context, req func() error, _ Acceptable) error {
 	return req()
 }
 
-func (b nopBreaker) DoWithFallbackAcceptable(req func() error, _ func(err error) error,
-	_ Acceptable) error {
+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()
 }
 

core/breaker/nopbreaker_test.go

@@ -1,6 +1,7 @@
 package breaker
 
 import (
+	"context"
 	"errors"
 	"testing"
 
@@ -8,10 +9,12 @@ import (
 )
 
 func TestNopBreaker(t *testing.T) {
-	b := newNopBreaker()
+	b := NopBreaker()
 	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/conf/readme.md

@@ -12,7 +12,7 @@ type RestfulConf struct {
 	MaxConns     int           `json:",default=10000"`
 	MaxBytes     int64         `json:",default=1048576"`
 	Timeout      time.Duration `json:",default=3s"`
-	CpuThreshold int64         `json:",default=900,range=[0:1000]"`
+	CpuThreshold int64         `json:",default=900,range=[0:1000)"`
 }
 ```
 

core/discov/internal/registry.go

@@ -222,7 +222,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/batcherror.go

@@ -1,11 +1,15 @@
 package errorx
 
-import "bytes"
+import (
+	"bytes"
+	"sync"
+)
 
 type (
 	// A BatchError is an error that can hold multiple errors.
 	BatchError struct {
 		errs errorArray
+		lock sync.Mutex
 	}
 
 	errorArray []error
@@ -13,6 +17,9 @@ type (
 
 // Add adds errs to be, nil errors are ignored.
 func (be *BatchError) Add(errs ...error) {
+	be.lock.Lock()
+	defer be.lock.Unlock()
+
 	for _, err := range errs {
 		if err != nil {
 			be.errs = append(be.errs, err)
@@ -22,6 +29,9 @@ func (be *BatchError) Add(errs ...error) {
 
 // Err returns an error that represents all errors.
 func (be *BatchError) Err() error {
+	be.lock.Lock()
+	defer be.lock.Unlock()
+
 	switch len(be.errs) {
 	case 0:
 		return nil
@@ -34,6 +44,9 @@ func (be *BatchError) Err() error {
 
 // NotNil checks if any error inside.
 func (be *BatchError) NotNil() bool {
+	be.lock.Lock()
+	defer be.lock.Unlock()
+
 	return len(be.errs) > 0
 }
 

core/errorx/batcherror_test.go

@@ -3,6 +3,7 @@ package errorx
 import (
 	"errors"
 	"fmt"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -33,7 +34,7 @@ func TestBatchErrorNilFromFunc(t *testing.T) {
 func TestBatchErrorOneError(t *testing.T) {
 	var batch BatchError
 	batch.Add(errors.New(err1))
-	assert.NotNil(t, batch)
+	assert.NotNil(t, batch.Err())
 	assert.Equal(t, err1, batch.Err().Error())
 	assert.True(t, batch.NotNil())
 }
@@ -42,7 +43,26 @@ func TestBatchErrorWithErrors(t *testing.T) {
 	var batch BatchError
 	batch.Add(errors.New(err1))
 	batch.Add(errors.New(err2))
-	assert.NotNil(t, batch)
+	assert.NotNil(t, batch.Err())
 	assert.Equal(t, fmt.Sprintf("%s\n%s", err1, err2), batch.Err().Error())
 	assert.True(t, batch.NotNil())
 }
+
+func TestBatchErrorConcurrentAdd(t *testing.T) {
+	const count = 10000
+	var batch BatchError
+	var wg sync.WaitGroup
+
+	wg.Add(count)
+	for i := 0; i < count; i++ {
+		go func() {
+			defer wg.Done()
+			batch.Add(errors.New(err1))
+		}()
+	}
+	wg.Wait()
+
+	assert.NotNil(t, batch.Err())
+	assert.Equal(t, count, len(batch.errs))
+	assert.True(t, batch.NotNil())
+}

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/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/retry.go

@@ -2,6 +2,7 @@ package fx
 
 import (
 	"context"
+	"errors"
 	"time"
 
 	"github.com/zeromicro/go-zero/core/errorx"
@@ -14,9 +15,10 @@ type (
 	RetryOption func(*retryOptions)
 
 	retryOptions struct {
-		times    int
-		interval time.Duration
-		timeout  time.Duration
+		times        int
+		interval     time.Duration
+		timeout      time.Duration
+		ignoreErrors []error
 	}
 )
 
@@ -62,6 +64,11 @@ func retry(ctx context.Context, fn func(errChan chan error, retryCount int), opt
 		select {
 		case err := <-errChan:
 			if err != nil {
+				for _, ignoreErr := range options.ignoreErrors {
+					if errors.Is(err, ignoreErr) {
+						return nil
+					}
+				}
 				berr.Add(err)
 			} else {
 				return nil
@@ -84,19 +91,28 @@ func retry(ctx context.Context, fn func(errChan chan error, retryCount int), opt
 	return berr.Err()
 }
 
-// WithRetry customize a DoWithRetry call with given retry times.
-func WithRetry(times int) RetryOption {
+// WithIgnoreErrors Ignore the specified errors
+func WithIgnoreErrors(ignoreErrors []error) RetryOption {
 	return func(options *retryOptions) {
-		options.times = times
+		options.ignoreErrors = ignoreErrors
 	}
 }
 
+// WithInterval customizes a DoWithRetry call with given interval.
 func WithInterval(interval time.Duration) RetryOption {
 	return func(options *retryOptions) {
 		options.interval = interval
 	}
 }
 
+// WithRetry customizes a DoWithRetry call with given retry times.
+func WithRetry(times int) RetryOption {
+	return func(options *retryOptions) {
+		options.times = times
+	}
+}
+
+// WithTimeout customizes a DoWithRetry call with given timeout.
 func WithTimeout(timeout time.Duration) RetryOption {
 	return func(options *retryOptions) {
 		options.timeout = timeout

core/fx/retry_test.go

@@ -97,6 +97,24 @@ func TestRetryWithInterval(t *testing.T) {
 
 }
 
+func TestRetryWithWithIgnoreErrors(t *testing.T) {
+	ignoreErr1 := errors.New("ignore error1")
+	ignoreErr2 := errors.New("ignore error2")
+	ignoreErrs := []error{ignoreErr1, ignoreErr2}
+
+	assert.Nil(t, DoWithRetry(func() error {
+		return ignoreErr1
+	}, WithIgnoreErrors(ignoreErrs)))
+
+	assert.Nil(t, DoWithRetry(func() error {
+		return ignoreErr2
+	}, WithIgnoreErrors(ignoreErrs)))
+
+	assert.NotNil(t, DoWithRetry(func() error {
+		return errors.New("any")
+	}))
+}
+
 func TestRetryCtx(t *testing.T) {
 	t.Run("with timeout", func(t *testing.T) {
 		assert.NotNil(t, DoWithRetryCtx(context.Background(), func(ctx context.Context, retryCount int) error {

core/fx/stream.go

@@ -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/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

@@ -9,6 +9,7 @@ import (
 
 	"github.com/zeromicro/go-zero/core/collection"
 	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/mathx"
 	"github.com/zeromicro/go-zero/core/stat"
 	"github.com/zeromicro/go-zero/core/syncx"
 	"github.com/zeromicro/go-zero/core/timex"
@@ -21,8 +22,11 @@ const (
 	defaultCpuThreshold = 900
 	defaultMinRt        = float64(time.Second / time.Millisecond)
 	// moving average hyperparameter beta for calculating requests on the fly
-	flyingBeta      = 0.9
-	coolOffDuration = time.Second
+	flyingBeta               = 0.9
+	coolOffDuration          = time.Second
+	cpuMax                   = 1000 // millicpu
+	millisecondsPerSecond    = 1000
+	overloadFactorLowerBound = 0.1
 )
 
 var (
@@ -66,14 +70,14 @@ type (
 
 	adaptiveShedder struct {
 		cpuThreshold    int64
-		windows         int64
+		windowScale     float64
 		flying          int64
 		avgFlying       float64
 		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]]
 	}
 )
 
@@ -103,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,
-		windows:         int64(time.Second / bucketDuration),
+		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]]()),
 	}
 }
 
@@ -134,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)
@@ -149,39 +154,42 @@ func (as *adaptiveShedder) highThru() bool {
 	as.avgFlyingLock.Lock()
 	avgFlying := as.avgFlying
 	as.avgFlyingLock.Unlock()
-	maxFlight := as.maxFlight()
-	return int64(avgFlying) > maxFlight && atomic.LoadInt64(&as.flying) > maxFlight
+	maxFlight := as.maxFlight() * as.overloadFactor()
+	return avgFlying > maxFlight && float64(atomic.LoadInt64(&as.flying)) > maxFlight
 }
 
-func (as *adaptiveShedder) maxFlight() int64 {
+func (as *adaptiveShedder) maxFlight() float64 {
 	// windows = buckets per second
 	// maxQPS = maxPASS * windows
 	// minRT = min average response time in milliseconds
-	// maxQPS * minRT / milliseconds_per_second
-	return int64(math.Max(1, float64(as.maxPass()*as.windows)*(as.minRt()/1e3)))
+	// allowedFlying = maxQPS * minRT / milliseconds_per_second
+	maxFlight := float64(as.maxPass()) * as.minRt() * as.windowScale
+	return mathx.AtLeast(maxFlight, 1)
 }
 
 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 {
+	// if no requests in previous windows, return defaultMinRt,
+	// 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
 		}
@@ -190,6 +198,13 @@ func (as *adaptiveShedder) minRt() float64 {
 	return result
 }
 
+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.
+	return mathx.Between(factor, overloadFactorLowerBound, 1)
+}
+
 func (as *adaptiveShedder) shouldDrop() bool {
 	if as.systemOverloaded() || as.stillHot() {
 		if as.highThru() {
@@ -236,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
@@ -269,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

@@ -19,6 +19,7 @@ import (
 const (
 	buckets        = 10
 	bucketDuration = time.Millisecond * 50
+	windowFactor   = 0.01
 )
 
 func init() {
@@ -57,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{
@@ -82,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{
@@ -106,18 +107,18 @@ 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{
 		passCounter:     passCounter,
 		rtCounter:       rtCounter,
-		windows:         buckets,
+		windowScale:     windowFactor,
 		droppedRecently: syncx.NewAtomicBool(),
 	}
-	assert.Equal(t, int64(54), shedder.maxFlight())
+	assert.Equal(t, float64(54), shedder.maxFlight())
 }
 
 func TestAdaptiveShedderShouldDrop(t *testing.T) {
@@ -128,15 +129,15 @@ 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{
 		passCounter:     passCounter,
 		rtCounter:       rtCounter,
-		windows:         buckets,
+		windowScale:     windowFactor,
 		overloadTime:    syncx.NewAtomicDuration(),
 		droppedRecently: syncx.NewAtomicBool(),
 	}
@@ -149,7 +150,8 @@ func TestAdaptiveShedderShouldDrop(t *testing.T) {
 
 	// cpu >=  800, inflight > maxPass
 	shedder.avgFlying = 80
-	shedder.flying = 50
+	// because of the overloadFactor, so we need to make sure maxFlight is greater than flying
+	shedder.flying = int64(shedder.maxFlight()*shedder.overloadFactor()) - 5
 	assert.False(t, shedder.shouldDrop())
 
 	// cpu >=  800, inflight > maxPass
@@ -182,15 +184,15 @@ 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{
 		passCounter:     passCounter,
 		rtCounter:       rtCounter,
-		windows:         buckets,
+		windowScale:     windowFactor,
 		overloadTime:    syncx.NewAtomicDuration(),
 		droppedRecently: syncx.ForAtomicBool(true),
 	}
@@ -239,6 +241,32 @@ func BenchmarkAdaptiveShedder_Allow(b *testing.B) {
 	b.Run("low load", bench)
 }
 
-func newRollingWindow() *collection.RollingWindow {
-	return collection.NewRollingWindow(buckets, bucketDuration, collection.IgnoreCurrentBucket())
+func BenchmarkMaxFlight(b *testing.B) {
+	passCounter := newRollingWindow()
+	rtCounter := newRollingWindow()
+	for i := 0; i < 10; i++ {
+		if i > 0 {
+			time.Sleep(bucketDuration)
+		}
+		passCounter.Add(int64((i + 1) * 100))
+		for j := i*10 + 1; j <= i*10+10; j++ {
+			rtCounter.Add(int64(j))
+		}
+	}
+	shedder := &adaptiveShedder{
+		passCounter:     passCounter,
+		rtCounter:       rtCounter,
+		windowScale:     windowFactor,
+		droppedRecently: syncx.NewAtomicBool(),
+	}
+
+	for i := 0; i < b.N; i++ {
+		_ = shedder.maxFlight()
+	}
+}
+
+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"
@@ -17,14 +18,13 @@ import (
 const callerDepth = 4
 
 var (
-	timeFormat = "2006-01-02T15:04:05.000Z07:00"
-	logLevel   uint32
+	timeFormat        = "2006-01-02T15:04:05.000Z07:00"
 	encoding   uint32 = jsonEncodingType
 	// maxContentLength is used to truncate the log content, 0 for not truncating.
 	maxContentLength uint32
 	// use uint32 for atomic operations
-	disableLog  uint32
 	disableStat uint32
+	logLevel    uint32
 	options     logOptions
 	writer      = new(atomicWriter)
 	setupOnce   sync.Once
@@ -87,7 +87,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...)
@@ -96,7 +96,7 @@ func Debugw(msg string, fields ...LogField) {
 
 // Disable disables the logging.
 func Disable() {
-	atomic.StoreUint32(&disableLog, 1)
+	atomic.StoreUint32(&logLevel, disableLevel)
 	writer.Store(nopWriter{})
 }
 
@@ -143,7 +143,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...)
@@ -154,11 +154,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:
@@ -176,11 +176,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:
@@ -209,7 +209,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...)
@@ -250,16 +250,17 @@ func SetLevel(level uint32) {
 
 // SetWriter sets the logging writer. It can be used to customize the logging.
 func SetWriter(w Writer) {
-	if atomic.LoadUint32(&disableLog) == 0 {
+	if atomic.LoadUint32(&logLevel) != disableLevel {
 		writer.Store(w)
 	}
 }
 
-// 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() {
@@ -414,6 +415,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 {
@@ -481,7 +508,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.
@@ -521,7 +548,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"))
 }
 
@@ -666,6 +687,7 @@ func TestDisable(t *testing.T) {
 	WithMaxSize(1024)(&opt)
 	assert.Nil(t, Close())
 	assert.Nil(t, Close())
+	assert.Equal(t, uint32(disableLevel), atomic.LoadUint32(&logLevel))
 }
 
 func TestDisableStat(t *testing.T) {
@@ -680,7 +702,7 @@ func TestDisableStat(t *testing.T) {
 }
 
 func TestSetWriter(t *testing.T) {
-	atomic.StoreUint32(&disableLog, 0)
+	atomic.StoreUint32(&logLevel, 0)
 	Reset()
 	SetWriter(nopWriter{})
 	assert.NotNil(t, writer.Load())
@@ -858,3 +880,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

@@ -15,6 +15,8 @@ const (
 	ErrorLevel
 	// SevereLevel only log severe messages
 	SevereLevel
+	// disableLevel doesn't log any messages
+	disableLevel = 0xff
 )
 
 const (
@@ -46,6 +48,7 @@ const (
 	levelDebug  = "debug"
 
 	backupFileDelimiter = "-"
+	nilAngleString      = "<nil>"
 	flags               = 0x0
 )
 

core/logx/writer.go

@@ -254,11 +254,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 +277,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 +302,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 +347,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

@@ -189,6 +189,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)

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
 	}
@@ -223,11 +225,11 @@ func (u *Unmarshaler) fillSliceFromString(fieldType reflect.Type, value reflect.
 	switch v := mapValue.(type) {
 	case fmt.Stringer:
 		if err := jsonx.UnmarshalFromString(v.String(), &slice); err != nil {
-			return err
+			return fmt.Errorf("fullName: `%s`, error: `%w`", fullName, err)
 		}
 	case string:
 		if err := jsonx.UnmarshalFromString(v, &slice); err != nil {
-			return err
+			return fmt.Errorf("fullName: `%s`, error: `%w`", fullName, err)
 		}
 	default:
 		return errUnsupportedType
@@ -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)
@@ -428,6 +466,10 @@ func (u *Unmarshaler) parseOptionsWithContext(field reflect.StructField, m Value
 		}
 	}
 
+	if u.opts.fillDefault {
+		return key, &options.fieldOptionsWithContext, nil
+	}
+
 	optsWithContext, err := options.toOptionsWithContext(key, m, fullName)
 	if err != nil {
 		return "", nil, err
@@ -572,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)
 	}
@@ -625,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)
 	}
 }
 
@@ -5411,6 +5414,15 @@ func TestFillDefaultUnmarshal(t *testing.T) {
 		assert.Equal(t, "c", st.C)
 	})
 
+	t.Run("optional !", func(t *testing.T) {
+		var st struct {
+			A string `json:",optional"`
+			B string `json:",optional=!A"`
+		}
+		err := fillDefaultUnmarshal.Unmarshal(map[string]any{}, &st)
+		assert.NoError(t, err)
+	})
+
 	t.Run("has value", func(t *testing.T) {
 		type St struct {
 			A string `json:",default=a"`
@@ -5751,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 {
@@ -5857,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/entropy_test.go

@@ -15,3 +15,17 @@ func TestCalcEntropy(t *testing.T) {
 	}
 	assert.True(t, CalcEntropy(m) > .99)
 }
+
+func TestCalcEmptyEntropy(t *testing.T) {
+	m := make(map[any]int)
+	assert.Equal(t, float64(1), CalcEntropy(m))
+}
+
+func TestCalcDiffEntropy(t *testing.T) {
+	const total = 1000
+	m := make(map[any]int, total)
+	for i := 0; i < total; i++ {
+		m[i] = i
+	}
+	assert.True(t, CalcEntropy(m) < .99)
+}

core/mathx/range.go

@@ -0,0 +1,34 @@
+package mathx
+
+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 {
+	if x < lower {
+		return lower
+	}
+	return x
+}
+
+// AtMost returns the smaller of x or upper.
+func AtMost[T Numerical](x, upper T) T {
+	if x > upper {
+		return upper
+	}
+	return x
+}
+
+// Between returns the value of x clamped to the range [lower, upper].
+func Between[T Numerical](x, lower, upper T) T {
+	if x < lower {
+		return lower
+	}
+	if x > upper {
+		return upper
+	}
+	return x
+}

core/mathx/range_test.go

@@ -0,0 +1,513 @@
+package mathx
+
+import "testing"
+
+func TestAtLeast(t *testing.T) {
+	t.Run("test int", func(t *testing.T) {
+		if got := AtLeast(10, 5); got != 10 {
+			t.Errorf("AtLeast() = %v, want 10", got)
+		}
+		if got := AtLeast(3, 5); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+		if got := AtLeast(5, 5); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test int8", func(t *testing.T) {
+		if got := AtLeast(int8(10), int8(5)); got != 10 {
+			t.Errorf("AtLeast() = %v, want 10", got)
+		}
+		if got := AtLeast(int8(3), int8(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+		if got := AtLeast(int8(5), int8(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test int16", func(t *testing.T) {
+		if got := AtLeast(int16(10), int16(5)); got != 10 {
+			t.Errorf("AtLeast() = %v, want 10", got)
+		}
+		if got := AtLeast(int16(3), int16(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+		if got := AtLeast(int16(5), int16(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test int32", func(t *testing.T) {
+		if got := AtLeast(int32(10), int32(5)); got != 10 {
+			t.Errorf("AtLeast() = %v, want 10", got)
+		}
+		if got := AtLeast(int32(3), int32(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+		if got := AtLeast(int32(5), int32(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test int64", func(t *testing.T) {
+		if got := AtLeast(int64(10), int64(5)); got != 10 {
+			t.Errorf("AtLeast() = %v, want 10", got)
+		}
+		if got := AtLeast(int64(3), int64(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+		if got := AtLeast(int64(5), int64(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test uint", func(t *testing.T) {
+		if got := AtLeast(uint(10), uint(5)); got != 10 {
+			t.Errorf("AtLeast() = %v, want 10", got)
+		}
+		if got := AtLeast(uint(3), uint(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+		if got := AtLeast(uint(5), uint(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test uint8", func(t *testing.T) {
+		if got := AtLeast(uint8(10), uint8(5)); got != 10 {
+			t.Errorf("AtLeast() = %v, want 10", got)
+		}
+		if got := AtLeast(uint8(3), uint8(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+		if got := AtLeast(uint8(5), uint8(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test uint16", func(t *testing.T) {
+		if got := AtLeast(uint16(10), uint16(5)); got != 10 {
+			t.Errorf("AtLeast() = %v, want 10", got)
+		}
+		if got := AtLeast(uint16(3), uint16(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+		if got := AtLeast(uint16(5), uint16(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test uint32", func(t *testing.T) {
+		if got := AtLeast(uint32(10), uint32(5)); got != 10 {
+			t.Errorf("AtLeast() = %v, want 10", got)
+		}
+		if got := AtLeast(uint32(3), uint32(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+		if got := AtLeast(uint32(5), uint32(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test uint64", func(t *testing.T) {
+		if got := AtLeast(uint64(10), uint64(5)); got != 10 {
+			t.Errorf("AtLeast() = %v, want 10", got)
+		}
+		if got := AtLeast(uint64(3), uint64(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+		if got := AtLeast(uint64(5), uint64(5)); got != 5 {
+			t.Errorf("AtLeast() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test float32", func(t *testing.T) {
+		if got := AtLeast(float32(10.0), float32(5.0)); got != 10.0 {
+			t.Errorf("AtLeast() = %v, want 10.0", got)
+		}
+		if got := AtLeast(float32(3.0), float32(5.0)); got != 5.0 {
+			t.Errorf("AtLeast() = %v, want 5.0", got)
+		}
+		if got := AtLeast(float32(5.0), float32(5.0)); got != 5.0 {
+			t.Errorf("AtLeast() = %v, want 5.0", got)
+		}
+	})
+
+	t.Run("test float64", func(t *testing.T) {
+		if got := AtLeast(10.0, 5.0); got != 10.0 {
+			t.Errorf("AtLeast() = %v, want 10.0", got)
+		}
+		if got := AtLeast(3.0, 5.0); got != 5.0 {
+			t.Errorf("AtLeast() = %v, want 5.0", got)
+		}
+		if got := AtLeast(5.0, 5.0); got != 5.0 {
+			t.Errorf("AtLeast() = %v, want 5.0", got)
+		}
+	})
+}
+
+func TestAtMost(t *testing.T) {
+	t.Run("test int", func(t *testing.T) {
+		if got := AtMost(10, 5); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+		if got := AtMost(3, 5); got != 3 {
+			t.Errorf("AtMost() = %v, want 3", got)
+		}
+		if got := AtMost(5, 5); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test int8", func(t *testing.T) {
+		if got := AtMost(int8(10), int8(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+		if got := AtMost(int8(3), int8(5)); got != 3 {
+			t.Errorf("AtMost() = %v, want 3", got)
+		}
+		if got := AtMost(int8(5), int8(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test int16", func(t *testing.T) {
+		if got := AtMost(int16(10), int16(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+		if got := AtMost(int16(3), int16(5)); got != 3 {
+			t.Errorf("AtMost() = %v, want 3", got)
+		}
+		if got := AtMost(int16(5), int16(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test int32", func(t *testing.T) {
+		if got := AtMost(int32(10), int32(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+		if got := AtMost(int32(3), int32(5)); got != 3 {
+			t.Errorf("AtMost() = %v, want 3", got)
+		}
+		if got := AtMost(int32(5), int32(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test int64", func(t *testing.T) {
+		if got := AtMost(int64(10), int64(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+		if got := AtMost(int64(3), int64(5)); got != 3 {
+			t.Errorf("AtMost() = %v, want 3", got)
+		}
+		if got := AtMost(int64(5), int64(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test uint", func(t *testing.T) {
+		if got := AtMost(uint(10), uint(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+		if got := AtMost(uint(3), uint(5)); got != 3 {
+			t.Errorf("AtMost() = %v, want 3", got)
+		}
+		if got := AtMost(uint(5), uint(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test uint8", func(t *testing.T) {
+		if got := AtMost(uint8(10), uint8(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+		if got := AtMost(uint8(3), uint8(5)); got != 3 {
+			t.Errorf("AtMost() = %v, want 3", got)
+		}
+		if got := AtMost(uint8(5), uint8(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test uint16", func(t *testing.T) {
+		if got := AtMost(uint16(10), uint16(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+		if got := AtMost(uint16(3), uint16(5)); got != 3 {
+			t.Errorf("AtMost() = %v, want 3", got)
+		}
+		if got := AtMost(uint16(5), uint16(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test uint32", func(t *testing.T) {
+		if got := AtMost(uint32(10), uint32(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+		if got := AtMost(uint32(3), uint32(5)); got != 3 {
+			t.Errorf("AtMost() = %v, want 3", got)
+		}
+		if got := AtMost(uint32(5), uint32(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test uint64", func(t *testing.T) {
+		if got := AtMost(uint64(10), uint64(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+		if got := AtMost(uint64(3), uint64(5)); got != 3 {
+			t.Errorf("AtMost() = %v, want 3", got)
+		}
+		if got := AtMost(uint64(5), uint64(5)); got != 5 {
+			t.Errorf("AtMost() = %v, want 5", got)
+		}
+	})
+
+	t.Run("test float32", func(t *testing.T) {
+		if got := AtMost(float32(10.0), float32(5.0)); got != 5.0 {
+			t.Errorf("AtMost() = %v, want 5.0", got)
+		}
+		if got := AtMost(float32(3.0), float32(5.0)); got != 3.0 {
+			t.Errorf("AtMost() = %v, want 3.0", got)
+		}
+		if got := AtMost(float32(5.0), float32(5.0)); got != 5.0 {
+			t.Errorf("AtMost() = %v, want 5.0", got)
+		}
+	})
+
+	t.Run("test float64", func(t *testing.T) {
+		if got := AtMost(10.0, 5.0); got != 5.0 {
+			t.Errorf("AtMost() = %v, want 5.0", got)
+		}
+		if got := AtMost(3.0, 5.0); got != 3.0 {
+			t.Errorf("AtMost() = %v, want 3.0", got)
+		}
+		if got := AtMost(5.0, 5.0); got != 5.0 {
+			t.Errorf("AtMost() = %v, want 5.0", got)
+		}
+	})
+}
+
+func TestBetween(t *testing.T) {
+	t.Run("test int", func(t *testing.T) {
+		if got := Between(10, 5, 15); got != 10 {
+			t.Errorf("Between() = %v, want 10", got)
+		}
+		if got := Between(3, 5, 15); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(20, 5, 15); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+		if got := Between(5, 5, 15); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(15, 5, 15); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+	})
+
+	t.Run("test int8", func(t *testing.T) {
+		if got := Between(int8(10), int8(5), int8(15)); got != 10 {
+			t.Errorf("Between() = %v, want 10", got)
+		}
+		if got := Between(int8(3), int8(5), int8(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(int8(20), int8(5), int8(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+		if got := Between(int8(5), int8(5), int8(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(int8(15), int8(5), int8(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+	})
+
+	t.Run("test int16", func(t *testing.T) {
+		if got := Between(int16(10), int16(5), int16(15)); got != 10 {
+			t.Errorf("Between() = %v, want 10", got)
+		}
+		if got := Between(int16(3), int16(5), int16(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(int16(20), int16(5), int16(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+		if got := Between(int16(5), int16(5), int16(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(int16(15), int16(5), int16(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+	})
+
+	t.Run("test int32", func(t *testing.T) {
+		if got := Between(int32(10), int32(5), int32(15)); got != 10 {
+			t.Errorf("Between() = %v, want 10", got)
+		}
+		if got := Between(int32(3), int32(5), int32(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(int32(20), int32(5), int32(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+		if got := Between(int32(5), int32(5), int32(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(int32(15), int32(5), int32(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+	})
+
+	t.Run("test int64", func(t *testing.T) {
+		if got := Between(int64(10), int64(5), int64(15)); got != 10 {
+			t.Errorf("Between() = %v, want 10", got)
+		}
+		if got := Between(int64(3), int64(5), int64(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(int64(20), int64(5), int64(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+		if got := Between(int64(5), int64(5), int64(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(int64(15), int64(5), int64(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+	})
+
+	t.Run("test uint", func(t *testing.T) {
+		if got := Between(uint(10), uint(5), uint(15)); got != 10 {
+			t.Errorf("Between() = %v, want 10", got)
+		}
+		if got := Between(uint(3), uint(5), uint(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(uint(20), uint(5), uint(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+		if got := Between(uint(5), uint(5), uint(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(uint(15), uint(5), uint(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+	})
+
+	t.Run("test uint8", func(t *testing.T) {
+		if got := Between(uint8(10), uint8(5), uint8(15)); got != 10 {
+			t.Errorf("Between() = %v, want 10", got)
+		}
+		if got := Between(uint8(3), uint8(5), uint8(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(uint8(20), uint8(5), uint8(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+		if got := Between(uint8(5), uint8(5), uint8(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(uint8(15), uint8(5), uint8(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+	})
+
+	t.Run("test uint16", func(t *testing.T) {
+		if got := Between(uint16(10), uint16(5), uint16(15)); got != 10 {
+			t.Errorf("Between() = %v, want 10", got)
+		}
+		if got := Between(uint16(3), uint16(5), uint16(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(uint16(20), uint16(5), uint16(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+		if got := Between(uint16(5), uint16(5), uint16(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(uint16(15), uint16(5), uint16(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+	})
+
+	t.Run("test uint32", func(t *testing.T) {
+		if got := Between(uint32(10), uint32(5), uint32(15)); got != 10 {
+			t.Errorf("Between() = %v, want 10", got)
+		}
+		if got := Between(uint32(3), uint32(5), uint32(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(uint32(20), uint32(5), uint32(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+		if got := Between(uint32(5), uint32(5), uint32(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(uint32(15), uint32(5), uint32(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+	})
+
+	t.Run("test uint64", func(t *testing.T) {
+		if got := Between(uint64(10), uint64(5), uint64(15)); got != 10 {
+			t.Errorf("Between() = %v, want 10", got)
+		}
+		if got := Between(uint64(3), uint64(5), uint64(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(uint64(20), uint64(5), uint64(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+		if got := Between(uint64(5), uint64(5), uint64(15)); got != 5 {
+			t.Errorf("Between() = %v, want 5", got)
+		}
+		if got := Between(uint64(15), uint64(5), uint64(15)); got != 15 {
+			t.Errorf("Between() = %v, want 15", got)
+		}
+	})
+
+	t.Run("test float32", func(t *testing.T) {
+		if got := Between(float32(10.0), float32(5.0), float32(15.0)); got != 10.0 {
+			t.Errorf("Between() = %v, want 10.0", got)
+		}
+		if got := Between(float32(3.0), float32(5.0), float32(15.0)); got != 5.0 {
+			t.Errorf("Between() = %v, want 5.0", got)
+		}
+		if got := Between(float32(20.0), float32(5.0), float32(15.0)); got != 15.0 {
+			t.Errorf("Between() = %v, want 15.0", got)
+		}
+		if got := Between(float32(5.0), float32(5.0), float32(15.0)); got != 5.0 {
+			t.Errorf("Between() = %v, want 5.0", got)
+		}
+		if got := Between(float32(15.0), float32(5.0), float32(15.0)); got != 15.0 {
+			t.Errorf("Between() = %v, want 15.0", got)
+		}
+	})
+
+	t.Run("test float64", func(t *testing.T) {
+		if got := Between(10.0, 5.0, 15.0); got != 10.0 {
+			t.Errorf("Between() = %v, want 10.0", got)
+		}
+		if got := Between(3.0, 5.0, 15.0); got != 5.0 {
+			t.Errorf("Between() = %v, want 5.0", got)
+		}
+		if got := Between(20.0, 5.0, 15.0); got != 15.0 {
+			t.Errorf("Between() = %v, want 15.0", got)
+		}
+		if got := Between(5.0, 5.0, 15.0); got != 5.0 {
+			t.Errorf("Between() = %v, want 5.0", got)
+		}
+		if got := Between(15.0, 5.0, 15.0); got != 15.0 {
+			t.Errorf("Between() = %v, want 15.0", got)
+		}
+	})
+}

core/mr/mapreduce_fuzz_test.go

@@ -13,7 +13,7 @@ import (
 )
 
 func FuzzMapReduce(f *testing.F) {
-	rand.Seed(time.Now().UnixNano())
+	rand.NewSource(time.Now().UnixNano())
 
 	f.Add(int64(10), runtime.NumCPU())
 	f.Fuzz(func(t *testing.T, n int64, workers int) {

core/mr/mapreduce_fuzzcase_test.go

@@ -20,7 +20,7 @@ import (
 // If Fuzz stuck, we don't know why, because it only returns hung or unexpected,
 // so we need to simulate the fuzz test in test mode.
 func TestMapReduceRandom(t *testing.T) {
-	rand.Seed(time.Now().UnixNano())
+	rand.NewSource(time.Now().UnixNano())
 
 	const (
 		times  = 10000

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/internal/cpu_linux.go

@@ -23,42 +23,17 @@ var (
 	preTotal  uint64
 	limit     float64
 	cores     uint64
+	noCgroup  bool
 	initOnce  sync.Once
 )
 
-// if /proc not present, ignore the cpu calculation, like wsl linux
-func initialize() {
-	cpus, err := effectiveCpus()
-	if err != nil {
-		logx.Error(err)
-		return
-	}
-
-	cores = uint64(cpus)
-	limit = float64(cpus)
-	quota, err := cpuQuota()
-	if err == nil && quota > 0 {
-		if quota < limit {
-			limit = quota
-		}
-	}
-
-	preSystem, err = systemCpuUsage()
-	if err != nil {
-		logx.Error(err)
-		return
-	}
-
-	preTotal, err = cpuUsage()
-	if err != nil {
-		logx.Error(err)
-		return
-	}
-}
-
 // RefreshCpu refreshes cpu usage and returns.
 func RefreshCpu() uint64 {
-	initOnce.Do(initialize)
+	initializeOnce()
+
+	if noCgroup {
+		return 0
+	}
 
 	total, err := cpuUsage()
 	if err != nil {
@@ -112,6 +87,47 @@ func effectiveCpus() (int, error) {
 	return cg.effectiveCpus()
 }
 
+// if /proc not present, ignore the cpu calculation, like wsl linux
+func initialize() error {
+	cpus, err := effectiveCpus()
+	if err != nil {
+		return err
+	}
+
+	cores = uint64(cpus)
+	limit = float64(cpus)
+	quota, err := cpuQuota()
+	if err == nil && quota > 0 {
+		if quota < limit {
+			limit = quota
+		}
+	}
+
+	preSystem, err = systemCpuUsage()
+	if err != nil {
+		return err
+	}
+
+	preTotal, err = cpuUsage()
+	return err
+}
+
+func initializeOnce() {
+	initOnce.Do(func() {
+		defer func() {
+			if p := recover(); p != nil {
+				noCgroup = true
+				logx.Error(p)
+			}
+		}()
+
+		if err := initialize(); err != nil {
+			noCgroup = true
+			logx.Error(err)
+		}
+	})
+}
+
 func systemCpuUsage() (uint64, error) {
 	lines, err := iox.ReadTextLines(statFile, iox.WithoutBlank())
 	if err != nil {

core/stat/metrics_test.go

@@ -1,11 +1,13 @@
 package stat
 
 import (
+	"errors"
 	"strconv"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/logx/logtest"
 )
 
 func TestMetrics(t *testing.T) {
@@ -30,6 +32,34 @@ func TestMetrics(t *testing.T) {
 	}
 }
 
+func TestTopDurationWithEmpty(t *testing.T) {
+	assert.Equal(t, float32(0), getTopDuration(nil))
+	assert.Equal(t, float32(0), getTopDuration([]Task{}))
+}
+
+func TestLogAndReport(t *testing.T) {
+	buf := logtest.NewCollector(t)
+	old := logEnabled.True()
+	logEnabled.Set(true)
+	t.Cleanup(func() {
+		logEnabled.Set(old)
+	})
+
+	log(&StatReport{})
+	assert.NotEmpty(t, buf.String())
+
+	writerLock.Lock()
+	writer := reportWriter
+	writerLock.Unlock()
+	buf = logtest.NewCollector(t)
+	t.Cleanup(func() {
+		SetReportWriter(writer)
+	})
+	SetReportWriter(&badWriter{})
+	writeReport(&StatReport{})
+	assert.NotEmpty(t, buf.String())
+}
+
 type mockedWriter struct {
 	report *StatReport
 }
@@ -38,3 +68,9 @@ func (m *mockedWriter) Write(report *StatReport) error {
 	m.report = report
 	return nil
 }
+
+type badWriter struct{}
+
+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/collection.go

@@ -2,10 +2,10 @@ package mon
 
 import (
 	"context"
-	"errors"
 	"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"
@@ -141,7 +141,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 +161,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 +181,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 +201,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 +221,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 +241,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 +261,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 +281,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 +301,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 +322,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 +344,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 +365,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 +386,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 +406,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 +426,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 +446,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 +466,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 +486,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 +527,10 @@ 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) ||
+	return err == nil || errorx.In(err, mongo.ErrNoDocuments, mongo.ErrNilValue,
+		mongo.ErrNilDocument, mongo.ErrNilCursor, 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)
+		session.ErrSessionEnded, session.ErrNoTransactStarted, session.ErrTransactInProgress,
+		session.ErrAbortAfterCommit, session.ErrAbortTwice, session.ErrCommitAfterAbort,
+		session.ErrUnackWCUnsupported, session.ErrSnapshotTransaction)
 }

core/stores/mon/collection_test.go

@@ -595,19 +595,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) DoWithFallback(_ func() error, _ func(err error) error) error {
+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) DoWithFallbackAcceptable(_ func() error, _ func(err error) error,
+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/options.go

@@ -1,9 +1,12 @@
 package mon
 
 import (
+	"reflect"
 	"time"
 
 	"github.com/zeromicro/go-zero/core/syncx"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/bsoncodec"
 	mopt "go.mongodb.org/mongo-driver/mongo/options"
 )
 
@@ -16,10 +19,17 @@ var (
 )
 
 type (
-	options = mopt.ClientOptions
-
 	// Option defines the method to customize a mongo model.
 	Option func(opts *options)
+
+	// TypeCodec is a struct that stores specific type Encoder/Decoder.
+	TypeCodec struct {
+		ValueType reflect.Type
+		Encoder   bsoncodec.ValueEncoder
+		Decoder   bsoncodec.ValueDecoder
+	}
+
+	options = mopt.ClientOptions
 )
 
 // DisableLog disables logging of mongo commands, includes info and slow logs.
@@ -38,15 +48,27 @@ func SetSlowThreshold(threshold time.Duration) {
 	slowThreshold.Set(threshold)
 }
 
-func defaultTimeoutOption() Option {
-	return func(opts *options) {
-		opts.SetTimeout(defaultTimeout)
-	}
-}
-
 // WithTimeout set the mon client operation timeout.
 func WithTimeout(timeout time.Duration) Option {
 	return func(opts *options) {
 		opts.SetTimeout(timeout)
 	}
 }
+
+// WithTypeCodec registers TypeCodecs to convert custom types.
+func WithTypeCodec(typeCodecs ...TypeCodec) Option {
+	return func(opts *options) {
+		registry := bson.NewRegistry()
+		for _, v := range typeCodecs {
+			registry.RegisterTypeEncoder(v.ValueType, v.Encoder)
+			registry.RegisterTypeDecoder(v.ValueType, v.Decoder)
+		}
+		opts.SetRegistry(registry)
+	}
+}
+
+func defaultTimeoutOption() Option {
+	return func(opts *options) {
+		opts.SetTimeout(defaultTimeout)
+	}
+}

core/stores/mon/options_test.go

@@ -1,10 +1,14 @@
 package mon
 
 import (
+	"fmt"
+	"reflect"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"go.mongodb.org/mongo-driver/bson/bsoncodec"
+	"go.mongodb.org/mongo-driver/bson/bsonrw"
 	mopt "go.mongodb.org/mongo-driver/mongo/options"
 )
 
@@ -51,3 +55,56 @@ func TestDisableInfoLog(t *testing.T) {
 	assert.False(t, logMon.True())
 	assert.True(t, logSlowMon.True())
 }
+
+func TestWithRegistryForTimestampRegisterType(t *testing.T) {
+	opts := mopt.Client()
+
+	// mongoDateTimeEncoder allow user convert time.Time to primitive.DateTime.
+	var mongoDateTimeEncoder bsoncodec.ValueEncoderFunc = func(ect bsoncodec.EncodeContext, w bsonrw.ValueWriter, value reflect.Value) error {
+		// Use reflect, determine if it can be converted to time.Time.
+		dec, ok := value.Interface().(time.Time)
+		if !ok {
+			return fmt.Errorf("value %v to encode is not of type time.Time", value)
+		}
+		return w.WriteDateTime(dec.Unix())
+	}
+
+	// mongoDateTimeEncoder allow user convert primitive.DateTime to time.Time.
+	var mongoDateTimeDecoder bsoncodec.ValueDecoderFunc = func(ect bsoncodec.DecodeContext, r bsonrw.ValueReader, value reflect.Value) error {
+		primTime, err := r.ReadDateTime()
+		if err != nil {
+			return fmt.Errorf("error reading primitive.DateTime from ValueReader: %v", err)
+		}
+		value.Set(reflect.ValueOf(time.Unix(primTime, 0)))
+		return nil
+	}
+
+	codecs := []TypeCodec{
+		{
+			ValueType: reflect.TypeOf(time.Time{}),
+			Encoder:   mongoDateTimeEncoder,
+			Decoder:   mongoDateTimeDecoder,
+		},
+	}
+	WithTypeCodec(codecs...)(opts)
+
+	for _, v := range codecs {
+		// Validate Encoder
+		enc, err := opts.Registry.LookupEncoder(v.ValueType)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if assert.ObjectsAreEqual(v.Encoder, enc) {
+			t.Errorf("Encoder got from Registry: %v, but want: %v", enc, v.Encoder)
+		}
+
+		// Validate Decoder
+		dec, err := opts.Registry.LookupDecoder(v.ValueType)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if assert.ObjectsAreEqual(v.Decoder, dec) {
+			t.Errorf("Decoder got from Registry: %v, but want: %v", dec, v.Decoder)
+		}
+	}
+}

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/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.
@@ -41,7 +39,7 @@ type RedisLock struct {
 }
 
 func init() {
-	rand.Seed(time.Now().UnixNano())
+	rand.NewSource(time.Now().UnixNano())
 }
 
 // NewRedisLock returns a RedisLock.
@@ -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/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)))