feat: improve mcp (#4828)

Signed-off-by: kevin <wanjunfeng@gmail.com>

.github/FUNDING.yml

@@ -1,13 +1,3 @@
 # These are supported funding model platforms
 
-github: # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]
+github: [zeromicro]
-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/workflows/go.yml

@@ -17,7 +17,7 @@ jobs:
       - name: Set up Go 1.x
         uses: actions/setup-go@v5
         with:
-          go-version: '1.19'
+          go-version-file: go.mod
           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@v4
+        uses: codecov/codecov-action@v5
 
   test-win:
     name: Windows
@@ -52,8 +52,8 @@ jobs:
       - name: Set up Go 1.x
         uses: actions/setup-go@v5
         with:
-          # use 1.19 to guarantee Go 1.19 compatibility
+          # make sure Go version compatible with go-zero
-          go-version: '1.19'
+          go-version-file: go.mod
           check-latest: true
           cache: true
 

.github/workflows/issue-translator.yml

@@ -1,18 +0,0 @@
-name: 'issue-translator'
-on: 
-  issue_comment: 
-    types: [created]
-  issues: 
-    types: [opened]
-
-jobs:
-  build:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: usthe/issues-translate-action@v2.7
-        with:
-          IS_MODIFY_TITLE: true
-          # not require, default false, . Decide whether to modify the issue title
-          # if true, the robot account @Issues-translate-bot must have modification permissions, invite @Issues-translate-bot to your project or use your custom bot.
-          CUSTOM_BOT_NOTE: Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿
-          # not require. Customize the translation robot prefix message.

.github/workflows/release.yaml

@@ -22,7 +22,7 @@ jobs:
           github_token: ${{ secrets.GITHUB_TOKEN }}
           goos: ${{ matrix.goos }}
           goarch: ${{ matrix.goarch }}
-          goversion: "https://dl.google.com/go/go1.19.13.linux-amd64.tar.gz"
+          goversion: "https://dl.google.com/go/go1.21.13.linux-amd64.tar.gz"
           project_path: "tools/goctl"
           binary_name: "goctl"
-          extra_files: tools/goctl/readme.md tools/goctl/readme-cn.md
+          extra_files: tools/goctl/readme.md tools/goctl/readme-cn.md

.github/workflows/reviewdog.yml

@@ -14,6 +14,6 @@ jobs:
           # Report all results.
           filter_mode: nofilter
           # Exit with 1 when it find at least one finding.
-          fail_on_error: true
+          fail_level: any
           # Set staticcheck flags
           staticcheck_flags: -checks=inherit,-SA1019,-SA1029,-SA5008

.github/workflows/version-check.yml

@@ -0,0 +1,42 @@
+name: Release Version Check
+
+on:
+  push:
+    tags:
+      - 'tools/goctl/v*'
+  workflow_dispatch:
+
+jobs:
+  version-check:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: '1.21'
+
+      - name: Extract tag version
+        id: get_version
+        run: |
+          # Extract version from tools/goctl/v* format
+          VERSION="${GITHUB_REF#refs/tags/tools/goctl/v}"
+          echo "VERSION=$VERSION" >> $GITHUB_ENV
+          echo "Extracted version: $VERSION"
+
+      - name: Check version in goctl source code
+        run: |
+          # Change to goctl directory
+          cd tools/goctl
+          
+          # Check version in BuildVersion constant
+          VERSION_IN_CODE=$(grep -r "const BuildVersion =" . | grep -o '".*"' | tr -d '"')
+          echo "Version in code: $VERSION_IN_CODE"
+          echo "Expected version: $VERSION"
+          
+          if [ "$VERSION_IN_CODE" != "$VERSION" ]; then
+            echo "Version mismatch: Version in code ($VERSION_IN_CODE) doesn't match tag version ($VERSION)"
+            exit 1
+          fi
+          echo "✅ Version check passed!"

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(`
+	//go:embed setscript.lua
-for _, offset in ipairs(ARGV) do
+	setLuaScript string
-	redis.call("setbit", KEYS[1], offset, 1)
+	setScript    = redis.NewScript(setLuaScript)
-end
+
-`)
+	//go:embed testscript.lua
-	testScript = redis.NewScript(`
+	testLuaScript string
-for _, offset in ipairs(ARGV) do
+	testScript    = redis.NewScript(testLuaScript)
-	if tonumber(redis.call("getbit", KEYS[1], offset)) == 0 then
-		return false
-	end
-end
-return true
-`)
 )
 
 type (
@@ -110,7 +105,7 @@ func newRedisBitSet(store *redis.Redis, key string, bits uint) *redisBitSet {
 }
 
 func (r *redisBitSet) buildOffsetArgs(offsets []uint) ([]string, error) {
-	var args []string
+	args := make([]string, 0, len(offsets))
 
 	for _, offset := range offsets {
 		if offset >= r.bits {

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"
@@ -36,12 +37,16 @@ type (
 		// 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.
@@ -49,12 +54,16 @@ type (
 		// and causes the same panic again.
 		// acceptable checks if it's a successful call, even if the error is not nil.
 		DoWithAcceptable(req func() error, acceptable Acceptable) error
+		// DoWithAcceptableCtx runs the given request if the Breaker accepts it when ctx isn't done.
+		DoWithAcceptableCtx(ctx context.Context, req func() error, acceptable Acceptable) error
 
 		// DoWithFallback runs the given request if the Breaker accepts it.
 		// DoWithFallback runs the fallback if the Breaker rejects the request.
 		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
 		DoWithFallback(req func() error, fallback Fallback) error
+		// DoWithFallbackCtx runs the given request if the Breaker accepts it when ctx isn't done.
+		DoWithFallbackCtx(ctx context.Context, req func() error, fallback Fallback) error
 
 		// DoWithFallbackAcceptable runs the given request if the Breaker accepts it.
 		// DoWithFallbackAcceptable runs the fallback if the Breaker rejects the request.
@@ -62,6 +71,9 @@ type (
 		// and causes the same panic again.
 		// acceptable checks if it's a successful call, even if the error is not nil.
 		DoWithFallbackAcceptable(req func() error, fallback Fallback, acceptable Acceptable) error
+		// DoWithFallbackAcceptableCtx runs the given request if the Breaker accepts it when ctx isn't done.
+		DoWithFallbackAcceptableCtx(ctx context.Context, req func() error, fallback Fallback,
+			acceptable Acceptable) error
 	}
 
 	// Fallback is the func to be called if the request is rejected.
@@ -118,23 +130,71 @@ func (cb *circuitBreaker) Allow() (Promise, error) {
 	return cb.throttle.allow()
 }
 
+func (cb *circuitBreaker) AllowCtx(ctx context.Context) (Promise, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+		return cb.Allow()
+	}
+}
+
 func (cb *circuitBreaker) Do(req func() error) error {
 	return cb.throttle.doReq(req, nil, defaultAcceptable)
 }
 
+func (cb *circuitBreaker) DoCtx(ctx context.Context, req func() error) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return cb.Do(req)
+	}
+}
+
 func (cb *circuitBreaker) DoWithAcceptable(req func() error, acceptable Acceptable) error {
 	return cb.throttle.doReq(req, nil, acceptable)
 }
 
+func (cb *circuitBreaker) DoWithAcceptableCtx(ctx context.Context, req func() error,
+	acceptable Acceptable) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return cb.DoWithAcceptable(req, acceptable)
+	}
+}
+
 func (cb *circuitBreaker) DoWithFallback(req func() error, fallback Fallback) error {
 	return cb.throttle.doReq(req, fallback, defaultAcceptable)
 }
 
+func (cb *circuitBreaker) DoWithFallbackCtx(ctx context.Context, req func() error,
+	fallback Fallback) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return cb.DoWithFallback(req, fallback)
+	}
+}
+
 func (cb *circuitBreaker) DoWithFallbackAcceptable(req func() error, fallback Fallback,
 	acceptable Acceptable) error {
 	return cb.throttle.doReq(req, fallback, acceptable)
 }
 
+func (cb *circuitBreaker) DoWithFallbackAcceptableCtx(ctx context.Context, req func() error,
+	fallback Fallback, acceptable Acceptable) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return cb.DoWithFallbackAcceptable(req, fallback, acceptable)
+	}
+}
+
 func (cb *circuitBreaker) Name() string {
 	return cb.name
 }
@@ -209,7 +269,7 @@ func (ew *errorWindow) add(reason string) {
 }
 
 func (ew *errorWindow) String() string {
-	var reasons []string
+	reasons := make([]string, 0, ew.count)
 
 	ew.lock.Lock()
 	// reverse order

core/breaker/breaker_test.go

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

core/breaker/breakers.go

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

core/breaker/breakers_test.go

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

core/breaker/bucket.go

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

core/breaker/bucket_test.go

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

core/breaker/googlebreaker.go

@@ -5,54 +5,83 @@ import (
 
 	"github.com/zeromicro/go-zero/core/collection"
 	"github.com/zeromicro/go-zero/core/mathx"
+	"github.com/zeromicro/go-zero/core/syncx"
+	"github.com/zeromicro/go-zero/core/timex"
 )
 
 const (
 	// 250ms for bucket duration
-	window     = time.Second * 10
+	window            = time.Second * 10
-	buckets    = 40
+	buckets           = 40
-	k          = 1.5
+	forcePassDuration = time.Second
-	protection = 5
+	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 {
+type (
-	k     float64
+	googleBreaker struct {
-	stat  *collection.RollingWindow
+		k        float64
-	proba *mathx.Proba
+		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,
+		stat:     st,
-		k:     k,
+		k:        k,
-		proba: mathx.NewProba(),
+		proba:    mathx.NewProba(),
+		lastPass: syncx.NewAtomicDuration(),
 	}
 }
 
 func (b *googleBreaker) accept() error {
-	accepts, total := b.history()
+	var w float64
-	weightedAccepts := b.k * float64(accepts)
+	history := b.history()
+	w = b.k - (b.k-minK)*float64(history.failingBuckets)/buckets
+	weightedAccepts := mathx.AtLeast(w, minK) * float64(history.accepts)
 	// https://landing.google.com/sre/sre-book/chapters/handling-overload/#eq2101
-	// for better performance, no need to care about negative ratio
+	// for better performance, no need to care about the negative ratio
-	dropRatio := (float64(total-protection) - weightedAccepts) / float64(total+1)
+	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.markFailure()
+		b.markDrop()
 		return nil, err
 	}
 
@@ -63,7 +92,7 @@ func (b *googleBreaker) allow() (internalPromise, error) {
 
 func (b *googleBreaker) doReq(req func() error, fallback Fallback, acceptable Acceptable) error {
 	if err := b.accept(); err != nil {
-		b.markFailure()
+		b.markDrop()
 		if fallback != nil {
 			return fallback(err)
 		}
@@ -71,10 +100,10 @@ func (b *googleBreaker) doReq(req func() error, fallback Fallback, acceptable Ac
 		return err
 	}
 
-	var success bool
+	var succ bool
 	defer func() {
 		// if req() panic, success is false, mark as failure
-		if success {
+		if succ {
 			b.markSuccess()
 		} else {
 			b.markFailure()
@@ -83,27 +112,43 @@ func (b *googleBreaker) doReq(req func() error, fallback Fallback, acceptable Ac
 
 	err := req()
 	if acceptable(err) {
-		success = true
+		succ = true
 	}
 
 	return err
 }
 
-func (b *googleBreaker) markSuccess() {
+func (b *googleBreaker) markDrop() {
-	b.stat.Add(1)
+	b.stat.Add(drop)
 }
 
 func (b *googleBreaker) markFailure() {
-	b.stat.Add(0)
+	b.stat.Add(fail)
 }
 
-func (b *googleBreaker) history() (accepts, total int64) {
+func (b *googleBreaker) markSuccess() {
-	b.stat.Reduce(func(b *collection.Bucket) {
+	b.stat.Add(success)
-		accepts += int64(b.Sum)
+}
-		total += b.Count
+
+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,
+		stat:     st,
-		k:     5,
+		k:        5,
-		proba: mathx.NewProba(),
+		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()
+		result := b.history()
-		assert.Equal(t, int64(10), accepts)
+		assert.Equal(t, int64(10), result.accepts)
-		assert.Equal(t, int64(10), total)
+		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()
+		result := b.history()
-		assert.Equal(t, int64(0), accepts)
+		assert.Equal(t, int64(0), result.accepts)
-		assert.Equal(t, int64(10), total)
+		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()
+		result := b.history()
-		assert.Equal(t, int64(5), accepts)
+		assert.Equal(t, int64(5), result.accepts)
-		assert.Equal(t, int64(10), total)
+		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()
+		result := b.history()
-		assert.Equal(t, int64(0), accepts)
+		assert.Equal(t, int64(0), result.accepts)
-		assert.Equal(t, int64(0), total)
+		assert.Equal(t, int64(0), result.total)
 	})
 }
 

core/breaker/nopbreaker.go

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

core/breaker/nopbreaker_test.go

@@ -1,6 +1,7 @@
 package breaker
 
 import (
+	"context"
 	"errors"
 	"testing"
 
@@ -10,7 +11,9 @@ import (
 func TestNopBreaker(t *testing.T) {
 	b := NopBreaker()
 	assert.Equal(t, nopBreakerName, b.Name())
-	p, err := b.Allow()
+	_, 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++ {
@@ -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/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.
+	// BucketInterface is the interface that defines the buckets.
-	RollingWindowOption func(rollingWindow *RollingWindow)
+	BucketInterface[T Numerical] interface {
+		Add(v T)
+		Reset()
+	}
 
-	// RollingWindow defines a rolling window to calculate the events in buckets with time interval.
+	// Numerical is the interface that restricts the numerical type.
-	RollingWindow struct {
+	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 {
+type Bucket[T Numerical] struct {
-	Sum   float64
+	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 {
+type window[T Numerical, B BucketInterface[T]] struct {
-	buckets []*Bucket
+	buckets []B
 	size    int
 }
 
-func newWindow(size int) *window {
+func newWindow[T Numerical, B BucketInterface[T]](newBucket func() B, size int) *window[T, B] {
-	buckets := make([]*Bucket, size)
+	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) {
+func (w *window[T, B]) add(offset int, v T) {
-	w.buckets[offset%w.size].add(v)
+	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) {
+func (w *window[T, B]) resetBucket(offset int) {
-	w.buckets[offset%w.size].reset()
+	w.buckets[offset%w.size].Reset()
 }
 
 // IgnoreCurrentBucket lets the Reduce call ignore current bucket.
-func IgnoreCurrentBucket() RollingWindowOption {
+func IgnoreCurrentBucket[T Numerical, B BucketInterface[T]]() RollingWindowOption[T, B] {
-	return func(w *RollingWindow) {
+	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

@@ -62,7 +62,11 @@ func Load(file string, v any, opts ...Option) error {
 		return loader([]byte(os.ExpandEnv(string(content))), v)
 	}
 
-	return loader(content, v)
+	if err = loader(content, v); err != nil {
+		return err
+	}
+
+	return validate(v)
 }
 
 // LoadConfig loads config into v from file, .json, .yaml and .yml are acceptable.
@@ -85,7 +89,12 @@ func LoadFromJsonBytes(content []byte, v any) error {
 
 	lowerCaseKeyMap := toLowerCaseKeyMap(m, info)
 
-	return mapping.UnmarshalJsonMap(lowerCaseKeyMap, v, mapping.WithCanonicalKeyFunc(toLowerCase))
+	if err = mapping.UnmarshalJsonMap(lowerCaseKeyMap, v,
+		mapping.WithCanonicalKeyFunc(toLowerCase)); err != nil {
+		return err
+	}
+
+	return validate(v)
 }
 
 // LoadConfigFromJsonBytes loads config into v from content json bytes.
@@ -189,10 +198,10 @@ func buildFieldsInfo(tp reflect.Type, fullName string) (*fieldInfo, error) {
 	switch tp.Kind() {
 	case reflect.Struct:
 		return buildStructFieldsInfo(tp, fullName)
-	case reflect.Array, reflect.Slice:
+	case reflect.Array, reflect.Slice, reflect.Map:
 		return buildFieldsInfo(mapping.Deref(tp.Elem()), fullName)
 	case reflect.Chan, reflect.Func:
-		return nil, fmt.Errorf("unsupported type: %s", tp.Kind())
+		return nil, fmt.Errorf("unsupported type: %s, fullName: %s", tp.Kind(), fullName)
 	default:
 		return &fieldInfo{
 			children: make(map[string]*fieldInfo),
@@ -332,6 +341,8 @@ func toLowerCaseKeyMap(m map[string]any, info *fieldInfo) map[string]any {
 			res[lk] = toLowerCaseInterface(v, ti)
 		} else if info.mapField != nil {
 			res[k] = toLowerCaseInterface(v, info.mapField)
+		} else if vv, ok := v.(map[string]any); ok {
+			res[k] = toLowerCaseKeyMap(vv, info)
 		} else {
 			res[k] = v
 		}

core/conf/config_test.go

@@ -1,6 +1,7 @@
 package conf
 
 import (
+	"errors"
 	"os"
 	"reflect"
 	"testing"
@@ -40,9 +41,8 @@ func TestConfigJson(t *testing.T) {
 	for _, test := range tests {
 		test := test
 		t.Run(test, func(t *testing.T) {
-			tmpfile, err := createTempFile(test, text)
+			tmpfile, err := createTempFile(t, test, text)
 			assert.Nil(t, err)
-			defer os.Remove(tmpfile)
 
 			var val struct {
 				A string `json:"a"`
@@ -82,9 +82,8 @@ c = "${FOO}"
 d = "abcd!@#$112"
 `
 	t.Setenv("FOO", "2")
-	tmpfile, err := createTempFile(".toml", text)
+	tmpfile, err := createTempFile(t, ".toml", text)
 	assert.Nil(t, err)
-	defer os.Remove(tmpfile)
 
 	var val struct {
 		A string `json:"a"`
@@ -105,9 +104,8 @@ b = 1
 c = "FOO"
 d = "abcd"
 `
-	tmpfile, err := createTempFile(".toml", text)
+	tmpfile, err := createTempFile(t, ".toml", text)
 	assert.Nil(t, err)
-	defer os.Remove(tmpfile)
 
 	var val struct {
 		A string `json:"a"`
@@ -127,9 +125,8 @@ func TestConfigWithLower(t *testing.T) {
 	text := `a = "foo"
 b = 1
 `
-	tmpfile, err := createTempFile(".toml", text)
+	tmpfile, err := createTempFile(t, ".toml", text)
 	assert.Nil(t, err)
-	defer os.Remove(tmpfile)
 
 	var val struct {
 		A string `json:"a"`
@@ -207,9 +204,8 @@ c = "${FOO}"
 d = "abcd!@#112"
 `
 	t.Setenv("FOO", "2")
-	tmpfile, err := createTempFile(".toml", text)
+	tmpfile, err := createTempFile(t, ".toml", text)
 	assert.Nil(t, err)
-	defer os.Remove(tmpfile)
 
 	var val struct {
 		A string `json:"a"`
@@ -241,9 +237,8 @@ func TestConfigJsonEnv(t *testing.T) {
 	for _, test := range tests {
 		test := test
 		t.Run(test, func(t *testing.T) {
-			tmpfile, err := createTempFile(test, text)
+			tmpfile, err := createTempFile(t, test, text)
 			assert.Nil(t, err)
-			defer os.Remove(tmpfile)
 
 			var val struct {
 				A string `json:"a"`
@@ -1192,6 +1187,42 @@ Email = "bar"`)
 			assert.Len(t, c.Value, 2)
 		}
 	})
+
+	t.Run("multi layer map", func(t *testing.T) {
+		type Value struct {
+			User struct {
+				Name string
+			}
+		}
+
+		type Config struct {
+			Value map[string]map[string]Value
+		}
+
+		var input = []byte(`
+[Value.first.User1.User]
+Name = "foo"
+[Value.second.User2.User]
+Name = "bar"
+`)
+		var c Config
+		if assert.NoError(t, LoadFromTomlBytes(input, &c)) {
+			assert.Len(t, c.Value, 2)
+		}
+	})
+}
+
+func Test_LoadBadConfig(t *testing.T) {
+	type Config struct {
+		Name string `json:"name,options=foo|bar"`
+	}
+
+	file, err := createTempFile(t, ".json", `{"name": "baz"}`)
+	assert.NoError(t, err)
+
+	var c Config
+	err = Load(file, &c)
+	assert.Error(t, err)
 }
 
 func Test_getFullName(t *testing.T) {
@@ -1199,6 +1230,26 @@ func Test_getFullName(t *testing.T) {
 	assert.Equal(t, "a", getFullName("", "a"))
 }
 
+func TestValidate(t *testing.T) {
+	t.Run("normal config", func(t *testing.T) {
+		var c mockConfig
+		err := LoadFromJsonBytes([]byte(`{"val": "hello", "number": 8}`), &c)
+		assert.NoError(t, err)
+	})
+
+	t.Run("error no int", func(t *testing.T) {
+		var c mockConfig
+		err := LoadFromJsonBytes([]byte(`{"val": "hello"}`), &c)
+		assert.Error(t, err)
+	})
+
+	t.Run("error no string", func(t *testing.T) {
+		var c mockConfig
+		err := LoadFromJsonBytes([]byte(`{"number": 8}`), &c)
+		assert.Error(t, err)
+	})
+}
+
 func Test_buildFieldsInfo(t *testing.T) {
 	type ParentSt struct {
 		Name string
@@ -1288,13 +1339,13 @@ func Test_buildFieldsInfo(t *testing.T) {
 	}
 }
 
-func createTempFile(ext, text string) (string, error) {
+func createTempFile(t *testing.T, ext, text string) (string, error) {
 	tmpFile, err := os.CreateTemp(os.TempDir(), hash.Md5Hex([]byte(text))+"*"+ext)
 	if err != nil {
 		return "", err
 	}
 
-	if err := os.WriteFile(tmpFile.Name(), []byte(text), os.ModeTemporary); err != nil {
+	if err = os.WriteFile(tmpFile.Name(), []byte(text), os.ModeTemporary); err != nil {
 		return "", err
 	}
 
@@ -1303,5 +1354,26 @@ func createTempFile(ext, text string) (string, error) {
 		return "", err
 	}
 
+	t.Cleanup(func() {
+		_ = os.Remove(filename)
+	})
+
 	return filename, nil
 }
+
+type mockConfig struct {
+	Val    string
+	Number int
+}
+
+func (m mockConfig) Validate() error {
+	if len(m.Val) == 0 {
+		return errors.New("val is empty")
+	}
+
+	if m.Number == 0 {
+		return errors.New("number is zero")
+	}
+
+	return nil
+}

core/conf/validate.go

@@ -0,0 +1,12 @@
+package conf
+
+import "github.com/zeromicro/go-zero/core/validation"
+
+// validate validates the value if it implements the Validator interface.
+func validate(v any) error {
+	if val, ok := v.(validation.Validator); ok {
+		return val.Validate()
+	}
+
+	return nil
+}

core/conf/validate_test.go

@@ -0,0 +1,81 @@
+package conf
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type mockType int
+
+func (m mockType) Validate() error {
+	if m < 10 {
+		return errors.New("invalid value")
+	}
+
+	return nil
+}
+
+type anotherMockType int
+
+func Test_validate(t *testing.T) {
+	tests := []struct {
+		name    string
+		v       any
+		wantErr bool
+	}{
+		{
+			name:    "invalid",
+			v:       mockType(5),
+			wantErr: true,
+		},
+		{
+			name:    "valid",
+			v:       mockType(10),
+			wantErr: false,
+		},
+		{
+			name:    "not validator",
+			v:       anotherMockType(5),
+			wantErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := validate(tt.v)
+			assert.Equal(t, tt.wantErr, err != nil)
+		})
+	}
+}
+
+type mockVal struct {
+}
+
+func (m mockVal) Validate() error {
+	return errors.New("invalid value")
+}
+
+func Test_validateValPtr(t *testing.T) {
+	tests := []struct {
+		name    string
+		v       any
+		wantErr bool
+	}{
+		{
+			name: "invalid",
+			v:    mockVal{},
+		},
+		{
+			name: "invalid value",
+			v:    &mockVal{},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Error(t, validate(tt.v))
+		})
+	}
+}

core/configcenter/configurator.go

@@ -0,0 +1,200 @@
+package configurator
+
+import (
+	"errors"
+	"fmt"
+	"reflect"
+	"strings"
+	"sync"
+	"sync/atomic"
+
+	"github.com/zeromicro/go-zero/core/configcenter/subscriber"
+	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/mapping"
+	"github.com/zeromicro/go-zero/core/threading"
+)
+
+var (
+	errEmptyConfig            = errors.New("empty config value")
+	errMissingUnmarshalerType = errors.New("missing unmarshaler type")
+)
+
+// Configurator is the interface for configuration center.
+type Configurator[T any] interface {
+	// GetConfig returns the subscription value.
+	GetConfig() (T, error)
+	// AddListener adds a listener to the subscriber.
+	AddListener(listener func())
+}
+
+type (
+	// Config is the configuration for Configurator.
+	Config struct {
+		// Type is the value type, yaml, json or toml.
+		Type string `json:",default=yaml,options=[yaml,json,toml]"`
+		// Log is the flag to control logging.
+		Log bool `json:",default=true"`
+	}
+
+	configCenter[T any] struct {
+		conf        Config
+		unmarshaler LoaderFn
+		subscriber  subscriber.Subscriber
+		listeners   []func()
+		lock        sync.Mutex
+		snapshot    atomic.Value
+	}
+
+	value[T any] struct {
+		data        string
+		marshalData T
+		err         error
+	}
+)
+
+// Configurator is the interface for configuration center.
+var _ Configurator[any] = (*configCenter[any])(nil)
+
+// MustNewConfigCenter returns a Configurator, exits on errors.
+func MustNewConfigCenter[T any](c Config, subscriber subscriber.Subscriber) Configurator[T] {
+	cc, err := NewConfigCenter[T](c, subscriber)
+	logx.Must(err)
+	return cc
+}
+
+// NewConfigCenter returns a Configurator.
+func NewConfigCenter[T any](c Config, subscriber subscriber.Subscriber) (Configurator[T], error) {
+	unmarshaler, ok := Unmarshaler(strings.ToLower(c.Type))
+	if !ok {
+		return nil, fmt.Errorf("unknown format: %s", c.Type)
+	}
+
+	cc := &configCenter[T]{
+		conf:        c,
+		unmarshaler: unmarshaler,
+		subscriber:  subscriber,
+	}
+
+	if err := cc.loadConfig(); err != nil {
+		return nil, err
+	}
+
+	if err := cc.subscriber.AddListener(cc.onChange); err != nil {
+		return nil, err
+	}
+
+	if _, err := cc.GetConfig(); err != nil {
+		return nil, err
+	}
+
+	return cc, nil
+}
+
+// AddListener adds listener to s.
+func (c *configCenter[T]) AddListener(listener func()) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	c.listeners = append(c.listeners, listener)
+}
+
+// GetConfig return structured config.
+func (c *configCenter[T]) GetConfig() (T, error) {
+	v := c.value()
+	if v == nil || len(v.data) == 0 {
+		var empty T
+		return empty, errEmptyConfig
+	}
+
+	return v.marshalData, v.err
+}
+
+// Value returns the subscription value.
+func (c *configCenter[T]) Value() string {
+	v := c.value()
+	if v == nil {
+		return ""
+	}
+	return v.data
+}
+
+func (c *configCenter[T]) loadConfig() error {
+	v, err := c.subscriber.Value()
+	if err != nil {
+		if c.conf.Log {
+			logx.Errorf("ConfigCenter loads changed configuration, error: %v", err)
+		}
+		return err
+	}
+
+	if c.conf.Log {
+		logx.Infof("ConfigCenter loads changed configuration, content [%s]", v)
+	}
+
+	c.snapshot.Store(c.genValue(v))
+	return nil
+}
+
+func (c *configCenter[T]) onChange() {
+	if err := c.loadConfig(); err != nil {
+		return
+	}
+
+	c.lock.Lock()
+	listeners := make([]func(), len(c.listeners))
+	copy(listeners, c.listeners)
+	c.lock.Unlock()
+
+	for _, l := range listeners {
+		threading.GoSafe(l)
+	}
+}
+
+func (c *configCenter[T]) value() *value[T] {
+	content := c.snapshot.Load()
+	if content == nil {
+		return nil
+	}
+	return content.(*value[T])
+}
+
+func (c *configCenter[T]) genValue(data string) *value[T] {
+	v := &value[T]{
+		data: data,
+	}
+	if len(data) == 0 {
+		return v
+	}
+
+	t := reflect.TypeOf(v.marshalData)
+	// if the type is nil, it means that the user has not set the type of the configuration.
+	if t == nil {
+		v.err = errMissingUnmarshalerType
+		return v
+	}
+
+	t = mapping.Deref(t)
+	switch t.Kind() {
+	case reflect.Struct, reflect.Array, reflect.Slice:
+		if err := c.unmarshaler([]byte(data), &v.marshalData); err != nil {
+			v.err = err
+			if c.conf.Log {
+				logx.Errorf("ConfigCenter unmarshal configuration failed, err: %+v, content [%s]",
+					err.Error(), data)
+			}
+		}
+	case reflect.String:
+		if str, ok := any(data).(T); ok {
+			v.marshalData = str
+		} else {
+			v.err = errMissingUnmarshalerType
+		}
+	default:
+		if c.conf.Log {
+			logx.Errorf("ConfigCenter unmarshal configuration missing unmarshaler for type: %s, content [%s]",
+				t.Kind(), data)
+		}
+		v.err = errMissingUnmarshalerType
+	}
+
+	return v
+}

core/configcenter/configurator_test.go

@@ -0,0 +1,233 @@
+package configurator
+
+import (
+	"errors"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewConfigCenter(t *testing.T) {
+	_, err := NewConfigCenter[any](Config{
+		Log: true,
+	}, &mockSubscriber{})
+	assert.Error(t, err)
+
+	_, err = NewConfigCenter[any](Config{
+		Type: "json",
+		Log:  true,
+	}, &mockSubscriber{})
+	assert.Error(t, err)
+}
+
+func TestConfigCenter_GetConfig(t *testing.T) {
+	mock := &mockSubscriber{}
+	type Data struct {
+		Name string `json:"name"`
+	}
+
+	mock.v = `{"name": "go-zero"}`
+	c1, err := NewConfigCenter[Data](Config{
+		Type: "json",
+		Log:  true,
+	}, mock)
+	assert.NoError(t, err)
+
+	data, err := c1.GetConfig()
+	assert.NoError(t, err)
+	assert.Equal(t, "go-zero", data.Name)
+
+	mock.v = `{"name": "111"}`
+	c2, err := NewConfigCenter[Data](Config{Type: "json"}, mock)
+	assert.NoError(t, err)
+
+	mock.v = `{}`
+	c3, err := NewConfigCenter[string](Config{
+		Type: "json",
+		Log:  true,
+	}, mock)
+	assert.NoError(t, err)
+	_, err = c3.GetConfig()
+	assert.NoError(t, err)
+
+	data, err = c2.GetConfig()
+	assert.NoError(t, err)
+
+	mock.lisErr = errors.New("mock error")
+	_, err = NewConfigCenter[Data](Config{
+		Type: "json",
+		Log:  true,
+	}, mock)
+	assert.Error(t, err)
+}
+
+func TestConfigCenter_onChange(t *testing.T) {
+	mock := &mockSubscriber{}
+	type Data struct {
+		Name string `json:"name"`
+	}
+
+	mock.v = `{"name": "go-zero"}`
+	c1, err := NewConfigCenter[Data](Config{Type: "json", Log: true}, mock)
+	assert.NoError(t, err)
+
+	data, err := c1.GetConfig()
+	assert.NoError(t, err)
+	assert.Equal(t, "go-zero", data.Name)
+
+	mock.v = `{"name": "go-zero2"}`
+	mock.change()
+
+	data, err = c1.GetConfig()
+	assert.NoError(t, err)
+	assert.Equal(t, "go-zero2", data.Name)
+
+	mock.valErr = errors.New("mock error")
+	_, err = NewConfigCenter[Data](Config{Type: "json", Log: false}, mock)
+	assert.Error(t, err)
+}
+
+func TestConfigCenter_Value(t *testing.T) {
+	mock := &mockSubscriber{}
+	mock.v = "1234"
+
+	c, err := NewConfigCenter[string](Config{
+		Type: "json",
+		Log:  true,
+	}, mock)
+	assert.NoError(t, err)
+
+	cc := c.(*configCenter[string])
+
+	assert.Equal(t, cc.Value(), "1234")
+
+	mock.valErr = errors.New("mock error")
+
+	_, err = NewConfigCenter[any](Config{
+		Type: "json",
+		Log:  true,
+	}, mock)
+	assert.Error(t, err)
+}
+
+func TestConfigCenter_AddListener(t *testing.T) {
+	mock := &mockSubscriber{}
+	mock.v = "1234"
+	c, err := NewConfigCenter[string](Config{
+		Type: "json",
+		Log:  true,
+	}, mock)
+	assert.NoError(t, err)
+
+	cc := c.(*configCenter[string])
+	var a, b int
+	var mutex sync.Mutex
+	cc.AddListener(func() {
+		mutex.Lock()
+		a = 1
+		mutex.Unlock()
+	})
+	cc.AddListener(func() {
+		mutex.Lock()
+		b = 2
+		mutex.Unlock()
+	})
+
+	assert.Equal(t, 2, len(cc.listeners))
+
+	mock.change()
+
+	time.Sleep(time.Millisecond * 100)
+
+	mutex.Lock()
+	assert.Equal(t, 1, a)
+	assert.Equal(t, 2, b)
+	mutex.Unlock()
+}
+
+func TestConfigCenter_genValue(t *testing.T) {
+	t.Run("data is empty", func(t *testing.T) {
+		c := &configCenter[string]{
+			unmarshaler: registry.unmarshalers["json"],
+			conf:        Config{Log: true},
+		}
+		v := c.genValue("")
+		assert.Equal(t, "", v.data)
+	})
+
+	t.Run("invalid template type", func(t *testing.T) {
+		c := &configCenter[any]{
+			unmarshaler: registry.unmarshalers["json"],
+			conf:        Config{Log: true},
+		}
+		v := c.genValue("xxxx")
+		assert.Equal(t, errMissingUnmarshalerType, v.err)
+	})
+
+	t.Run("unsupported template type", func(t *testing.T) {
+		c := &configCenter[int]{
+			unmarshaler: registry.unmarshalers["json"],
+			conf:        Config{Log: true},
+		}
+		v := c.genValue("1")
+		assert.Equal(t, errMissingUnmarshalerType, v.err)
+	})
+
+	t.Run("supported template string type", func(t *testing.T) {
+		c := &configCenter[string]{
+			unmarshaler: registry.unmarshalers["json"],
+			conf:        Config{Log: true},
+		}
+		v := c.genValue("12345")
+		assert.NoError(t, v.err)
+		assert.Equal(t, "12345", v.data)
+	})
+
+	t.Run("unmarshal fail", func(t *testing.T) {
+		c := &configCenter[struct {
+			Name string `json:"name"`
+		}]{
+			unmarshaler: registry.unmarshalers["json"],
+			conf:        Config{Log: true},
+		}
+		v := c.genValue(`{"name":"new name}`)
+		assert.Equal(t, `{"name":"new name}`, v.data)
+		assert.Error(t, v.err)
+	})
+
+	t.Run("success", func(t *testing.T) {
+		c := &configCenter[struct {
+			Name string `json:"name"`
+		}]{
+			unmarshaler: registry.unmarshalers["json"],
+			conf:        Config{Log: true},
+		}
+		v := c.genValue(`{"name":"new name"}`)
+		assert.Equal(t, `{"name":"new name"}`, v.data)
+		assert.Equal(t, "new name", v.marshalData.Name)
+		assert.NoError(t, v.err)
+	})
+}
+
+type mockSubscriber struct {
+	v              string
+	lisErr, valErr error
+	listener       func()
+}
+
+func (m *mockSubscriber) AddListener(listener func()) error {
+	m.listener = listener
+	return m.lisErr
+}
+
+func (m *mockSubscriber) Value() (string, error) {
+	return m.v, m.valErr
+}
+
+func (m *mockSubscriber) change() {
+	if m.listener != nil {
+		m.listener()
+	}
+}

core/configcenter/subscriber/etcd.go

@@ -0,0 +1,67 @@
+package subscriber
+
+import (
+	"github.com/zeromicro/go-zero/core/discov"
+	"github.com/zeromicro/go-zero/core/logx"
+)
+
+type (
+	// etcdSubscriber is a subscriber that subscribes to etcd.
+	etcdSubscriber struct {
+		*discov.Subscriber
+	}
+
+	// EtcdConf is the configuration for etcd.
+	EtcdConf = discov.EtcdConf
+)
+
+// MustNewEtcdSubscriber returns an etcd Subscriber, exits on errors.
+func MustNewEtcdSubscriber(conf EtcdConf) Subscriber {
+	s, err := NewEtcdSubscriber(conf)
+	logx.Must(err)
+	return s
+}
+
+// NewEtcdSubscriber returns an etcd Subscriber.
+func NewEtcdSubscriber(conf EtcdConf) (Subscriber, error) {
+	opts := buildSubOptions(conf)
+	s, err := discov.NewSubscriber(conf.Hosts, conf.Key, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	return &etcdSubscriber{Subscriber: s}, nil
+}
+
+// buildSubOptions constructs the options for creating a new etcd subscriber.
+func buildSubOptions(conf EtcdConf) []discov.SubOption {
+	opts := []discov.SubOption{
+		discov.WithExactMatch(),
+	}
+
+	if len(conf.User) > 0 {
+		opts = append(opts, discov.WithSubEtcdAccount(conf.User, conf.Pass))
+	}
+	if len(conf.CertFile) > 0 || len(conf.CertKeyFile) > 0 || len(conf.CACertFile) > 0 {
+		opts = append(opts, discov.WithSubEtcdTLS(conf.CertFile, conf.CertKeyFile,
+			conf.CACertFile, conf.InsecureSkipVerify))
+	}
+
+	return opts
+}
+
+// AddListener adds a listener to the subscriber.
+func (s *etcdSubscriber) AddListener(listener func()) error {
+	s.Subscriber.AddListener(listener)
+	return nil
+}
+
+// Value returns the value of the subscriber.
+func (s *etcdSubscriber) Value() (string, error) {
+	vs := s.Subscriber.Values()
+	if len(vs) > 0 {
+		return vs[len(vs)-1], nil
+	}
+
+	return "", nil
+}

core/configcenter/subscriber/subscriber.go

@@ -0,0 +1,9 @@
+package subscriber
+
+// Subscriber is the interface for configcenter subscribers.
+type Subscriber interface {
+	// AddListener adds a listener to the subscriber.
+	AddListener(listener func()) error
+	// Value returns the value of the subscriber.
+	Value() (string, error)
+}

core/configcenter/unmarshaler.go

@@ -0,0 +1,41 @@
+package configurator
+
+import (
+	"sync"
+
+	"github.com/zeromicro/go-zero/core/conf"
+)
+
+var registry = &unmarshalerRegistry{
+	unmarshalers: map[string]LoaderFn{
+		"json": conf.LoadFromJsonBytes,
+		"toml": conf.LoadFromTomlBytes,
+		"yaml": conf.LoadFromYamlBytes,
+	},
+}
+
+type (
+	// LoaderFn is the function type for loading configuration.
+	LoaderFn func([]byte, any) error
+
+	// unmarshalerRegistry is the registry for unmarshalers.
+	unmarshalerRegistry struct {
+		unmarshalers map[string]LoaderFn
+		mu           sync.RWMutex
+	}
+)
+
+// RegisterUnmarshaler registers an unmarshaler.
+func RegisterUnmarshaler(name string, fn LoaderFn) {
+	registry.mu.Lock()
+	defer registry.mu.Unlock()
+	registry.unmarshalers[name] = fn
+}
+
+// Unmarshaler returns the unmarshaler by name.
+func Unmarshaler(name string) (LoaderFn, bool) {
+	registry.mu.RLock()
+	defer registry.mu.RUnlock()
+	fn, ok := registry.unmarshalers[name]
+	return fn, ok
+}

core/configcenter/unmarshaler_test.go

@@ -0,0 +1,28 @@
+package configurator
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRegisterUnmarshaler(t *testing.T) {
+	RegisterUnmarshaler("test", func(data []byte, v interface{}) error {
+		return nil
+	})
+
+	_, ok := Unmarshaler("test")
+	assert.True(t, ok)
+
+	_, ok = Unmarshaler("test2")
+	assert.False(t, ok)
+
+	_, ok = Unmarshaler("json")
+	assert.True(t, ok)
+
+	_, ok = Unmarshaler("toml")
+	assert.True(t, ok)
+
+	_, ok = Unmarshaler("yaml")
+	assert.True(t, ok)
+}

core/discov/internal/registry.go

@@ -10,27 +10,30 @@ import (
 	"sync"
 	"time"
 
-	"github.com/zeromicro/go-zero/core/contextx"
 	"github.com/zeromicro/go-zero/core/lang"
-	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logc"
+	"github.com/zeromicro/go-zero/core/mathx"
 	"github.com/zeromicro/go-zero/core/syncx"
 	"github.com/zeromicro/go-zero/core/threading"
 	"go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
 	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
+const coolDownDeviation = 0.05
+
 var (
 	registry = Registry{
 		clusters: make(map[string]*cluster),
 	}
-	connManager = syncx.NewResourceManager()
+	connManager      = syncx.NewResourceManager()
-	errClosed   = errors.New("etcd monitor chan has been closed")
+	coolDownUnstable = mathx.NewUnstable(coolDownDeviation)
+	errClosed        = errors.New("etcd monitor chan has been closed")
 )
 
 // A Registry is a registry that manages the etcd client connections.
 type Registry struct {
 	clusters map[string]*cluster
-	lock     sync.Mutex
+	lock     sync.RWMutex
 }
 
 // GetRegistry returns a global Registry.
@@ -40,60 +43,148 @@ func GetRegistry() *Registry {
 
 // GetConn returns an etcd client connection associated with given endpoints.
 func (r *Registry) GetConn(endpoints []string) (EtcdClient, error) {
-	c, _ := r.getCluster(endpoints)
+	c, _ := r.getOrCreateCluster(endpoints)
 	return c.getClient()
 }
 
 // Monitor monitors the key on given etcd endpoints, notify with the given UpdateListener.
-func (r *Registry) Monitor(endpoints []string, key string, l UpdateListener) error {
+func (r *Registry) Monitor(endpoints []string, key string, exactMatch bool, l UpdateListener) error {
-	c, exists := r.getCluster(endpoints)
+	wkey := watchKey{
+		key:        key,
+		exactMatch: exactMatch,
+	}
+
+	c, exists := r.getOrCreateCluster(endpoints)
 	// if exists, the existing values should be updated to the listener.
 	if exists {
-		kvs := c.getCurrent(key)
+		c.lock.Lock()
-		for _, kv := range kvs {
+		watcher, ok := c.watchers[wkey]
-			l.OnAdd(kv)
+		if ok {
+			watcher.listeners = append(watcher.listeners, l)
+		}
+		c.lock.Unlock()
+
+		if ok {
+			kvs := c.getCurrent(wkey)
+			for _, kv := range kvs {
+				l.OnAdd(kv)
+			}
+
+			return nil
 		}
 	}
 
-	return c.monitor(key, l)
+	return c.monitor(wkey, l)
 }
 
-func (r *Registry) getCluster(endpoints []string) (c *cluster, exists bool) {
+func (r *Registry) Unmonitor(endpoints []string, key string, exactMatch bool, l UpdateListener) {
-	clusterKey := getClusterKey(endpoints)
+	c, exists := r.getCluster(endpoints)
-	r.lock.Lock()
-	defer r.lock.Unlock()
-	c, exists = r.clusters[clusterKey]
 	if !exists {
-		c = newCluster(endpoints)
+		return
-		r.clusters[clusterKey] = c
+	}
+
+	wkey := watchKey{
+		key:        key,
+		exactMatch: exactMatch,
+	}
+
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	watcher, ok := c.watchers[wkey]
+	if !ok {
+		return
+	}
+
+	for i, listener := range watcher.listeners {
+		if listener == l {
+			watcher.listeners = append(watcher.listeners[:i], watcher.listeners[i+1:]...)
+			break
+		}
+	}
+
+	if len(watcher.listeners) == 0 {
+		if watcher.cancel != nil {
+			watcher.cancel()
+		}
+		delete(c.watchers, wkey)
+	}
+}
+
+func (r *Registry) getCluster(endpoints []string) (*cluster, bool) {
+	clusterKey := getClusterKey(endpoints)
+
+	r.lock.RLock()
+	c, ok := r.clusters[clusterKey]
+	r.lock.RUnlock()
+
+	return c, ok
+}
+
+func (r *Registry) getOrCreateCluster(endpoints []string) (c *cluster, exists bool) {
+	c, exists = r.getCluster(endpoints)
+	if !exists {
+		clusterKey := getClusterKey(endpoints)
+
+		r.lock.Lock()
+		defer r.lock.Unlock()
+
+		// double-check locking
+		c, exists = r.clusters[clusterKey]
+		if !exists {
+			c = newCluster(endpoints)
+			r.clusters[clusterKey] = c
+		}
 	}
 
 	return
 }
 
-type cluster struct {
+type (
-	endpoints  []string
+	watchKey struct {
-	key        string
+		key        string
-	values     map[string]map[string]string
+		exactMatch bool
-	listeners  map[string][]UpdateListener
+	}
-	watchGroup *threading.RoutineGroup
+
-	done       chan lang.PlaceholderType
+	watchValue struct {
-	lock       sync.Mutex
+		listeners []UpdateListener
-}
+		values    map[string]string
+		cancel    context.CancelFunc
+	}
+
+	cluster struct {
+		endpoints  []string
+		key        string
+		watchers   map[watchKey]*watchValue
+		watchGroup *threading.RoutineGroup
+		done       chan lang.PlaceholderType
+		lock       sync.RWMutex
+	}
+)
 
 func newCluster(endpoints []string) *cluster {
 	return &cluster{
 		endpoints:  endpoints,
 		key:        getClusterKey(endpoints),
-		values:     make(map[string]map[string]string),
+		watchers:   make(map[watchKey]*watchValue),
-		listeners:  make(map[string][]UpdateListener),
 		watchGroup: threading.NewRoutineGroup(),
 		done:       make(chan lang.PlaceholderType),
 	}
 }
 
-func (c *cluster) context(cli EtcdClient) context.Context {
+func (c *cluster) addListener(key watchKey, l UpdateListener) {
-	return contextx.ValueOnlyFrom(cli.Ctx())
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	watcher, ok := c.watchers[key]
+	if ok {
+		watcher.listeners = append(watcher.listeners, l)
+		return
+	}
+
+	val := newWatchValue()
+	val.listeners = []UpdateListener{l}
+	c.watchers[key] = val
 }
 
 func (c *cluster) getClient() (EtcdClient, error) {
@@ -107,12 +198,17 @@ func (c *cluster) getClient() (EtcdClient, error) {
 	return val.(EtcdClient), nil
 }
 
-func (c *cluster) getCurrent(key string) []KV {
+func (c *cluster) getCurrent(key watchKey) []KV {
-	c.lock.Lock()
+	c.lock.RLock()
-	defer c.lock.Unlock()
+	defer c.lock.RUnlock()
+
+	watcher, ok := c.watchers[key]
+	if !ok {
+		return nil
+	}
 
 	var kvs []KV
-	for k, v := range c.values[key] {
+	for k, v := range watcher.values {
 		kvs = append(kvs, KV{
 			Key: k,
 			Val: v,
@@ -122,42 +218,23 @@ func (c *cluster) getCurrent(key string) []KV {
 	return kvs
 }
 
-func (c *cluster) handleChanges(key string, kvs []KV) {
+func (c *cluster) handleChanges(key watchKey, kvs []KV) {
-	var add []KV
-	var remove []KV
 	c.lock.Lock()
-	listeners := append([]UpdateListener(nil), c.listeners[key]...)
+	watcher, ok := c.watchers[key]
-	vals, ok := c.values[key]
 	if !ok {
-		add = kvs
+		c.lock.Unlock()
-		vals = make(map[string]string)
+		return
-		for _, kv := range kvs {
-			vals[kv.Key] = kv.Val
-		}
-		c.values[key] = vals
-	} else {
-		m := make(map[string]string)
-		for _, kv := range kvs {
-			m[kv.Key] = kv.Val
-		}
-		for k, v := range vals {
-			if val, ok := m[k]; !ok || v != val {
-				remove = append(remove, KV{
-					Key: k,
-					Val: v,
-				})
-			}
-		}
-		for k, v := range m {
-			if val, ok := vals[k]; !ok || v != val {
-				add = append(add, KV{
-					Key: k,
-					Val: v,
-				})
-			}
-		}
-		c.values[key] = m
 	}
+
+	listeners := append([]UpdateListener(nil), watcher.listeners...)
+	// watcher.values cannot be nil
+	vals := watcher.values
+	newVals := make(map[string]string, len(kvs)+len(vals))
+	for _, kv := range kvs {
+		newVals[kv.Key] = kv.Val
+	}
+	add, remove := calculateChanges(vals, newVals)
+	watcher.values = newVals
 	c.lock.Unlock()
 
 	for _, kv := range add {
@@ -172,20 +249,22 @@ func (c *cluster) handleChanges(key string, kvs []KV) {
 	}
 }
 
-func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
+func (c *cluster) handleWatchEvents(ctx context.Context, key watchKey, events []*clientv3.Event) {
-	c.lock.Lock()
+	c.lock.RLock()
-	listeners := append([]UpdateListener(nil), c.listeners[key]...)
+	watcher, ok := c.watchers[key]
-	c.lock.Unlock()
+	if !ok {
+		c.lock.RUnlock()
+		return
+	}
+
+	listeners := append([]UpdateListener(nil), watcher.listeners...)
+	c.lock.RUnlock()
 
 	for _, ev := range events {
 		switch ev.Type {
 		case clientv3.EventTypePut:
 			c.lock.Lock()
-			if vals, ok := c.values[key]; ok {
+			watcher.values[string(ev.Kv.Key)] = string(ev.Kv.Value)
-				vals[string(ev.Kv.Key)] = string(ev.Kv.Value)
-			} else {
-				c.values[key] = map[string]string{string(ev.Kv.Key): string(ev.Kv.Value)}
-			}
 			c.lock.Unlock()
 			for _, l := range listeners {
 				l.OnAdd(KV{
@@ -195,9 +274,7 @@ func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
 			}
 		case clientv3.EventTypeDelete:
 			c.lock.Lock()
-			if vals, ok := c.values[key]; ok {
+			delete(watcher.values, string(ev.Kv.Key))
-				delete(vals, string(ev.Kv.Key))
-			}
 			c.lock.Unlock()
 			for _, l := range listeners {
 				l.OnDelete(KV{
@@ -206,24 +283,29 @@ func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
 				})
 			}
 		default:
-			logx.Errorf("Unknown event type: %v", ev.Type)
+			logc.Errorf(ctx, "Unknown event type: %v", ev.Type)
 		}
 	}
 }
 
-func (c *cluster) load(cli EtcdClient, key string) int64 {
+func (c *cluster) load(cli EtcdClient, key watchKey) int64 {
 	var resp *clientv3.GetResponse
 	for {
 		var err error
-		ctx, cancel := context.WithTimeout(c.context(cli), RequestTimeout)
+		ctx, cancel := context.WithTimeout(cli.Ctx(), RequestTimeout)
-		resp, err = cli.Get(ctx, makeKeyPrefix(key), clientv3.WithPrefix())
+		if key.exactMatch {
+			resp, err = cli.Get(ctx, key.key)
+		} else {
+			resp, err = cli.Get(ctx, makeKeyPrefix(key.key), clientv3.WithPrefix())
+		}
+
 		cancel()
 		if err == nil {
 			break
 		}
 
-		logx.Errorf("%s, key is %s", err.Error(), key)
+		logc.Errorf(cli.Ctx(), "%s, key: %s, exactMatch: %t", err.Error(), key.key, key.exactMatch)
-		time.Sleep(coolDownInterval)
+		time.Sleep(coolDownUnstable.AroundDuration(coolDownInterval))
 	}
 
 	var kvs []KV
@@ -239,16 +321,13 @@ func (c *cluster) load(cli EtcdClient, key string) int64 {
 	return resp.Header.Revision
 }
 
-func (c *cluster) monitor(key string, l UpdateListener) error {
+func (c *cluster) monitor(key watchKey, l UpdateListener) error {
-	c.lock.Lock()
-	c.listeners[key] = append(c.listeners[key], l)
-	c.lock.Unlock()
-
 	cli, err := c.getClient()
 	if err != nil {
 		return err
 	}
 
+	c.addListener(key, l)
 	rev := c.load(cli, key)
 	c.watchGroup.Run(func() {
 		c.watch(cli, key, rev)
@@ -270,16 +349,22 @@ func (c *cluster) newClient() (EtcdClient, error) {
 
 func (c *cluster) reload(cli EtcdClient) {
 	c.lock.Lock()
+	// cancel the previous watches
 	close(c.done)
 	c.watchGroup.Wait()
+	var keys []watchKey
+	for wk, wval := range c.watchers {
+		keys = append(keys, wk)
+		if wval.cancel != nil {
+			wval.cancel()
+		}
+	}
+
 	c.done = make(chan lang.PlaceholderType)
 	c.watchGroup = threading.NewRoutineGroup()
-	var keys []string
-	for k := range c.listeners {
-		keys = append(keys, k)
-	}
 	c.lock.Unlock()
 
+	// start new watches
 	for _, key := range keys {
 		k := key
 		c.watchGroup.Run(func() {
@@ -289,7 +374,7 @@ func (c *cluster) reload(cli EtcdClient) {
 	}
 }
 
-func (c *cluster) watch(cli EtcdClient, key string, rev int64) {
+func (c *cluster) watch(cli EtcdClient, key watchKey, rev int64) {
 	for {
 		err := c.watchStream(cli, key, rev)
 		if err == nil {
@@ -297,24 +382,17 @@ func (c *cluster) watch(cli EtcdClient, key string, rev int64) {
 		}
 
 		if rev != 0 && errors.Is(err, rpctypes.ErrCompacted) {
-			logx.Errorf("etcd watch stream has been compacted, try to reload, rev %d", rev)
+			logc.Errorf(cli.Ctx(), "etcd watch stream has been compacted, try to reload, rev %d", rev)
 			rev = c.load(cli, key)
 		}
 
 		// log the error and retry
-		logx.Error(err)
+		logc.Error(cli.Ctx(), err)
 	}
 }
 
-func (c *cluster) watchStream(cli EtcdClient, key string, rev int64) error {
+func (c *cluster) watchStream(cli EtcdClient, key watchKey, rev int64) error {
-	var rch clientv3.WatchChan
+	ctx, rch := c.setupWatch(cli, key, rev)
-	if rev != 0 {
-		rch = cli.Watch(clientv3.WithRequireLeader(c.context(cli)), makeKeyPrefix(key),
-			clientv3.WithPrefix(), clientv3.WithRev(rev+1))
-	} else {
-		rch = cli.Watch(clientv3.WithRequireLeader(c.context(cli)), makeKeyPrefix(key),
-			clientv3.WithPrefix())
-	}
 
 	for {
 		select {
@@ -329,13 +407,47 @@ func (c *cluster) watchStream(cli EtcdClient, key string, rev int64) error {
 				return fmt.Errorf("etcd monitor chan error: %w", wresp.Err())
 			}
 
-			c.handleWatchEvents(key, wresp.Events)
+			c.handleWatchEvents(ctx, key, wresp.Events)
+		case <-ctx.Done():
+			return nil
 		case <-c.done:
 			return nil
 		}
 	}
 }
 
+func (c *cluster) setupWatch(cli EtcdClient, key watchKey, rev int64) (context.Context, clientv3.WatchChan) {
+	var (
+		rch  clientv3.WatchChan
+		ops  []clientv3.OpOption
+		wkey = key.key
+	)
+
+	if !key.exactMatch {
+		wkey = makeKeyPrefix(key.key)
+		ops = append(ops, clientv3.WithPrefix())
+	}
+	if rev != 0 {
+		ops = append(ops, clientv3.WithRev(rev+1))
+	}
+
+	ctx, cancel := context.WithCancel(cli.Ctx())
+	if watcher, ok := c.watchers[key]; ok {
+		watcher.cancel = cancel
+	} else {
+		val := newWatchValue()
+		val.cancel = cancel
+
+		c.lock.Lock()
+		c.watchers[key] = val
+		c.lock.Unlock()
+	}
+
+	rch = cli.Watch(clientv3.WithRequireLeader(ctx), wkey, ops...)
+
+	return ctx, rch
+}
+
 func (c *cluster) watchConnState(cli EtcdClient) {
 	watcher := newStateWatcher()
 	watcher.addListener(func() {
@@ -364,6 +476,28 @@ func DialClient(endpoints []string) (EtcdClient, error) {
 	return clientv3.New(cfg)
 }
 
+func calculateChanges(oldVals, newVals map[string]string) (add, remove []KV) {
+	for k, v := range newVals {
+		if val, ok := oldVals[k]; !ok || v != val {
+			add = append(add, KV{
+				Key: k,
+				Val: v,
+			})
+		}
+	}
+
+	for k, v := range oldVals {
+		if val, ok := newVals[k]; !ok || v != val {
+			remove = append(remove, KV{
+				Key: k,
+				Val: v,
+			})
+		}
+	}
+
+	return add, remove
+}
+
 func getClusterKey(endpoints []string) string {
 	sort.Strings(endpoints)
 	return strings.Join(endpoints, endpointsSeparator)
@@ -372,3 +506,10 @@ func getClusterKey(endpoints []string) string {
 func makeKeyPrefix(key string) string {
 	return fmt.Sprintf("%s%c", key, Delimiter)
 }
+
+// NewClient returns a watchValue that make sure values are not nil.
+func newWatchValue() *watchValue {
+	return &watchValue{
+		values: make(map[string]string),
+	}
+}

core/discov/internal/registry_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/zeromicro/go-zero/core/lang"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stringx"
+	"github.com/zeromicro/go-zero/core/threading"
 	"go.etcd.io/etcd/api/v3/etcdserverpb"
 	"go.etcd.io/etcd/api/v3/mvccpb"
 	clientv3 "go.etcd.io/etcd/client/v3"
@@ -38,9 +39,9 @@ func setMockClient(cli EtcdClient) func() {
 
 func TestGetCluster(t *testing.T) {
 	AddAccount([]string{"first"}, "foo", "bar")
-	c1, _ := GetRegistry().getCluster([]string{"first"})
+	c1, _ := GetRegistry().getOrCreateCluster([]string{"first"})
-	c2, _ := GetRegistry().getCluster([]string{"second"})
+	c2, _ := GetRegistry().getOrCreateCluster([]string{"second"})
-	c3, _ := GetRegistry().getCluster([]string{"first"})
+	c3, _ := GetRegistry().getOrCreateCluster([]string{"first"})
 	assert.Equal(t, c1, c3)
 	assert.NotEqual(t, c1, c2)
 }
@@ -50,6 +51,36 @@ func TestGetClusterKey(t *testing.T) {
 		getClusterKey([]string{"remotehost:5678", "localhost:1234"}))
 }
 
+func TestUnmonitor(t *testing.T) {
+	t.Run("no listener", func(t *testing.T) {
+		reg := &Registry{
+			clusters: map[string]*cluster{},
+		}
+		assert.NotPanics(t, func() {
+			reg.Unmonitor([]string{"any"}, "any", false, nil)
+		})
+	})
+
+	t.Run("no value", func(t *testing.T) {
+		reg := &Registry{
+			clusters: map[string]*cluster{
+				"any": {
+					watchers: map[watchKey]*watchValue{
+						{
+							key: "any",
+						}: {
+							values: map[string]string{},
+						},
+					},
+				},
+			},
+		}
+		assert.NotPanics(t, func() {
+			reg.Unmonitor([]string{"any"}, "another", false, nil)
+		})
+	})
+}
+
 func TestCluster_HandleChanges(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	l := NewMockUpdateListener(ctrl)
@@ -78,8 +109,14 @@ func TestCluster_HandleChanges(t *testing.T) {
 		Val: "4",
 	})
 	c := newCluster([]string{"any"})
-	c.listeners["any"] = []UpdateListener{l}
+	key := watchKey{
-	c.handleChanges("any", []KV{
+		key:        "any",
+		exactMatch: false,
+	}
+	c.watchers[key] = &watchValue{
+		listeners: []UpdateListener{l},
+	}
+	c.handleChanges(key, []KV{
 		{
 			Key: "first",
 			Val: "1",
@@ -92,8 +129,8 @@ func TestCluster_HandleChanges(t *testing.T) {
 	assert.EqualValues(t, map[string]string{
 		"first":  "1",
 		"second": "2",
-	}, c.values["any"])
+	}, c.watchers[key].values)
-	c.handleChanges("any", []KV{
+	c.handleChanges(key, []KV{
 		{
 			Key: "third",
 			Val: "3",
@@ -106,7 +143,7 @@ func TestCluster_HandleChanges(t *testing.T) {
 	assert.EqualValues(t, map[string]string{
 		"third":  "3",
 		"fourth": "4",
-	}, c.values["any"])
+	}, c.watchers[key].values)
 }
 
 func TestCluster_Load(t *testing.T) {
@@ -126,9 +163,11 @@ func TestCluster_Load(t *testing.T) {
 	}, nil)
 	cli.EXPECT().Ctx().Return(context.Background())
 	c := &cluster{
-		values: make(map[string]map[string]string),
+		watchers: make(map[watchKey]*watchValue),
 	}
-	c.load(cli, "any")
+	c.load(cli, watchKey{
+		key: "any",
+	})
 }
 
 func TestCluster_Watch(t *testing.T) {
@@ -160,11 +199,16 @@ func TestCluster_Watch(t *testing.T) {
 			var wg sync.WaitGroup
 			wg.Add(1)
 			c := &cluster{
-				listeners: make(map[string][]UpdateListener),
+				watchers: make(map[watchKey]*watchValue),
-				values:    make(map[string]map[string]string),
+			}
+			key := watchKey{
+				key: "any",
 			}
 			listener := NewMockUpdateListener(ctrl)
-			c.listeners["any"] = []UpdateListener{listener}
+			c.watchers[key] = &watchValue{
+				listeners: []UpdateListener{listener},
+				values:    make(map[string]string),
+			}
 			listener.EXPECT().OnAdd(gomock.Any()).Do(func(kv KV) {
 				assert.Equal(t, "hello", kv.Key)
 				assert.Equal(t, "world", kv.Val)
@@ -173,7 +217,7 @@ func TestCluster_Watch(t *testing.T) {
 			listener.EXPECT().OnDelete(gomock.Any()).Do(func(_ any) {
 				wg.Done()
 			}).MaxTimes(1)
-			go c.watch(cli, "any", 0)
+			go c.watch(cli, key, 0)
 			ch <- clientv3.WatchResponse{
 				Events: []*clientv3.Event{
 					{
@@ -211,17 +255,111 @@ func TestClusterWatch_RespFailures(t *testing.T) {
 			ch := make(chan clientv3.WatchResponse)
 			cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch).AnyTimes()
 			cli.EXPECT().Ctx().Return(context.Background()).AnyTimes()
-			c := new(cluster)
+			c := &cluster{
+				watchers: make(map[watchKey]*watchValue),
+			}
 			c.done = make(chan lang.PlaceholderType)
 			go func() {
 				ch <- resp
 				close(c.done)
 			}()
-			c.watch(cli, "any", 0)
+			key := watchKey{
+				key: "any",
+			}
+			c.watch(cli, key, 0)
 		})
 	}
 }
 
+func TestCluster_getCurrent(t *testing.T) {
+	t.Run("no value", func(t *testing.T) {
+		c := &cluster{
+			watchers: map[watchKey]*watchValue{
+				{
+					key: "any",
+				}: {
+					values: map[string]string{},
+				},
+			},
+		}
+		assert.Nil(t, c.getCurrent(watchKey{
+			key: "another",
+		}))
+	})
+}
+
+func TestCluster_handleWatchEvents(t *testing.T) {
+	t.Run("no value", func(t *testing.T) {
+		c := &cluster{
+			watchers: map[watchKey]*watchValue{
+				{
+					key: "any",
+				}: {
+					values: map[string]string{},
+				},
+			},
+		}
+		assert.NotPanics(t, func() {
+			c.handleWatchEvents(context.Background(), watchKey{
+				key: "another",
+			}, nil)
+		})
+	})
+}
+
+func TestCluster_addListener(t *testing.T) {
+	t.Run("has listener", func(t *testing.T) {
+		c := &cluster{
+			watchers: map[watchKey]*watchValue{
+				{
+					key: "any",
+				}: {
+					listeners: make([]UpdateListener, 0),
+				},
+			},
+		}
+		assert.NotPanics(t, func() {
+			c.addListener(watchKey{
+				key: "any",
+			}, nil)
+		})
+	})
+
+	t.Run("no listener", func(t *testing.T) {
+		c := &cluster{
+			watchers: map[watchKey]*watchValue{
+				{
+					key: "any",
+				}: {
+					listeners: make([]UpdateListener, 0),
+				},
+			},
+		}
+		assert.NotPanics(t, func() {
+			c.addListener(watchKey{
+				key: "another",
+			}, nil)
+		})
+	})
+}
+
+func TestCluster_reload(t *testing.T) {
+	c := &cluster{
+		watchers:   map[watchKey]*watchValue{},
+		watchGroup: threading.NewRoutineGroup(),
+		done:       make(chan lang.PlaceholderType),
+	}
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	cli := NewMockEtcdClient(ctrl)
+	restore := setMockClient(cli)
+	defer restore()
+	assert.NotPanics(t, func() {
+		c.reload(cli)
+	})
+}
+
 func TestClusterWatch_CloseChan(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
@@ -231,13 +369,17 @@ func TestClusterWatch_CloseChan(t *testing.T) {
 	ch := make(chan clientv3.WatchResponse)
 	cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch).AnyTimes()
 	cli.EXPECT().Ctx().Return(context.Background()).AnyTimes()
-	c := new(cluster)
+	c := &cluster{
+		watchers: make(map[watchKey]*watchValue),
+	}
 	c.done = make(chan lang.PlaceholderType)
 	go func() {
 		close(ch)
 		close(c.done)
 	}()
-	c.watch(cli, "any", 0)
+	c.watch(cli, watchKey{
+		key: "any",
+	}, 0)
 }
 
 func TestValueOnlyContext(t *testing.T) {
@@ -280,16 +422,59 @@ func TestRegistry_Monitor(t *testing.T) {
 	GetRegistry().lock.Lock()
 	GetRegistry().clusters = map[string]*cluster{
 		getClusterKey(endpoints): {
-			listeners: map[string][]UpdateListener{},
+			watchers: map[watchKey]*watchValue{
-			values: map[string]map[string]string{
+				watchKey{
-				"foo": {
+					key:        "foo",
-					"bar": "baz",
+					exactMatch: true,
+				}: {
+					values: map[string]string{
+						"bar": "baz",
+					},
 				},
 			},
 		},
 	}
 	GetRegistry().lock.Unlock()
-	assert.Error(t, GetRegistry().Monitor(endpoints, "foo", new(mockListener)))
+	assert.Error(t, GetRegistry().Monitor(endpoints, "foo", false, new(mockListener)))
+}
+
+func TestRegistry_Unmonitor(t *testing.T) {
+	svr, err := mockserver.StartMockServers(1)
+	assert.NoError(t, err)
+	svr.StartAt(0)
+
+	_, cancel := context.WithCancel(context.Background())
+	endpoints := []string{svr.Servers[0].Address}
+	GetRegistry().lock.Lock()
+	GetRegistry().clusters = map[string]*cluster{
+		getClusterKey(endpoints): {
+			watchers: map[watchKey]*watchValue{
+				watchKey{
+					key:        "foo",
+					exactMatch: true,
+				}: {
+					values: map[string]string{
+						"bar": "baz",
+					},
+					cancel: cancel,
+				},
+			},
+		},
+	}
+	GetRegistry().lock.Unlock()
+	l := new(mockListener)
+	assert.NoError(t, GetRegistry().Monitor(endpoints, "foo", true, l))
+	watchVals := GetRegistry().clusters[getClusterKey(endpoints)].watchers[watchKey{
+		key:        "foo",
+		exactMatch: true,
+	}]
+	assert.Equal(t, 1, len(watchVals.listeners))
+	GetRegistry().Unmonitor(endpoints, "foo", true, l)
+	watchVals = GetRegistry().clusters[getClusterKey(endpoints)].watchers[watchKey{
+		key:        "foo",
+		exactMatch: true,
+	}]
+	assert.Nil(t, watchVals)
 }
 
 type mockListener struct {

core/discov/internal/updatelistener.go

@@ -10,6 +10,7 @@ type (
 	}
 
 	// UpdateListener wraps the OnAdd and OnDelete methods.
+	// The implementation should be thread-safe and idempotent.
 	UpdateListener interface {
 		OnAdd(kv KV)
 		OnDelete(kv KV)

core/discov/publisher.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/zeromicro/go-zero/core/discov/internal"
 	"github.com/zeromicro/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/logc"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/proc"
 	"github.com/zeromicro/go-zero/core/syncx"
@@ -91,12 +92,12 @@ func (p *Publisher) doKeepAlive() error {
 		default:
 			cli, err := p.doRegister()
 			if err != nil {
-				logx.Errorf("etcd publisher doRegister: %s", err.Error())
+				logc.Errorf(cli.Ctx(), "etcd publisher doRegister: %s", err.Error())
 				break
 			}
 
 			if err := p.keepAliveAsync(cli); err != nil {
-				logx.Errorf("etcd publisher keepAliveAsync: %s", err.Error())
+				logc.Errorf(cli.Ctx(), "etcd publisher keepAliveAsync: %s", err.Error())
 				break
 			}
 
@@ -130,17 +131,17 @@ func (p *Publisher) keepAliveAsync(cli internal.EtcdClient) error {
 				if !ok {
 					p.revoke(cli)
 					if err := p.doKeepAlive(); err != nil {
-						logx.Errorf("etcd publisher KeepAlive: %s", err.Error())
+						logc.Errorf(cli.Ctx(), "etcd publisher KeepAlive: %s", err.Error())
 					}
 					return
 				}
 			case <-p.pauseChan:
-				logx.Infof("paused etcd renew, key: %s, value: %s", p.key, p.value)
+				logc.Infof(cli.Ctx(), "paused etcd renew, key: %s, value: %s", p.key, p.value)
 				p.revoke(cli)
 				select {
 				case <-p.resumeChan:
 					if err := p.doKeepAlive(); err != nil {
-						logx.Errorf("etcd publisher KeepAlive: %s", err.Error())
+						logc.Errorf(cli.Ctx(), "etcd publisher KeepAlive: %s", err.Error())
 					}
 					return
 				case <-p.quit.Done():
@@ -175,7 +176,7 @@ func (p *Publisher) register(client internal.EtcdClient) (clientv3.LeaseID, erro
 
 func (p *Publisher) revoke(cli internal.EtcdClient) {
 	if _, err := cli.Revoke(cli.Ctx(), p.lease); err != nil {
-		logx.Errorf("etcd publisher revoke: %s", err.Error())
+		logc.Errorf(cli.Ctx(), "etcd publisher revoke: %s", err.Error())
 	}
 }
 

core/discov/subscriber.go

@@ -15,9 +15,11 @@ type (
 
 	// A Subscriber is used to subscribe the given key on an etcd cluster.
 	Subscriber struct {
-		endpoints []string
+		endpoints  []string
-		exclusive bool
+		exclusive  bool
-		items     *container
+		key        string
+		exactMatch bool
+		items      *container
 	}
 )
 
@@ -28,13 +30,14 @@ type (
 func NewSubscriber(endpoints []string, key string, opts ...SubOption) (*Subscriber, error) {
 	sub := &Subscriber{
 		endpoints: endpoints,
+		key:       key,
 	}
 	for _, opt := range opts {
 		opt(sub)
 	}
 	sub.items = newContainer(sub.exclusive)
 
-	if err := internal.GetRegistry().Monitor(endpoints, key, sub.items); err != nil {
+	if err := internal.GetRegistry().Monitor(endpoints, key, sub.exactMatch, sub.items); err != nil {
 		return nil, err
 	}
 
@@ -46,6 +49,11 @@ func (s *Subscriber) AddListener(listener func()) {
 	s.items.addListener(listener)
 }
 
+// Close closes the subscriber.
+func (s *Subscriber) Close() {
+	internal.GetRegistry().Unmonitor(s.endpoints, s.key, s.exactMatch, s.items)
+}
+
 // Values returns all the subscription values.
 func (s *Subscriber) Values() []string {
 	return s.items.getValues()
@@ -59,6 +67,13 @@ func Exclusive() SubOption {
 	}
 }
 
+// WithExactMatch turn off querying using key prefixes.
+func WithExactMatch() SubOption {
+	return func(sub *Subscriber) {
+		sub.exactMatch = true
+	}
+}
+
 // WithSubEtcdAccount provides the etcd username/password.
 func WithSubEtcdAccount(user, pass string) SubOption {
 	return func(sub *Subscriber) {

core/discov/subscriber_test.go

@@ -225,3 +225,28 @@ func TestWithSubEtcdAccount(t *testing.T) {
 	assert.Equal(t, user, account.User)
 	assert.Equal(t, "bar", account.Pass)
 }
+
+func TestWithExactMatch(t *testing.T) {
+	sub := new(Subscriber)
+	WithExactMatch()(sub)
+	sub.items = newContainer(sub.exclusive)
+	var count int32
+	sub.AddListener(func() {
+		atomic.AddInt32(&count, 1)
+	})
+	sub.items.notifyChange()
+	assert.Empty(t, sub.Values())
+	assert.Equal(t, int32(1), atomic.LoadInt32(&count))
+}
+
+func TestSubscriberClose(t *testing.T) {
+	l := newContainer(false)
+	sub := &Subscriber{
+		endpoints: []string{"localhost:12379"},
+		key:       "foo",
+		items:     l,
+	}
+	assert.NotPanics(t, func() {
+		sub.Close()
+	})
+}

core/errorx/batcherror.go

@@ -1,21 +1,17 @@
 package errorx
 
 import (
-	"bytes"
+	"errors"
 	"sync"
 )
 
-type (
+// BatchError is an error that can hold multiple errors.
-	// A BatchError is an error that can hold multiple errors.
+type BatchError struct {
-	BatchError struct {
+	errs []error
-		errs errorArray
+	lock sync.RWMutex
-		lock sync.Mutex
+}
-	}
 
-	errorArray []error
+// Add adds one or more non-nil errors to the BatchError instance.
-)
-
-// Add adds errs to be, nil errors are ignored.
 func (be *BatchError) Add(errs ...error) {
 	be.lock.Lock()
 	defer be.lock.Unlock()
@@ -27,39 +23,20 @@ func (be *BatchError) Add(errs ...error) {
 	}
 }
 
-// Err returns an error that represents all errors.
+// Err returns an error that represents all accumulated errors.
+// It returns nil if there are no errors.
 func (be *BatchError) Err() error {
-	be.lock.Lock()
+	be.lock.RLock()
-	defer be.lock.Unlock()
+	defer be.lock.RUnlock()
 
-	switch len(be.errs) {
+	// If there are no non-nil errors, errors.Join(...) returns nil.
-	case 0:
+	return errors.Join(be.errs...)
-		return nil
-	case 1:
-		return be.errs[0]
-	default:
-		return be.errs
-	}
 }
 
-// NotNil checks if any error inside.
+// NotNil checks if there is at least one error inside the BatchError.
 func (be *BatchError) NotNil() bool {
-	be.lock.Lock()
+	be.lock.RLock()
-	defer be.lock.Unlock()
+	defer be.lock.RUnlock()
 
 	return len(be.errs) > 0
 }
-
-// Error returns a string that represents inside errors.
-func (ea errorArray) Error() string {
-	var buf bytes.Buffer
-
-	for i := range ea {
-		if i > 0 {
-			buf.WriteByte('\n')
-		}
-		buf.WriteString(ea[i].Error())
-	}
-
-	return buf.String()
-}

core/errorx/batcherror_test.go

@@ -66,3 +66,82 @@ func TestBatchErrorConcurrentAdd(t *testing.T) {
 	assert.Equal(t, count, len(batch.errs))
 	assert.True(t, batch.NotNil())
 }
+
+func TestBatchError_Unwrap(t *testing.T) {
+	t.Run("nil", func(t *testing.T) {
+		var be BatchError
+		assert.Nil(t, be.Err())
+		assert.True(t, errors.Is(be.Err(), nil))
+	})
+
+	t.Run("one error", func(t *testing.T) {
+		var errFoo = errors.New("foo")
+		var errBar = errors.New("bar")
+		var be BatchError
+		be.Add(errFoo)
+		assert.True(t, errors.Is(be.Err(), errFoo))
+		assert.False(t, errors.Is(be.Err(), errBar))
+	})
+
+	t.Run("two errors", func(t *testing.T) {
+		var errFoo = errors.New("foo")
+		var errBar = errors.New("bar")
+		var errBaz = errors.New("baz")
+		var be BatchError
+		be.Add(errFoo)
+		be.Add(errBar)
+		assert.True(t, errors.Is(be.Err(), errFoo))
+		assert.True(t, errors.Is(be.Err(), errBar))
+		assert.False(t, errors.Is(be.Err(), errBaz))
+	})
+}
+
+func TestBatchError_Add(t *testing.T) {
+	var be BatchError
+
+	// Test adding nil errors
+	be.Add(nil, nil)
+	assert.False(t, be.NotNil(), "Expected BatchError to be empty after adding nil errors")
+
+	// Test adding non-nil errors
+	err1 := errors.New("error 1")
+	err2 := errors.New("error 2")
+	be.Add(err1, err2)
+	assert.True(t, be.NotNil(), "Expected BatchError to be non-empty after adding errors")
+
+	// Test adding a mix of nil and non-nil errors
+	err3 := errors.New("error 3")
+	be.Add(nil, err3, nil)
+	assert.True(t, be.NotNil(), "Expected BatchError to be non-empty after adding a mix of nil and non-nil errors")
+}
+
+func TestBatchError_Err(t *testing.T) {
+	var be BatchError
+
+	// Test Err() on empty BatchError
+	assert.Nil(t, be.Err(), "Expected nil error for empty BatchError")
+
+	// Test Err() with multiple errors
+	err1 := errors.New("error 1")
+	err2 := errors.New("error 2")
+	be.Add(err1, err2)
+
+	combinedErr := be.Err()
+	assert.NotNil(t, combinedErr, "Expected nil error for BatchError with multiple errors")
+
+	// Check if the combined error contains both error messages
+	errString := combinedErr.Error()
+	assert.Truef(t, errors.Is(combinedErr, err1), "Combined error doesn't contain first error: %s", errString)
+	assert.Truef(t, errors.Is(combinedErr, err2), "Combined error doesn't contain second error: %s", errString)
+}
+
+func TestBatchError_NotNil(t *testing.T) {
+	var be BatchError
+
+	// Test NotNil() on empty BatchError
+	assert.Nil(t, be.Err(), "Expected nil error for empty BatchError")
+
+	// Test NotNil() after adding an error
+	be.Add(errors.New("test error"))
+	assert.NotNil(t, be.Err(), "Expected non-nil error after adding an error")
+}

core/errorx/check.go

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

core/errorx/check_test.go

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

core/filex/file.go

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

core/filex/file_test.go

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

core/fs/files.go

@@ -1,4 +1,4 @@
-//go:build linux || darwin
+//go:build linux || darwin || freebsd
 
 package fs
 

core/fx/parallel.go

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

core/fx/parallel_test.go

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

core/fx/stream.go

@@ -84,10 +84,10 @@ func Range(source <-chan any) Stream {
 	}
 }
 
-// AllMach returns whether all elements of this stream match the provided predicate.
+// AllMatch returns whether all elements of this stream match the provided predicate.
 // May not evaluate the predicate on all elements if not necessary for determining the result.
 // If the stream is empty then true is returned and the predicate is not evaluated.
-func (s Stream) AllMach(predicate func(item any) bool) bool {
+func (s Stream) AllMatch(predicate func(item any) bool) bool {
 	for item := range s.source {
 		if !predicate(item) {
 			// make sure the former goroutine not block, and current func returns fast.
@@ -99,10 +99,10 @@ func (s Stream) AllMach(predicate func(item any) bool) bool {
 	return true
 }
 
-// AnyMach returns whether any elements of this stream match the provided predicate.
+// AnyMatch returns whether any elements of this stream match the provided predicate.
 // May not evaluate the predicate on all elements if not necessary for determining the result.
 // If the stream is empty then false is returned and the predicate is not evaluated.
-func (s Stream) AnyMach(predicate func(item any) bool) bool {
+func (s Stream) AnyMatch(predicate func(item any) bool) bool {
 	for item := range s.source {
 		if predicate(item) {
 			// make sure the former goroutine not block, and current func returns fast.

core/fx/stream_test.go

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

core/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
+	//go:embed periodscript.lua
-	periodScript = redis.NewScript(`local limit = tonumber(ARGV[1])
+	periodLuaScript string
-local window = tonumber(ARGV[2])
+	periodScript    = redis.NewScript(periodLuaScript)
-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`)
 )
 
 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
+var (
-// KEYS[1] as tokens_key
+	//go:embed tokenscript.lua
-// KEYS[2] as timestamp_key
+	tokenLuaScript string
-var script = redis.NewScript(`local rate = tonumber(ARGV[1])
+	tokenScript    = redis.NewScript(tokenLuaScript)
-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`)
 
 // 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,
@@ -128,7 +104,7 @@ func (lim *TokenLimiter) reserveN(ctx context.Context, now time.Time, n int) boo
 	if errors.Is(err, redis.Nil) {
 		return false
 	}
-	if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
+	if errorx.In(err, context.DeadlineExceeded, context.Canceled) {
 		logx.Errorf("fail to use rate limiter: %s", err)
 		return false
 	}

core/limit/tokenscript.lua

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

core/load/adaptiveshedder.go

@@ -76,8 +76,8 @@ type (
 		avgFlyingLock   syncx.SpinLock
 		overloadTime    *syncx.AtomicDuration
 		droppedRecently *syncx.AtomicBool
-		passCounter     *collection.RollingWindow
+		passCounter     *collection.RollingWindow[int64, *collection.Bucket[int64]]
-		rtCounter       *collection.RollingWindow
+		rtCounter       *collection.RollingWindow[int64, *collection.Bucket[int64]]
 	}
 )
 
@@ -107,15 +107,16 @@ func NewAdaptiveShedder(opts ...ShedderOption) Shedder {
 		opt(&options)
 	}
 	bucketDuration := options.window / time.Duration(options.buckets)
+	newBucket := func() *collection.Bucket[int64] {
+		return new(collection.Bucket[int64])
+	}
 	return &adaptiveShedder{
 		cpuThreshold:    options.cpuThreshold,
 		windowScale:     float64(time.Second) / float64(bucketDuration) / millisecondsPerSecond,
 		overloadTime:    syncx.NewAtomicDuration(),
 		droppedRecently: syncx.NewAtomicBool(),
-		passCounter: collection.NewRollingWindow(options.buckets, bucketDuration,
+		passCounter:     collection.NewRollingWindow[int64, *collection.Bucket[int64]](newBucket, options.buckets, bucketDuration, collection.IgnoreCurrentBucket[int64, *collection.Bucket[int64]]()),
-			collection.IgnoreCurrentBucket()),
+		rtCounter:       collection.NewRollingWindow[int64, *collection.Bucket[int64]](newBucket, options.buckets, bucketDuration, collection.IgnoreCurrentBucket[int64, *collection.Bucket[int64]]()),
-		rtCounter: collection.NewRollingWindow(options.buckets, bucketDuration,
-			collection.IgnoreCurrentBucket()),
 	}
 }
 
@@ -167,15 +168,15 @@ func (as *adaptiveShedder) maxFlight() float64 {
 }
 
 func (as *adaptiveShedder) maxPass() int64 {
-	var result float64 = 1
+	var result int64 = 1
 
-	as.passCounter.Reduce(func(b *collection.Bucket) {
+	as.passCounter.Reduce(func(b *collection.Bucket[int64]) {
 		if b.Sum > result {
 			result = b.Sum
 		}
 	})
 
-	return int64(result)
+	return result
 }
 
 func (as *adaptiveShedder) minRt() float64 {
@@ -183,12 +184,12 @@ func (as *adaptiveShedder) minRt() float64 {
 	// its a reasonable large value to avoid dropping requests.
 	result := defaultMinRt
 
-	as.rtCounter.Reduce(func(b *collection.Bucket) {
+	as.rtCounter.Reduce(func(b *collection.Bucket[int64]) {
 		if b.Count <= 0 {
 			return
 		}
 
-		avg := math.Round(b.Sum / float64(b.Count))
+		avg := math.Round(float64(b.Sum) / float64(b.Count))
 		if avg < result {
 			result = avg
 		}
@@ -283,6 +284,6 @@ func (p *promise) Fail() {
 func (p *promise) Pass() {
 	rt := float64(timex.Since(p.start)) / float64(time.Millisecond)
 	p.shedder.addFlying(-1)
-	p.shedder.rtCounter.Add(math.Ceil(rt))
+	p.shedder.rtCounter.Add(int64(math.Ceil(rt)))
 	p.shedder.passCounter.Add(1)
 }

core/load/adaptiveshedder_test.go

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

@@ -37,6 +37,13 @@ func Debugf(ctx context.Context, format string, v ...interface{}) {
 	getLogger(ctx).Debugf(format, v...)
 }
 
+// Debugfn writes fn result into access log.
+// This is useful when the function is expensive to compute,
+// and we want to log it only when necessary.
+func Debugfn(ctx context.Context, fn func() any) {
+	getLogger(ctx).Debugfn(fn)
+}
+
 // Debugv writes v into access log with json content.
 func Debugv(ctx context.Context, v interface{}) {
 	getLogger(ctx).Debugv(v)
@@ -57,6 +64,13 @@ func Errorf(ctx context.Context, format string, v ...any) {
 	getLogger(ctx).Errorf(fmt.Errorf(format, v...).Error())
 }
 
+// Errorfn writes fn result into error log.
+// This is useful when the function is expensive to compute,
+// and we want to log it only when necessary.
+func Errorfn(ctx context.Context, fn func() any) {
+	getLogger(ctx).Errorfn(fn)
+}
+
 // Errorv writes v into error log with json content.
 // No call stack attached, because not elegant to pack the messages.
 func Errorv(ctx context.Context, v any) {
@@ -83,6 +97,13 @@ func Infof(ctx context.Context, format string, v ...any) {
 	getLogger(ctx).Infof(format, v...)
 }
 
+// Infofn writes fn result into access log.
+// This is useful when the function is expensive to compute,
+// and we want to log it only when necessary.
+func Infofn(ctx context.Context, fn func() any) {
+	getLogger(ctx).Infofn(fn)
+}
+
 // Infov writes v into access log with json content.
 func Infov(ctx context.Context, v any) {
 	getLogger(ctx).Infov(v)
@@ -127,6 +148,13 @@ func Slowf(ctx context.Context, format string, v ...any) {
 	getLogger(ctx).Slowf(format, v...)
 }
 
+// Slowfn writes fn result into slow log.
+// This is useful when the function is expensive to compute,
+// and we want to log it only when necessary.
+func Slowfn(ctx context.Context, fn func() any) {
+	getLogger(ctx).Slowfn(fn)
+}
+
 // Slowv writes v into slow log with json content.
 func Slowv(ctx context.Context, v any) {
 	getLogger(ctx).Slowv(v)

core/logc/logs_test.go

@@ -49,6 +49,15 @@ func TestErrorf(t *testing.T) {
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
+func TestErrorfn(t *testing.T) {
+	buf := logtest.NewCollector(t)
+	file, line := getFileLine()
+	Errorfn(context.Background(), func() any {
+		return fmt.Sprintf("foo %s", "bar")
+	})
+	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
+}
+
 func TestErrorv(t *testing.T) {
 	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
@@ -77,6 +86,15 @@ func TestInfof(t *testing.T) {
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
+func TestInfofn(t *testing.T) {
+	buf := logtest.NewCollector(t)
+	file, line := getFileLine()
+	Infofn(context.Background(), func() any {
+		return fmt.Sprintf("foo %s", "bar")
+	})
+	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
+}
+
 func TestInfov(t *testing.T) {
 	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
@@ -105,6 +123,15 @@ func TestDebugf(t *testing.T) {
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
+func TestDebugfn(t *testing.T) {
+	buf := logtest.NewCollector(t)
+	file, line := getFileLine()
+	Debugfn(context.Background(), func() any {
+		return fmt.Sprintf("foo %s", "bar")
+	})
+	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
+}
+
 func TestDebugv(t *testing.T) {
 	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
@@ -148,6 +175,15 @@ func TestSlowf(t *testing.T) {
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)), buf.String())
 }
 
+func TestSlowfn(t *testing.T) {
+	buf := logtest.NewCollector(t)
+	file, line := getFileLine()
+	Slowfn(context.Background(), func() any {
+		return fmt.Sprintf("foo %s", "bar")
+	})
+	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)), buf.String())
+}
+
 func TestSlowv(t *testing.T) {
 	buf := logtest.NewCollector(t)
 	file, line := getFileLine()

core/logx/config.go

@@ -42,4 +42,6 @@ type LogConf struct {
 	// daily: daily rotation.
 	// size: size limited rotation.
 	Rotation string `json:",default=daily,options=[daily,size]"`
+	// FileTimeFormat represents the time format for file name, default is `2006-01-02T15:04:05.000Z07:00`.
+	FileTimeFormat string `json:",optional"`
 }

core/logx/logger.go

@@ -7,18 +7,22 @@ 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.
+	// Debugfn logs a message at debug level.
+	Debugfn(func() any)
+	// 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)
 	// Errorf logs a message at error level.
 	Errorf(string, ...any)
+	// Errorfn logs a message at error level.
+	Errorfn(func() any)
 	// Errorv logs a message at error level.
 	Errorv(any)
 	// Errorw logs a message at error level.
@@ -27,6 +31,8 @@ type Logger interface {
 	Info(...any)
 	// Infof logs a message at info level.
 	Infof(string, ...any)
+	// Infofn logs a message at info level.
+	Infofn(func() any)
 	// Infov logs a message at info level.
 	Infov(any)
 	// Infow logs a message at info level.
@@ -35,6 +41,8 @@ type Logger interface {
 	Slow(...any)
 	// Slowf logs a message at slow level.
 	Slowf(string, ...any)
+	// Slowfn logs a message at slow level.
+	Slowfn(func() any)
 	// Slowv logs a message at slow level.
 	Slowv(any)
 	// Sloww logs a message at slow level.

core/logx/logs.go

@@ -6,6 +6,7 @@ import (
 	"log"
 	"os"
 	"path"
+	"reflect"
 	"runtime/debug"
 	"sync"
 	"sync/atomic"
@@ -51,6 +52,26 @@ type (
 	}
 )
 
+// AddWriter adds a new writer.
+// If there is already a writer, the new writer will be added to the writer chain.
+// For example, to write logs to both file and console, if there is already a file writer,
+// ```go
+// logx.AddWriter(logx.NewWriter(os.Stdout))
+// ```
+func AddWriter(w Writer) {
+	ow := Reset()
+	if ow == nil {
+		SetWriter(w)
+	} else {
+		// no need to check if the existing writer is a comboWriter,
+		// because it is not common to add more than one writer.
+		// even more than one writer, the behavior is the same.
+		SetWriter(comboWriter{
+			writers: []Writer{ow, w},
+		})
+	}
+}
+
 // Alert alerts v in alert level, and the message is written to error log.
 func Alert(v string) {
 	getWriter().Alert(v)
@@ -79,6 +100,14 @@ func Debugf(format string, v ...any) {
 	}
 }
 
+// Debugfn writes function result into access log if debug level enabled.
+// This is useful when the function is expensive to call and debug level disabled.
+func Debugfn(fn func() any) {
+	if shallLog(DebugLevel) {
+		writeDebug(fn())
+	}
+}
+
 // Debugv writes v into access log with json content.
 func Debugv(v any) {
 	if shallLog(DebugLevel) {
@@ -118,6 +147,13 @@ func Errorf(format string, v ...any) {
 	}
 }
 
+// Errorfn writes function result into error log.
+func Errorfn(fn func() any) {
+	if shallLog(ErrorLevel) {
+		writeError(fn())
+	}
+}
+
 // ErrorStack writes v along with call stack into error log.
 func ErrorStack(v ...any) {
 	if shallLog(ErrorLevel) {
@@ -153,11 +189,11 @@ func Errorw(msg string, fields ...LogField) {
 func Field(key string, value any) LogField {
 	switch val := value.(type) {
 	case error:
-		return LogField{Key: key, Value: val.Error()}
+		return LogField{Key: key, Value: encodeError(val)}
 	case []error:
 		var errs []string
 		for _, err := range val {
-			errs = append(errs, err.Error())
+			errs = append(errs, encodeError(err))
 		}
 		return LogField{Key: key, Value: errs}
 	case time.Duration:
@@ -175,11 +211,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:
@@ -201,6 +237,14 @@ func Infof(format string, v ...any) {
 	}
 }
 
+// Infofn writes function result into access log.
+// This is useful when the function is expensive to call and info level disabled.
+func Infofn(fn func() any) {
+	if shallLog(InfoLevel) {
+		writeInfo(fn())
+	}
+}
+
 // Infov writes v into access log with json content.
 func Infov(v any) {
 	if shallLog(InfoLevel) {
@@ -273,6 +317,10 @@ func SetUp(c LogConf) (err error) {
 			timeFormat = c.TimeFormat
 		}
 
+		if len(c.FileTimeFormat) > 0 {
+			fileTimeFormat = c.FileTimeFormat
+		}
+
 		atomic.StoreUint32(&maxContentLength, c.MaxContentLength)
 
 		switch c.Encoding {
@@ -323,6 +371,14 @@ func Slowf(format string, v ...any) {
 	}
 }
 
+// Slowfn writes function result into slow log.
+// This is useful when the function is expensive to call and slow level disabled.
+func Slowfn(fn func() any) {
+	if shallLog(ErrorLevel) {
+		writeSlow(fn())
+	}
+}
+
 // Slowv writes v into slow log with json content.
 func Slowv(v any) {
 	if shallLog(ErrorLevel) {
@@ -414,6 +470,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 {
@@ -478,7 +560,7 @@ func shallLogStat() bool {
 // 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.
 func writeDebug(val any, fields ...LogField) {
-	getWriter().Debug(val, addCaller(fields...)...)
+	getWriter().Debug(val, mergeGlobalFields(addCaller(fields...))...)
 }
 
 // writeError writes v into the error log.
@@ -486,7 +568,7 @@ func writeDebug(val any, fields ...LogField) {
 // 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.
 func writeError(val any, fields ...LogField) {
-	getWriter().Error(val, addCaller(fields...)...)
+	getWriter().Error(val, mergeGlobalFields(addCaller(fields...))...)
 }
 
 // writeInfo writes v into info log.
@@ -494,7 +576,7 @@ func writeError(val any, fields ...LogField) {
 // 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.
 func writeInfo(val any, fields ...LogField) {
-	getWriter().Info(val, addCaller(fields...)...)
+	getWriter().Info(val, mergeGlobalFields(addCaller(fields...))...)
 }
 
 // writeSevere writes v into severe log.
@@ -510,7 +592,7 @@ func writeSevere(msg string) {
 // 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.
 func writeSlow(val any, fields ...LogField) {
-	getWriter().Slow(val, addCaller(fields...)...)
+	getWriter().Slow(val, mergeGlobalFields(addCaller(fields...))...)
 }
 
 // writeStack writes v into stack log.
@@ -526,5 +608,5 @@ func writeStack(msg string) {
 // 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.
 func writeStat(msg string) {
-	getWriter().Stat(msg, addCaller()...)
+	getWriter().Stat(msg, mergeGlobalFields(addCaller())...)
 }

core/logx/logs_test.go

@@ -248,6 +248,32 @@ func TestStructedLogDebugf(t *testing.T) {
 	})
 }
 
+func TestStructedLogDebugfn(t *testing.T) {
+	t.Run("debugfn with output", func(t *testing.T) {
+		w := new(mockWriter)
+		old := writer.Swap(w)
+		defer writer.Store(old)
+
+		doTestStructedLog(t, levelDebug, w, func(v ...any) {
+			Debugfn(func() any {
+				return fmt.Sprint(v...)
+			})
+		})
+	})
+
+	t.Run("debugfn without output", func(t *testing.T) {
+		w := new(mockWriter)
+		old := writer.Swap(w)
+		defer writer.Store(old)
+
+		doTestStructedLogEmpty(t, w, InfoLevel, func(v ...any) {
+			Debugfn(func() any {
+				return fmt.Sprint(v...)
+			})
+		})
+	})
+}
+
 func TestStructedLogDebugv(t *testing.T) {
 	w := new(mockWriter)
 	old := writer.Swap(w)
@@ -288,6 +314,32 @@ func TestStructedLogErrorf(t *testing.T) {
 	})
 }
 
+func TestStructedLogErrorfn(t *testing.T) {
+	t.Run("errorfn with output", func(t *testing.T) {
+		w := new(mockWriter)
+		old := writer.Swap(w)
+		defer writer.Store(old)
+
+		doTestStructedLog(t, levelError, w, func(v ...any) {
+			Errorfn(func() any {
+				return fmt.Sprint(v...)
+			})
+		})
+	})
+
+	t.Run("errorfn without output", func(t *testing.T) {
+		w := new(mockWriter)
+		old := writer.Swap(w)
+		defer writer.Store(old)
+
+		doTestStructedLogEmpty(t, w, SevereLevel, func(v ...any) {
+			Errorfn(func() any {
+				return fmt.Sprint(v...)
+			})
+		})
+	})
+}
+
 func TestStructedLogErrorv(t *testing.T) {
 	w := new(mockWriter)
 	old := writer.Swap(w)
@@ -328,6 +380,32 @@ func TestStructedLogInfof(t *testing.T) {
 	})
 }
 
+func TestStructedInfofn(t *testing.T) {
+	t.Run("infofn with output", func(t *testing.T) {
+		w := new(mockWriter)
+		old := writer.Swap(w)
+		defer writer.Store(old)
+
+		doTestStructedLog(t, levelInfo, w, func(v ...any) {
+			Infofn(func() any {
+				return fmt.Sprint(v...)
+			})
+		})
+	})
+
+	t.Run("infofn without output", func(t *testing.T) {
+		w := new(mockWriter)
+		old := writer.Swap(w)
+		defer writer.Store(old)
+
+		doTestStructedLogEmpty(t, w, ErrorLevel, func(v ...any) {
+			Infofn(func() any {
+				return fmt.Sprint(v...)
+			})
+		})
+	})
+}
+
 func TestStructedLogInfov(t *testing.T) {
 	w := new(mockWriter)
 	old := writer.Swap(w)
@@ -348,6 +426,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)
@@ -430,6 +529,17 @@ func TestStructedLogInfoConsoleText(t *testing.T) {
 	})
 }
 
+func TestInfofnWithErrorLevel(t *testing.T) {
+	called := false
+	SetLevel(ErrorLevel)
+	defer SetLevel(DebugLevel)
+	Infofn(func() any {
+		called = true
+		return "info log"
+	})
+	assert.False(t, called)
+}
+
 func TestStructedLogSlow(t *testing.T) {
 	w := new(mockWriter)
 	old := writer.Swap(w)
@@ -450,6 +560,32 @@ func TestStructedLogSlowf(t *testing.T) {
 	})
 }
 
+func TestStructedLogSlowfn(t *testing.T) {
+	t.Run("slowfn with output", func(t *testing.T) {
+		w := new(mockWriter)
+		old := writer.Swap(w)
+		defer writer.Store(old)
+
+		doTestStructedLog(t, levelSlow, w, func(v ...any) {
+			Slowfn(func() any {
+				return fmt.Sprint(v...)
+			})
+		})
+	})
+
+	t.Run("slowfn without output", func(t *testing.T) {
+		w := new(mockWriter)
+		old := writer.Swap(w)
+		defer writer.Store(old)
+
+		doTestStructedLogEmpty(t, w, SevereLevel, func(v ...any) {
+			Slowfn(func() any {
+				return fmt.Sprint(v...)
+			})
+		})
+	})
+}
+
 func TestStructedLogSlowv(t *testing.T) {
 	w := new(mockWriter)
 	old := writer.Swap(w)
@@ -658,6 +794,10 @@ func TestSetup(t *testing.T) {
 
 func TestDisable(t *testing.T) {
 	Disable()
+	defer func() {
+		SetLevel(InfoLevel)
+		atomic.StoreUint32(&encoding, jsonEncodingType)
+	}()
 
 	var opt logOptions
 	WithKeepDays(1)(&opt)
@@ -680,6 +820,17 @@ func TestDisableStat(t *testing.T) {
 	assert.Equal(t, 0, w.builder.Len())
 }
 
+func TestAddWriter(t *testing.T) {
+	const message = "hello there"
+	w := new(mockWriter)
+	AddWriter(w)
+	w1 := new(mockWriter)
+	AddWriter(w1)
+	Error(message)
+	assert.Contains(t, w.String(), message)
+	assert.Contains(t, w1.String(), message)
+}
+
 func TestSetWriter(t *testing.T) {
 	atomic.StoreUint32(&logLevel, 0)
 	Reset()
@@ -811,15 +962,26 @@ func doTestStructedLogConsole(t *testing.T, w *mockWriter, write func(...any)) {
 	assert.True(t, strings.Contains(w.String(), message))
 }
 
+func doTestStructedLogEmpty(t *testing.T, w *mockWriter, level uint32, write func(...any)) {
+	olevel := atomic.LoadUint32(&logLevel)
+	SetLevel(level)
+	defer SetLevel(olevel)
+
+	const message = "hello there"
+	write(message)
+	assert.Empty(t, w.String())
+}
+
 func testSetLevelTwiceWithMode(t *testing.T, mode string, w *mockWriter) {
 	writer.Store(nil)
 	SetUp(LogConf{
-		Mode:       mode,
+		Mode:           mode,
-		Level:      "debug",
+		Level:          "debug",
-		Path:       "/dev/null",
+		Path:           "/dev/null",
-		Encoding:   plainEncoding,
+		Encoding:       plainEncoding,
-		Stat:       false,
+		Stat:           false,
-		TimeFormat: time.RFC3339,
+		TimeFormat:     time.RFC3339,
+		FileTimeFormat: time.DateTime,
 	})
 	SetUp(LogConf{
 		Mode:  mode,
@@ -859,3 +1021,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

@@ -52,6 +52,12 @@ func (l *richLogger) Debugf(format string, v ...any) {
 	}
 }
 
+func (l *richLogger) Debugfn(fn func() any) {
+	if shallLog(DebugLevel) {
+		l.debug(fn())
+	}
+}
+
 func (l *richLogger) Debugv(v any) {
 	if shallLog(DebugLevel) {
 		l.debug(v)
@@ -76,6 +82,12 @@ func (l *richLogger) Errorf(format string, v ...any) {
 	}
 }
 
+func (l *richLogger) Errorfn(fn func() any) {
+	if shallLog(ErrorLevel) {
+		l.err(fn())
+	}
+}
+
 func (l *richLogger) Errorv(v any) {
 	if shallLog(ErrorLevel) {
 		l.err(v)
@@ -100,6 +112,12 @@ func (l *richLogger) Infof(format string, v ...any) {
 	}
 }
 
+func (l *richLogger) Infofn(fn func() any) {
+	if shallLog(InfoLevel) {
+		l.info(fn())
+	}
+}
+
 func (l *richLogger) Infov(v any) {
 	if shallLog(InfoLevel) {
 		l.info(v)
@@ -124,6 +142,12 @@ func (l *richLogger) Slowf(format string, v ...any) {
 	}
 }
 
+func (l *richLogger) Slowfn(fn func() any) {
+	if shallLog(ErrorLevel) {
+		l.slow(fn())
+	}
+}
+
 func (l *richLogger) Slowv(v any) {
 	if shallLog(ErrorLevel) {
 		l.slow(v)
@@ -141,28 +165,50 @@ func (l *richLogger) WithCallerSkip(skip int) Logger {
 		return l
 	}
 
-	l.callerSkip = skip
+	return &richLogger{
-	return l
+		ctx:        l.ctx,
+		callerSkip: skip,
+		fields:     l.fields,
+	}
 }
 
 func (l *richLogger) WithContext(ctx context.Context) Logger {
-	l.ctx = ctx
+	return &richLogger{
-	return l
+		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)))
+	fields := append(l.fields, Field(durationKey, timex.ReprOfDuration(duration)))
-	return l
+
+	return &richLogger{
+		ctx:        l.ctx,
+		callerSkip: l.callerSkip,
+		fields:     fields,
+	}
 }
 
 func (l *richLogger) WithFields(fields ...LogField) Logger {
-	l.fields = append(l.fields, fields...)
+	if len(fields) == 0 {
-	return l
+		return l
+	}
+
+	f := append(l.fields, fields...)
+
+	return &richLogger{
+		ctx:        l.ctx,
+		callerSkip: l.callerSkip,
+		fields:     f,
+	}
 }
 
 func (l *richLogger) buildFields(fields ...LogField) []LogField {
 	fields = append(l.fields, fields...)
+	// caller field should always appear together with global fields
 	fields = append(fields, Field(callerKey, getCaller(callerDepth+l.callerSkip)))
+	fields = mergeGlobalFields(fields)
 
 	if l.ctx == nil {
 		return fields

core/logx/richlogger_test.go

@@ -63,6 +63,11 @@ func TestTraceDebug(t *testing.T) {
 	l.WithDuration(time.Second).Debugf(testlog)
 	validate(t, w.String(), true, true)
 	w.Reset()
+	l.WithDuration(time.Second).Debugfn(func() any {
+		return testlog
+	})
+	validate(t, w.String(), true, true)
+	w.Reset()
 	l.WithDuration(time.Second).Debugv(testlog)
 	validate(t, w.String(), true, true)
 	w.Reset()
@@ -103,6 +108,11 @@ func TestTraceError(t *testing.T) {
 	l.WithDuration(time.Second).Errorf(testlog)
 	validate(t, w.String(), true, true)
 	w.Reset()
+	l.WithDuration(time.Second).Errorfn(func() any {
+		return testlog
+	})
+	validate(t, w.String(), true, true)
+	w.Reset()
 	l.WithDuration(time.Second).Errorv(testlog)
 	validate(t, w.String(), true, true)
 	w.Reset()
@@ -140,6 +150,11 @@ func TestTraceInfo(t *testing.T) {
 	l.WithDuration(time.Second).Infof(testlog)
 	validate(t, w.String(), true, true)
 	w.Reset()
+	l.WithDuration(time.Second).Infofn(func() any {
+		return testlog
+	})
+	validate(t, w.String(), true, true)
+	w.Reset()
 	l.WithDuration(time.Second).Infov(testlog)
 	validate(t, w.String(), true, true)
 	w.Reset()
@@ -213,6 +228,11 @@ func TestTraceSlow(t *testing.T) {
 	l.WithDuration(time.Second).Slowf(testlog)
 	validate(t, w.String(), true, true)
 	w.Reset()
+	l.WithDuration(time.Second).Slowfn(func() any {
+		return testlog
+	})
+	validate(t, w.String(), true, true)
+	w.Reset()
 	l.WithDuration(time.Second).Slowv(testlog)
 	validate(t, w.String(), true, true)
 	w.Reset()
@@ -287,6 +307,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

@@ -19,7 +19,6 @@ import (
 
 const (
 	dateFormat      = "2006-01-02"
-	fileTimeFormat  = time.RFC3339
 	hoursPerDay     = 24
 	bufferSize      = 100
 	defaultDirMode  = 0o755
@@ -28,8 +27,12 @@ const (
 	megaBytes       = 1 << 20
 )
 
-// ErrLogFileClosed is an error that indicates the log file is already closed.
+var (
-var ErrLogFileClosed = errors.New("error: log file closed")
+	// ErrLogFileClosed is an error that indicates the log file is already closed.
+	ErrLogFileClosed = errors.New("error: log file closed")
+
+	fileTimeFormat = time.RFC3339
+)
 
 type (
 	// A RotateRule interface is used to define the log rotating rules.

core/logx/vars.go

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

core/logx/writer.go

@@ -13,18 +13,31 @@ import (
 
 	fatihcolor "github.com/fatih/color"
 	"github.com/zeromicro/go-zero/core/color"
+	"github.com/zeromicro/go-zero/core/errorx"
 )
 
 type (
+	// Writer is the interface for writing logs.
+	// It's designed to let users customize their own log writer,
+	// such as writing logs to a kafka, a database, or using third-party loggers.
 	Writer interface {
+		// Alert sends an alert message, if your writer implemented alerting functionality.
 		Alert(v any)
+		// Close closes the writer.
 		Close() error
+		// Debug logs a message at debug level.
 		Debug(v any, fields ...LogField)
+		// Error logs a message at error level.
 		Error(v any, fields ...LogField)
+		// Info logs a message at info level.
 		Info(v any, fields ...LogField)
+		// Severe logs a message at severe level.
 		Severe(v any)
+		// Slow logs a message at slow level.
 		Slow(v any, fields ...LogField)
+		// Stack logs a message at error level.
 		Stack(v any)
+		// Stat logs a message at stat level.
 		Stat(v any, fields ...LogField)
 	}
 
@@ -33,6 +46,10 @@ type (
 		lock   sync.RWMutex
 	}
 
+	comboWriter struct {
+		writers []Writer
+	}
+
 	concreteWriter struct {
 		infoLog   io.WriteCloser
 		errorLog  io.WriteCloser
@@ -88,6 +105,62 @@ func (w *atomicWriter) Swap(v Writer) Writer {
 	return old
 }
 
+func (c comboWriter) Alert(v any) {
+	for _, w := range c.writers {
+		w.Alert(v)
+	}
+}
+
+func (c comboWriter) Close() error {
+	var be errorx.BatchError
+	for _, w := range c.writers {
+		be.Add(w.Close())
+	}
+	return be.Err()
+}
+
+func (c comboWriter) Debug(v any, fields ...LogField) {
+	for _, w := range c.writers {
+		w.Debug(v, fields...)
+	}
+}
+
+func (c comboWriter) Error(v any, fields ...LogField) {
+	for _, w := range c.writers {
+		w.Error(v, fields...)
+	}
+}
+
+func (c comboWriter) Info(v any, fields ...LogField) {
+	for _, w := range c.writers {
+		w.Info(v, fields...)
+	}
+}
+
+func (c comboWriter) Severe(v any) {
+	for _, w := range c.writers {
+		w.Severe(v)
+	}
+}
+
+func (c comboWriter) Slow(v any, fields ...LogField) {
+	for _, w := range c.writers {
+		w.Slow(v, fields...)
+	}
+}
+
+func (c comboWriter) Stack(v any) {
+	for _, w := range c.writers {
+		w.Stack(v)
+	}
+}
+
+func (c comboWriter) Stat(v any, fields ...LogField) {
+	for _, w := range c.writers {
+		w.Stat(v, fields...)
+	}
+}
+
 func newConsoleWriter() Writer {
 	outLog := newLogWriter(log.New(fatihcolor.Output, "", flags))
 	errLog := newLogWriter(log.New(fatihcolor.Error, "", flags))
@@ -254,17 +327,30 @@ func (n nopWriter) Stack(_ any) {
 func (n nopWriter) Stat(_ any, _ ...LogField) {
 }
 
-func buildPlainFields(fields ...LogField) []string {
+func buildPlainFields(fields logEntry) []string {
-	var items []string
+	items := make([]string, 0, len(fields))
-
+	for k, v := range fields {
-	for _, field := range fields {
+		items = append(items, fmt.Sprintf("%s=%+v", k, v))
-		items = append(items, fmt.Sprintf("%s=%+v", field.Key, field.Value))
 	}
 
 	return items
 }
 
-func combineGlobalFields(fields []LogField) []LogField {
+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 mergeGlobalFields(fields []LogField) []LogField {
 	globals := globalFields.Load()
 	if globals == nil {
 		return fields
@@ -288,16 +374,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 +419,7 @@ func wrapLevelWithColor(level string) string {
 }
 
 func writeJson(writer io.Writer, info any) {
-	if content, err := json.Marshal(info); err != nil {
+	if content, err := marshalJson(info); err != nil {
 		log.Printf("err: %s\n\n%s", err.Error(), debug.Stack())
 	} else if writer == nil {
 		log.Println(string(content))

core/logx/writer_test.go

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

core/mapping/marshaler.go

@@ -13,6 +13,15 @@ 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 commonly used.
+// support anonymous field, e.g.:
+//
+//	type Foo struct {
+//		Token string `header:"token"`
+//	}
+//	type FooB struct {
+//		Foo
+//		Bar string  `json:"bar"`
+//	}
 func Marshal(val any) (map[string]map[string]any, error) {
 	ret := make(map[string]map[string]any)
 	tp := reflect.TypeOf(val)
@@ -44,6 +53,16 @@ func getTag(field reflect.StructField) (string, bool) {
 	return strings.TrimSpace(tag), false
 }
 
+func insertValue(collector map[string]map[string]any, tag string, key string, val any) {
+	if m, ok := collector[tag]; ok {
+		m[key] = val
+	} else {
+		collector[tag] = map[string]any{
+			key: val,
+		}
+	}
+}
+
 func processMember(field reflect.StructField, value reflect.Value,
 	collector map[string]map[string]any) error {
 	var key string
@@ -69,15 +88,20 @@ func processMember(field reflect.StructField, value reflect.Value,
 		val = fmt.Sprint(val)
 	}
 
-	m, ok := collector[tag]
+	if field.Anonymous {
-	if ok {
+		anonCollector, err := Marshal(val)
-		m[key] = val
+		if err != nil {
-	} else {
+			return err
-		m = map[string]any{
-			key: val,
 		}
+
+		for anonTag, anonMap := range anonCollector {
+			for anonKey, anonVal := range anonMap {
+				insertValue(collector, anonTag, anonKey, anonVal)
+			}
+		}
+	} else {
+		insertValue(collector, tag, key, val)
 	}
-	collector[tag] = m
 
 	return nil
 }
@@ -118,7 +142,7 @@ func validateOptional(field reflect.StructField, value reflect.Value) error {
 		if value.IsNil() {
 			return fmt.Errorf("field %q is nil", field.Name)
 		}
-	case reflect.Array, reflect.Slice, reflect.Map:
+	case reflect.Slice, reflect.Map:
 		if value.IsNil() || value.Len() == 0 {
 			return fmt.Errorf("field %q is empty", field.Name)
 		}

core/mapping/marshaler_test.go

@@ -27,6 +27,124 @@ func TestMarshal(t *testing.T) {
 	assert.True(t, m[emptyTag]["Anonymous"].(bool))
 }
 
+func TestMarshal_Anonymous(t *testing.T) {
+	t.Run("anonymous", func(t *testing.T) {
+		type BaseHeader struct {
+			Token string `header:"token"`
+		}
+		v := struct {
+			Name    string `json:"name"`
+			Address string `json:"address,options=[beijing,shanghai]"`
+			Age     int    `json:"age"`
+			BaseHeader
+		}{
+			Name:    "kevin",
+			Address: "shanghai",
+			Age:     20,
+			BaseHeader: BaseHeader{
+				Token: "token_xxx",
+			},
+		}
+		m, err := Marshal(v)
+		assert.Nil(t, err)
+		assert.Equal(t, "kevin", m["json"]["name"])
+		assert.Equal(t, "shanghai", m["json"]["address"])
+		assert.Equal(t, 20, m["json"]["age"].(int))
+		assert.Equal(t, "token_xxx", m["header"]["token"])
+
+		v1 := struct {
+			Name    string `json:"name"`
+			Address string `json:"address,options=[beijing,shanghai]"`
+			Age     int    `json:"age"`
+			BaseHeader
+		}{
+			Name:    "kevin",
+			Address: "shanghai",
+			Age:     20,
+		}
+		m1, err1 := Marshal(v1)
+		assert.Nil(t, err1)
+		assert.Equal(t, "kevin", m1["json"]["name"])
+		assert.Equal(t, "shanghai", m1["json"]["address"])
+		assert.Equal(t, 20, m1["json"]["age"].(int))
+
+		type AnotherHeader struct {
+			Version string `header:"version"`
+		}
+		v2 := struct {
+			Name    string `json:"name"`
+			Address string `json:"address,options=[beijing,shanghai]"`
+			Age     int    `json:"age"`
+			BaseHeader
+			AnotherHeader
+		}{
+			Name:    "kevin",
+			Address: "shanghai",
+			Age:     20,
+			BaseHeader: BaseHeader{
+				Token: "token_xxx",
+			},
+			AnotherHeader: AnotherHeader{
+				Version: "v1.0",
+			},
+		}
+		m2, err2 := Marshal(v2)
+		assert.Nil(t, err2)
+		assert.Equal(t, "kevin", m2["json"]["name"])
+		assert.Equal(t, "shanghai", m2["json"]["address"])
+		assert.Equal(t, 20, m2["json"]["age"].(int))
+		assert.Equal(t, "token_xxx", m2["header"]["token"])
+		assert.Equal(t, "v1.0", m2["header"]["version"])
+
+		type PointerHeader struct {
+			Ref *string `header:"ref"`
+		}
+		ref := "reference"
+		v3 := struct {
+			Name    string `json:"name"`
+			Address string `json:"address,options=[beijing,shanghai]"`
+			Age     int    `json:"age"`
+			PointerHeader
+		}{
+			Name:    "kevin",
+			Address: "shanghai",
+			Age:     20,
+			PointerHeader: PointerHeader{
+				Ref: &ref,
+			},
+		}
+		m3, err3 := Marshal(v3)
+		assert.Nil(t, err3)
+		assert.Equal(t, "kevin", m3["json"]["name"])
+		assert.Equal(t, "shanghai", m3["json"]["address"])
+		assert.Equal(t, 20, m3["json"]["age"].(int))
+		assert.Equal(t, "reference", *m3["header"]["ref"].(*string))
+	})
+
+	t.Run("bad anonymous", func(t *testing.T) {
+		type BaseHeader struct {
+			Token string `json:"token,options=[a,b]"`
+		}
+
+		v := struct {
+			Name    string `json:"name"`
+			Address string `json:"address,options=[beijing,shanghai]"`
+			Age     int    `json:"age"`
+			BaseHeader
+		}{
+			Name:    "kevin",
+			Address: "shanghai",
+			Age:     20,
+			BaseHeader: BaseHeader{
+				Token: "c",
+			},
+		}
+
+		_, err := Marshal(v)
+		assert.NotNil(t, err)
+	})
+}
+
 func TestMarshal_Ptr(t *testing.T) {
 	v := &struct {
 		Name      string `path:"name"`
@@ -344,3 +462,15 @@ func TestMarshal_FromString(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, "10", m["json"]["age"].(string))
 }
+
+func TestMarshal_Array(t *testing.T) {
+	v := struct {
+		H [1]int `json:"h,string"`
+	}{
+		H: [1]int{1},
+	}
+
+	m, err := Marshal(v)
+	assert.Nil(t, err)
+	assert.Equal(t, "[1]", m["json"]["h"].(string))
+}

core/mapping/unmarshaler.go

@@ -2,6 +2,7 @@ package mapping
 
 import (
 	"encoding"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -50,6 +51,7 @@ type (
 
 	unmarshalOptions struct {
 		fillDefault  bool
+		fromArray    bool
 		fromString   bool
 		opaqueKeys   bool
 		canonicalKey func(key string) string
@@ -79,41 +81,13 @@ func (u *Unmarshaler) Unmarshal(i, v any) error {
 	return u.unmarshal(i, v, "")
 }
 
-func (u *Unmarshaler) unmarshal(i, v any, fullName string) error {
-	valueType := reflect.TypeOf(v)
-	if valueType.Kind() != reflect.Ptr {
-		return errValueNotSettable
-	}
-
-	elemType := Deref(valueType)
-	switch iv := i.(type) {
-	case map[string]any:
-		if elemType.Kind() != reflect.Struct {
-			return errTypeMismatch
-		}
-
-		return u.unmarshalValuer(mapValuer(iv), v, fullName)
-	case []any:
-		if elemType.Kind() != reflect.Slice {
-			return errTypeMismatch
-		}
-
-		return u.fillSlice(elemType, reflect.ValueOf(v).Elem(), iv, fullName)
-	default:
-		return errUnsupportedType
-	}
-}
-
 // UnmarshalValuer unmarshals m into v.
 func (u *Unmarshaler) UnmarshalValuer(m Valuer, v any) error {
 	return u.unmarshalValuer(simpleValuer{current: m}, v, "")
 }
 
-func (u *Unmarshaler) unmarshalValuer(m Valuer, v any, fullName string) error {
+func (u *Unmarshaler) fillMap(fieldType reflect.Type, value reflect.Value,
-	return u.unmarshalWithFullName(simpleValuer{current: m}, v, fullName)
+	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 +128,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
 	}
@@ -170,13 +145,14 @@ func (u *Unmarshaler) fillSlice(fieldType reflect.Type, value reflect.Value, map
 	baseType := fieldType.Elem()
 	dereffedBaseType := Deref(baseType)
 	dereffedBaseKind := dereffedBaseType.Kind()
-	conv := reflect.MakeSlice(reflect.SliceOf(baseType), refValue.Len(), refValue.Cap())
 	if refValue.Len() == 0 {
-		value.Set(conv)
+		value.Set(reflect.MakeSlice(reflect.SliceOf(baseType), 0, 0))
 		return nil
 	}
 
 	var valid bool
+	conv := reflect.MakeSlice(reflect.SliceOf(baseType), refValue.Len(), refValue.Cap())
+
 	for i := 0; i < refValue.Len(); i++ {
 		ithValue := refValue.Index(i).Interface()
 		if ithValue == nil {
@@ -188,17 +164,9 @@ func (u *Unmarshaler) fillSlice(fieldType reflect.Type, value reflect.Value, map
 
 		switch dereffedBaseKind {
 		case reflect.Struct:
-			target := reflect.New(dereffedBaseType)
+			if err := u.fillStructElement(baseType, conv.Index(i), ithValue, sliceFullName); err != nil {
-			val, ok := ithValue.(map[string]any)
-			if !ok {
-				return errTypeMismatch
-			}
-
-			if err := u.unmarshal(val, target.Interface(), sliceFullName); err != nil {
 				return err
 			}
-
-			SetValue(fieldType.Elem(), conv.Index(i), target.Elem())
 		case reflect.Slice:
 			if err := u.fillSlice(dereffedBaseType, conv.Index(i), ithValue, sliceFullName); err != nil {
 				return err
@@ -233,7 +201,7 @@ func (u *Unmarshaler) fillSliceFromString(fieldType reflect.Type, value reflect.
 		return errUnsupportedType
 	}
 
-	baseFieldType := Deref(fieldType.Elem())
+	baseFieldType := fieldType.Elem()
 	baseFieldKind := baseFieldType.Kind()
 	conv := reflect.MakeSlice(reflect.SliceOf(baseFieldType), len(slice), cap(slice))
 
@@ -254,29 +222,39 @@ func (u *Unmarshaler) fillSliceValue(slice reflect.Value, index int,
 	}
 
 	ithVal := slice.Index(index)
+	ithValType := ithVal.Type()
+
 	switch v := value.(type) {
 	case fmt.Stringer:
 		return setValueFromString(baseKind, ithVal, v.String())
 	case string:
 		return setValueFromString(baseKind, ithVal, v)
 	case map[string]any:
-		return u.fillMap(ithVal.Type(), ithVal, value, fullName)
+		// deref to handle both pointer and non-pointer types.
+		switch Deref(ithValType).Kind() {
+		case reflect.Struct:
+			return u.fillStructElement(ithValType, ithVal, v, fullName)
+		case reflect.Map:
+			return u.fillMap(ithValType, ithVal, value, fullName)
+		default:
+			return errTypeMismatch
+		}
 	default:
 		// don't need to consider the difference between int, int8, int16, int32, int64,
 		// uint, uint8, uint16, uint32, uint64, because they're handled as json.Number.
 		if ithVal.Kind() == reflect.Ptr {
-			baseType := Deref(ithVal.Type())
+			baseType := Deref(ithValType)
 			if !reflect.TypeOf(value).AssignableTo(baseType) {
 				return errTypeMismatch
 			}
 
 			target := reflect.New(baseType).Elem()
 			target.Set(reflect.ValueOf(value))
-			SetValue(ithVal.Type(), ithVal, target)
+			SetValue(ithValType, ithVal, target)
 			return nil
 		}
 
-		if !reflect.TypeOf(value).AssignableTo(ithVal.Type()) {
+		if !reflect.TypeOf(value).AssignableTo(ithValType) {
 			return errTypeMismatch
 		}
 
@@ -307,7 +285,51 @@ 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) fillStructElement(baseType reflect.Type, target reflect.Value,
+	value any, fullName string) error {
+	val, ok := value.(map[string]any)
+	if !ok {
+		return errTypeMismatch
+	}
+
+	// use Deref(baseType) to get the base type in case the type is a pointer type.
+	ptr := reflect.New(Deref(baseType))
+	if err := u.unmarshal(val, ptr.Interface(), fullName); err != nil {
+		return err
+	}
+
+	SetValue(baseType, target, ptr.Elem())
+	return nil
+}
+
+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 +421,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)
@@ -573,9 +604,27 @@ func (u *Unmarshaler) processFieldNotFromString(fieldType reflect.Type, value re
 	case valueKind == reflect.String && typeKind == reflect.Map:
 		return u.fillMapFromString(value, mapValue)
 	case valueKind == reflect.String && typeKind == reflect.Slice:
+		// try to find out if it's a byte slice,
+		// more details https://pkg.go.dev/encoding/json#Marshal
+		// array and slice values encode as JSON arrays,
+		// except that []byte encodes as a base64-encoded string,
+		// and a nil slice encoded as the null JSON value.
+		// https://stackoverflow.com/questions/34089750/marshal-byte-to-json-giving-a-strange-string
+		if fieldType.Elem().Kind() == reflect.Uint8 {
+			// check whether string type, because the kind of some other types can be string
+			if strVal, ok := mapValue.(string); ok {
+				if decodedBytes, err := base64.StdEncoding.DecodeString(strVal); err == nil {
+					value.Set(reflect.ValueOf(decodedBytes))
+					return nil
+				}
+			}
+		}
+
 		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)
 	}
@@ -771,6 +820,19 @@ func (u *Unmarshaler) processNamedField(field reflect.StructField, value reflect
 		return u.processNamedFieldWithoutValue(field.Type, value, opts, fullName)
 	}
 
+	if u.opts.fromArray {
+		fieldKind := field.Type.Kind()
+		if fieldKind != reflect.Slice && fieldKind != reflect.Array {
+			valueKind := reflect.TypeOf(mapValue).Kind()
+			if valueKind == reflect.Slice || valueKind == reflect.Array {
+				val := reflect.ValueOf(mapValue)
+				if val.Len() > 0 {
+					mapValue = val.Index(0).Interface()
+				}
+			}
+		}
+	}
+
 	return u.processNamedFieldWithValue(field.Type, value, valueWithParent{
 		value:  mapValue,
 		parent: valuer,
@@ -898,6 +960,35 @@ func (u *Unmarshaler) processNamedFieldWithoutValue(fieldType reflect.Type, valu
 	return nil
 }
 
+func (u *Unmarshaler) unmarshal(i, v any, fullName string) error {
+	valueType := reflect.TypeOf(v)
+	if valueType.Kind() != reflect.Ptr {
+		return errValueNotSettable
+	}
+
+	elemType := Deref(valueType)
+	switch iv := i.(type) {
+	case map[string]any:
+		if elemType.Kind() != reflect.Struct {
+			return errTypeMismatch
+		}
+
+		return u.unmarshalValuer(mapValuer(iv), v, fullName)
+	case []any:
+		if elemType.Kind() != reflect.Slice {
+			return errTypeMismatch
+		}
+
+		return u.fillSlice(elemType, reflect.ValueOf(v).Elem(), iv, fullName)
+	default:
+		return errUnsupportedType
+	}
+}
+
+func (u *Unmarshaler) unmarshalValuer(m Valuer, v any, fullName string) error {
+	return u.unmarshalWithFullName(simpleValuer{current: m}, v, fullName)
+}
+
 func (u *Unmarshaler) unmarshalWithFullName(m valuerWithParent, v any, fullName string) error {
 	rv := reflect.ValueOf(v)
 	if err := ValidatePtr(rv); err != nil {
@@ -950,6 +1041,16 @@ func WithDefault() UnmarshalOption {
 	}
 }
 
+// WithFromArray customizes an Unmarshaler with converting array values to non-array types.
+// For example, if the field type is []string, and the value is [hello],
+// the field type can be `string`, instead of `[]string`.
+// Typically, this option is used for unmarshaling from form values.
+func WithFromArray() UnmarshalOption {
+	return func(opt *unmarshalOptions) {
+		opt.fromArray = true
+	}
+}
+
 // WithOpaqueKeys customizes an Unmarshaler with opaque keys.
 // Opaque keys are keys that are not processed by the unmarshaler.
 func WithOpaqueKeys() UnmarshalOption {

core/mapping/unmarshaler_test.go

@@ -2,6 +2,7 @@ package mapping
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"reflect"
 	"strconv"
@@ -12,6 +13,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/jsonx"
 	"github.com/zeromicro/go-zero/core/stringx"
 )
 
@@ -260,6 +262,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 +291,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)
 	}
 }
 
@@ -348,7 +352,7 @@ func TestUnmarshalIntSliceOfPtr(t *testing.T) {
 		assert.Error(t, UnmarshalKey(m, &in))
 	})
 
-	t.Run("int slice with nil", func(t *testing.T) {
+	t.Run("int slice with nil element", func(t *testing.T) {
 		type inner struct {
 			Ints []int `key:"ints"`
 		}
@@ -362,6 +366,21 @@ func TestUnmarshalIntSliceOfPtr(t *testing.T) {
 			assert.Empty(t, in.Ints)
 		}
 	})
+
+	t.Run("int slice with nil", func(t *testing.T) {
+		type inner struct {
+			Ints []int `key:"ints"`
+		}
+
+		m := map[string]any{
+			"ints": []any(nil),
+		}
+
+		var in inner
+		if assert.NoError(t, UnmarshalKey(m, &in)) {
+			assert.Empty(t, in.Ints)
+		}
+	})
 }
 
 func TestUnmarshalIntWithDefault(t *testing.T) {
@@ -1371,20 +1390,80 @@ func TestUnmarshalWithFloatPtr(t *testing.T) {
 }
 
 func TestUnmarshalIntSlice(t *testing.T) {
-	var v struct {
+	t.Run("int slice from int", func(t *testing.T) {
-		Ages  []int `key:"ages"`
+		var v struct {
-		Slice []int `key:"slice"`
+			Ages  []int `key:"ages"`
-	}
+			Slice []int `key:"slice"`
-	m := map[string]any{
+		}
-		"ages":  []int{1, 2},
+		m := map[string]any{
-		"slice": []any{},
+			"ages":  []int{1, 2},
-	}
+			"slice": []any{},
+		}
 
-	ast := assert.New(t)
+		ast := assert.New(t)
-	if ast.NoError(UnmarshalKey(m, &v)) {
+		if ast.NoError(UnmarshalKey(m, &v)) {
-		ast.ElementsMatch([]int{1, 2}, v.Ages)
+			ast.ElementsMatch([]int{1, 2}, v.Ages)
-		ast.Equal([]int{}, v.Slice)
+			ast.Equal([]int{}, v.Slice)
-	}
+		}
+	})
+
+	t.Run("int slice from one int", func(t *testing.T) {
+		var v struct {
+			Ages []int `key:"ages"`
+		}
+		m := map[string]any{
+			"ages": []int{2},
+		}
+
+		ast := assert.New(t)
+		unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
+		if ast.NoError(unmarshaler.Unmarshal(m, &v)) {
+			ast.ElementsMatch([]int{2}, v.Ages)
+		}
+	})
+
+	t.Run("int slice from one int string", func(t *testing.T) {
+		var v struct {
+			Ages []int `key:"ages"`
+		}
+		m := map[string]any{
+			"ages": []string{"2"},
+		}
+
+		ast := assert.New(t)
+		unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
+		if ast.NoError(unmarshaler.Unmarshal(m, &v)) {
+			ast.ElementsMatch([]int{2}, v.Ages)
+		}
+	})
+
+	t.Run("int slice from one json.Number", func(t *testing.T) {
+		var v struct {
+			Ages []int `key:"ages"`
+		}
+		m := map[string]any{
+			"ages": []json.Number{"2"},
+		}
+
+		ast := assert.New(t)
+		unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
+		if ast.NoError(unmarshaler.Unmarshal(m, &v)) {
+			ast.ElementsMatch([]int{2}, v.Ages)
+		}
+	})
+
+	t.Run("int slice from one int strings", func(t *testing.T) {
+		var v struct {
+			Ages []int `key:"ages"`
+		}
+		m := map[string]any{
+			"ages": []string{"1,2"},
+		}
+
+		ast := assert.New(t)
+		unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
+		ast.Error(unmarshaler.Unmarshal(m, &v))
+	})
 }
 
 func TestUnmarshalString(t *testing.T) {
@@ -1439,6 +1518,51 @@ func TestUnmarshalStringSliceFromString(t *testing.T) {
 		}
 	})
 
+	t.Run("slice from empty string", func(t *testing.T) {
+		var v struct {
+			Names []string `key:"names"`
+		}
+		m := map[string]any{
+			"names": []string{""},
+		}
+
+		ast := assert.New(t)
+		unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
+		if ast.NoError(unmarshaler.Unmarshal(m, &v)) {
+			ast.ElementsMatch([]string{""}, v.Names)
+		}
+	})
+
+	t.Run("slice from empty and valid string", func(t *testing.T) {
+		var v struct {
+			Names []string `key:"names"`
+		}
+		m := map[string]any{
+			"names": []string{","},
+		}
+
+		ast := assert.New(t)
+		unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
+		if ast.NoError(unmarshaler.Unmarshal(m, &v)) {
+			ast.ElementsMatch([]string{","}, v.Names)
+		}
+	})
+
+	t.Run("slice from valid strings with comma", func(t *testing.T) {
+		var v struct {
+			Names []string `key:"names"`
+		}
+		m := map[string]any{
+			"names": []string{"aa,bb"},
+		}
+
+		ast := assert.New(t)
+		unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
+		if ast.NoError(unmarshaler.Unmarshal(m, &v)) {
+			ast.ElementsMatch([]string{"aa,bb"}, v.Names)
+		}
+	})
+
 	t.Run("slice from string with slice error", func(t *testing.T) {
 		var v struct {
 			Names []int `key:"names"`
@@ -4758,14 +4882,28 @@ func TestUnmarshal_EnvWithOptionsWrongValueString(t *testing.T) {
 
 func TestUnmarshalJsonReaderMultiArray(t *testing.T) {
 	t.Run("reader multi array", func(t *testing.T) {
-		var res struct {
+		type testRes struct {
 			A string     `json:"a"`
 			B [][]string `json:"b"`
+			C []byte     `json:"c"`
 		}
-		payload := `{"a": "133", "b": [["add", "cccd"], ["eeee"]]}`
+
+		var res testRes
+		marshal := testRes{
+			A: "133",
+			B: [][]string{
+				{"add", "cccd"},
+				{"eeee"},
+			},
+			C: []byte("11122344wsss"),
+		}
+		bytes, err := jsonx.Marshal(marshal)
+		assert.NoError(t, err)
+		payload := string(bytes)
 		reader := strings.NewReader(payload)
 		if assert.NoError(t, UnmarshalJsonReader(reader, &res)) {
 			assert.Equal(t, 2, len(res.B))
+			assert.Equal(t, string(marshal.C), string(res.C))
 		}
 	})
 
@@ -5636,6 +5774,62 @@ func TestUnmarshalFromStringSliceForTypeMismatch(t *testing.T) {
 	}, &v))
 }
 
+func TestUnmarshalWithFromArray(t *testing.T) {
+	t.Run("array", func(t *testing.T) {
+		var v struct {
+			Value []string `key:"value"`
+		}
+		unmarshaler := NewUnmarshaler("key", WithFromArray())
+		if assert.NoError(t, unmarshaler.Unmarshal(map[string]any{
+			"value": []string{"foo", "bar"},
+		}, &v)) {
+			assert.ElementsMatch(t, []string{"foo", "bar"}, v.Value)
+		}
+	})
+
+	t.Run("not array", func(t *testing.T) {
+		var v struct {
+			Value string `key:"value"`
+		}
+		unmarshaler := NewUnmarshaler("key", WithFromArray())
+		if assert.NoError(t, unmarshaler.Unmarshal(map[string]any{
+			"value": []string{"foo"},
+		}, &v)) {
+			assert.Equal(t, "foo", v.Value)
+		}
+	})
+
+	t.Run("not array and empty", func(t *testing.T) {
+		var v struct {
+			Value string `key:"value"`
+		}
+		unmarshaler := NewUnmarshaler("key", WithFromArray())
+		if assert.NoError(t, unmarshaler.Unmarshal(map[string]any{
+			"value": []string{""},
+		}, &v)) {
+			assert.Empty(t, v.Value)
+		}
+	})
+
+	t.Run("not array and no value", func(t *testing.T) {
+		var v struct {
+			Value string `key:"value"`
+		}
+		unmarshaler := NewUnmarshaler("key", WithFromArray())
+		assert.Error(t, unmarshaler.Unmarshal(map[string]any{}, &v))
+	})
+
+	t.Run("not array and no value and optional", func(t *testing.T) {
+		var v struct {
+			Value string `key:"value,optional"`
+		}
+		unmarshaler := NewUnmarshaler("key", WithFromArray())
+		if assert.NoError(t, unmarshaler.Unmarshal(map[string]any{}, &v)) {
+			assert.Empty(t, v.Value)
+		}
+	})
+}
+
 func TestUnmarshalWithOpaqueKeys(t *testing.T) {
 	var v struct {
 		Opaque string `key:"opaque.key"`
@@ -5760,6 +5954,81 @@ 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 TestParseJsonStringValue(t *testing.T) {
+	t.Run("string", func(t *testing.T) {
+		type GoodsInfo struct {
+			Sku int64 `json:"sku,optional"`
+		}
+
+		type GetReq struct {
+			GoodsList []*GoodsInfo `json:"goods_list"`
+		}
+
+		input := map[string]any{"goods_list": "[{\"sku\":11},{\"sku\":22}]"}
+		var v GetReq
+		assert.NotPanics(t, func() {
+			assert.NoError(t, UnmarshalJsonMap(input, &v))
+			assert.Equal(t, 2, len(v.GoodsList))
+			assert.ElementsMatch(t, []int64{11, 22}, []int64{v.GoodsList[0].Sku, v.GoodsList[1].Sku})
+		})
+	})
+
+	t.Run("string with invalid type", func(t *testing.T) {
+		type GetReq struct {
+			GoodsList []*int `json:"goods_list"`
+		}
+
+		input := map[string]any{"goods_list": "[{\"sku\":11},{\"sku\":22}]"}
+		var v GetReq
+		assert.NotPanics(t, func() {
+			assert.Error(t, UnmarshalJsonMap(input, &v))
+		})
+	})
+}
+
 func BenchmarkDefaultValue(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		var a struct {
@@ -5873,3 +6142,20 @@ func (m mockValuerWithParent) Value(_ string) (any, bool) {
 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/mathx/range.go

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

core/mr/mapreduce.go

@@ -142,89 +142,6 @@ func MapReduceChan[T, U, V any](source <-chan T, mapper MapperFunc[T, U], reduce
 	return mapReduceWithPanicChan(source, panicChan, mapper, reducer, opts...)
 }
 
-// mapReduceWithPanicChan maps all elements from source, and reduce the output elements with given reducer.
-func mapReduceWithPanicChan[T, U, V any](source <-chan T, panicChan *onceChan, mapper MapperFunc[T, U],
-	reducer ReducerFunc[U, V], opts ...Option) (val V, err error) {
-	options := buildOptions(opts...)
-	// output is used to write the final result
-	output := make(chan V)
-	defer func() {
-		// reducer can only write once, if more, panic
-		for range output {
-			panic("more than one element written in reducer")
-		}
-	}()
-
-	// collector is used to collect data from mapper, and consume in reducer
-	collector := make(chan U, options.workers)
-	// if done is closed, all mappers and reducer should stop processing
-	done := make(chan struct{})
-	writer := newGuardedWriter(options.ctx, output, done)
-	var closeOnce sync.Once
-	// use atomic type to avoid data race
-	var retErr errorx.AtomicError
-	finish := func() {
-		closeOnce.Do(func() {
-			close(done)
-			close(output)
-		})
-	}
-	cancel := once(func(err error) {
-		if err != nil {
-			retErr.Set(err)
-		} else {
-			retErr.Set(ErrCancelWithNil)
-		}
-
-		drain(source)
-		finish()
-	})
-
-	go func() {
-		defer func() {
-			drain(collector)
-			if r := recover(); r != nil {
-				panicChan.write(r)
-			}
-			finish()
-		}()
-
-		reducer(collector, writer, cancel)
-	}()
-
-	go executeMappers(mapperContext[T, U]{
-		ctx: options.ctx,
-		mapper: func(item T, w Writer[U]) {
-			mapper(item, w, cancel)
-		},
-		source:    source,
-		panicChan: panicChan,
-		collector: collector,
-		doneChan:  done,
-		workers:   options.workers,
-	})
-
-	select {
-	case <-options.ctx.Done():
-		cancel(context.DeadlineExceeded)
-		err = context.DeadlineExceeded
-	case v := <-panicChan.channel:
-		// drain output here, otherwise for loop panic in defer
-		drain(output)
-		panic(v)
-	case v, ok := <-output:
-		if e := retErr.Load(); e != nil {
-			err = e
-		} else if ok {
-			val = v
-		} else {
-			err = ErrReduceNoOutput
-		}
-	}
-
-	return
-}
-
 // MapReduceVoid maps all elements generated from given generate,
 // and reduce the output elements with given reducer.
 func MapReduceVoid[T, U any](generate GenerateFunc[T], mapper MapperFunc[T, U],
@@ -330,6 +247,89 @@ func executeMappers[T, U any](mCtx mapperContext[T, U]) {
 	}
 }
 
+// mapReduceWithPanicChan maps all elements from source, and reduce the output elements with given reducer.
+func mapReduceWithPanicChan[T, U, V any](source <-chan T, panicChan *onceChan, mapper MapperFunc[T, U],
+	reducer ReducerFunc[U, V], opts ...Option) (val V, err error) {
+	options := buildOptions(opts...)
+	// output is used to write the final result
+	output := make(chan V)
+	defer func() {
+		// reducer can only write once, if more, panic
+		for range output {
+			panic("more than one element written in reducer")
+		}
+	}()
+
+	// collector is used to collect data from mapper, and consume in reducer
+	collector := make(chan U, options.workers)
+	// if done is closed, all mappers and reducer should stop processing
+	done := make(chan struct{})
+	writer := newGuardedWriter(options.ctx, output, done)
+	var closeOnce sync.Once
+	// use atomic type to avoid data race
+	var retErr errorx.AtomicError
+	finish := func() {
+		closeOnce.Do(func() {
+			close(done)
+			close(output)
+		})
+	}
+	cancel := once(func(err error) {
+		if err != nil {
+			retErr.Set(err)
+		} else {
+			retErr.Set(ErrCancelWithNil)
+		}
+
+		drain(source)
+		finish()
+	})
+
+	go func() {
+		defer func() {
+			drain(collector)
+			if r := recover(); r != nil {
+				panicChan.write(r)
+			}
+			finish()
+		}()
+
+		reducer(collector, writer, cancel)
+	}()
+
+	go executeMappers(mapperContext[T, U]{
+		ctx: options.ctx,
+		mapper: func(item T, w Writer[U]) {
+			mapper(item, w, cancel)
+		},
+		source:    source,
+		panicChan: panicChan,
+		collector: collector,
+		doneChan:  done,
+		workers:   options.workers,
+	})
+
+	select {
+	case <-options.ctx.Done():
+		cancel(context.DeadlineExceeded)
+		err = context.DeadlineExceeded
+	case v := <-panicChan.channel:
+		// drain output here, otherwise for loop panic in defer
+		drain(output)
+		panic(v)
+	case v, ok := <-output:
+		if e := retErr.Load(); e != nil {
+			err = e
+		} else if ok {
+			val = v
+		} else {
+			err = ErrReduceNoOutput
+		}
+	}
+
+	return
+}
+
 func newOptions() *mapReduceOptions {
 	return &mapReduceOptions{
 		ctx:     context.Background(),
@@ -363,9 +363,7 @@ func newGuardedWriter[T any](ctx context.Context, channel chan<- T, done <-chan
 func (gw guardedWriter[T]) Write(v T) {
 	select {
 	case <-gw.ctx.Done():
-		return
 	case <-gw.done:
-		return
 	default:
 		gw.channel <- v
 	}

core/proc/goroutines.go

@@ -1,4 +1,4 @@
-//go:build linux || darwin
+//go:build linux || darwin || freebsd
 
 package proc
 

core/proc/goroutines_test.go

@@ -1,4 +1,4 @@
-//go:build linux || darwin
+//go:build linux || darwin || freebsd
 
 package proc
 

core/proc/profile.go

@@ -1,4 +1,4 @@
-//go:build linux || darwin
+//go:build linux || darwin || freebsd
 
 package proc
 

core/proc/shutdown+polyfill.go

@@ -4,6 +4,9 @@ package proc
 
 import "time"
 
+// ShutdownConf is empty on windows.
+type ShutdownConf struct{}
+
 // AddShutdownListener returns fn itself on windows, lets callers call fn on their own.
 func AddShutdownListener(fn func()) func() {
 	return fn
@@ -18,6 +21,10 @@ func AddWrapUpListener(fn func()) func() {
 func SetTimeToForceQuit(duration time.Duration) {
 }
 
+// Setup does nothing on windows.
+func Setup(conf ShutdownConf) {
+}
+
 // Shutdown does nothing on windows.
 func Shutdown() {
 }

core/proc/shutdown.go

@@ -1,4 +1,4 @@
-//go:build linux || darwin
+//go:build linux || darwin || freebsd
 
 package proc
 
@@ -14,17 +14,29 @@ import (
 )
 
 const (
-	wrapUpTime = time.Second
+	// defaultWrapUpTime is the default time to wait before calling wrap up listeners.
-	// why we use 5500 milliseconds is because most of our queue are blocking mode with 5 seconds
+	defaultWrapUpTime = time.Second
-	waitTime = 5500 * time.Millisecond
+	// defaultWaitTime is the default time to wait before force quitting.
+	// why we use 5500 milliseconds is because most of our queues are blocking mode with 5 seconds
+	defaultWaitTime = 5500 * time.Millisecond
 )
 
 var (
-	wrapUpListeners          = new(listenerManager)
+	wrapUpListeners   = new(listenerManager)
-	shutdownListeners        = new(listenerManager)
+	shutdownListeners = new(listenerManager)
-	delayTimeBeforeForceQuit = waitTime
+	wrapUpTime        = defaultWrapUpTime
+	waitTime          = defaultWaitTime
+	shutdownLock      sync.Mutex
 )
 
+// ShutdownConf defines the shutdown configuration for the process.
+type ShutdownConf struct {
+	// WrapUpTime is the time to wait before calling shutdown listeners.
+	WrapUpTime time.Duration `json:",default=1s"`
+	// WaitTime is the time to wait before force quitting.
+	WaitTime time.Duration `json:",default=5.5s"`
+}
+
 // AddShutdownListener adds fn as a shutdown listener.
 // The returned func can be used to wait for fn getting called.
 func AddShutdownListener(fn func()) (waitForCalled func()) {
@@ -39,7 +51,21 @@ func AddWrapUpListener(fn func()) (waitForCalled func()) {
 
 // SetTimeToForceQuit sets the waiting time before force quitting.
 func SetTimeToForceQuit(duration time.Duration) {
-	delayTimeBeforeForceQuit = duration
+	shutdownLock.Lock()
+	defer shutdownLock.Unlock()
+	waitTime = duration
+}
+
+func Setup(conf ShutdownConf) {
+	shutdownLock.Lock()
+	defer shutdownLock.Unlock()
+
+	if conf.WrapUpTime > 0 {
+		wrapUpTime = conf.WrapUpTime
+	}
+	if conf.WaitTime > 0 {
+		waitTime = conf.WaitTime
+	}
 }
 
 // Shutdown calls the registered shutdown listeners, only for test purpose.
@@ -61,8 +87,12 @@ func gracefulStop(signals chan os.Signal, sig syscall.Signal) {
 	time.Sleep(wrapUpTime)
 	go shutdownListeners.notifyListeners()
 
-	time.Sleep(delayTimeBeforeForceQuit - wrapUpTime)
+	shutdownLock.Lock()
-	logx.Infof("Still alive after %v, going to force kill the process...", delayTimeBeforeForceQuit)
+	remainingTime := waitTime - wrapUpTime
+	shutdownLock.Unlock()
+
+	time.Sleep(remainingTime)
+	logx.Infof("Still alive after %v, going to force kill the process...", waitTime)
 	_ = syscall.Kill(syscall.Getpid(), sig)
 }
 
@@ -82,6 +112,9 @@ func (lm *listenerManager) addListener(fn func()) (waitForCalled func()) {
 	})
 	lm.lock.Unlock()
 
+	// we can return lm.waitGroup.Wait directly,
+	// but we want to make the returned func more readable.
+	// creating an extra closure would be negligible in practice.
 	return func() {
 		lm.waitGroup.Wait()
 	}

core/proc/shutdown_test.go

@@ -1,8 +1,9 @@
-//go:build linux || darwin
+//go:build linux || darwin || freebsd
 
 package proc
 
 import (
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -10,8 +11,12 @@ import (
 )
 
 func TestShutdown(t *testing.T) {
+	t.Cleanup(restoreSettings)
+
 	SetTimeToForceQuit(time.Hour)
-	assert.Equal(t, time.Hour, delayTimeBeforeForceQuit)
+	shutdownLock.Lock()
+	assert.Equal(t, time.Hour, waitTime)
+	shutdownLock.Unlock()
 
 	var val int
 	called := AddWrapUpListener(func() {
@@ -29,7 +34,53 @@ func TestShutdown(t *testing.T) {
 	assert.Equal(t, 3, val)
 }
 
+func TestShutdownWithMultipleServices(t *testing.T) {
+	t.Cleanup(restoreSettings)
+
+	SetTimeToForceQuit(time.Hour)
+	shutdownLock.Lock()
+	assert.Equal(t, time.Hour, waitTime)
+	shutdownLock.Unlock()
+
+	var val int32
+	called1 := AddShutdownListener(func() {
+		atomic.AddInt32(&val, 1)
+	})
+	called2 := AddShutdownListener(func() {
+		atomic.AddInt32(&val, 2)
+	})
+	Shutdown()
+	called1()
+	called2()
+
+	assert.Equal(t, int32(3), atomic.LoadInt32(&val))
+}
+
+func TestWrapUpWithMultipleServices(t *testing.T) {
+	t.Cleanup(restoreSettings)
+
+	SetTimeToForceQuit(time.Hour)
+	shutdownLock.Lock()
+	assert.Equal(t, time.Hour, waitTime)
+	shutdownLock.Unlock()
+
+	var val int32
+	called1 := AddWrapUpListener(func() {
+		atomic.AddInt32(&val, 1)
+	})
+	called2 := AddWrapUpListener(func() {
+		atomic.AddInt32(&val, 2)
+	})
+	WrapUp()
+	called1()
+	called2()
+
+	assert.Equal(t, int32(3), atomic.LoadInt32(&val))
+}
+
 func TestNotifyMoreThanOnce(t *testing.T) {
+	t.Cleanup(restoreSettings)
+
 	ch := make(chan struct{}, 1)
 
 	go func() {
@@ -58,3 +109,38 @@ func TestNotifyMoreThanOnce(t *testing.T) {
 		t.Fatal("timeout, check error logs")
 	}
 }
+
+func TestSetup(t *testing.T) {
+	t.Run("valid time", func(t *testing.T) {
+		defer restoreSettings()
+
+		Setup(ShutdownConf{
+			WrapUpTime: time.Second * 2,
+			WaitTime:   time.Second * 30,
+		})
+
+		shutdownLock.Lock()
+		assert.Equal(t, time.Second*2, wrapUpTime)
+		assert.Equal(t, time.Second*30, waitTime)
+		shutdownLock.Unlock()
+	})
+
+	t.Run("valid time", func(t *testing.T) {
+		defer restoreSettings()
+
+		Setup(ShutdownConf{})
+
+		shutdownLock.Lock()
+		assert.Equal(t, defaultWrapUpTime, wrapUpTime)
+		assert.Equal(t, defaultWaitTime, waitTime)
+		shutdownLock.Unlock()
+	})
+}
+
+func restoreSettings() {
+	shutdownLock.Lock()
+	defer shutdownLock.Unlock()
+
+	wrapUpTime = defaultWrapUpTime
+	waitTime = defaultWaitTime
+}

core/proc/signals.go

@@ -1,4 +1,4 @@
-//go:build linux || darwin
+//go:build linux || darwin || freebsd
 
 package proc
 

core/proc/signals_test.go

@@ -1,4 +1,4 @@
-//go:build linux || darwin
+//go:build linux || darwin || freebsd
 
 package proc
 

core/service/serviceconf.go

@@ -37,6 +37,7 @@ type (
 		Prometheus prometheus.Config `json:",optional"`
 		Telemetry  trace.Config      `json:",optional"`
 		DevServer  DevServerConfig   `json:",optional"`
+		Shutdown   proc.ShutdownConf `json:",optional"`
 	}
 )
 
@@ -61,6 +62,7 @@ func (sc ServiceConf) SetUp() error {
 		sc.Telemetry.Name = sc.Name
 	}
 	trace.StartAgent(sc.Telemetry)
+	proc.Setup(sc.Shutdown)
 	proc.AddShutdownListener(func() {
 		trace.StopAgent()
 	})

core/service/servicegroup.go

@@ -76,9 +76,14 @@ func (sg *ServiceGroup) doStart() {
 }
 
 func (sg *ServiceGroup) doStop() {
+	group := threading.NewRoutineGroup()
 	for _, service := range sg.services {
-		service.Stop()
+		// new variable to avoid closure problems, can be removed after go 1.22
+		// see https://golang.org/doc/faq#closures_and_goroutines
+		service := service
+		group.Run(service.Stop)
 	}
+	group.Wait()
 }
 
 // WithStart wraps a start func as a Service.

core/stores/mon/clientmanager_test.go

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

core/stores/mon/collection.go

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

core/stores/mon/collection_test.go

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

core/stores/mon/model.go

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

core/stores/mon/trace.go

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

core/stores/redis/breakerhook.go

@@ -26,7 +26,7 @@ func (h breakerHook) ProcessHook(next red.ProcessHook) red.ProcessHook {
 			return next(ctx, cmd)
 		}
 
-		return h.brk.DoWithAcceptable(func() error {
+		return h.brk.DoWithAcceptableCtx(ctx, func() error {
 			return next(ctx, cmd)
 		}, acceptable)
 	}
@@ -34,7 +34,7 @@ func (h breakerHook) ProcessHook(next red.ProcessHook) red.ProcessHook {
 
 func (h breakerHook) ProcessPipelineHook(next red.ProcessPipelineHook) red.ProcessPipelineHook {
 	return func(ctx context.Context, cmds []red.Cmder) error {
-		return h.brk.DoWithAcceptable(func() error {
+		return h.brk.DoWithAcceptableCtx(ctx, func() error {
 			return next(ctx, cmds)
 		}, acceptable)
 	}

core/stores/redis/conf.go

@@ -19,6 +19,7 @@ type (
 	RedisConf struct {
 		Host     string
 		Type     string `json:",default=node,options=node|cluster"`
+		User     string `json:",optional"`
 		Pass     string `json:",optional"`
 		Tls      bool   `json:",optional"`
 		NonBlock bool   `json:",default=true"`
@@ -40,6 +41,9 @@ func (rc RedisConf) NewRedis() *Redis {
 	if rc.Type == ClusterType {
 		opts = append(opts, Cluster())
 	}
+	if len(rc.User) > 0 {
+		opts = append(opts, WithUser(rc.User))
+	}
 	if len(rc.Pass) > 0 {
 		opts = append(opts, WithPass(rc.Pass))
 	}

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

@@ -55,6 +55,7 @@ type (
 	Redis struct {
 		Addr  string
 		Type  string
+		User  string
 		Pass  string
 		tls   bool
 		brk   breaker.Breaker
@@ -90,6 +91,18 @@ type (
 	StringCmd = red.StringCmd
 	// Script is an alias of redis.Script.
 	Script = red.Script
+
+	// Hook is an alias of redis.Hook.
+	Hook = red.Hook
+	// DialHook is an alias of redis.DialHook.
+	DialHook = red.DialHook
+	// ProcessHook is an alias of redis.ProcessHook.
+	ProcessHook = red.ProcessHook
+	// ProcessPipelineHook is an alias of redis.ProcessPipelineHook.
+	ProcessPipelineHook = red.ProcessPipelineHook
+
+	// Cmder is an alias of redis.Cmder.
+	Cmder = red.Cmder
 )
 
 // MustNewRedis returns a Redis with given options.
@@ -114,6 +127,9 @@ func NewRedis(conf RedisConf, opts ...Option) (*Redis, error) {
 	if conf.Type == ClusterType {
 		opts = append([]Option{Cluster()}, opts...)
 	}
+	if len(conf.User) > 0 {
+		opts = append([]Option{WithUser(conf.User)}, opts...)
+	}
 	if len(conf.Pass) > 0 {
 		opts = append([]Option{WithPass(conf.Pass)}, opts...)
 	}
@@ -593,6 +609,28 @@ func (s *Redis) GetBitCtx(ctx context.Context, key string, offset int64) (int, e
 	return int(v), nil
 }
 
+// GetDel is the implementation of redis getdel command.
+// Available since: redis version 6.2.0
+func (s *Redis) GetDel(key string) (string, error) {
+	return s.GetDelCtx(context.Background(), key)
+}
+
+// GetDelCtx is the implementation of redis getdel command.
+// Available since: redis version 6.2.0
+func (s *Redis) GetDelCtx(ctx context.Context, key string) (string, error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return "", err
+	}
+
+	val, err := conn.GetDel(ctx, key).Result()
+	if errors.Is(err, red.Nil) {
+		return "", nil
+	}
+
+	return val, err
+}
+
 // GetSet is the implementation of redis getset command.
 func (s *Redis) GetSet(key, value string) (string, error) {
 	return s.GetSetCtx(context.Background(), key, value)
@@ -1185,6 +1223,18 @@ func (s *Redis) PipelinedCtx(ctx context.Context, fn func(Pipeliner) error) erro
 	return err
 }
 
+func (s *Redis) Publish(channel string, message interface{}) (int64, error) {
+	return s.PublishCtx(context.Background(), channel, message)
+}
+
+func (s *Redis) PublishCtx(ctx context.Context, channel string, message interface{}) (int64, error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return 0, err
+	}
+	return conn.Publish(ctx, channel, message).Result()
+}
+
 // Rpop is the implementation of redis rpop command.
 func (s *Redis) Rpop(key string) (string, error) {
 	return s.RpopCtx(context.Background(), key)
@@ -1235,6 +1285,18 @@ func (s *Redis) RpushCtx(ctx context.Context, key string, values ...any) (int, e
 	return int(v), nil
 }
 
+func (s *Redis) RPopLPush(source string, destination string) (string, error) {
+	return s.RPopLPushCtx(context.Background(), source, destination)
+}
+
+func (s *Redis) RPopLPushCtx(ctx context.Context, source string, destination string) (string, error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return "", err
+	}
+	return conn.RPopLPush(ctx, source, destination).Result()
+}
+
 // Sadd is the implementation of redis sadd command.
 func (s *Redis) Sadd(key string, values ...any) (int, error) {
 	return s.SaddCtx(context.Background(), key, values...)
@@ -1633,6 +1695,26 @@ func (s *Redis) TtlCtx(ctx context.Context, key string) (int, error) {
 	return int(duration), nil
 }
 
+func (s *Redis) TxPipeline() (pipe Pipeliner, err error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return nil, err
+	}
+	return conn.TxPipeline(), nil
+}
+
+func (s *Redis) Unlink(keys ...string) (int64, error) {
+	return s.UnlinkCtx(context.Background(), keys...)
+}
+
+func (s *Redis) UnlinkCtx(ctx context.Context, keys ...string) (int64, error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return 0, err
+	}
+	return conn.Unlink(ctx, keys...).Result()
+}
+
 // Zadd is the implementation of redis zadd command.
 func (s *Redis) Zadd(key string, score int64, value string) (bool, error) {
 	return s.ZaddCtx(context.Background(), key, score, value)
@@ -2349,6 +2431,13 @@ func SetSlowThreshold(threshold time.Duration) {
 	slowThreshold.Set(threshold)
 }
 
+// WithHook customizes the given Redis with given durationHook.
+func WithHook(hook Hook) Option {
+	return func(r *Redis) {
+		r.hooks = append(r.hooks, hook)
+	}
+}
+
 // WithPass customizes the given Redis with given password.
 func WithPass(pass string) Option {
 	return func(r *Redis) {
@@ -2363,16 +2452,15 @@ func WithTLS() Option {
 	}
 }
 
-// withHook customizes the given Redis with given durationHook, only for private use now,
+// WithUser customizes the given Redis with given username.
-// maybe expose later.
+func WithUser(user string) Option {
-func withHook(hook red.Hook) Option {
 	return func(r *Redis) {
-		r.hooks = append(r.hooks, hook)
+		r.User = user
 	}
 }
 
 func acceptable(err error) bool {
-	return err == nil || errors.Is(err, red.Nil) || errors.Is(err, context.Canceled)
+	return err == nil || errorx.In(err, red.Nil, context.Canceled)
 }
 
 func getRedis(r *Redis) (RedisNode, error) {

core/stores/redis/redis_test.go

@@ -160,7 +160,7 @@ func TestRedis_NonBlock(t *testing.T) {
 			Host:     s.Addr(),
 			NonBlock: true,
 			Type:     NodeType,
-		}, withHook(myHook{includePing: true}))
+		}, WithHook(myHook{includePing: true}))
 		assert.NoError(t, err)
 	})
 
@@ -170,7 +170,7 @@ func TestRedis_NonBlock(t *testing.T) {
 			Host:     s.Addr(),
 			NonBlock: false,
 			Type:     NodeType,
-		}, withHook(myHook{includePing: true}))
+		}, WithHook(myHook{includePing: true}))
 		assert.ErrorContains(t, err, "redis connect error")
 	})
 }
@@ -1071,6 +1071,34 @@ func TestRedis_Set(t *testing.T) {
 	})
 }
 
+func TestRedis_GetDel(t *testing.T) {
+	t.Run("get_del", func(t *testing.T) {
+		runOnRedis(t, func(client *Redis) {
+			val, err := newRedis(client.Addr).GetDel("hello")
+			assert.Equal(t, "", val)
+			assert.Nil(t, err)
+			err = client.Set("hello", "world")
+			assert.Nil(t, err)
+			val, err = client.Get("hello")
+			assert.Nil(t, err)
+			assert.Equal(t, "world", val)
+			val, err = client.GetDel("hello")
+			assert.Nil(t, err)
+			assert.Equal(t, "world", val)
+			val, err = client.Get("hello")
+			assert.Nil(t, err)
+			assert.Equal(t, "", val)
+		})
+	})
+
+	t.Run("get_del_with_error", func(t *testing.T) {
+		runOnRedisWithError(t, func(client *Redis) {
+			_, err := newRedis(client.Addr, badType()).GetDel("hello")
+			assert.Error(t, err)
+		})
+	})
+}
+
 func TestRedis_GetSet(t *testing.T) {
 	t.Run("set_get", func(t *testing.T) {
 		runOnRedis(t, func(client *Redis) {
@@ -1996,9 +2024,9 @@ func TestSetSlowThreshold(t *testing.T) {
 	assert.Equal(t, time.Second, slowThreshold.Load())
 }
 
-func TestRedis_WithPass(t *testing.T) {
+func TestRedis_WithUserPass(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
-		err := newRedis(client.Addr, WithPass("any")).Ping()
+		err := newRedis(client.Addr, WithUser("any"), WithPass("any")).Ping()
 		assert.NotNil(t, err)
 	})
 }
@@ -2080,3 +2108,70 @@ func (n mockedNode) BLPop(_ context.Context, _ time.Duration, _ ...string) *red.
 
 	return cmd
 }
+
+func TestRedisPublish(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		_, err := newRedis(client.Addr, badType()).Publish("Test", "message")
+		assert.NotNil(t, err)
+		_, err = client.Publish("Test", "message")
+		assert.Nil(t, err)
+	})
+}
+
+func TestRedisRPopLPush(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		_, err := newRedis(client.Addr, badType()).RPopLPush("Source", "Destination")
+		assert.NotNil(t, err)
+		_, err = client.Rpush("Source", "Destination")
+		assert.Nil(t, err)
+		_, err = client.RPopLPush("Source", "Destination")
+		assert.Nil(t, err)
+	})
+}
+
+func TestRedisUnlink(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		_, err := newRedis(client.Addr, badType()).Unlink("Key1", "Key2")
+		assert.NotNil(t, err)
+		err = client.Set("Key1", "Key2")
+		assert.Nil(t, err)
+		get, err := client.Get("Key1")
+		assert.Nil(t, err)
+		assert.Equal(t, "Key2", get)
+		res, err := client.Unlink("Key1")
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), res)
+	})
+}
+
+func TestRedisTxPipeline(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		ctx := context.Background()
+		_, err := newRedis(client.Addr, badType()).TxPipeline()
+		assert.NotNil(t, err)
+		pipe, err := client.TxPipeline()
+		assert.Nil(t, err)
+		key := "key"
+		hashKey := "field"
+		hashValue := "value"
+
+		// setting value
+		pipe.HSet(ctx, key, hashKey, hashValue)
+
+		existsCmd := pipe.Exists(ctx, key)
+		getCmd := pipe.HGet(ctx, key, hashKey)
+
+		// execution
+		_, err = pipe.Exec(ctx)
+		assert.Nil(t, err)
+
+		// verification results
+		exists, err := existsCmd.Result()
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), exists)
+
+		value, err := getCmd.Result()
+		assert.Nil(t, err)
+		assert.Equal(t, hashValue, value)
+	})
+}

core/stores/redis/redisblockingnode.go

@@ -21,6 +21,7 @@ func CreateBlockingNode(r *Redis) (ClosableNode, error) {
 	case NodeType:
 		client := red.NewClient(&red.Options{
 			Addr:         r.Addr,
+			Username:     r.User,
 			Password:     r.Pass,
 			DB:           defaultDatabase,
 			MaxRetries:   maxRetries,
@@ -32,6 +33,7 @@ func CreateBlockingNode(r *Redis) (ClosableNode, error) {
 	case ClusterType:
 		client := red.NewClusterClient(&red.ClusterOptions{
 			Addrs:        splitClusterAddrs(r.Addr),
+			Username:     r.User,
 			Password:     r.Pass,
 			MaxRetries:   maxRetries,
 			PoolSize:     1,

core/stores/redis/redisclientmanager.go

@@ -31,6 +31,7 @@ func getClient(r *Redis) (*red.Client, error) {
 		}
 		store := red.NewClient(&red.Options{
 			Addr:         r.Addr,
+			Username:     r.User,
 			Password:     r.Pass,
 			DB:           defaultDatabase,
 			MaxRetries:   maxRetries,

core/stores/redis/redisclustermanager.go

@@ -28,6 +28,7 @@ func getCluster(r *Redis) (*red.ClusterClient, error) {
 		}
 		store := red.NewClusterClient(&red.ClusterOptions{
 			Addrs:        splitClusterAddrs(r.Addr),
+			Username:     r.User,
 			Password:     r.Pass,
 			MaxRetries:   maxRetries,
 			MinIdleConns: idleConns,

core/stores/redis/redislock.go

@@ -2,6 +2,7 @@ package redis
 
 import (
 	"context"
+	_ "embed"
 	"errors"
 	"math/rand"
 	"strconv"
@@ -20,17 +21,13 @@ const (
 )
 
 var (
-	lockScript = NewScript(`if redis.call("GET", KEYS[1]) == ARGV[1] then
+	//go:embed lockscript.lua
-    redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2])
+	lockLuaScript string
-    return "OK"
+	lockScript    = NewScript(lockLuaScript)
-else
+
-    return redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2])
+	//go:embed delscript.lua
-end`)
+	delLuaScript string
-	delScript = NewScript(`if redis.call("GET", KEYS[1]) == ARGV[1] then
+	delScript    = NewScript(delLuaScript)
-    return redis.call("DEL", KEYS[1])
-else
-    return 0
-end`)
 )
 
 // A RedisLock is a redis lock.

core/stores/sqlc/cachedsql.go

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

core/stores/sqlc/cachedsql_test.go

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

core/stores/sqlx/metrics.go

@@ -41,6 +41,7 @@ var (
 
 type (
 	statGetter struct {
+		host      string
 		dbName    string
 		hash      string
 		poolStats func() sql.DBStats