fix(zrpc): fix slow threshold priority in stat interceptor (#5310)

Co-authored-by: qiuwenhao <qiushaotest@qq.com>

.codecov.yml

@@ -1,13 +0,0 @@
-coverage:
-  status:
-    patch: true
-    project: true
-comment:
-  layout: "flags, files"
-  behavior: once
-  require_changes: true
-ignore:
-  - "tools"
-  - "**/mock"
-  - "**/*_mock.go"
-  - "**/*test"

.github/copilot-instructions.md

@@ -0,0 +1,241 @@
+# GitHub Copilot Instructions for go-zero
+
+This document provides guidelines for GitHub Copilot when assisting with development in the go-zero project.
+
+## Project Overview
+
+go-zero is a web and RPC framework with lots of built-in engineering practices designed to ensure the stability of busy services with resilience design. It has been serving sites with tens of millions of users for years.
+
+### Key Architecture Components
+
+- **REST API framework** (`rest/`) - HTTP service framework with middleware chain support
+- **RPC framework** (`zrpc/`) - gRPC-based RPC framework with etcd service discovery and p2c_ewma load balancing
+- **Gateway** (`gateway/`) - API gateway supporting both HTTP and gRPC upstreams with proto-based routing
+- **MCP Server** (`mcp/`) - Model Context Protocol server for AI agent integration via SSE
+- **Core utilities** (`core/`) - Production-grade components:
+  - Resilience: circuit breakers (`breaker/`), rate limiters (`limit/`), adaptive load shedding (`load/`)
+  - Storage: SQL with cache (`stores/sqlc/`), Redis (`stores/redis/`), MongoDB (`stores/mongo/`)
+  - Concurrency: MapReduce (`mr/`), worker pools (`executors/`), sync primitives (`syncx/`)
+  - Observability: metrics (`metric/`), tracing (`trace/`), structured logging (`logx/`)
+- **Code generation tool** (`tools/goctl/`) - CLI tool for generating Go code from `.api` and `.proto` files
+
+## Coding Standards and Conventions
+
+### Code Style
+
+1. **Follow Go conventions**: Use `gofmt` for formatting, follow effective Go practices
+2. **Package naming**: Use lowercase, single-word package names when possible
+3. **Error handling**: Always handle errors explicitly, use `errorx.BatchError` for multiple errors
+4. **Context propagation**: Always pass `context.Context` as the first parameter for functions that may block
+5. **Configuration structures**: Use struct tags with JSON annotations, defaults, and validation
+
+**Pattern**: All service configs embed `service.ServiceConf` for common fields (Name, Log, Mode, Telemetry)
+```go
+type Config struct {
+    service.ServiceConf              // Always embed for services
+    Host     string `json:",default=0.0.0.0"`
+    Port     int    // Required field (no default)
+    Timeout  int64  `json:",default=3000"`  // Timeouts in milliseconds
+    Optional string `json:",optional"`      // Optional field
+    Mode     string `json:",default=pro,options=dev|test|rt|pre|pro"`  // Validated options
+}
+```
+
+**Service modes**: `dev`/`test`/`rt` disable load shedding and stats; `pre`/`pro` enable all resilience features
+
+### Interface Design
+
+1. **Small interfaces**: Follow Go's preference for small, focused interfaces
+2. **Context methods**: Provide both context and non-context versions of methods
+3. **Options pattern**: Use functional options for complex configuration
+
+Example:
+```go
+func (c *Client) Get(key string, val any) error {
+    return c.GetCtx(context.Background(), key, val)
+}
+
+func (c *Client) GetCtx(ctx context.Context, key string, val any) error {
+    // implementation
+}
+```
+
+### Testing Patterns
+
+1. **Test file naming**: Use `*_test.go` suffix
+2. **Test function naming**: Use `TestFunctionName` pattern
+3. **Use testify/assert**: Prefer `assert` package for assertions
+4. **Table-driven tests**: Use table-driven tests for multiple scenarios
+5. **Mock interfaces**: Use `go.uber.org/mock` for mocking
+6. **Test helpers**: Use `redistest`, `mongtest` helpers for database testing
+
+Example test pattern:
+```go
+func TestSomething(t *testing.T) {
+    tests := []struct {
+        name     string
+        input    string
+        expected string
+        wantErr  bool
+    }{
+        {"valid case", "input", "output", false},
+        {"error case", "bad", "", true},
+    }
+
+    for _, tt := range tests {
+        t.Run(tt.name, func(t *testing.T) {
+            result, err := SomeFunction(tt.input)
+            if tt.wantErr {
+                assert.Error(t, err)
+                return
+            }
+            assert.NoError(t, err)
+            assert.Equal(t, tt.expected, result)
+        })
+    }
+}
+```
+
+## Framework-Specific Guidelines
+
+### REST API Development
+
+1. **API Definition**: Use `.api` files to define REST APIs with goctl codegen
+2. **Handler pattern**: Separate business logic into logic packages (handlers call logic layer)
+3. **Middleware chain**: Middlewares wrap via `chain.Chain` interface - use `Append()` or `Prepend()` to control order
+   - Built-in middlewares (all in `rest/handler/`): tracing, logging, metrics, recovery, breaker, shedding, timeout, maxconns, maxbytes, gunzip
+   - Custom middleware: `func(http.Handler) http.Handler` - call `next.ServeHTTP(w, r)` to continue chain
+4. **Response handling**: Use `httpx.WriteJson(w, code, v)` for JSON responses
+5. **Error handling**: Use `httpx.Error(w, err)` or `httpx.ErrorCtx(ctx, w, err)` for HTTP error responses
+6. **Route registration**: Routes defined with `Method`, `Path`, and `Handler` - wildcards use `:param` syntax
+
+### RPC Development
+
+1. **Protocol Buffers**: Use protobuf for service definitions, generate code with goctl
+2. **Service discovery**: Use etcd for dynamic service registration/discovery, or direct endpoints for static routing
+3. **Load balancing**: Default is `p2c_ewma` (power of 2 choices with EWMA), configurable via `BalancerName`
+4. **Client configuration**: Support `Etcd`, `Endpoints`, or `Target` - use `BuildTarget()` to construct connection string
+5. **Interceptors**: Implement gRPC interceptors for cross-cutting concerns (auth, logging, metrics)
+6. **Health checks**: gRPC health checks enabled by default (`Health: true`)
+
+### Database Operations
+
+1. **SQL operations**: Use `sqlx.SqlConn` interface - methods always end with `Ctx` for context support
+2. **Caching pattern**: `stores/sqlc` provides `CachedConn` for automatic cache-aside pattern
+   - `QueryRowCtx`: Query with cache key, auto-populate on cache miss
+   - `ExecCtx`: Execute and delete cache keys
+3. **Transactions**: Use `sqlx.SqlConn.TransactCtx()` to get transaction session
+4. **Connection pooling**: Managed automatically (64 max idle/open, 1min lifetime)
+5. **Test helpers**: Use `redistest.CreateRedis(t)` for Redis, SQL mocks for DB testing
+
+Example cache pattern:
+```go
+err := c.QueryRowCtx(ctx, &dest, key, func(ctx context.Context, conn sqlx.SqlConn) error {
+    return conn.QueryRowCtx(ctx, &dest, query, args...)
+})
+```
+
+### Configuration Management
+
+1. **YAML configuration**: Use YAML for configuration files
+2. **Environment variables**: Support environment variable overrides
+3. **Validation**: Include proper validation for configuration parameters
+4. **Sensible defaults**: Provide reasonable default values
+
+## Error Handling Best Practices
+
+1. **Wrap errors**: Use `fmt.Errorf` with `%w` verb to wrap errors
+2. **Custom errors**: Define custom error types when needed
+3. **Error logging**: Log errors appropriately with context
+4. **Graceful degradation**: Implement fallback mechanisms
+
+## Performance Considerations
+
+1. **Resource pools**: Use connection pools and worker pools
+2. **Circuit breakers**: Implement circuit breaker patterns for external calls
+3. **Rate limiting**: Apply rate limiting to protect services
+4. **Load shedding**: Implement adaptive load shedding
+5. **Metrics**: Add appropriate metrics and monitoring
+
+## Security Guidelines
+
+1. **Input validation**: Validate all input parameters
+2. **SQL injection prevention**: Use parameterized queries
+3. **Authentication**: Implement proper JWT token handling
+4. **HTTPS**: Support TLS/HTTPS configurations
+5. **CORS**: Configure CORS appropriately for web APIs
+
+## Documentation Standards
+
+1. **Package documentation**: Include package-level documentation
+2. **Function documentation**: Document exported functions with examples
+3. **API documentation**: Maintain API documentation in sync
+4. **README updates**: Update README for significant changes
+
+## Common Patterns to Follow
+
+### Service Configuration
+```go
+type ServiceConf struct {
+    Name string
+    Log  logx.LogConf
+    Mode string `json:",default=pro,options=[dev,test,pre,pro]"`
+    // ... other common fields
+}
+```
+
+### Middleware Implementation
+```go
+func SomeMiddleware() rest.Middleware {
+    return func(next http.HandlerFunc) http.HandlerFunc {
+        return func(w http.ResponseWriter, r *http.Request) {
+            // Pre-processing
+            next.ServeHTTP(w, r)
+            // Post-processing
+        }
+    }
+}
+```
+
+### Resource Management
+Always implement proper resource cleanup using defer and context cancellation.
+
+## Build and Test Commands
+
+- Build: `go build ./...`
+- Test: `go test ./...`  
+- Test with race detection: `go test -race ./...`
+- Format: `gofmt -w .`
+- Code generation:
+  - REST API: `goctl api go -api *.api -dir .`
+  - RPC: `goctl rpc protoc *.proto --go_out=. --go-grpc_out=. --zrpc_out=.`
+  - Model from SQL: `goctl model mysql datasource -url="user:pass@tcp(host:port)/db" -table="*" -dir="./model"`
+
+## Critical Architecture Patterns
+
+### Resilience Design Philosophy
+go-zero implements defense-in-depth with multiple protection layers:
+1. **Circuit Breaker** (`core/breaker`): Google SRE breaker - tracks success/failure, opens on error threshold
+2. **Adaptive Load Shedding** (`core/load`): CPU-based auto-rejection when system overloaded (disabled in dev/test/rt modes)
+3. **Rate Limiting** (`core/limit`): Token bucket (Redis-based) and period limiters
+4. **Timeout Control**: Cascading timeouts via context - set at multiple levels (client, server, handler)
+
+### Middleware Chain Architecture
+`rest/chain` provides middleware composition:
+```go
+// Middleware signature
+type Middleware func(http.Handler) http.Handler
+
+// Chain operations
+chain := chain.New(m1, m2)
+chain.Append(m3)    // Adds to end: m1 -> m2 -> m3
+chain.Prepend(m0)   // Adds to start: m0 -> m1 -> m2 -> m3
+handler := chain.Then(finalHandler)
+```
+
+### Concurrency Patterns
+- **MapReduce** (`core/mr`): Parallel processing with worker pools - use for batch operations
+- **Executors** (`core/executors`): Bulk/period executors for batching operations
+- **SingleFlight** (`core/syncx`): Deduplicates concurrent identical requests
+
+Remember to run tests and ensure all checks pass before submitting changes. The project emphasizes high quality, performance, and reliability, so these should be primary considerations in all development work.

.github/workflows/codeql-analysis.yml

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

.github/workflows/go.yml

@@ -12,10 +12,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Check out code into the Go module directory
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Set up Go 1.x
-        uses: actions/setup-go@v5
+        uses: actions/setup-go@v6
         with:
           go-version-file: go.mod
           check-latest: true
@@ -41,16 +41,21 @@ jobs:
 
       - name: Codecov
         uses: codecov/codecov-action@v5
+        with:
+          files: ./coverage.txt
+          flags: unittests
+          name: codecov-umbrella
+          fail_ci_if_error: false
 
   test-win:
     name: Windows
     runs-on: windows-latest
     steps:
       - name: Checkout codebase
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Set up Go 1.x
-        uses: actions/setup-go@v5
+        uses: actions/setup-go@v6
         with:
           # make sure Go version compatible with go-zero
           go-version-file: go.mod

.github/workflows/issues.yml

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

.github/workflows/release.yaml

@@ -16,7 +16,7 @@ jobs:
           - goarch: "386"
             goos: darwin
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
       - uses: zeromicro/go-zero-release-action@master
         with:
           github_token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/reviewdog.yml

@@ -5,7 +5,12 @@ jobs:
     name: runner / staticcheck
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
+      - uses: actions/setup-go@v6
+        with:
+          go-version-file: go.mod
+          check-latest: true
+          cache: true
       - uses: reviewdog/action-staticcheck@v1
         with:
           github_token: ${{ secrets.github_token }}

.github/workflows/version-check.yml

@@ -10,10 +10,10 @@ jobs:
   version-check:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
 
       - name: Set up Go
-        uses: actions/setup-go@v5
+        uses: actions/setup-go@v6
         with:
           go-version: '1.21'
 

.gitignore

@@ -17,6 +17,7 @@
 **/logs
 **/adhoc
 **/coverage.txt
+**/WARP.md
 
 # for test purpose
 go.work

core/bloom/bloom.go

@@ -40,7 +40,7 @@ type (
 	}
 )
 
-// New create a Filter, store is the backed redis, key is the key for the bloom filter,
+// New creates a Filter, store is the backed redis, key is the key for the bloom filter,
 // bits is how many bits will be used, maps is how many hashes for each addition.
 // best practices:
 // elements - means how many actual elements

core/collection/cache.go

@@ -81,6 +81,10 @@ func (c *Cache) Del(key string) {
 	delete(c.data, key)
 	c.lruCache.remove(key)
 	c.lock.Unlock()
+
+	// RemoveTimer is called outside the lock to avoid performance impact from this
+	// potentially time-consuming operation. Data integrity is maintained by lruCache,
+	// which will eventually evict any remaining entries when capacity is exceeded.
 	c.timingWheel.RemoveTimer(key)
 }
 

core/collection/set.go

@@ -1,235 +1,53 @@
 package collection
 
-import (
-	"github.com/zeromicro/go-zero/core/lang"
-	"github.com/zeromicro/go-zero/core/logx"
-)
+import "github.com/zeromicro/go-zero/core/lang"
 
-const (
-	unmanaged = iota
-	untyped
-	intType
-	int64Type
-	uintType
-	uint64Type
-	stringType
-)
-
-// Set is not thread-safe, for concurrent use, make sure to use it with synchronization.
-type Set struct {
-	data map[any]lang.PlaceholderType
-	tp   int
+// Set is a type-safe generic set collection.
+// It's not thread-safe, use with synchronization for concurrent access.
+type Set[T comparable] struct {
+	data map[T]lang.PlaceholderType
 }
 
-// NewSet returns a managed Set, can only put the values with the same type.
-func NewSet() *Set {
-	return &Set{
-		data: make(map[any]lang.PlaceholderType),
-		tp:   untyped,
+// NewSet returns a new type-safe set.
+func NewSet[T comparable]() *Set[T] {
+	return &Set[T]{
+		data: make(map[T]lang.PlaceholderType),
 	}
 }
 
-// NewUnmanagedSet returns an unmanaged Set, which can put values with different types.
-func NewUnmanagedSet() *Set {
-	return &Set{
-		data: make(map[any]lang.PlaceholderType),
-		tp:   unmanaged,
+// Add adds items to the set. Duplicates are automatically ignored.
+func (s *Set[T]) Add(items ...T) {
+	for _, item := range items {
+		s.data[item] = lang.Placeholder
 	}
 }
 
-// Add adds i into s.
-func (s *Set) Add(i ...any) {
-	for _, each := range i {
-		s.add(each)
-	}
+// Clear removes all items from the set.
+func (s *Set[T]) Clear() {
+	clear(s.data)
 }
 
-// AddInt adds int values ii into s.
-func (s *Set) AddInt(ii ...int) {
-	for _, each := range ii {
-		s.add(each)
-	}
-}
-
-// AddInt64 adds int64 values ii into s.
-func (s *Set) AddInt64(ii ...int64) {
-	for _, each := range ii {
-		s.add(each)
-	}
-}
-
-// AddUint adds uint values ii into s.
-func (s *Set) AddUint(ii ...uint) {
-	for _, each := range ii {
-		s.add(each)
-	}
-}
-
-// AddUint64 adds uint64 values ii into s.
-func (s *Set) AddUint64(ii ...uint64) {
-	for _, each := range ii {
-		s.add(each)
-	}
-}
-
-// AddStr adds string values ss into s.
-func (s *Set) AddStr(ss ...string) {
-	for _, each := range ss {
-		s.add(each)
-	}
-}
-
-// Contains checks if i is in s.
-func (s *Set) Contains(i any) bool {
-	if len(s.data) == 0 {
-		return false
-	}
-
-	s.validate(i)
-	_, ok := s.data[i]
+// Contains checks if an item exists in the set.
+func (s *Set[T]) Contains(item T) bool {
+	_, ok := s.data[item]
 	return ok
 }
 
-// Keys returns the keys in s.
-func (s *Set) Keys() []any {
-	var keys []any
-
-	for key := range s.data {
-		keys = append(keys, key)
-	}
-
-	return keys
-}
-
-// KeysInt returns the int keys in s.
-func (s *Set) KeysInt() []int {
-	var keys []int
-
-	for key := range s.data {
-		if intKey, ok := key.(int); ok {
-			keys = append(keys, intKey)
-		}
-	}
-
-	return keys
-}
-
-// KeysInt64 returns int64 keys in s.
-func (s *Set) KeysInt64() []int64 {
-	var keys []int64
-
-	for key := range s.data {
-		if intKey, ok := key.(int64); ok {
-			keys = append(keys, intKey)
-		}
-	}
-
-	return keys
-}
-
-// KeysUint returns uint keys in s.
-func (s *Set) KeysUint() []uint {
-	var keys []uint
-
-	for key := range s.data {
-		if intKey, ok := key.(uint); ok {
-			keys = append(keys, intKey)
-		}
-	}
-
-	return keys
-}
-
-// KeysUint64 returns uint64 keys in s.
-func (s *Set) KeysUint64() []uint64 {
-	var keys []uint64
-
-	for key := range s.data {
-		if intKey, ok := key.(uint64); ok {
-			keys = append(keys, intKey)
-		}
-	}
-
-	return keys
-}
-
-// KeysStr returns string keys in s.
-func (s *Set) KeysStr() []string {
-	var keys []string
-
-	for key := range s.data {
-		if strKey, ok := key.(string); ok {
-			keys = append(keys, strKey)
-		}
-	}
-
-	return keys
-}
-
-// Remove removes i from s.
-func (s *Set) Remove(i any) {
-	s.validate(i)
-	delete(s.data, i)
-}
-
-// Count returns the number of items in s.
-func (s *Set) Count() int {
+// Count returns the number of items in the set.
+func (s *Set[T]) Count() int {
 	return len(s.data)
 }
 
-func (s *Set) add(i any) {
-	switch s.tp {
-	case unmanaged:
-		// do nothing
-	case untyped:
-		s.setType(i)
-	default:
-		s.validate(i)
+// Keys returns all elements in the set as a slice.
+func (s *Set[T]) Keys() []T {
+	keys := make([]T, 0, len(s.data))
+	for key := range s.data {
+		keys = append(keys, key)
 	}
-	s.data[i] = lang.Placeholder
+	return keys
 }
 
-func (s *Set) setType(i any) {
-	// s.tp can only be untyped here
-	switch i.(type) {
-	case int:
-		s.tp = intType
-	case int64:
-		s.tp = int64Type
-	case uint:
-		s.tp = uintType
-	case uint64:
-		s.tp = uint64Type
-	case string:
-		s.tp = stringType
-	}
-}
-
-func (s *Set) validate(i any) {
-	if s.tp == unmanaged {
-		return
-	}
-
-	switch i.(type) {
-	case int:
-		if s.tp != intType {
-			logx.Errorf("element is int, but set contains elements with type %d", s.tp)
-		}
-	case int64:
-		if s.tp != int64Type {
-			logx.Errorf("element is int64, but set contains elements with type %d", s.tp)
-		}
-	case uint:
-		if s.tp != uintType {
-			logx.Errorf("element is uint, but set contains elements with type %d", s.tp)
-		}
-	case uint64:
-		if s.tp != uint64Type {
-			logx.Errorf("element is uint64, but set contains elements with type %d", s.tp)
-		}
-	case string:
-		if s.tp != stringType {
-			logx.Errorf("element is string, but set contains elements with type %d", s.tp)
-		}
-	}
+// Remove removes an item from the set.
+func (s *Set[T]) Remove(item T) {
+	delete(s.data, item)
 }

core/collection/set_test.go

@@ -12,6 +12,105 @@ func init() {
 	logx.Disable()
 }
 
+// Set functionality tests
+func TestTypedSetInt(t *testing.T) {
+	set := NewSet[int]()
+	values := []int{1, 2, 3, 2, 1} // Contains duplicates
+
+	// Test adding
+	set.Add(values...)
+	assert.Equal(t, 3, set.Count()) // Should only have 3 elements after deduplication
+
+	// Test contains
+	assert.True(t, set.Contains(1))
+	assert.True(t, set.Contains(2))
+	assert.True(t, set.Contains(3))
+	assert.False(t, set.Contains(4))
+
+	// Test getting all keys
+	keys := set.Keys()
+	sort.Ints(keys)
+	assert.EqualValues(t, []int{1, 2, 3}, keys)
+
+	// Test removal
+	set.Remove(2)
+	assert.False(t, set.Contains(2))
+	assert.Equal(t, 2, set.Count())
+}
+
+func TestTypedSetStringOps(t *testing.T) {
+	set := NewSet[string]()
+	values := []string{"a", "b", "c", "b", "a"}
+
+	set.Add(values...)
+	assert.Equal(t, 3, set.Count())
+
+	assert.True(t, set.Contains("a"))
+	assert.True(t, set.Contains("b"))
+	assert.True(t, set.Contains("c"))
+	assert.False(t, set.Contains("d"))
+
+	keys := set.Keys()
+	sort.Strings(keys)
+	assert.EqualValues(t, []string{"a", "b", "c"}, keys)
+}
+
+func TestTypedSetClear(t *testing.T) {
+	set := NewSet[int]()
+	set.Add(1, 2, 3)
+	assert.Equal(t, 3, set.Count())
+
+	set.Clear()
+	assert.Equal(t, 0, set.Count())
+	assert.False(t, set.Contains(1))
+}
+
+func TestTypedSetEmpty(t *testing.T) {
+	set := NewSet[int]()
+	assert.Equal(t, 0, set.Count())
+	assert.False(t, set.Contains(1))
+	assert.Empty(t, set.Keys())
+}
+
+func TestTypedSetMultipleTypes(t *testing.T) {
+	// Test different typed generic sets
+	intSet := NewSet[int]()
+	int64Set := NewSet[int64]()
+	uintSet := NewSet[uint]()
+	uint64Set := NewSet[uint64]()
+	stringSet := NewSet[string]()
+
+	intSet.Add(1, 2, 3)
+	int64Set.Add(1, 2, 3)
+	uintSet.Add(1, 2, 3)
+	uint64Set.Add(1, 2, 3)
+	stringSet.Add("1", "2", "3")
+
+	assert.Equal(t, 3, intSet.Count())
+	assert.Equal(t, 3, int64Set.Count())
+	assert.Equal(t, 3, uintSet.Count())
+	assert.Equal(t, 3, uint64Set.Count())
+	assert.Equal(t, 3, stringSet.Count())
+}
+
+// Set benchmarks
+func BenchmarkTypedIntSet(b *testing.B) {
+	s := NewSet[int]()
+	for i := 0; i < b.N; i++ {
+		s.Add(i)
+		_ = s.Contains(i)
+	}
+}
+
+func BenchmarkTypedStringSet(b *testing.B) {
+	s := NewSet[string]()
+	for i := 0; i < b.N; i++ {
+		s.Add(string(rune(i)))
+		_ = s.Contains(string(rune(i)))
+	}
+}
+
+// Legacy tests remain unchanged for backward compatibility
 func BenchmarkRawSet(b *testing.B) {
 	m := make(map[any]struct{})
 	for i := 0; i < b.N; i++ {
@@ -20,26 +119,10 @@ func BenchmarkRawSet(b *testing.B) {
 	}
 }
 
-func BenchmarkUnmanagedSet(b *testing.B) {
-	s := NewUnmanagedSet()
-	for i := 0; i < b.N; i++ {
-		s.Add(i)
-		_ = s.Contains(i)
-	}
-}
-
-func BenchmarkSet(b *testing.B) {
-	s := NewSet()
-	for i := 0; i < b.N; i++ {
-		s.AddInt(i)
-		_ = s.Contains(i)
-	}
-}
-
 func TestAdd(t *testing.T) {
 	// given
-	set := NewUnmanagedSet()
-	values := []any{1, 2, 3}
+	set := NewSet[int]()
+	values := []int{1, 2, 3}
 
 	// when
 	set.Add(values...)
@@ -51,82 +134,74 @@ func TestAdd(t *testing.T) {
 
 func TestAddInt(t *testing.T) {
 	// given
-	set := NewSet()
+	set := NewSet[int]()
 	values := []int{1, 2, 3}
 
 	// when
-	set.AddInt(values...)
+	set.Add(values...)
 
 	// then
 	assert.True(t, set.Contains(1) && set.Contains(2) && set.Contains(3))
-	keys := set.KeysInt()
+	keys := set.Keys()
 	sort.Ints(keys)
 	assert.EqualValues(t, values, keys)
 }
 
 func TestAddInt64(t *testing.T) {
 	// given
-	set := NewSet()
+	set := NewSet[int64]()
 	values := []int64{1, 2, 3}
 
 	// when
-	set.AddInt64(values...)
+	set.Add(values...)
 
 	// then
-	assert.True(t, set.Contains(int64(1)) && set.Contains(int64(2)) && set.Contains(int64(3)))
-	assert.Equal(t, len(values), len(set.KeysInt64()))
+	assert.True(t, set.Contains(1) && set.Contains(2) && set.Contains(3))
+	assert.Equal(t, len(values), len(set.Keys()))
 }
 
 func TestAddUint(t *testing.T) {
 	// given
-	set := NewSet()
+	set := NewSet[uint]()
 	values := []uint{1, 2, 3}
 
 	// when
-	set.AddUint(values...)
+	set.Add(values...)
 
 	// then
-	assert.True(t, set.Contains(uint(1)) && set.Contains(uint(2)) && set.Contains(uint(3)))
-	assert.Equal(t, len(values), len(set.KeysUint()))
+	assert.True(t, set.Contains(1) && set.Contains(2) && set.Contains(3))
+	assert.Equal(t, len(values), len(set.Keys()))
 }
 
 func TestAddUint64(t *testing.T) {
 	// given
-	set := NewSet()
+	set := NewSet[uint64]()
 	values := []uint64{1, 2, 3}
 
 	// when
-	set.AddUint64(values...)
+	set.Add(values...)
 
 	// then
-	assert.True(t, set.Contains(uint64(1)) && set.Contains(uint64(2)) && set.Contains(uint64(3)))
-	assert.Equal(t, len(values), len(set.KeysUint64()))
+	assert.True(t, set.Contains(1) && set.Contains(2) && set.Contains(3))
+	assert.Equal(t, len(values), len(set.Keys()))
 }
 
 func TestAddStr(t *testing.T) {
 	// given
-	set := NewSet()
+	set := NewSet[string]()
 	values := []string{"1", "2", "3"}
 
 	// when
-	set.AddStr(values...)
+	set.Add(values...)
 
 	// then
 	assert.True(t, set.Contains("1") && set.Contains("2") && set.Contains("3"))
-	assert.Equal(t, len(values), len(set.KeysStr()))
+	assert.Equal(t, len(values), len(set.Keys()))
 }
 
 func TestContainsWithoutElements(t *testing.T) {
 	// given
-	set := NewSet()
-
-	// then
-	assert.False(t, set.Contains(1))
-}
-
-func TestContainsUnmanagedWithoutElements(t *testing.T) {
-	// given
-	set := NewUnmanagedSet()
+	set := NewSet[int]()
 
 	// then
 	assert.False(t, set.Contains(1))
@@ -134,8 +209,8 @@ func TestContainsUnmanagedWithoutElements(t *testing.T) {
 
 func TestRemove(t *testing.T) {
 	// given
-	set := NewSet()
-	set.Add([]any{1, 2, 3}...)
+	set := NewSet[int]()
+	set.Add([]int{1, 2, 3}...)
 
 	// when
 	set.Remove(2)
@@ -146,57 +221,9 @@ func TestRemove(t *testing.T) {
 
 func TestCount(t *testing.T) {
 	// given
-	set := NewSet()
-	set.Add([]any{1, 2, 3}...)
+	set := NewSet[int]()
+	set.Add([]int{1, 2, 3}...)
 
 	// then
 	assert.Equal(t, set.Count(), 3)
 }
-
-func TestKeysIntMismatch(t *testing.T) {
-	set := NewSet()
-	set.add(int64(1))
-	set.add(2)
-	vals := set.KeysInt()
-	assert.EqualValues(t, []int{2}, vals)
-}
-
-func TestKeysInt64Mismatch(t *testing.T) {
-	set := NewSet()
-	set.add(1)
-	set.add(int64(2))
-	vals := set.KeysInt64()
-	assert.EqualValues(t, []int64{2}, vals)
-}
-
-func TestKeysUintMismatch(t *testing.T) {
-	set := NewSet()
-	set.add(1)
-	set.add(uint(2))
-	vals := set.KeysUint()
-	assert.EqualValues(t, []uint{2}, vals)
-}
-
-func TestKeysUint64Mismatch(t *testing.T) {
-	set := NewSet()
-	set.add(1)
-	set.add(uint64(2))
-	vals := set.KeysUint64()
-	assert.EqualValues(t, []uint64{2}, vals)
-}
-
-func TestKeysStrMismatch(t *testing.T) {
-	set := NewSet()
-	set.add(1)
-	set.add("2")
-	vals := set.KeysStr()
-	assert.EqualValues(t, []string{"2"}, vals)
-}
-
-func TestSetType(t *testing.T) {
-	set := NewUnmanagedSet()
-	set.add(1)
-	set.add("2")
-	vals := set.Keys()
-	assert.ElementsMatch(t, []any{1, "2"}, vals)
-}

core/collection/timingwheel.go

@@ -164,6 +164,7 @@ func (tw *TimingWheel) Stop() {
 
 func (tw *TimingWheel) drainAll(fn func(key, value any)) {
 	runner := threading.NewTaskRunner(drainWorkers)
+
 	for _, slot := range tw.slots {
 		for e := slot.Front(); e != nil; {
 			task := e.Value.(*timingEntry)
@@ -177,6 +178,8 @@ func (tw *TimingWheel) drainAll(fn func(key, value any)) {
 			}
 		}
 	}
+
+	runner.Wait()
 }
 
 func (tw *TimingWheel) getPositionAndCircle(d time.Duration) (pos, circle int) {

core/collection/timingwheel_test.go

@@ -629,6 +629,157 @@ func TestMoveAndRemoveTask(t *testing.T) {
 	assert.Equal(t, 0, len(keys))
 }
 
+// TestTimingWheel_DrainClosureBug tests the closure capture bug in drainAll
+// Issue: https://github.com/zeromicro/go-zero/issues/5314
+func TestTimingWheel_DrainClosureBug(t *testing.T) {
+	ticker := timex.NewFakeTicker()
+	tw, _ := NewTimingWheelWithTicker(testStep, 10, func(k, v any) {}, ticker)
+	defer tw.Stop()
+
+	// Set multiple timers with different values
+	for i := 0; i < 10; i++ {
+		tw.SetTimer(i, i*10, testStep*5)
+	}
+
+	// Give time for timers to be set
+	time.Sleep(time.Millisecond * 100)
+
+	var mu sync.Mutex
+	received := make(map[int]int)
+	var wg sync.WaitGroup
+	wg.Add(10)
+
+	tw.Drain(func(key, value any) {
+		mu.Lock()
+		defer mu.Unlock()
+		k := key.(int)
+		v := value.(int)
+		received[k] = v
+		wg.Done()
+	})
+
+	wg.Wait()
+
+	// Check if all values match their keys
+	for k, v := range received {
+		expected := k * 10
+		assert.Equal(t, expected, v, "key %d should have value %d, got %d", k, expected, v)
+	}
+}
+
+// TestTimingWheel_RunTasksClosureBug tests the closure capture bug in runTasks
+// Issue: https://github.com/zeromicro/go-zero/issues/5314
+func TestTimingWheel_RunTasksClosureBug(t *testing.T) {
+	ticker := timex.NewFakeTicker()
+	var mu sync.Mutex
+	executed := make(map[int]int)
+	var wg sync.WaitGroup
+
+	tw, _ := NewTimingWheelWithTicker(testStep, 10, func(k, v any) {
+		mu.Lock()
+		defer mu.Unlock()
+		key := k.(int)
+		val := v.(int)
+		executed[key] = val
+		wg.Done()
+	}, ticker)
+	defer tw.Stop()
+
+	// Set multiple timers that should fire in the same tick
+	count := 10
+	wg.Add(count)
+	for i := 0; i < count; i++ {
+		tw.SetTimer(i, i*10, testStep)
+	}
+
+	// Advance ticker to trigger tasks
+	ticker.Tick()
+
+	// Wait for execution with timeout
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		// Success
+	case <-time.After(2 * time.Second):
+		t.Fatal("timeout waiting for tasks to execute")
+	}
+
+	// Verify all tasks executed with correct values
+	assert.Equal(t, count, len(executed), "should have executed all tasks")
+	for k, v := range executed {
+		expected := k * 10
+		assert.Equal(t, expected, v, "key %d should have value %d, got %d", k, expected, v)
+	}
+}
+
+// TestTimingWheel_RunTasksRaceCondition tests for race conditions in runTasks
+// This test specifically targets the loop variable capture bug
+func TestTimingWheel_RunTasksRaceCondition(t *testing.T) {
+	// Run multiple times to increase likelihood of catching the bug
+	for attempt := 0; attempt < 10; attempt++ {
+		t.Run("", func(t *testing.T) {
+			ticker := timex.NewFakeTicker()
+			var mu sync.Mutex
+			keyValues := make(map[int][]int)
+			var wg sync.WaitGroup
+
+			tw, _ := NewTimingWheelWithTicker(testStep, 10, func(k, v any) {
+				// Add small delay to increase chance of race
+				time.Sleep(time.Microsecond)
+				mu.Lock()
+				defer mu.Unlock()
+				key := k.(int)
+				val := v.(int)
+				keyValues[key] = append(keyValues[key], val)
+				wg.Done()
+			}, ticker)
+			defer tw.Stop()
+
+			// Set many timers rapidly to increase chance of race
+			count := 50
+			wg.Add(count)
+			for i := 0; i < count; i++ {
+				tw.SetTimer(i, i*100, testStep)
+			}
+
+			ticker.Tick()
+
+			done := make(chan struct{})
+			go func() {
+				wg.Wait()
+				close(done)
+			}()
+
+			select {
+			case <-done:
+			case <-time.After(5 * time.Second):
+				t.Fatal("timeout waiting for tasks")
+			}
+
+			// Check for duplicates or wrong values
+			wrongCount := 0
+			for key, values := range keyValues {
+				assert.Equal(t, 1, len(values), "key %d should only execute once, got %v", key, values)
+				if len(values) > 0 {
+					expected := key * 100
+					if values[0] != expected {
+						wrongCount++
+						t.Logf("BUG DETECTED: key %d should have value %d, got %d", key, expected, values[0])
+					}
+				}
+			}
+			if wrongCount > 0 {
+				t.Errorf("Found %d tasks with wrong values due to closure bug", wrongCount)
+			}
+		})
+	}
+}
+
 func BenchmarkTimingWheel(b *testing.B) {
 	b.ReportAllocs()
 

core/conf/config.go

@@ -316,7 +316,7 @@ func toLowerCaseInterface(v any, info *fieldInfo) any {
 	case map[string]any:
 		return toLowerCaseKeyMap(vv, info)
 	case []any:
-		var arr []any
+		arr := make([]any, 0, len(vv))
 		for _, vvv := range vv {
 			arr = append(arr, toLowerCaseInterface(vvv, info))
 		}
@@ -368,5 +368,5 @@ func getFullName(parent, child string) string {
 		return child
 	}
 
-	return strings.Join([]string{parent, child}, ".")
+	return parent + "." + child
 }

core/conf/config_test.go

@@ -1377,3 +1377,23 @@ func (m mockConfig) Validate() error {
 
 	return nil
 }
+
+func TestGetFullName(t *testing.T) {
+	tests := []struct {
+		parent string
+		child  string
+		want   string
+	}{
+		{"", "child", "child"},
+		{"parent", "child", "parent.child"},
+		{"a.b", "c", "a.b.c"},
+		{"root", "nested.field", "root.nested.field"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.parent+"."+tt.child, func(t *testing.T) {
+			got := getFullName(tt.parent, tt.child)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}

core/configcenter/subscriber/etcd.go

@@ -1,6 +1,9 @@
 package subscriber
 
 import (
+	"sync"
+	"sync/atomic"
+
 	"github.com/zeromicro/go-zero/core/discov"
 	"github.com/zeromicro/go-zero/core/logx"
 )
@@ -37,6 +40,7 @@ func NewEtcdSubscriber(conf EtcdConf) (Subscriber, error) {
 func buildSubOptions(conf EtcdConf) []discov.SubOption {
 	opts := []discov.SubOption{
 		discov.WithExactMatch(),
+		discov.WithContainer(newContainer()),
 	}
 
 	if len(conf.User) > 0 {
@@ -65,3 +69,47 @@ func (s *etcdSubscriber) Value() (string, error) {
 
 	return "", nil
 }
+
+type container struct {
+	value     atomic.Value
+	listeners []func()
+	lock      sync.Mutex
+}
+
+func newContainer() *container {
+	return &container{}
+}
+
+func (c *container) OnAdd(kv discov.KV) {
+	c.value.Store([]string{kv.Val})
+	c.notifyChange()
+}
+
+func (c *container) OnDelete(_ discov.KV) {
+	c.value.Store([]string(nil))
+	c.notifyChange()
+}
+
+func (c *container) AddListener(listener func()) {
+	c.lock.Lock()
+	c.listeners = append(c.listeners, listener)
+	c.lock.Unlock()
+}
+
+func (c *container) GetValues() []string {
+	if vals, ok := c.value.Load().([]string); ok {
+		return vals
+	}
+
+	return []string(nil)
+}
+
+func (c *container) notifyChange() {
+	c.lock.Lock()
+	listeners := append(([]func())(nil), c.listeners...)
+	c.lock.Unlock()
+
+	for _, listener := range listeners {
+		listener()
+	}
+}

core/configcenter/subscriber/etcd_test.go

@@ -0,0 +1,186 @@
+package subscriber
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/discov"
+)
+
+const (
+	actionAdd = iota
+	actionDel
+)
+
+func TestConfigCenterContainer(t *testing.T) {
+	type action struct {
+		act int
+		key string
+		val string
+	}
+	tests := []struct {
+		name   string
+		do     []action
+		expect []string
+	}{
+		{
+			name: "add one",
+			do: []action{
+				{
+					act: actionAdd,
+					key: "first",
+					val: "a",
+				},
+			},
+			expect: []string{
+				"a",
+			},
+		},
+		{
+			name: "add two",
+			do: []action{
+				{
+					act: actionAdd,
+					key: "first",
+					val: "a",
+				},
+				{
+					act: actionAdd,
+					key: "second",
+					val: "b",
+				},
+			},
+			expect: []string{
+				"b",
+			},
+		},
+		{
+			name: "add two, delete one",
+			do: []action{
+				{
+					act: actionAdd,
+					key: "first",
+					val: "a",
+				},
+				{
+					act: actionAdd,
+					key: "second",
+					val: "b",
+				},
+				{
+					act: actionDel,
+					key: "first",
+				},
+			},
+			expect: []string(nil),
+		},
+		{
+			name: "add two, delete two",
+			do: []action{
+				{
+					act: actionAdd,
+					key: "first",
+					val: "a",
+				},
+				{
+					act: actionAdd,
+					key: "second",
+					val: "b",
+				},
+				{
+					act: actionDel,
+					key: "first",
+				},
+				{
+					act: actionDel,
+					key: "second",
+				},
+			},
+			expect: []string(nil),
+		},
+		{
+			name: "add two, dup values",
+			do: []action{
+				{
+					act: actionAdd,
+					key: "first",
+					val: "a",
+				},
+				{
+					act: actionAdd,
+					key: "second",
+					val: "b",
+				},
+				{
+					act: actionAdd,
+					key: "third",
+					val: "a",
+				},
+			},
+			expect: []string{"a"},
+		},
+		{
+			name: "add three, dup values, delete two, add one",
+			do: []action{
+				{
+					act: actionAdd,
+					key: "first",
+					val: "a",
+				},
+				{
+					act: actionAdd,
+					key: "second",
+					val: "b",
+				},
+				{
+					act: actionAdd,
+					key: "third",
+					val: "a",
+				},
+				{
+					act: actionDel,
+					key: "first",
+				},
+				{
+					act: actionDel,
+					key: "second",
+				},
+				{
+					act: actionAdd,
+					key: "forth",
+					val: "c",
+				},
+			},
+			expect: []string{"c"},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var changed bool
+			c := newContainer()
+			c.AddListener(func() {
+				changed = true
+			})
+			assert.Nil(t, c.GetValues())
+			assert.False(t, changed)
+
+			for _, order := range test.do {
+				if order.act == actionAdd {
+					c.OnAdd(discov.KV{
+						Key: order.key,
+						Val: order.val,
+					})
+				} else {
+					c.OnDelete(discov.KV{
+						Key: order.key,
+						Val: order.val,
+					})
+				}
+			}
+
+			assert.True(t, changed)
+			assert.ElementsMatch(t, test.expect, c.GetValues())
+		})
+	}
+}

core/discov/internal/etcdclient_mock.go

@@ -1,5 +1,10 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: etcdclient.go
+//
+// Generated by this command:
+//
+//	mockgen -package internal -destination etcdclient_mock.go -source etcdclient.go EtcdClient
+//
 
 // Package internal is a generated GoMock package.
 package internal
@@ -8,35 +13,36 @@ import (
 	context "context"
 	reflect "reflect"
 
-	gomock "github.com/golang/mock/gomock"
 	clientv3 "go.etcd.io/etcd/client/v3"
+	gomock "go.uber.org/mock/gomock"
 	grpc "google.golang.org/grpc"
 )
 
-// MockEtcdClient is a mock of EtcdClient interface
+// MockEtcdClient is a mock of EtcdClient interface.
 type MockEtcdClient struct {
 	ctrl     *gomock.Controller
 	recorder *MockEtcdClientMockRecorder
+	isgomock struct{}
 }
 
-// MockEtcdClientMockRecorder is the mock recorder for MockEtcdClient
+// MockEtcdClientMockRecorder is the mock recorder for MockEtcdClient.
 type MockEtcdClientMockRecorder struct {
 	mock *MockEtcdClient
 }
 
-// NewMockEtcdClient creates a new mock instance
+// NewMockEtcdClient creates a new mock instance.
 func NewMockEtcdClient(ctrl *gomock.Controller) *MockEtcdClient {
 	mock := &MockEtcdClient{ctrl: ctrl}
 	mock.recorder = &MockEtcdClientMockRecorder{mock}
 	return mock
 }
 
-// EXPECT returns an object that allows the caller to indicate expected use
+// EXPECT returns an object that allows the caller to indicate expected use.
 func (m *MockEtcdClient) EXPECT() *MockEtcdClientMockRecorder {
 	return m.recorder
 }
 
-// ActiveConnection mocks base method
+// ActiveConnection mocks base method.
 func (m *MockEtcdClient) ActiveConnection() *grpc.ClientConn {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "ActiveConnection")
@@ -44,13 +50,13 @@ func (m *MockEtcdClient) ActiveConnection() *grpc.ClientConn {
 	return ret0
 }
 
-// ActiveConnection indicates an expected call of ActiveConnection
+// ActiveConnection indicates an expected call of ActiveConnection.
 func (mr *MockEtcdClientMockRecorder) ActiveConnection() *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ActiveConnection", reflect.TypeOf((*MockEtcdClient)(nil).ActiveConnection))
 }
 
-// Close mocks base method
+// Close mocks base method.
 func (m *MockEtcdClient) Close() error {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "Close")
@@ -58,13 +64,13 @@ func (m *MockEtcdClient) Close() error {
 	return ret0
 }
 
-// Close indicates an expected call of Close
+// Close indicates an expected call of Close.
 func (mr *MockEtcdClientMockRecorder) Close() *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockEtcdClient)(nil).Close))
 }
 
-// Ctx mocks base method
+// Ctx mocks base method.
 func (m *MockEtcdClient) Ctx() context.Context {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "Ctx")
@@ -72,13 +78,13 @@ func (m *MockEtcdClient) Ctx() context.Context {
 	return ret0
 }
 
-// Ctx indicates an expected call of Ctx
+// Ctx indicates an expected call of Ctx.
 func (mr *MockEtcdClientMockRecorder) Ctx() *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Ctx", reflect.TypeOf((*MockEtcdClient)(nil).Ctx))
 }
 
-// Get mocks base method
+// Get mocks base method.
 func (m *MockEtcdClient) Get(ctx context.Context, key string, opts ...clientv3.OpOption) (*clientv3.GetResponse, error) {
 	m.ctrl.T.Helper()
 	varargs := []any{ctx, key}
@@ -91,14 +97,14 @@ func (m *MockEtcdClient) Get(ctx context.Context, key string, opts ...clientv3.O
 	return ret0, ret1
 }
 
-// Get indicates an expected call of Get
+// Get indicates an expected call of Get.
 func (mr *MockEtcdClientMockRecorder) Get(ctx, key any, opts ...any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	varargs := append([]any{ctx, key}, opts...)
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockEtcdClient)(nil).Get), varargs...)
 }
 
-// Grant mocks base method
+// Grant mocks base method.
 func (m *MockEtcdClient) Grant(ctx context.Context, ttl int64) (*clientv3.LeaseGrantResponse, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "Grant", ctx, ttl)
@@ -107,13 +113,13 @@ func (m *MockEtcdClient) Grant(ctx context.Context, ttl int64) (*clientv3.LeaseG
 	return ret0, ret1
 }
 
-// Grant indicates an expected call of Grant
+// Grant indicates an expected call of Grant.
 func (mr *MockEtcdClientMockRecorder) Grant(ctx, ttl any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Grant", reflect.TypeOf((*MockEtcdClient)(nil).Grant), ctx, ttl)
 }
 
-// KeepAlive mocks base method
+// KeepAlive mocks base method.
 func (m *MockEtcdClient) KeepAlive(ctx context.Context, id clientv3.LeaseID) (<-chan *clientv3.LeaseKeepAliveResponse, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "KeepAlive", ctx, id)
@@ -122,13 +128,13 @@ func (m *MockEtcdClient) KeepAlive(ctx context.Context, id clientv3.LeaseID) (<-
 	return ret0, ret1
 }
 
-// KeepAlive indicates an expected call of KeepAlive
+// KeepAlive indicates an expected call of KeepAlive.
 func (mr *MockEtcdClientMockRecorder) KeepAlive(ctx, id any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "KeepAlive", reflect.TypeOf((*MockEtcdClient)(nil).KeepAlive), ctx, id)
 }
 
-// Put mocks base method
+// Put mocks base method.
 func (m *MockEtcdClient) Put(ctx context.Context, key, val string, opts ...clientv3.OpOption) (*clientv3.PutResponse, error) {
 	m.ctrl.T.Helper()
 	varargs := []any{ctx, key, val}
@@ -141,14 +147,14 @@ func (m *MockEtcdClient) Put(ctx context.Context, key, val string, opts ...clien
 	return ret0, ret1
 }
 
-// Put indicates an expected call of Put
+// Put indicates an expected call of Put.
 func (mr *MockEtcdClientMockRecorder) Put(ctx, key, val any, opts ...any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	varargs := append([]any{ctx, key, val}, opts...)
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Put", reflect.TypeOf((*MockEtcdClient)(nil).Put), varargs...)
 }
 
-// Revoke mocks base method
+// Revoke mocks base method.
 func (m *MockEtcdClient) Revoke(ctx context.Context, id clientv3.LeaseID) (*clientv3.LeaseRevokeResponse, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "Revoke", ctx, id)
@@ -157,13 +163,13 @@ func (m *MockEtcdClient) Revoke(ctx context.Context, id clientv3.LeaseID) (*clie
 	return ret0, ret1
 }
 
-// Revoke indicates an expected call of Revoke
+// Revoke indicates an expected call of Revoke.
 func (mr *MockEtcdClientMockRecorder) Revoke(ctx, id any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Revoke", reflect.TypeOf((*MockEtcdClient)(nil).Revoke), ctx, id)
 }
 
-// Watch mocks base method
+// Watch mocks base method.
 func (m *MockEtcdClient) Watch(ctx context.Context, key string, opts ...clientv3.OpOption) clientv3.WatchChan {
 	m.ctrl.T.Helper()
 	varargs := []any{ctx, key}
@@ -175,7 +181,7 @@ func (m *MockEtcdClient) Watch(ctx context.Context, key string, opts ...clientv3
 	return ret0
 }
 
-// Watch indicates an expected call of Watch
+// Watch indicates an expected call of Watch.
 func (mr *MockEtcdClientMockRecorder) Watch(ctx, key any, opts ...any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	varargs := append([]any{ctx, key}, opts...)

core/discov/internal/registry.go

@@ -207,7 +207,7 @@ func (c *cluster) getCurrent(key watchKey) []KV {
 		return nil
 	}
 
-	var kvs []KV
+	kvs := make([]KV, 0, len(watcher.values))
 	for k, v := range watcher.values {
 		kvs = append(kvs, KV{
 			Key: k,
@@ -308,7 +308,7 @@ func (c *cluster) load(cli EtcdClient, key watchKey) int64 {
 		time.Sleep(coolDownUnstable.AroundDuration(coolDownInterval))
 	}
 
-	var kvs []KV
+	kvs := make([]KV, 0, len(resp.Kvs))
 	for _, ev := range resp.Kvs {
 		kvs = append(kvs, KV{
 			Key: string(ev.Key),
@@ -352,7 +352,7 @@ func (c *cluster) reload(cli EtcdClient) {
 	// cancel the previous watches
 	close(c.done)
 	c.watchGroup.Wait()
-	var keys []watchKey
+	keys := make([]watchKey, 0, len(c.watchers))
 	for wk, wval := range c.watchers {
 		keys = append(keys, wk)
 		if wval.cancel != nil {
@@ -386,8 +386,9 @@ func (c *cluster) watch(cli EtcdClient, key watchKey, rev int64) {
 			rev = c.load(cli, key)
 		}
 
-		// log the error and retry
+		// log the error and retry with cooldown to prevent CPU/disk exhaustion
 		logc.Error(cli.Ctx(), err)
+		time.Sleep(coolDownUnstable.AroundDuration(coolDownInterval))
 	}
 }
 

core/discov/internal/registry_test.go

@@ -7,7 +7,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/contextx"
 	"github.com/zeromicro/go-zero/core/lang"
@@ -18,6 +17,7 @@ import (
 	"go.etcd.io/etcd/api/v3/mvccpb"
 	clientv3 "go.etcd.io/etcd/client/v3"
 	"go.etcd.io/etcd/client/v3/mock/mockserver"
+	"go.uber.org/mock/gomock"
 )
 
 var mockLock sync.Mutex

core/discov/internal/statewatcher_mock.go

@@ -1,5 +1,10 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: statewatcher.go
+//
+// Generated by this command:
+//
+//	mockgen -package internal -destination statewatcher_mock.go -source statewatcher.go etcdConn
+//
 
 // Package internal is a generated GoMock package.
 package internal
@@ -8,34 +13,35 @@ import (
 	context "context"
 	reflect "reflect"
 
-	gomock "github.com/golang/mock/gomock"
+	gomock "go.uber.org/mock/gomock"
 	connectivity "google.golang.org/grpc/connectivity"
 )
 
-// MocketcdConn is a mock of etcdConn interface
+// MocketcdConn is a mock of etcdConn interface.
 type MocketcdConn struct {
 	ctrl     *gomock.Controller
 	recorder *MocketcdConnMockRecorder
+	isgomock struct{}
 }
 
-// MocketcdConnMockRecorder is the mock recorder for MocketcdConn
+// MocketcdConnMockRecorder is the mock recorder for MocketcdConn.
 type MocketcdConnMockRecorder struct {
 	mock *MocketcdConn
 }
 
-// NewMocketcdConn creates a new mock instance
+// NewMocketcdConn creates a new mock instance.
 func NewMocketcdConn(ctrl *gomock.Controller) *MocketcdConn {
 	mock := &MocketcdConn{ctrl: ctrl}
 	mock.recorder = &MocketcdConnMockRecorder{mock}
 	return mock
 }
 
-// EXPECT returns an object that allows the caller to indicate expected use
+// EXPECT returns an object that allows the caller to indicate expected use.
 func (m *MocketcdConn) EXPECT() *MocketcdConnMockRecorder {
 	return m.recorder
 }
 
-// GetState mocks base method
+// GetState mocks base method.
 func (m *MocketcdConn) GetState() connectivity.State {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "GetState")
@@ -43,13 +49,13 @@ func (m *MocketcdConn) GetState() connectivity.State {
 	return ret0
 }
 
-// GetState indicates an expected call of GetState
+// GetState indicates an expected call of GetState.
 func (mr *MocketcdConnMockRecorder) GetState() *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetState", reflect.TypeOf((*MocketcdConn)(nil).GetState))
 }
 
-// WaitForStateChange mocks base method
+// WaitForStateChange mocks base method.
 func (m *MocketcdConn) WaitForStateChange(ctx context.Context, sourceState connectivity.State) bool {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "WaitForStateChange", ctx, sourceState)
@@ -57,7 +63,7 @@ func (m *MocketcdConn) WaitForStateChange(ctx context.Context, sourceState conne
 	return ret0
 }
 
-// WaitForStateChange indicates an expected call of WaitForStateChange
+// WaitForStateChange indicates an expected call of WaitForStateChange.
 func (mr *MocketcdConnMockRecorder) WaitForStateChange(ctx, sourceState any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WaitForStateChange", reflect.TypeOf((*MocketcdConn)(nil).WaitForStateChange), ctx, sourceState)

core/discov/internal/statewatcher_test.go

@@ -4,7 +4,7 @@ import (
 	"sync"
 	"testing"
 
-	"github.com/golang/mock/gomock"
+	"go.uber.org/mock/gomock"
 	"google.golang.org/grpc/connectivity"
 )
 

core/discov/internal/updatelistener_mock.go

@@ -1,5 +1,10 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: updatelistener.go
+//
+// Generated by this command:
+//
+//	mockgen -package internal -destination updatelistener_mock.go -source updatelistener.go UpdateListener
+//
 
 // Package internal is a generated GoMock package.
 package internal
@@ -7,51 +12,52 @@ package internal
 import (
 	reflect "reflect"
 
-	gomock "github.com/golang/mock/gomock"
+	gomock "go.uber.org/mock/gomock"
 )
 
-// MockUpdateListener is a mock of UpdateListener interface
+// MockUpdateListener is a mock of UpdateListener interface.
 type MockUpdateListener struct {
 	ctrl     *gomock.Controller
 	recorder *MockUpdateListenerMockRecorder
+	isgomock struct{}
 }
 
-// MockUpdateListenerMockRecorder is the mock recorder for MockUpdateListener
+// MockUpdateListenerMockRecorder is the mock recorder for MockUpdateListener.
 type MockUpdateListenerMockRecorder struct {
 	mock *MockUpdateListener
 }
 
-// NewMockUpdateListener creates a new mock instance
+// NewMockUpdateListener creates a new mock instance.
 func NewMockUpdateListener(ctrl *gomock.Controller) *MockUpdateListener {
 	mock := &MockUpdateListener{ctrl: ctrl}
 	mock.recorder = &MockUpdateListenerMockRecorder{mock}
 	return mock
 }
 
-// EXPECT returns an object that allows the caller to indicate expected use
+// EXPECT returns an object that allows the caller to indicate expected use.
 func (m *MockUpdateListener) EXPECT() *MockUpdateListenerMockRecorder {
 	return m.recorder
 }
 
-// OnAdd mocks base method
+// OnAdd mocks base method.
 func (m *MockUpdateListener) OnAdd(kv KV) {
 	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "OnAdd", kv)
 }
 
-// OnAdd indicates an expected call of OnAdd
+// OnAdd indicates an expected call of OnAdd.
 func (mr *MockUpdateListenerMockRecorder) OnAdd(kv any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnAdd", reflect.TypeOf((*MockUpdateListener)(nil).OnAdd), kv)
 }
 
-// OnDelete mocks base method
+// OnDelete mocks base method.
 func (m *MockUpdateListener) OnDelete(kv KV) {
 	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "OnDelete", kv)
 }
 
-// OnDelete indicates an expected call of OnDelete
+// OnDelete indicates an expected call of OnDelete.
 func (mr *MockUpdateListenerMockRecorder) OnDelete(kv any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnDelete", reflect.TypeOf((*MockUpdateListener)(nil).OnDelete), kv)

core/discov/publisher.go

@@ -92,12 +92,12 @@ func (p *Publisher) doKeepAlive() error {
 		default:
 			cli, err := p.doRegister()
 			if err != nil {
-				logc.Errorf(cli.Ctx(), "etcd publisher doRegister: %s", err.Error())
+				logc.Errorf(cli.Ctx(), "etcd publisher doRegister: %v", err)
 				break
 			}
 
 			if err := p.keepAliveAsync(cli); err != nil {
-				logc.Errorf(cli.Ctx(), "etcd publisher keepAliveAsync: %s", err.Error())
+				logc.Errorf(cli.Ctx(), "etcd publisher keepAliveAsync: %v", err)
 				break
 			}
 
@@ -125,23 +125,48 @@ func (p *Publisher) keepAliveAsync(cli internal.EtcdClient) error {
 	}
 
 	threading.GoSafe(func() {
+		wch := cli.Watch(cli.Ctx(), p.fullKey, clientv3.WithFilterPut())
+
 		for {
 			select {
 			case _, ok := <-ch:
 				if !ok {
 					p.revoke(cli)
 					if err := p.doKeepAlive(); err != nil {
-						logc.Errorf(cli.Ctx(), "etcd publisher KeepAlive: %s", err.Error())
+						logc.Errorf(cli.Ctx(), "etcd publisher KeepAlive: %v", err)
 					}
 					return
 				}
+
+			case c := <-wch:
+				if c.Err() != nil {
+					logc.Errorf(cli.Ctx(), "etcd publisher watch: %v", c.Err())
+					if err := p.doKeepAlive(); err != nil {
+						logc.Errorf(cli.Ctx(), "etcd publisher KeepAlive: %v", err)
+					}
+					return
+				}
+
+				for _, evt := range c.Events {
+					if evt.Type == clientv3.EventTypeDelete {
+						logc.Infof(cli.Ctx(), "etcd publisher watch: %s, event: %v",
+							evt.Kv.Key, evt.Type)
+						_, err := cli.Put(cli.Ctx(), p.fullKey, p.value, clientv3.WithLease(p.lease))
+						if err != nil {
+							logc.Errorf(cli.Ctx(), "etcd publisher re-put key: %v", err)
+						} else {
+							logc.Infof(cli.Ctx(), "etcd publisher re-put key: %s, value: %s",
+								p.fullKey, p.value)
+						}
+					}
+				}
 			case <-p.pauseChan:
 				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 {
-						logc.Errorf(cli.Ctx(), "etcd publisher KeepAlive: %s", err.Error())
+						logc.Errorf(cli.Ctx(), "etcd publisher KeepAlive: %v", err)
 					}
 					return
 				case <-p.quit.Done():
@@ -176,7 +201,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 {
-		logc.Errorf(cli.Ctx(), "etcd publisher revoke: %s", err.Error())
+		logc.Errorf(cli.Ctx(), "etcd publisher revoke: %v", err)
 	}
 }
 

core/discov/publisher_test.go

@@ -9,13 +9,14 @@ import (
 	"testing"
 	"time"
 
-	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/discov/internal"
 	"github.com/zeromicro/go-zero/core/lang"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stringx"
+	"go.etcd.io/etcd/api/v3/mvccpb"
 	clientv3 "go.etcd.io/etcd/client/v3"
+	"go.uber.org/mock/gomock"
 	"golang.org/x/net/http2"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
@@ -211,6 +212,9 @@ func TestPublisher_keepAliveAsyncQuit(t *testing.T) {
 	defer restore()
 	cli.EXPECT().Ctx().AnyTimes()
 	cli.EXPECT().KeepAlive(gomock.Any(), id)
+	// Add Watch expectation for the new watch mechanism
+	watchChan := make(<-chan clientv3.WatchResponse)
+	cli.EXPECT().Watch(gomock.Any(), gomock.Any(), gomock.Any()).Return(watchChan)
 	var wg sync.WaitGroup
 	wg.Add(1)
 	cli.EXPECT().Revoke(gomock.Any(), id).Do(func(_, _ any) {
@@ -232,6 +236,9 @@ func TestPublisher_keepAliveAsyncPause(t *testing.T) {
 	defer restore()
 	cli.EXPECT().Ctx().AnyTimes()
 	cli.EXPECT().KeepAlive(gomock.Any(), id)
+	// Add Watch expectation for the new watch mechanism
+	watchChan := make(<-chan clientv3.WatchResponse)
+	cli.EXPECT().Watch(gomock.Any(), gomock.Any(), gomock.Any()).Return(watchChan)
 	pub := NewPublisher(nil, "thekey", "thevalue")
 	var wg sync.WaitGroup
 	wg.Add(1)
@@ -245,6 +252,112 @@ func TestPublisher_keepAliveAsyncPause(t *testing.T) {
 	wg.Wait()
 }
 
+// Test case for key deletion and re-registration (covers lines 148-155)
+func TestPublisher_keepAliveAsyncKeyDeletion(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	const id clientv3.LeaseID = 1
+	cli := internal.NewMockEtcdClient(ctrl)
+	restore := setMockClient(cli)
+	defer restore()
+	cli.EXPECT().Ctx().AnyTimes()
+	cli.EXPECT().KeepAlive(gomock.Any(), id)
+
+	// Create a watch channel that will send a delete event
+	watchChan := make(chan clientv3.WatchResponse, 1)
+	watchResp := clientv3.WatchResponse{
+		Events: []*clientv3.Event{{
+			Type: clientv3.EventTypeDelete,
+			Kv: &mvccpb.KeyValue{
+				Key: []byte("thekey"),
+			},
+		}},
+	}
+	watchChan <- watchResp
+
+	cli.EXPECT().Watch(gomock.Any(), gomock.Any(), gomock.Any()).Return((<-chan clientv3.WatchResponse)(watchChan))
+
+	var wg sync.WaitGroup
+	wg.Add(1) // Only wait for Revoke call
+
+	// Use a channel to signal when Put has been called
+	putCalled := make(chan struct{})
+
+	// Expect the re-put operation when key is deleted
+	cli.EXPECT().Put(gomock.Any(), "thekey", "thevalue", gomock.Any()).Do(func(_, _, _, _ any) {
+		close(putCalled) // Signal that Put has been called
+	}).Return(nil, nil)
+
+	// Expect revoke when Stop is called
+	cli.EXPECT().Revoke(gomock.Any(), id).Do(func(_, _ any) {
+		wg.Done()
+	})
+
+	pub := NewPublisher(nil, "thekey", "thevalue")
+	pub.lease = id
+	pub.fullKey = "thekey"
+
+	assert.Nil(t, pub.keepAliveAsync(cli))
+
+	// Wait for Put to be called, then stop
+	<-putCalled
+	pub.Stop()
+	wg.Wait()
+}
+
+// Test case for key deletion with re-put error (covers error branch in lines 151-152)
+func TestPublisher_keepAliveAsyncKeyDeletionPutError(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	const id clientv3.LeaseID = 1
+	cli := internal.NewMockEtcdClient(ctrl)
+	restore := setMockClient(cli)
+	defer restore()
+	cli.EXPECT().Ctx().AnyTimes()
+	cli.EXPECT().KeepAlive(gomock.Any(), id)
+
+	// Create a watch channel that will send a delete event
+	watchChan := make(chan clientv3.WatchResponse, 1)
+	watchResp := clientv3.WatchResponse{
+		Events: []*clientv3.Event{{
+			Type: clientv3.EventTypeDelete,
+			Kv: &mvccpb.KeyValue{
+				Key: []byte("thekey"),
+			},
+		}},
+	}
+	watchChan <- watchResp
+
+	cli.EXPECT().Watch(gomock.Any(), gomock.Any(), gomock.Any()).Return((<-chan clientv3.WatchResponse)(watchChan))
+
+	var wg sync.WaitGroup
+	wg.Add(1) // Only wait for Revoke call
+
+	// Use a channel to signal when Put has been called
+	putCalled := make(chan struct{})
+
+	// Expect the re-put operation to fail
+	cli.EXPECT().Put(gomock.Any(), "thekey", "thevalue", gomock.Any()).Do(func(_, _, _, _ any) {
+		close(putCalled) // Signal that Put has been called
+	}).Return(nil, errors.New("put error"))
+
+	// Expect revoke when Stop is called
+	cli.EXPECT().Revoke(gomock.Any(), id).Do(func(_, _ any) {
+		wg.Done()
+	})
+
+	pub := NewPublisher(nil, "thekey", "thevalue")
+	pub.lease = id
+	pub.fullKey = "thekey"
+
+	assert.Nil(t, pub.keepAliveAsync(cli))
+
+	// Wait for Put to be called, then stop
+	<-putCalled
+	pub.Stop()
+	wg.Wait()
+}
+
 func TestPublisher_Resume(t *testing.T) {
 	publisher := new(Publisher)
 	publisher.resumeChan = make(chan lang.PlaceholderType)
@@ -273,6 +386,9 @@ func TestPublisher_keepAliveAsync(t *testing.T) {
 	defer restore()
 	cli.EXPECT().Ctx().AnyTimes()
 	cli.EXPECT().KeepAlive(gomock.Any(), id)
+	// Add Watch expectation for the new watch mechanism
+	watchChan := make(<-chan clientv3.WatchResponse)
+	cli.EXPECT().Watch(gomock.Any(), gomock.Any(), gomock.Any()).Return(watchChan)
 	cli.EXPECT().Grant(gomock.Any(), timeToLive).Return(&clientv3.LeaseGrantResponse{
 		ID: 1,
 	}, nil)

core/discov/subscriber.go

@@ -19,8 +19,9 @@ type (
 		exclusive  bool
 		key        string
 		exactMatch bool
-		items      *container
+		items      Container
 	}
+	KV = internal.KV
 )
 
 // NewSubscriber returns a Subscriber.
@@ -35,7 +36,9 @@ func NewSubscriber(endpoints []string, key string, opts ...SubOption) (*Subscrib
 	for _, opt := range opts {
 		opt(sub)
 	}
-	sub.items = newContainer(sub.exclusive)
+	if sub.items == nil {
+		sub.items = newContainer(sub.exclusive)
+	}
 
 	if err := internal.GetRegistry().Monitor(endpoints, key, sub.exactMatch, sub.items); err != nil {
 		return nil, err
@@ -46,7 +49,7 @@ func NewSubscriber(endpoints []string, key string, opts ...SubOption) (*Subscrib
 
 // AddListener adds listener to s.
 func (s *Subscriber) AddListener(listener func()) {
-	s.items.addListener(listener)
+	s.items.AddListener(listener)
 }
 
 // Close closes the subscriber.
@@ -56,7 +59,7 @@ func (s *Subscriber) Close() {
 
 // Values returns all the subscription values.
 func (s *Subscriber) Values() []string {
-	return s.items.getValues()
+	return s.items.GetValues()
 }
 
 // Exclusive means that key value can only be 1:1,
@@ -88,16 +91,32 @@ func WithSubEtcdTLS(certFile, certKeyFile, caFile string, insecureSkipVerify boo
 	}
 }
 
-type container struct {
-	exclusive bool
-	values    map[string][]string
-	mapping   map[string]string
-	snapshot  atomic.Value
-	dirty     *syncx.AtomicBool
-	listeners []func()
-	lock      sync.Mutex
+// WithContainer provides a custom container to the subscriber.
+func WithContainer(container Container) SubOption {
+	return func(sub *Subscriber) {
+		sub.items = container
+	}
 }
 
+type (
+	Container interface {
+		OnAdd(kv internal.KV)
+		OnDelete(kv internal.KV)
+		AddListener(listener func())
+		GetValues() []string
+	}
+
+	container struct {
+		exclusive bool
+		values    map[string][]string
+		mapping   map[string]string
+		snapshot  atomic.Value
+		dirty     *syncx.AtomicBool
+		listeners []func()
+		lock      sync.Mutex
+	}
+)
+
 func newContainer(exclusive bool) *container {
 	return &container{
 		exclusive: exclusive,
@@ -141,7 +160,7 @@ func (c *container) addKv(key, value string) ([]string, bool) {
 	return nil, false
 }
 
-func (c *container) addListener(listener func()) {
+func (c *container) AddListener(listener func()) {
 	c.lock.Lock()
 	c.listeners = append(c.listeners, listener)
 	c.lock.Unlock()
@@ -170,7 +189,7 @@ func (c *container) doRemoveKey(key string) {
 	}
 }
 
-func (c *container) getValues() []string {
+func (c *container) GetValues() []string {
 	if !c.dirty.True() {
 		return c.snapshot.Load().([]string)
 	}

core/discov/subscriber_test.go

@@ -171,10 +171,10 @@ func TestContainer(t *testing.T) {
 			t.Run(test.name, func(t *testing.T) {
 				var changed bool
 				c := newContainer(exclusive)
-				c.addListener(func() {
+				c.AddListener(func() {
 					changed = true
 				})
-				assert.Nil(t, c.getValues())
+				assert.Nil(t, c.GetValues())
 				assert.False(t, changed)
 
 				for _, order := range test.do {
@@ -193,9 +193,9 @@ func TestContainer(t *testing.T) {
 
 				assert.True(t, changed)
 				assert.True(t, c.dirty.True())
-				assert.ElementsMatch(t, test.expect, c.getValues())
+				assert.ElementsMatch(t, test.expect, c.GetValues())
 				assert.False(t, c.dirty.True())
-				assert.ElementsMatch(t, test.expect, c.getValues())
+				assert.ElementsMatch(t, test.expect, c.GetValues())
 			})
 		}
 	}
@@ -204,12 +204,14 @@ func TestContainer(t *testing.T) {
 func TestSubscriber(t *testing.T) {
 	sub := new(Subscriber)
 	Exclusive()(sub)
-	sub.items = newContainer(sub.exclusive)
+	c := newContainer(sub.exclusive)
+	WithContainer(c)(sub)
+	sub.items = c
 	var count int32
 	sub.AddListener(func() {
 		atomic.AddInt32(&count, 1)
 	})
-	sub.items.notifyChange()
+	c.notifyChange()
 	assert.Empty(t, sub.Values())
 	assert.Equal(t, int32(1), atomic.LoadInt32(&count))
 }
@@ -229,12 +231,13 @@ func TestWithSubEtcdAccount(t *testing.T) {
 func TestWithExactMatch(t *testing.T) {
 	sub := new(Subscriber)
 	WithExactMatch()(sub)
-	sub.items = newContainer(sub.exclusive)
+	c := newContainer(sub.exclusive)
+	sub.items = c
 	var count int32
 	sub.AddListener(func() {
 		atomic.AddInt32(&count, 1)
 	})
-	sub.items.notifyChange()
+	c.notifyChange()
 	assert.Empty(t, sub.Values())
 	assert.Equal(t, int32(1), atomic.LoadInt32(&count))
 }

core/fx/stream.go

@@ -168,7 +168,7 @@ func (s Stream) Count() (count int) {
 	return
 }
 
-// Distinct removes the duplicated items base on the given KeyFunc.
+// Distinct removes the duplicated items based on the given KeyFunc.
 func (s Stream) Distinct(fn KeyFunc) Stream {
 	source := make(chan any)
 
@@ -459,7 +459,7 @@ func (s Stream) Tail(n int64) Stream {
 	return Range(source)
 }
 
-// Walk lets the callers handle each item, the caller may write zero, one or more items base on the given item.
+// Walk lets the callers handle each item, the caller may write zero, one or more items based on the given item.
 func (s Stream) Walk(fn WalkFunc, opts ...Option) Stream {
 	option := buildOptions(opts...)
 	if option.unlimitedWorkers {

core/fx/stream_test.go

@@ -1,8 +1,6 @@
 package fx
 
 import (
-	"io"
-	"log"
 	"math/rand"
 	"reflect"
 	"runtime"
@@ -13,6 +11,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/logx/logtest"
 	"github.com/zeromicro/go-zero/core/stringx"
 	"go.uber.org/goleak"
 )
@@ -238,7 +237,7 @@ func TestLast(t *testing.T) {
 
 func TestMap(t *testing.T) {
 	runCheckedTest(t, func(t *testing.T) {
-		log.SetOutput(io.Discard)
+		logtest.Discard(t)
 
 		tests := []struct {
 			mapper MapFunc

core/hash/consistenthash.go

@@ -96,7 +96,7 @@ func (h *ConsistentHash) AddWithWeight(node any, weight int) {
 	h.AddWithReplicas(node, replicas)
 }
 
-// Get returns the corresponding node from h base on the given v.
+// Get returns the corresponding node from h based on the given v.
 func (h *ConsistentHash) Get(v any) (any, bool) {
 	h.lock.RLock()
 	defer h.lock.RUnlock()

core/jsonx/json.go

@@ -8,9 +8,25 @@ import (
 	"strings"
 )
 
-// Marshal marshals v into json bytes.
+// Marshal marshals v into json bytes, without escaping HTML and removes the trailing newline.
 func Marshal(v any) ([]byte, error) {
-	return json.Marshal(v)
+	// why not use json.Marshal?  https://github.com/golang/go/issues/28453
+	// it changes the behavior of json.Marshal, like & -> \u0026, < -> \u003c, > -> \u003e
+	// which is not what we want in API responses
+	var buf bytes.Buffer
+	enc := json.NewEncoder(&buf)
+	enc.SetEscapeHTML(false)
+	if err := enc.Encode(v); err != nil {
+		return nil, err
+	}
+
+	bs := buf.Bytes()
+	// Remove trailing newline added by json.Encoder.Encode
+	if len(bs) > 0 && bs[len(bs)-1] == '\n' {
+		bs = bs[:len(bs)-1]
+	}
+
+	return bs, nil
 }
 
 // MarshalToString marshals v into a string.

core/jsonx/json_test.go

@@ -1,6 +1,7 @@
 package jsonx
 
 import (
+	"fmt"
 	"strings"
 	"testing"
 
@@ -101,3 +102,105 @@ func TestUnmarshalFromReaderError(t *testing.T) {
 	err := UnmarshalFromReader(strings.NewReader(s), &v)
 	assert.NotNil(t, err)
 }
+
+func Test_doMarshalJson(t *testing.T) {
+	type args struct {
+		v any
+	}
+
+	tests := []struct {
+		name    string
+		args    args
+		want    []byte
+		wantErr assert.ErrorAssertionFunc
+	}{
+		{
+			name:    "nil",
+			args:    args{nil},
+			want:    []byte("null"),
+			wantErr: assert.NoError,
+		},
+		{
+			name:    "string",
+			args:    args{"hello"},
+			want:    []byte(`"hello"`),
+			wantErr: assert.NoError,
+		},
+		{
+			name:    "int",
+			args:    args{42},
+			want:    []byte("42"),
+			wantErr: assert.NoError,
+		},
+		{
+			name:    "bool",
+			args:    args{true},
+			want:    []byte("true"),
+			wantErr: assert.NoError,
+		},
+		{
+			name: "struct",
+			args: args{
+				struct {
+					Name string `json:"name"`
+				}{Name: "test"},
+			},
+			want:    []byte(`{"name":"test"}`),
+			wantErr: assert.NoError,
+		},
+		{
+			name:    "slice",
+			args:    args{[]int{1, 2, 3}},
+			want:    []byte("[1,2,3]"),
+			wantErr: assert.NoError,
+		},
+		{
+			name:    "map",
+			args:    args{map[string]int{"a": 1, "b": 2}},
+			want:    []byte(`{"a":1,"b":2}`),
+			wantErr: assert.NoError,
+		},
+		{
+			name:    "unmarshalable type",
+			args:    args{complex(1, 2)},
+			want:    nil,
+			wantErr: assert.Error,
+		},
+		{
+			name:    "channel type",
+			args:    args{make(chan int)},
+			want:    nil,
+			wantErr: assert.Error,
+		},
+		{
+			name:    "url with query params",
+			args:    args{"https://example.com/api?name=test&age=25"},
+			want:    []byte(`"https://example.com/api?name=test&age=25"`),
+			wantErr: assert.NoError,
+		},
+		{
+			name:    "url with encoded query params",
+			args:    args{"https://example.com/api?data=hello%20world&special=%26%3D"},
+			want:    []byte(`"https://example.com/api?data=hello%20world&special=%26%3D"`),
+			wantErr: assert.NoError,
+		},
+		{
+			name:    "url with multiple query params",
+			args:    args{"http://localhost:8080/users?page=1&limit=10&sort=name&order=asc"},
+			want:    []byte(`"http://localhost:8080/users?page=1&limit=10&sort=name&order=asc"`),
+			wantErr: assert.NoError,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := Marshal(tt.args.v)
+			if !tt.wantErr(t, err, fmt.Sprintf("Marshal(%v)", tt.args.v)) {
+				return
+			}
+
+			assert.Equalf(t, string(tt.want), string(got), "Marshal(%v)", tt.args.v)
+		})
+	}
+}

core/logx/config.go

@@ -1,47 +1,70 @@
 package logx
 
-// A LogConf is a logging config.
-type LogConf struct {
-	// ServiceName represents the service name.
-	ServiceName string `json:",optional"`
-	// Mode represents the logging mode, default is `console`.
-	// console: log to console.
-	// file: log to file.
-	// volume: used in k8s, prepend the hostname to the log file name.
-	Mode string `json:",default=console,options=[console,file,volume]"`
-	// Encoding represents the encoding type, default is `json`.
-	// json: json encoding.
-	// plain: plain text encoding, typically used in development.
-	Encoding string `json:",default=json,options=[json,plain]"`
-	// TimeFormat represents the time format, default is `2006-01-02T15:04:05.000Z07:00`.
-	TimeFormat string `json:",optional"`
-	// Path represents the log file path, default is `logs`.
-	Path string `json:",default=logs"`
-	// Level represents the log level, default is `info`.
-	Level string `json:",default=info,options=[debug,info,error,severe]"`
-	// MaxContentLength represents the max content bytes, default is no limit.
-	MaxContentLength uint32 `json:",optional"`
-	// Compress represents whether to compress the log file, default is `false`.
-	Compress bool `json:",optional"`
-	// Stat represents whether to log statistics, default is `true`.
-	Stat bool `json:",default=true"`
-	// KeepDays represents how many days the log files will be kept. Default to keep all files.
-	// Only take effect when Mode is `file` or `volume`, both work when Rotation is `daily` or `size`.
-	KeepDays int `json:",optional"`
-	// StackCooldownMillis represents the cooldown time for stack logging, default is 100ms.
-	StackCooldownMillis int `json:",default=100"`
-	// MaxBackups represents how many backup log files will be kept. 0 means all files will be kept forever.
-	// Only take effect when RotationRuleType is `size`.
-	// Even though `MaxBackups` sets 0, log files will still be removed
-	// if the `KeepDays` limitation is reached.
-	MaxBackups int `json:",default=0"`
-	// MaxSize represents how much space the writing log file takes up. 0 means no limit. The unit is `MB`.
-	// Only take effect when RotationRuleType is `size`
-	MaxSize int `json:",default=0"`
-	// Rotation represents the type of log rotation rule. Default is `daily`.
-	// 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"`
-}
+type (
+	// A LogConf is a logging config.
+	LogConf struct {
+		// ServiceName represents the service name.
+		ServiceName string `json:",optional"`
+		// Mode represents the logging mode, default is `console`.
+		// console: log to console.
+		// file: log to file.
+		// volume: used in k8s, prepend the hostname to the log file name.
+		Mode string `json:",default=console,options=[console,file,volume]"`
+		// Encoding represents the encoding type, default is `json`.
+		// json: json encoding.
+		// plain: plain text encoding, typically used in development.
+		Encoding string `json:",default=json,options=[json,plain]"`
+		// TimeFormat represents the time format, default is `2006-01-02T15:04:05.000Z07:00`.
+		TimeFormat string `json:",optional"`
+		// Path represents the log file path, default is `logs`.
+		Path string `json:",default=logs"`
+		// Level represents the log level, default is `info`.
+		Level string `json:",default=info,options=[debug,info,error,severe]"`
+		// MaxContentLength represents the max content bytes, default is no limit.
+		MaxContentLength uint32 `json:",optional"`
+		// Compress represents whether to compress the log file, default is `false`.
+		Compress bool `json:",optional"`
+		// Stat represents whether to log statistics, default is `true`.
+		Stat bool `json:",default=true"`
+		// KeepDays represents how many days the log files will be kept. Default to keep all files.
+		// Only take effect when Mode is `file` or `volume`, both work when Rotation is `daily` or `size`.
+		KeepDays int `json:",optional"`
+		// StackCooldownMillis represents the cooldown time for stack logging, default is 100ms.
+		StackCooldownMillis int `json:",default=100"`
+		// MaxBackups represents how many backup log files will be kept. 0 means all files will be kept forever.
+		// Only take effect when RotationRuleType is `size`.
+		// Even though `MaxBackups` sets 0, log files will still be removed
+		// if the `KeepDays` limitation is reached.
+		MaxBackups int `json:",default=0"`
+		// MaxSize represents how much space the writing log file takes up. 0 means no limit. The unit is `MB`.
+		// Only take effect when RotationRuleType is `size`
+		MaxSize int `json:",default=0"`
+		// Rotation represents the type of log rotation rule. Default is `daily`.
+		// 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"`
+		// FieldKeys represents the field keys.
+		FieldKeys fieldKeyConf `json:",optional"`
+	}
+
+	fieldKeyConf struct {
+		// CallerKey represents the caller key.
+		CallerKey string `json:",default=caller"`
+		// ContentKey represents the content key.
+		ContentKey string `json:",default=content"`
+		// DurationKey represents the duration key.
+		DurationKey string `json:",default=duration"`
+		// LevelKey represents the level key.
+		LevelKey string `json:",default=level"`
+		// SpanKey represents the span key.
+		SpanKey string `json:",default=span"`
+		// TimestampKey represents the timestamp key.
+		TimestampKey string `json:",default=@timestamp"`
+		// TraceKey represents the trace key.
+		TraceKey string `json:",default=trace"`
+		// TruncatedKey represents the truncated key.
+		TruncatedKey string `json:",default=truncated"`
+	}
+)

core/logx/logs.go

@@ -10,7 +10,6 @@ import (
 	"runtime/debug"
 	"sync"
 	"sync/atomic"
-	"time"
 
 	"github.com/zeromicro/go-zero/core/sysx"
 )
@@ -187,39 +186,9 @@ func Errorw(msg string, fields ...LogField) {
 
 // Field returns a LogField for the given key and value.
 func Field(key string, value any) LogField {
-	switch val := value.(type) {
-	case error:
-		return LogField{Key: key, Value: encodeError(val)}
-	case []error:
-		var errs []string
-		for _, err := range val {
-			errs = append(errs, encodeError(err))
-		}
-		return LogField{Key: key, Value: errs}
-	case time.Duration:
-		return LogField{Key: key, Value: fmt.Sprint(val)}
-	case []time.Duration:
-		var durs []string
-		for _, dur := range val {
-			durs = append(durs, fmt.Sprint(dur))
-		}
-		return LogField{Key: key, Value: durs}
-	case []time.Time:
-		var times []string
-		for _, t := range val {
-			times = append(times, fmt.Sprint(t))
-		}
-		return LogField{Key: key, Value: times}
-	case fmt.Stringer:
-		return LogField{Key: key, Value: encodeStringer(val)}
-	case []fmt.Stringer:
-		var strs []string
-		for _, str := range val {
-			strs = append(strs, encodeStringer(str))
-		}
-		return LogField{Key: key, Value: strs}
-	default:
-		return LogField{Key: key, Value: val}
+	return LogField{
+		Key:   key,
+		Value: value,
 	}
 }
 
@@ -307,7 +276,8 @@ func SetUp(c LogConf) (err error) {
 	// Because multiple services in one process might call SetUp respectively.
 	// Need to wait for the first caller to complete the execution.
 	setupOnce.Do(func() {
-		setupLogLevel(c)
+		setupLogLevel(c.Level)
+		setupFieldKeys(c.FieldKeys)
 
 		if !c.Stat {
 			DisableStat()
@@ -511,8 +481,35 @@ func handleOptions(opts []LogOption) {
 	}
 }
 
-func setupLogLevel(c LogConf) {
-	switch c.Level {
+func setupFieldKeys(c fieldKeyConf) {
+	if len(c.CallerKey) > 0 {
+		callerKey = c.CallerKey
+	}
+	if len(c.ContentKey) > 0 {
+		contentKey = c.ContentKey
+	}
+	if len(c.DurationKey) > 0 {
+		durationKey = c.DurationKey
+	}
+	if len(c.LevelKey) > 0 {
+		levelKey = c.LevelKey
+	}
+	if len(c.SpanKey) > 0 {
+		spanKey = c.SpanKey
+	}
+	if len(c.TimestampKey) > 0 {
+		timestampKey = c.TimestampKey
+	}
+	if len(c.TraceKey) > 0 {
+		traceKey = c.TraceKey
+	}
+	if len(c.TruncatedKey) > 0 {
+		truncatedKey = c.TruncatedKey
+	}
+}
+
+func setupLogLevel(level string) {
+	switch level {
 	case levelDebug:
 		SetLevel(DebugLevel)
 	case levelInfo:

core/logx/logs_test.go

@@ -1,6 +1,7 @@
 package logx
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -16,6 +17,8 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/sdk/trace"
 )
 
 var (
@@ -776,15 +779,9 @@ func TestSetup(t *testing.T) {
 		MaxBackups:  3,
 		MaxSize:     1024 * 1024,
 	}))
-	setupLogLevel(LogConf{
-		Level: levelInfo,
-	})
-	setupLogLevel(LogConf{
-		Level: levelError,
-	})
-	setupLogLevel(LogConf{
-		Level: levelSevere,
-	})
+	setupLogLevel(levelInfo)
+	setupLogLevel(levelError)
+	setupLogLevel(levelSevere)
 	_, err := createOutput("")
 	assert.NotNil(t, err)
 	Disable()
@@ -856,6 +853,95 @@ func TestWithKeepDays(t *testing.T) {
 	assert.Equal(t, 1, opt.keepDays)
 }
 
+func TestWithField_LogLevel(t *testing.T) {
+	tests := []struct {
+		name  string
+		level uint32
+		fn    func(string, ...LogField)
+		count int32
+	}{
+		{
+			name:  "debug/info",
+			level: DebugLevel,
+			fn:    Infow,
+			count: 1,
+		},
+		{
+			name:  "info/error",
+			level: InfoLevel,
+			fn:    Errorw,
+			count: 1,
+		},
+		{
+			name:  "info/info",
+			level: InfoLevel,
+			fn:    Infow,
+			count: 1,
+		},
+		{
+			name:  "info/severe",
+			level: InfoLevel,
+			fn:    Errorw,
+			count: 1,
+		},
+		{
+			name:  "error/info",
+			level: ErrorLevel,
+			fn:    Infow,
+			count: 0,
+		},
+		{
+			name:  "error/debug",
+			level: ErrorLevel,
+			fn:    Debugw,
+			count: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			olevel := atomic.LoadUint32(&logLevel)
+			SetLevel(tt.level)
+			defer SetLevel(olevel)
+
+			var val countingStringer
+			tt.fn("hello there", Field("foo", &val))
+			assert.Equal(t, tt.count, val.Count())
+		})
+	}
+}
+
+func TestWithField_LogLevelWithContext(t *testing.T) {
+	t.Run("context more than once with info/info", func(t *testing.T) {
+		olevel := atomic.LoadUint32(&logLevel)
+		SetLevel(InfoLevel)
+		defer SetLevel(olevel)
+
+		var val countingStringer
+		ctx := ContextWithFields(context.Background(), Field("foo", &val))
+		logger := WithContext(ctx)
+		logger.Info("hello there")
+		logger.Info("hello there")
+		logger.Info("hello there")
+		assert.True(t, val.Count() > 0)
+	})
+
+	t.Run("context more than once with error/info", func(t *testing.T) {
+		olevel := atomic.LoadUint32(&logLevel)
+		SetLevel(ErrorLevel)
+		defer SetLevel(olevel)
+
+		var val countingStringer
+		ctx := ContextWithFields(context.Background(), Field("foo", &val))
+		logger := WithContext(ctx)
+		logger.Info("hello there")
+		logger.Info("hello there")
+		logger.Info("hello there")
+		assert.Equal(t, int32(0), val.Count())
+	})
+}
+
 func BenchmarkCopyByteSliceAppend(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		var buf []byte
@@ -1054,3 +1140,79 @@ type panicStringer struct {
 func (s panicStringer) String() string {
 	panic("panic")
 }
+
+type countingStringer struct {
+	count int32
+}
+
+func (s *countingStringer) Count() int32 {
+	return atomic.LoadInt32(&s.count)
+}
+
+func (s *countingStringer) String() string {
+	atomic.AddInt32(&s.count, 1)
+	return "countingStringer"
+}
+
+func TestLogKey(t *testing.T) {
+	setupOnce = sync.Once{}
+	MustSetup(LogConf{
+		ServiceName: "any",
+		Mode:        "console",
+		Encoding:    "json",
+		TimeFormat:  timeFormat,
+		FieldKeys: fieldKeyConf{
+			CallerKey:    "_caller",
+			ContentKey:   "_content",
+			DurationKey:  "_duration",
+			LevelKey:     "_level",
+			SpanKey:      "_span",
+			TimestampKey: "_timestamp",
+			TraceKey:     "_trace",
+			TruncatedKey: "_truncated",
+		},
+	})
+
+	t.Cleanup(func() {
+		setupFieldKeys(fieldKeyConf{
+			CallerKey:    defaultCallerKey,
+			ContentKey:   defaultContentKey,
+			DurationKey:  defaultDurationKey,
+			LevelKey:     defaultLevelKey,
+			SpanKey:      defaultSpanKey,
+			TimestampKey: defaultTimestampKey,
+			TraceKey:     defaultTraceKey,
+			TruncatedKey: defaultTruncatedKey,
+		})
+	})
+
+	const message = "hello there"
+	w := new(mockWriter)
+	old := writer.Swap(w)
+	defer writer.Store(old)
+
+	otp := otel.GetTracerProvider()
+	tp := trace.NewTracerProvider(trace.WithSampler(trace.AlwaysSample()))
+	otel.SetTracerProvider(tp)
+	defer otel.SetTracerProvider(otp)
+
+	ctx, span := tp.Tracer("trace-id").Start(context.Background(), "span-id")
+	defer span.End()
+
+	WithContext(ctx).WithDuration(time.Second).Info(message)
+	now := time.Now()
+
+	var m map[string]string
+	if err := json.Unmarshal([]byte(w.String()), &m); err != nil {
+		t.Error(err)
+	}
+	assert.Equal(t, "info", m["_level"])
+	assert.Equal(t, message, m["_content"])
+	assert.Equal(t, "1000.0ms", m["_duration"])
+	assert.Regexp(t, `logx/logs_test.go:\d+`, m["_caller"])
+	assert.NotEmpty(t, m["_trace"])
+	assert.NotEmpty(t, m["_span"])
+	parsedTime, err := time.Parse(timeFormat, m["_timestamp"])
+	assert.True(t, err == nil)
+	assert.Equal(t, now.Minute(), parsedTime.Minute())
+}

core/logx/richlogger_test.go

@@ -423,3 +423,49 @@ type mockValue struct {
 	Foo     string `json:"foo"`
 	Content any    `json:"content"`
 }
+
+type testJson struct {
+	Name  string  `json:"name"`
+	Age   int     `json:"age"`
+	Score float64 `json:"score"`
+}
+
+func (t testJson) MarshalJSON() ([]byte, error) {
+	type testJsonImpl testJson
+	return json.Marshal(testJsonImpl(t))
+}
+
+func (t testJson) String() string {
+	return fmt.Sprintf("%s %d %f", t.Name, t.Age, t.Score)
+}
+
+func TestLogWithJson(t *testing.T) {
+	w := new(mockWriter)
+	old := writer.Swap(w)
+	writer.lock.RLock()
+	defer func() {
+		writer.lock.RUnlock()
+		writer.Store(old)
+	}()
+
+	l := WithContext(context.Background()).WithFields(Field("bar", testJson{
+		Name:  "foo",
+		Age:   1,
+		Score: 1.0,
+	}))
+	l.Info(testlog)
+
+	type mockValue2 struct {
+		mockValue
+		Bar testJson `json:"bar"`
+	}
+
+	var val mockValue2
+	err := json.Unmarshal([]byte(w.String()), &val)
+	assert.NoError(t, err)
+
+	assert.Equal(t, testlog, val.Content)
+	assert.Equal(t, "foo", val.Bar.Name)
+	assert.Equal(t, 1, val.Bar.Age)
+	assert.Equal(t, 1.0, val.Bar.Score)
+}

core/logx/rotatelogger.go

@@ -66,7 +66,7 @@ type (
 		gzip        bool
 	}
 
-	// SizeLimitRotateRule a rotation rule that make the log file rotated base on size
+	// SizeLimitRotateRule a rotation rule that makes the log file rotated based on size
 	SizeLimitRotateRule struct {
 		DailyRotateRule
 		maxSize    int64
@@ -211,7 +211,7 @@ func (r *SizeLimitRotateRule) OutdatedFiles() []string {
 		}
 	}
 
-	var result []string
+	result := make([]string, 0, len(outdated))
 	for k := range outdated {
 		result = append(result, k)
 	}

core/logx/sensitive.go

@@ -0,0 +1,21 @@
+package logx
+
+// Sensitive is an interface that defines a method for masking sensitive information in logs.
+// It is typically implemented by types that contain sensitive data,
+// such as passwords or personal information.
+// Infov, Errorv, Debugv, and Slowv methods will call this method to mask sensitive data.
+// The values in LogField will also be masked if they implement the Sensitive interface.
+type Sensitive interface {
+	// MaskSensitive masks sensitive information in the log.
+	MaskSensitive() any
+}
+
+// maskSensitive returns the value returned by MaskSensitive method,
+// if the value implements Sensitive interface.
+func maskSensitive(v any) any {
+	if s, ok := v.(Sensitive); ok {
+		return s.MaskSensitive()
+	}
+
+	return v
+}

core/logx/sensitive_test.go

@@ -0,0 +1,50 @@
+package logx
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const maskedContent = "******"
+
+type User struct {
+	Name string
+	Pass string
+}
+
+func (u User) MaskSensitive() any {
+	return User{
+		Name: u.Name,
+		Pass: maskedContent,
+	}
+}
+
+type NonSensitiveUser struct {
+	Name string
+	Pass string
+}
+
+func TestMaskSensitive(t *testing.T) {
+	t.Run("sensitive", func(t *testing.T) {
+		user := User{
+			Name: "kevin",
+			Pass: "123",
+		}
+
+		mu := maskSensitive(user)
+		assert.Equal(t, user.Name, mu.(User).Name)
+		assert.Equal(t, maskedContent, mu.(User).Pass)
+	})
+
+	t.Run("non-sensitive", func(t *testing.T) {
+		user := NonSensitiveUser{
+			Name: "kevin",
+			Pass: "123",
+		}
+
+		mu := maskSensitive(user)
+		assert.Equal(t, user.Name, mu.(NonSensitiveUser).Name)
+		assert.Equal(t, user.Pass, mu.(NonSensitiveUser).Pass)
+	})
+}

core/logx/vars.go

@@ -53,14 +53,14 @@ const (
 )
 
 const (
-	callerKey    = "caller"
-	contentKey   = "content"
-	durationKey  = "duration"
-	levelKey     = "level"
-	spanKey      = "span"
-	timestampKey = "@timestamp"
-	traceKey     = "trace"
-	truncatedKey = "truncated"
+	defaultCallerKey    = "caller"
+	defaultContentKey   = "content"
+	defaultDurationKey  = "duration"
+	defaultLevelKey     = "level"
+	defaultSpanKey      = "span"
+	defaultTimestampKey = "@timestamp"
+	defaultTraceKey     = "trace"
+	defaultTruncatedKey = "truncated"
 )
 
 var (
@@ -73,3 +73,14 @@ var (
 
 	truncatedField = Field(truncatedKey, true)
 )
+
+var (
+	callerKey    = defaultCallerKey
+	contentKey   = defaultContentKey
+	durationKey  = defaultDurationKey
+	levelKey     = defaultLevelKey
+	spanKey      = defaultSpanKey
+	timestampKey = defaultTimestampKey
+	traceKey     = defaultTraceKey
+	truncatedKey = defaultTruncatedKey
+)

core/logx/writer.go

@@ -10,6 +10,7 @@ import (
 	"runtime/debug"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	fatihcolor "github.com/fatih/color"
 	"github.com/zeromicro/go-zero/core/color"
@@ -211,7 +212,6 @@ func newFileWriter(c LogConf) (Writer, error) {
 	statFile := path.Join(c.Path, statFilename)
 
 	handleOptions(opts)
-	setupLogLevel(c)
 
 	if infoLog, err = createOutput(accessFile); err != nil {
 		return nil, err
@@ -365,19 +365,25 @@ func mergeGlobalFields(fields []LogField) []LogField {
 }
 
 func output(writer io.Writer, level string, val any, fields ...LogField) {
-	// only truncate string content, don't know how to truncate the values of other types.
-	if v, ok := val.(string); ok {
+	switch v := val.(type) {
+	case string:
+		// only truncate string content, don't know how to truncate the values of other types.
 		maxLen := atomic.LoadUint32(&maxContentLength)
 		if maxLen > 0 && len(v) > int(maxLen) {
 			val = v[:maxLen]
 			fields = append(fields, truncatedField)
 		}
+	case Sensitive:
+		val = v.MaskSensitive()
 	}
 
 	// +3 for timestamp, level and content
 	entry := make(logEntry, len(fields)+3)
 	for _, field := range fields {
-		entry[field.Key] = field.Value
+		// mask sensitive data before processing types,
+		// in case field.Value is a sensitive type and also implemented fmt.Stringer.
+		mval := maskSensitive(field.Value)
+		entry[field.Key] = processFieldValue(mval)
 	}
 
 	switch atomic.LoadUint32(&encoding) {
@@ -392,6 +398,45 @@ func output(writer io.Writer, level string, val any, fields ...LogField) {
 	}
 }
 
+func processFieldValue(value any) any {
+	switch val := value.(type) {
+	case error:
+		return encodeError(val)
+	case []error:
+		var errs []string
+		for _, err := range val {
+			errs = append(errs, encodeError(err))
+		}
+		return errs
+	case time.Duration:
+		return fmt.Sprint(val)
+	case []time.Duration:
+		var durs []string
+		for _, dur := range val {
+			durs = append(durs, fmt.Sprint(dur))
+		}
+		return durs
+	case []time.Time:
+		var times []string
+		for _, t := range val {
+			times = append(times, fmt.Sprint(t))
+		}
+		return times
+	case json.Marshaler:
+		return val
+	case fmt.Stringer:
+		return encodeStringer(val)
+	case []fmt.Stringer:
+		var strs []string
+		for _, str := range val {
+			strs = append(strs, encodeStringer(str))
+		}
+		return strs
+	default:
+		return val
+	}
+}
+
 func wrapLevelWithColor(level string) string {
 	var colour color.Color
 	switch level {
@@ -399,6 +444,8 @@ func wrapLevelWithColor(level string) string {
 		colour = color.FgRed
 	case levelError:
 		colour = color.FgRed
+	case levelSevere:
+		colour = color.FgRed
 	case levelFatal:
 		colour = color.FgRed
 	case levelInfo:

core/logx/writer_test.go

@@ -225,6 +225,48 @@ func TestWritePlainDuplicate(t *testing.T) {
 	assert.Contains(t, buf.String(), "second=c")
 }
 
+func TestLogWithSensitive(t *testing.T) {
+	old := atomic.SwapUint32(&encoding, plainEncodingType)
+	t.Cleanup(func() {
+		atomic.StoreUint32(&encoding, old)
+	})
+
+	t.Run("sensitive", func(t *testing.T) {
+		var buf bytes.Buffer
+		output(&buf, levelInfo, User{
+			Name: "kevin",
+			Pass: "123",
+		}, LogField{
+			Key:   "first",
+			Value: "a",
+		}, LogField{
+			Key:   "first",
+			Value: "b",
+		})
+		assert.Contains(t, buf.String(), maskedContent)
+		assert.NotContains(t, buf.String(), "first=a")
+		assert.Contains(t, buf.String(), "first=b")
+	})
+
+	t.Run("sensitive fields", func(t *testing.T) {
+		var buf bytes.Buffer
+		output(&buf, levelInfo, "foo", LogField{
+			Key: "first",
+			Value: User{
+				Name: "kevin",
+				Pass: "123",
+			},
+		}, LogField{
+			Key:   "second",
+			Value: "b",
+		})
+		assert.Contains(t, buf.String(), "foo")
+		assert.Contains(t, buf.String(), "first")
+		assert.Contains(t, buf.String(), maskedContent)
+		assert.Contains(t, buf.String(), "second=b")
+	})
+}
+
 func TestLogWithLimitContentLength(t *testing.T) {
 	maxLen := atomic.LoadUint32(&maxContentLength)
 	atomic.StoreUint32(&maxContentLength, 10)

core/mapping/marshaler.go

@@ -3,6 +3,7 @@ package mapping
 import (
 	"fmt"
 	"reflect"
+	"slices"
 	"strings"
 )
 
@@ -152,15 +153,8 @@ func validateOptional(field reflect.StructField, value reflect.Value) error {
 }
 
 func validateOptions(value reflect.Value, opt *fieldOptions) error {
-	var found bool
 	val := fmt.Sprint(value.Interface())
-	for i := range opt.Options {
-		if opt.Options[i] == val {
-			found = true
-			break
-		}
-	}
-	if !found {
+	if !slices.Contains(opt.Options, val) {
 		return fmt.Errorf("field %q not in options", val)
 	}
 

core/mapping/unmarshaler.go

@@ -30,9 +30,7 @@ var (
 	errValueNotSettable = errors.New("value is not settable")
 	errValueNotStruct   = errors.New("value type is not struct")
 	keyUnmarshaler      = NewUnmarshaler(defaultKeyName)
-	boolType            = reflect.TypeOf(false)
 	durationType        = reflect.TypeOf(time.Duration(0))
-	stringType          = reflect.TypeOf("")
 	cacheKeys           = make(map[string][]string)
 	cacheKeysLock       sync.Mutex
 	defaultCache        = make(map[string]any)
@@ -768,23 +766,25 @@ func (u *Unmarshaler) processFieldWithEnvValue(fieldType reflect.Type, value ref
 	}
 
 	derefType := Deref(fieldType)
-	switch derefType {
-	case boolType:
+	derefKind := derefType.Kind()
+	switch {
+	case derefKind == reflect.String:
+		SetValue(fieldType, value, toReflectValue(derefType, envVal))
+		return nil
+	case derefKind == reflect.Bool:
 		val, err := strconv.ParseBool(envVal)
 		if err != nil {
 			return fmt.Errorf("unmarshal field %q with environment variable, %w", fullName, err)
 		}
 
-		SetValue(fieldType, value, reflect.ValueOf(val))
+		SetValue(fieldType, value, toReflectValue(derefType, val))
 		return nil
-	case durationType:
+	case derefType == durationType:
+		// time.Duration is a special case, its derefKind is reflect.Int64.
 		if err := fillDurationValue(fieldType, value, envVal); err != nil {
 			return fmt.Errorf("unmarshal field %q with environment variable, %w", fullName, err)
 		}
 
-		return nil
-	case stringType:
-		SetValue(fieldType, value, reflect.ValueOf(envVal))
 		return nil
 	default:
 		return u.processFieldPrimitiveWithJSONNumber(fieldType, value, json.Number(envVal), opts, fullName)

core/mapping/unmarshaler_test.go

@@ -6083,6 +6083,105 @@ func TestParseJsonStringValue(t *testing.T) {
 	})
 }
 
+// issue #5033, string type
+func TestUnmarshalFromEnvString(t *testing.T) {
+	t.Setenv("STRING_ENV", "dev")
+
+	t.Run("by value", func(t *testing.T) {
+		type (
+			Env    string
+			Config struct {
+				Env Env `json:",env=STRING_ENV,default=prod"`
+			}
+		)
+
+		var c Config
+		if assert.NoError(t, UnmarshalJsonMap(map[string]any{}, &c)) {
+			assert.Equal(t, Env("dev"), c.Env)
+		}
+	})
+
+	t.Run("by ptr", func(t *testing.T) {
+		type (
+			Env    string
+			Config struct {
+				Env *Env `json:",env=STRING_ENV,default=prod"`
+			}
+		)
+
+		var c Config
+		if assert.NoError(t, UnmarshalJsonMap(map[string]any{}, &c)) {
+			assert.Equal(t, Env("dev"), *c.Env)
+		}
+	})
+}
+
+// issue #5033, bool type
+func TestUnmarshalFromEnvBool(t *testing.T) {
+	t.Setenv("BOOL_ENV", "true")
+
+	t.Run("by value", func(t *testing.T) {
+		type (
+			Env    bool
+			Config struct {
+				Env Env `json:",env=BOOL_ENV,default=false"`
+			}
+		)
+
+		var c Config
+		if assert.NoError(t, UnmarshalJsonMap(map[string]any{}, &c)) {
+			assert.Equal(t, Env(true), c.Env)
+		}
+	})
+
+	t.Run("by ptr", func(t *testing.T) {
+		type (
+			Env    bool
+			Config struct {
+				Env *Env `json:",env=BOOL_ENV,default=false"`
+			}
+		)
+
+		var c Config
+		if assert.NoError(t, UnmarshalJsonMap(map[string]any{}, &c)) {
+			assert.Equal(t, Env(true), *c.Env)
+		}
+	})
+}
+
+// issue #5033, customized int type
+func TestUnmarshalFromEnvInt(t *testing.T) {
+	t.Setenv("INT_ENV", "2")
+
+	t.Run("by value", func(t *testing.T) {
+		type (
+			Env    int
+			Config struct {
+				Env Env `json:",env=INT_ENV,default=0"`
+			}
+		)
+
+		var c Config
+		if assert.NoError(t, UnmarshalJsonMap(map[string]any{}, &c)) {
+			assert.Equal(t, Env(2), c.Env)
+		}
+	})
+
+	t.Run("by ptr", func(t *testing.T) {
+		type (
+			Env    int
+			Config struct {
+				Env *Env `json:",env=INT_ENV,default=0"`
+			}
+		)
+
+		var c Config
+		if assert.NoError(t, UnmarshalJsonMap(map[string]any{}, &c)) {
+			assert.Equal(t, Env(2), *c.Env)
+		}
+	})
+}
+
 func BenchmarkDefaultValue(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		var a struct {

core/mapping/utils.go

@@ -104,14 +104,13 @@ func convertToString(val any, fullName string) (string, error) {
 func convertTypeFromString(kind reflect.Kind, str string) (any, error) {
 	switch kind {
 	case reflect.Bool:
-		switch strings.ToLower(str) {
-		case "1", "true":
+		if str == "1" || strings.EqualFold(str, "true") {
 			return true, nil
-		case "0", "false":
-			return false, nil
-		default:
-			return false, errTypeMismatch
 		}
+		if str == "0" || strings.EqualFold(str, "false") {
+			return false, nil
+		}
+		return false, errTypeMismatch
 	case reflect.Int:
 		return strconv.ParseInt(str, 10, intSize)
 	case reflect.Int8:
@@ -583,6 +582,10 @@ func toFloat64(v any) (float64, bool) {
 	}
 }
 
+func toReflectValue(tp reflect.Type, v any) reflect.Value {
+	return reflect.ValueOf(v).Convert(Deref(tp))
+}
+
 func usingDifferentKeys(key string, field reflect.StructField) bool {
 	if len(field.Tag) > 0 {
 		if _, ok := field.Tag.Lookup(key); !ok {

core/mapping/utils_test.go

@@ -334,3 +334,43 @@ func TestValidateValueRange(t *testing.T) {
 func TestSetMatchedPrimitiveValue(t *testing.T) {
 	assert.Error(t, setMatchedPrimitiveValue(reflect.Func, reflect.ValueOf(2), "1"))
 }
+
+func TestConvertTypeFromString_Bool(t *testing.T) {
+	tests := []struct {
+		name    string
+		input   string
+		want    bool
+		wantErr bool
+	}{
+		// true cases
+		{name: "1", input: "1", want: true, wantErr: false},
+		{name: "true lowercase", input: "true", want: true, wantErr: false},
+		{name: "True mixed", input: "True", want: true, wantErr: false},
+		{name: "TRUE uppercase", input: "TRUE", want: true, wantErr: false},
+		{name: "TrUe mixed", input: "TrUe", want: true, wantErr: false},
+		// false cases
+		{name: "0", input: "0", want: false, wantErr: false},
+		{name: "false lowercase", input: "false", want: false, wantErr: false},
+		{name: "False mixed", input: "False", want: false, wantErr: false},
+		{name: "FALSE uppercase", input: "FALSE", want: false, wantErr: false},
+		{name: "FaLsE mixed", input: "FaLsE", want: false, wantErr: false},
+		// error cases
+		{name: "invalid yes", input: "yes", want: false, wantErr: true},
+		{name: "invalid no", input: "no", want: false, wantErr: true},
+		{name: "invalid empty", input: "", want: false, wantErr: true},
+		{name: "invalid 2", input: "2", want: false, wantErr: true},
+		{name: "invalid truee", input: "truee", want: false, wantErr: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := convertTypeFromString(reflect.Bool, tt.input)
+			if tt.wantErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tt.want, got)
+			}
+		})
+	}
+}

core/mathx/unstable.go

@@ -6,7 +6,7 @@ import (
 	"time"
 )
 
-// An Unstable is used to generate random value around the mean value base on given deviation.
+// An Unstable is used to generate random value around the mean value based on given deviation.
 type Unstable struct {
 	deviation float64
 	r         *rand.Rand

core/mr/mapreduce.go

@@ -3,6 +3,9 @@ package mr
 import (
 	"context"
 	"errors"
+	"fmt"
+	"runtime/debug"
+	"strings"
 	"sync"
 	"sync/atomic"
 
@@ -183,12 +186,16 @@ func buildOptions(opts ...Option) *mapReduceOptions {
 	return options
 }
 
+func buildPanicInfo(r any, stack []byte) string {
+	return fmt.Sprintf("%+v\n\n%s", r, strings.TrimSpace(string(stack)))
+}
+
 func buildSource[T any](generate GenerateFunc[T], panicChan *onceChan) chan T {
 	source := make(chan T)
 	go func() {
 		defer func() {
 			if r := recover(); r != nil {
-				panicChan.write(r)
+				panicChan.write(buildPanicInfo(r, debug.Stack()))
 			}
 			close(source)
 		}()
@@ -235,7 +242,7 @@ func executeMappers[T, U any](mCtx mapperContext[T, U]) {
 				defer func() {
 					if r := recover(); r != nil {
 						atomic.AddInt32(&failed, 1)
-						mCtx.panicChan.write(r)
+						mCtx.panicChan.write(buildPanicInfo(r, debug.Stack()))
 					}
 					wg.Done()
 					<-pool
@@ -289,7 +296,7 @@ func mapReduceWithPanicChan[T, U, V any](source <-chan T, panicChan *onceChan, m
 		defer func() {
 			drain(collector)
 			if r := recover(); r != nil {
-				panicChan.write(r)
+				panicChan.write(buildPanicInfo(r, debug.Stack()))
 			}
 			finish()
 		}()

core/mr/mapreduce_test.go

@@ -3,8 +3,7 @@ package mr
 import (
 	"context"
 	"errors"
-	"io"
-	"log"
+	"fmt"
 	"runtime"
 	"sync/atomic"
 	"testing"
@@ -16,9 +15,6 @@ import (
 
 var errDummy = errors.New("dummy")
 
-func init() {
-	log.SetOutput(io.Discard)
-}
 
 func TestFinish(t *testing.T) {
 	defer goleak.VerifyNone(t)
@@ -39,6 +35,36 @@ func TestFinish(t *testing.T) {
 	assert.Nil(t, err)
 }
 
+func TestFinishWithPartialErrors(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
+	errDummy := errors.New("dummy")
+
+	t.Run("one error", func(t *testing.T) {
+		err := Finish(func() error {
+			return errDummy
+		}, func() error {
+			return nil
+		}, func() error {
+			return nil
+		})
+
+		assert.Equal(t, errDummy, err)
+	})
+
+	t.Run("two errors", func(t *testing.T) {
+		err := Finish(func() error {
+			return errDummy
+		}, func() error {
+			return errDummy
+		}, func() error {
+			return nil
+		})
+
+		assert.Equal(t, errDummy, err)
+	})
+}
+
 func TestFinishNone(t *testing.T) {
 	defer goleak.VerifyNone(t)
 
@@ -118,11 +144,28 @@ func TestForEach(t *testing.T) {
 
 		assert.Equal(t, tasks/2, int(count))
 	})
+}
 
-	t.Run("all", func(t *testing.T) {
-		defer goleak.VerifyNone(t)
+func TestPanics(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
+	const tasks = 1000
+	verify := func(t *testing.T, r any) {
+		panicStr := fmt.Sprintf("%v", r)
+		assert.Contains(t, panicStr, "foo")
+		assert.Contains(t, panicStr, "goroutine")
+		assert.Contains(t, panicStr, "runtime/debug.Stack")
+		panic(r)
+	}
+
+	t.Run("ForEach run panics", func(t *testing.T) {
+		assert.Panics(t, func() {
+			defer func() {
+				if r := recover(); r != nil {
+					verify(t, r)
+				}
+			}()
 
-		assert.PanicsWithValue(t, "foo", func() {
 			ForEach(func(source chan<- int) {
 				for i := 0; i < tasks; i++ {
 					source <- i
@@ -132,28 +175,31 @@ func TestForEach(t *testing.T) {
 			})
 		})
 	})
-}
 
-func TestGeneratePanic(t *testing.T) {
-	defer goleak.VerifyNone(t)
+	t.Run("ForEach generate panics", func(t *testing.T) {
+		assert.Panics(t, func() {
+			defer func() {
+				if r := recover(); r != nil {
+					verify(t, r)
+				}
+			}()
 
-	t.Run("all", func(t *testing.T) {
-		assert.PanicsWithValue(t, "foo", func() {
 			ForEach(func(source chan<- int) {
 				panic("foo")
 			}, func(item int) {
 			})
 		})
 	})
-}
 
-func TestMapperPanic(t *testing.T) {
-	defer goleak.VerifyNone(t)
-
-	const tasks = 1000
 	var run int32
-	t.Run("all", func(t *testing.T) {
-		assert.PanicsWithValue(t, "foo", func() {
+	t.Run("Mapper panics", func(t *testing.T) {
+		assert.Panics(t, func() {
+			defer func() {
+				if r := recover(); r != nil {
+					verify(t, r)
+				}
+			}()
+
 			_, _ = MapReduce(func(source chan<- int) {
 				for i := 0; i < tasks; i++ {
 					source <- i

core/prof/runtime.go

@@ -5,6 +5,8 @@ import (
 	"io"
 	"os"
 	"runtime"
+	"runtime/debug"
+	"runtime/metrics"
 	"time"
 )
 
@@ -28,10 +30,29 @@ func displayStatsWithWriter(writer io.Writer, interval ...time.Duration) {
 		ticker := time.NewTicker(duration)
 		defer ticker.Stop()
 		for range ticker.C {
-			var m runtime.MemStats
-			runtime.ReadMemStats(&m)
+			var (
+				alloc, totalAlloc, sys uint64
+				samples                = []metrics.Sample{
+					{Name: "/memory/classes/heap/objects:bytes"},
+					{Name: "/gc/heap/allocs:bytes"},
+					{Name: "/memory/classes/total:bytes"},
+				}
+			)
+			metrics.Read(samples)
+
+			if samples[0].Value.Kind() == metrics.KindUint64 {
+				alloc = samples[0].Value.Uint64()
+			}
+			if samples[1].Value.Kind() == metrics.KindUint64 {
+				totalAlloc = samples[1].Value.Uint64()
+			}
+			if samples[2].Value.Kind() == metrics.KindUint64 {
+				sys = samples[2].Value.Uint64()
+			}
+			var stats debug.GCStats
+			debug.ReadGCStats(&stats)
 			fmt.Fprintf(writer, "Goroutines: %d, Alloc: %vm, TotalAlloc: %vm, Sys: %vm, NumGC: %v\n",
-				runtime.NumGoroutine(), m.Alloc/mega, m.TotalAlloc/mega, m.Sys/mega, m.NumGC)
+				runtime.NumGoroutine(), alloc/mega, totalAlloc/mega, sys/mega, stats.NumGC)
 		}
 	}()
 }

core/stat/usage.go

@@ -1,7 +1,8 @@
 package stat
 
 import (
-	"runtime"
+	"runtime/debug"
+	"runtime/metrics"
 	"sync/atomic"
 	"time"
 
@@ -56,8 +57,28 @@ func bToMb(b uint64) float32 {
 }
 
 func printUsage() {
-	var m runtime.MemStats
-	runtime.ReadMemStats(&m)
+	var (
+		alloc, totalAlloc, sys uint64
+		samples                = []metrics.Sample{
+			{Name: "/memory/classes/heap/objects:bytes"},
+			{Name: "/gc/heap/allocs:bytes"},
+			{Name: "/memory/classes/total:bytes"},
+		}
+		stats debug.GCStats
+	)
+	metrics.Read(samples)
+
+	if samples[0].Value.Kind() == metrics.KindUint64 {
+		alloc = samples[0].Value.Uint64()
+	}
+	if samples[1].Value.Kind() == metrics.KindUint64 {
+		totalAlloc = samples[1].Value.Uint64()
+	}
+	if samples[2].Value.Kind() == metrics.KindUint64 {
+		sys = samples[2].Value.Uint64()
+	}
+	debug.ReadGCStats(&stats)
+
 	logx.Statf("CPU: %dm, MEMORY: Alloc=%.1fMi, TotalAlloc=%.1fMi, Sys=%.1fMi, NumGC=%d",
-		CpuUsage(), bToMb(m.Alloc), bToMb(m.TotalAlloc), bToMb(m.Sys), m.NumGC)
+		CpuUsage(), bToMb(alloc), bToMb(totalAlloc), bToMb(sys), stats.NumGC)
 }

core/stores/mon/bulkinserter.go

@@ -1,3 +1,4 @@
+//go:generate mockgen -package mon -destination collectioninserter_mock.go -source bulkinserter.go collectionInserter
 package mon
 
 import (
@@ -6,7 +7,8 @@ import (
 
 	"github.com/zeromicro/go-zero/core/executors"
 	"github.com/zeromicro/go-zero/core/logx"
-	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
 )
 
 const (
@@ -27,10 +29,7 @@ type (
 
 // NewBulkInserter returns a BulkInserter.
 func NewBulkInserter(coll Collection, interval ...time.Duration) (*BulkInserter, error) {
-	cloneColl, err := coll.Clone()
-	if err != nil {
-		return nil, err
-	}
+	cloneColl := coll.Clone()
 
 	inserter := &dbInserter{
 		collection: cloneColl,
@@ -64,8 +63,16 @@ func (bi *BulkInserter) SetResultHandler(handler ResultHandler) {
 	})
 }
 
+type collectionInserter interface {
+	InsertMany(
+		ctx context.Context,
+		documents interface{},
+		opts ...options.Lister[options.InsertManyOptions],
+	) (*mongo.InsertManyResult, error)
+}
+
 type dbInserter struct {
-	collection    *mongo.Collection
+	collection    collectionInserter
 	documents     []any
 	resultHandler ResultHandler
 }

core/stores/mon/bulkinserter_test.go

@@ -1,26 +1,131 @@
 package mon
 
 import (
+	"errors"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
-	"go.mongodb.org/mongo-driver/bson"
-	"go.mongodb.org/mongo-driver/mongo"
-	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.uber.org/mock/gomock"
 )
 
-func TestBulkInserter(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
-		bulk, err := NewBulkInserter(createModel(mt).Collection)
-		assert.Equal(t, err, nil)
-		bulk.SetResultHandler(func(result *mongo.InsertManyResult, err error) {
-			assert.Nil(t, err)
-			assert.Equal(t, 2, len(result.InsertedIDs))
-		})
-		bulk.Insert(bson.D{{Key: "foo", Value: "bar"}})
-		bulk.Insert(bson.D{{Key: "foo", Value: "baz"}})
-		bulk.Flush()
+func TestBulkInserter_InsertAndFlush(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockCollection(ctrl)
+	mockCollection.EXPECT().Clone().Return(&mongo.Collection{})
+	bulkInserter, err := NewBulkInserter(mockCollection, time.Second)
+	assert.NoError(t, err)
+	bulkInserter.SetResultHandler(func(result *mongo.InsertManyResult, err error) {
+		assert.Nil(t, err)
+		assert.Equal(t, 2, len(result.InsertedIDs))
 	})
+	doc := map[string]interface{}{"name": "test"}
+	bulkInserter.Insert(doc)
+	bulkInserter.Flush()
+}
+
+func TestBulkInserter_SetResultHandler(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockCollection(ctrl)
+	mockCollection.EXPECT().Clone().Return(nil)
+	bulkInserter, err := NewBulkInserter(mockCollection)
+	assert.NoError(t, err)
+	mockHandler := func(result *mongo.InsertManyResult, err error) {}
+	bulkInserter.SetResultHandler(mockHandler)
+}
+
+func TestDbInserter_RemoveAll(t *testing.T) {
+	inserter := &dbInserter{}
+	inserter.documents = []interface{}{}
+	docs := inserter.RemoveAll()
+	assert.NotNil(t, docs)
+	assert.Empty(t, inserter.documents)
+}
+
+func Test_dbInserter_Execute(t *testing.T) {
+	type fields struct {
+		collection    collectionInserter
+		documents     []any
+		resultHandler ResultHandler
+	}
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockcollectionInserter(ctrl)
+	type args struct {
+		objs any
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+		mock   func()
+	}{
+		{
+			name: "empty doc",
+			fields: fields{
+				collection:    nil,
+				documents:     nil,
+				resultHandler: nil,
+			},
+			args: args{
+				objs: make([]any, 0),
+			},
+			mock: func() {},
+		},
+		{
+			name: "result handler",
+			fields: fields{
+				collection: mockCollection,
+				resultHandler: func(result *mongo.InsertManyResult, err error) {
+					assert.NotNil(t, err)
+				},
+			},
+			args: args{
+				objs: make([]any, 1),
+			},
+			mock: func() {
+				mockCollection.EXPECT().InsertMany(gomock.Any(), gomock.Any()).Return(&mongo.InsertManyResult{}, errors.New("error"))
+			},
+		},
+		{
+			name: "normal error handler",
+			fields: fields{
+				collection:    mockCollection,
+				resultHandler: nil,
+			},
+			args: args{
+				objs: make([]any, 1),
+			},
+			mock: func() {
+				mockCollection.EXPECT().InsertMany(gomock.Any(), gomock.Any()).Return(&mongo.InsertManyResult{}, errors.New("error"))
+			},
+		},
+		{
+			name: "no error",
+			fields: fields{
+				collection:    mockCollection,
+				resultHandler: nil,
+			},
+			args: args{
+				objs: make([]any, 1),
+			},
+			mock: func() {
+				mockCollection.EXPECT().InsertMany(gomock.Any(), gomock.Any()).Return(&mongo.InsertManyResult{}, nil)
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.mock()
+			in := &dbInserter{
+				collection:    tt.fields.collection,
+				documents:     tt.fields.documents,
+				resultHandler: tt.fields.resultHandler,
+			}
+			in.Execute(tt.args.objs)
+		})
+	}
 }

core/stores/mon/clientmanager.go

@@ -5,8 +5,8 @@ import (
 	"io"
 
 	"github.com/zeromicro/go-zero/core/syncx"
-	"go.mongodb.org/mongo-driver/mongo"
-	mopt "go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
 )
 
 var clientManager = syncx.NewResourceManager()
@@ -29,13 +29,13 @@ func Inject(key string, client *mongo.Client) {
 
 func getClient(url string, opts ...Option) (*mongo.Client, error) {
 	val, err := clientManager.GetResource(url, func() (io.Closer, error) {
-		o := mopt.Client().ApplyURI(url)
+		o := options.Client().ApplyURI(url)
 		opts = append([]Option{defaultTimeoutOption()}, opts...)
 		for _, opt := range opts {
 			opt(o)
 		}
 
-		cli, err := mongo.Connect(context.Background(), o)
+		cli, err := mongo.Connect(o)
 		if err != nil {
 			return nil, err
 		}

core/stores/mon/clientmanager_test.go

@@ -4,19 +4,13 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
+	"go.mongodb.org/mongo-driver/v2/mongo"
 )
 
-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) {
-		Inject(mtest.ClusterURI(), mt.Client)
-		cli, err := getClient(mtest.ClusterURI())
-		assert.Nil(t, err)
-		assert.Equal(t, mt.Client, cli)
-	})
+	c := &mongo.Client{}
+	Inject("foo", c)
+	cli, err := getClient("foo")
+	assert.Nil(t, err)
+	assert.Equal(t, c, cli)
 }

core/stores/mon/collection.go

@@ -1,3 +1,4 @@
+//go:generate mockgen -package mon -destination collection_mock.go -source collection.go Collection,monCollection
 package mon
 
 import (
@@ -8,9 +9,9 @@ import (
 	"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"
-	"go.mongodb.org/mongo-driver/x/mongo/driver/session"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
+	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/session"
 )
 
 const (
@@ -47,79 +48,79 @@ type (
 	// Collection defines a MongoDB collection.
 	Collection interface {
 		// Aggregate executes an aggregation pipeline.
-		Aggregate(ctx context.Context, pipeline any, opts ...*mopt.AggregateOptions) (
+		Aggregate(ctx context.Context, pipeline any, opts ...options.Lister[options.AggregateOptions]) (
 			*mongo.Cursor, error)
 		// BulkWrite performs a bulk write operation.
-		BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...*mopt.BulkWriteOptions) (
+		BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...options.Lister[options.BulkWriteOptions]) (
 			*mongo.BulkWriteResult, error)
 		// Clone creates a copy of this collection with the same settings.
-		Clone(opts ...*mopt.CollectionOptions) (*mongo.Collection, error)
+		Clone(opts ...options.Lister[options.CollectionOptions]) *mongo.Collection
 		// CountDocuments returns the number of documents in the collection that match the filter.
-		CountDocuments(ctx context.Context, filter any, opts ...*mopt.CountOptions) (int64, error)
+		CountDocuments(ctx context.Context, filter any, opts ...options.Lister[options.CountOptions]) (int64, error)
 		// Database returns the database that this collection is a part of.
 		Database() *mongo.Database
 		// DeleteMany deletes documents from the collection that match the filter.
-		DeleteMany(ctx context.Context, filter any, opts ...*mopt.DeleteOptions) (
+		DeleteMany(ctx context.Context, filter any, opts ...options.Lister[options.DeleteManyOptions]) (
 			*mongo.DeleteResult, error)
 		// DeleteOne deletes at most one document from the collection that matches the filter.
-		DeleteOne(ctx context.Context, filter any, opts ...*mopt.DeleteOptions) (
+		DeleteOne(ctx context.Context, filter any, opts ...options.Lister[options.DeleteOneOptions]) (
 			*mongo.DeleteResult, error)
 		// Distinct returns a list of distinct values for the given key across the collection.
 		Distinct(ctx context.Context, fieldName string, filter any,
-			opts ...*mopt.DistinctOptions) ([]any, error)
+			opts ...options.Lister[options.DistinctOptions]) (*mongo.DistinctResult, error)
 		// Drop drops this collection from database.
-		Drop(ctx context.Context) error
+		Drop(ctx context.Context, opts ...options.Lister[options.DropCollectionOptions]) error
 		// EstimatedDocumentCount returns an estimate of the count of documents in a collection
 		// using collection metadata.
-		EstimatedDocumentCount(ctx context.Context, opts ...*mopt.EstimatedDocumentCountOptions) (int64, error)
+		EstimatedDocumentCount(ctx context.Context, opts ...options.Lister[options.EstimatedDocumentCountOptions]) (int64, error)
 		// Find finds the documents matching the provided filter.
-		Find(ctx context.Context, filter any, opts ...*mopt.FindOptions) (*mongo.Cursor, error)
+		Find(ctx context.Context, filter any, opts ...options.Lister[options.FindOptions]) (*mongo.Cursor, error)
 		// FindOne returns up to one document that matches the provided filter.
-		FindOne(ctx context.Context, filter any, opts ...*mopt.FindOneOptions) (
+		FindOne(ctx context.Context, filter any, opts ...options.Lister[options.FindOneOptions]) (
 			*mongo.SingleResult, error)
 		// FindOneAndDelete returns at most one document that matches the filter. If the filter
 		// matches multiple documents, only the first document is deleted.
-		FindOneAndDelete(ctx context.Context, filter any, opts ...*mopt.FindOneAndDeleteOptions) (
+		FindOneAndDelete(ctx context.Context, filter any, opts ...options.Lister[options.FindOneAndDeleteOptions]) (
 			*mongo.SingleResult, error)
 		// FindOneAndReplace returns at most one document that matches the filter. If the filter
 		// matches multiple documents, FindOneAndReplace returns the first document in the
 		// collection that matches the filter.
 		FindOneAndReplace(ctx context.Context, filter, replacement any,
-			opts ...*mopt.FindOneAndReplaceOptions) (*mongo.SingleResult, error)
+			opts ...options.Lister[options.FindOneAndReplaceOptions]) (*mongo.SingleResult, error)
 		// FindOneAndUpdate returns at most one document that matches the filter. If the filter
 		// matches multiple documents, FindOneAndUpdate returns the first document in the
 		// collection that matches the filter.
 		FindOneAndUpdate(ctx context.Context, filter, update any,
-			opts ...*mopt.FindOneAndUpdateOptions) (*mongo.SingleResult, error)
+			opts ...options.Lister[options.FindOneAndUpdateOptions]) (*mongo.SingleResult, error)
 		// Indexes returns the index view for this collection.
 		Indexes() mongo.IndexView
 		// InsertMany inserts the provided documents.
-		InsertMany(ctx context.Context, documents []any, opts ...*mopt.InsertManyOptions) (
+		InsertMany(ctx context.Context, documents []any, opts ...options.Lister[options.InsertManyOptions]) (
 			*mongo.InsertManyResult, error)
 		// InsertOne inserts the provided document.
-		InsertOne(ctx context.Context, document any, opts ...*mopt.InsertOneOptions) (
+		InsertOne(ctx context.Context, document any, opts ...options.Lister[options.InsertOneOptions]) (
 			*mongo.InsertOneResult, error)
 		// ReplaceOne replaces at most one document that matches the filter.
 		ReplaceOne(ctx context.Context, filter, replacement any,
-			opts ...*mopt.ReplaceOptions) (*mongo.UpdateResult, error)
+			opts ...options.Lister[options.ReplaceOptions]) (*mongo.UpdateResult, error)
 		// UpdateByID updates a single document matching the provided filter.
 		UpdateByID(ctx context.Context, id, update any,
-			opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error)
+			opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error)
 		// UpdateMany updates the provided documents.
 		UpdateMany(ctx context.Context, filter, update any,
-			opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error)
+			opts ...options.Lister[options.UpdateManyOptions]) (*mongo.UpdateResult, error)
 		// UpdateOne updates a single document matching the provided filter.
 		UpdateOne(ctx context.Context, filter, update any,
-			opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error)
+			opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error)
 		// Watch returns a change stream cursor used to receive notifications of changes to the collection.
-		Watch(ctx context.Context, pipeline any, opts ...*mopt.ChangeStreamOptions) (
+		Watch(ctx context.Context, pipeline any, opts ...options.Lister[options.ChangeStreamOptions]) (
 			*mongo.ChangeStream, error)
 	}
 
 	decoratedCollection struct {
-		*mongo.Collection
-		name string
-		brk  breaker.Breaker
+		Collection monCollection
+		name       string
+		brk        breaker.Breaker
 	}
 
 	keepablePromise struct {
@@ -137,7 +138,7 @@ func newCollection(collection *mongo.Collection, brk breaker.Breaker) Collection
 }
 
 func (c *decoratedCollection) Aggregate(ctx context.Context, pipeline any,
-	opts ...*mopt.AggregateOptions) (cur *mongo.Cursor, err error) {
+	opts ...options.Lister[options.AggregateOptions]) (cur *mongo.Cursor, err error) {
 	ctx, span := startSpan(ctx, aggregate)
 	defer func() {
 		endSpan(span, err)
@@ -157,7 +158,7 @@ func (c *decoratedCollection) Aggregate(ctx context.Context, pipeline any,
 }
 
 func (c *decoratedCollection) BulkWrite(ctx context.Context, models []mongo.WriteModel,
-	opts ...*mopt.BulkWriteOptions) (res *mongo.BulkWriteResult, err error) {
+	opts ...options.Lister[options.BulkWriteOptions]) (res *mongo.BulkWriteResult, err error) {
 	ctx, span := startSpan(ctx, bulkWrite)
 	defer func() {
 		endSpan(span, err)
@@ -176,8 +177,12 @@ func (c *decoratedCollection) BulkWrite(ctx context.Context, models []mongo.Writ
 	return
 }
 
+func (c *decoratedCollection) Clone(opts ...options.Lister[options.CollectionOptions]) *mongo.Collection {
+	return c.Collection.Clone(opts...)
+}
+
 func (c *decoratedCollection) CountDocuments(ctx context.Context, filter any,
-	opts ...*mopt.CountOptions) (count int64, err error) {
+	opts ...options.Lister[options.CountOptions]) (count int64, err error) {
 	ctx, span := startSpan(ctx, countDocuments)
 	defer func() {
 		endSpan(span, err)
@@ -196,8 +201,12 @@ func (c *decoratedCollection) CountDocuments(ctx context.Context, filter any,
 	return
 }
 
+func (c *decoratedCollection) Database() *mongo.Database {
+	return c.Collection.Database()
+}
+
 func (c *decoratedCollection) DeleteMany(ctx context.Context, filter any,
-	opts ...*mopt.DeleteOptions) (res *mongo.DeleteResult, err error) {
+	opts ...options.Lister[options.DeleteManyOptions]) (res *mongo.DeleteResult, err error) {
 	ctx, span := startSpan(ctx, deleteMany)
 	defer func() {
 		endSpan(span, err)
@@ -217,7 +226,7 @@ func (c *decoratedCollection) DeleteMany(ctx context.Context, filter any,
 }
 
 func (c *decoratedCollection) DeleteOne(ctx context.Context, filter any,
-	opts ...*mopt.DeleteOptions) (res *mongo.DeleteResult, err error) {
+	opts ...options.Lister[options.DeleteOneOptions]) (res *mongo.DeleteResult, err error) {
 	ctx, span := startSpan(ctx, deleteOne)
 	defer func() {
 		endSpan(span, err)
@@ -237,7 +246,7 @@ func (c *decoratedCollection) DeleteOne(ctx context.Context, filter any,
 }
 
 func (c *decoratedCollection) Distinct(ctx context.Context, fieldName string, filter any,
-	opts ...*mopt.DistinctOptions) (val []any, err error) {
+	opts ...options.Lister[options.DistinctOptions]) (res *mongo.DistinctResult, err error) {
 	ctx, span := startSpan(ctx, distinct)
 	defer func() {
 		endSpan(span, err)
@@ -249,15 +258,20 @@ func (c *decoratedCollection) Distinct(ctx context.Context, fieldName string, fi
 			c.logDurationSimple(ctx, distinct, startTime, err)
 		}()
 
-		val, err = c.Collection.Distinct(ctx, fieldName, filter, opts...)
+		res = c.Collection.Distinct(ctx, fieldName, filter, opts...)
+		err = res.Err()
 		return err
 	}, acceptable)
 
 	return
 }
 
+func (c *decoratedCollection) Drop(ctx context.Context, opts ...options.Lister[options.DropCollectionOptions]) error {
+	return c.Collection.Drop(ctx, opts...)
+}
+
 func (c *decoratedCollection) EstimatedDocumentCount(ctx context.Context,
-	opts ...*mopt.EstimatedDocumentCountOptions) (val int64, err error) {
+	opts ...options.Lister[options.EstimatedDocumentCountOptions]) (val int64, err error) {
 	ctx, span := startSpan(ctx, estimatedDocumentCount)
 	defer func() {
 		endSpan(span, err)
@@ -277,7 +291,7 @@ func (c *decoratedCollection) EstimatedDocumentCount(ctx context.Context,
 }
 
 func (c *decoratedCollection) Find(ctx context.Context, filter any,
-	opts ...*mopt.FindOptions) (cur *mongo.Cursor, err error) {
+	opts ...options.Lister[options.FindOptions]) (cur *mongo.Cursor, err error) {
 	ctx, span := startSpan(ctx, find)
 	defer func() {
 		endSpan(span, err)
@@ -297,7 +311,7 @@ func (c *decoratedCollection) Find(ctx context.Context, filter any,
 }
 
 func (c *decoratedCollection) FindOne(ctx context.Context, filter any,
-	opts ...*mopt.FindOneOptions) (res *mongo.SingleResult, err error) {
+	opts ...options.Lister[options.FindOneOptions]) (res *mongo.SingleResult, err error) {
 	ctx, span := startSpan(ctx, findOne)
 	defer func() {
 		endSpan(span, err)
@@ -318,7 +332,7 @@ func (c *decoratedCollection) FindOne(ctx context.Context, filter any,
 }
 
 func (c *decoratedCollection) FindOneAndDelete(ctx context.Context, filter any,
-	opts ...*mopt.FindOneAndDeleteOptions) (res *mongo.SingleResult, err error) {
+	opts ...options.Lister[options.FindOneAndDeleteOptions]) (res *mongo.SingleResult, err error) {
 	ctx, span := startSpan(ctx, findOneAndDelete)
 	defer func() {
 		endSpan(span, err)
@@ -339,7 +353,7 @@ func (c *decoratedCollection) FindOneAndDelete(ctx context.Context, filter any,
 }
 
 func (c *decoratedCollection) FindOneAndReplace(ctx context.Context, filter any,
-	replacement any, opts ...*mopt.FindOneAndReplaceOptions) (
+	replacement any, opts ...options.Lister[options.FindOneAndReplaceOptions]) (
 	res *mongo.SingleResult, err error) {
 	ctx, span := startSpan(ctx, findOneAndReplace)
 	defer func() {
@@ -361,7 +375,7 @@ func (c *decoratedCollection) FindOneAndReplace(ctx context.Context, filter any,
 }
 
 func (c *decoratedCollection) FindOneAndUpdate(ctx context.Context, filter, update any,
-	opts ...*mopt.FindOneAndUpdateOptions) (res *mongo.SingleResult, err error) {
+	opts ...options.Lister[options.FindOneAndUpdateOptions]) (res *mongo.SingleResult, err error) {
 	ctx, span := startSpan(ctx, findOneAndUpdate)
 	defer func() {
 		endSpan(span, err)
@@ -381,8 +395,12 @@ func (c *decoratedCollection) FindOneAndUpdate(ctx context.Context, filter, upda
 	return
 }
 
+func (c *decoratedCollection) Indexes() mongo.IndexView {
+	return c.Collection.Indexes()
+}
+
 func (c *decoratedCollection) InsertMany(ctx context.Context, documents []any,
-	opts ...*mopt.InsertManyOptions) (res *mongo.InsertManyResult, err error) {
+	opts ...options.Lister[options.InsertManyOptions]) (res *mongo.InsertManyResult, err error) {
 	ctx, span := startSpan(ctx, insertMany)
 	defer func() {
 		endSpan(span, err)
@@ -402,7 +420,7 @@ func (c *decoratedCollection) InsertMany(ctx context.Context, documents []any,
 }
 
 func (c *decoratedCollection) InsertOne(ctx context.Context, document any,
-	opts ...*mopt.InsertOneOptions) (res *mongo.InsertOneResult, err error) {
+	opts ...options.Lister[options.InsertOneOptions]) (res *mongo.InsertOneResult, err error) {
 	ctx, span := startSpan(ctx, insertOne)
 	defer func() {
 		endSpan(span, err)
@@ -422,7 +440,7 @@ func (c *decoratedCollection) InsertOne(ctx context.Context, document any,
 }
 
 func (c *decoratedCollection) ReplaceOne(ctx context.Context, filter, replacement any,
-	opts ...*mopt.ReplaceOptions) (res *mongo.UpdateResult, err error) {
+	opts ...options.Lister[options.ReplaceOptions]) (res *mongo.UpdateResult, err error) {
 	ctx, span := startSpan(ctx, replaceOne)
 	defer func() {
 		endSpan(span, err)
@@ -442,7 +460,7 @@ func (c *decoratedCollection) ReplaceOne(ctx context.Context, filter, replacemen
 }
 
 func (c *decoratedCollection) UpdateByID(ctx context.Context, id, update any,
-	opts ...*mopt.UpdateOptions) (res *mongo.UpdateResult, err error) {
+	opts ...options.Lister[options.UpdateOneOptions]) (res *mongo.UpdateResult, err error) {
 	ctx, span := startSpan(ctx, updateByID)
 	defer func() {
 		endSpan(span, err)
@@ -462,7 +480,7 @@ func (c *decoratedCollection) UpdateByID(ctx context.Context, id, update any,
 }
 
 func (c *decoratedCollection) UpdateMany(ctx context.Context, filter, update any,
-	opts ...*mopt.UpdateOptions) (res *mongo.UpdateResult, err error) {
+	opts ...options.Lister[options.UpdateManyOptions]) (res *mongo.UpdateResult, err error) {
 	ctx, span := startSpan(ctx, updateMany)
 	defer func() {
 		endSpan(span, err)
@@ -482,7 +500,7 @@ func (c *decoratedCollection) UpdateMany(ctx context.Context, filter, update any
 }
 
 func (c *decoratedCollection) UpdateOne(ctx context.Context, filter, update any,
-	opts ...*mopt.UpdateOptions) (res *mongo.UpdateResult, err error) {
+	opts ...options.Lister[options.UpdateOneOptions]) (res *mongo.UpdateResult, err error) {
 	ctx, span := startSpan(ctx, updateOne)
 	defer func() {
 		endSpan(span, err)
@@ -501,6 +519,11 @@ func (c *decoratedCollection) UpdateOne(ctx context.Context, filter, update any,
 	return
 }
 
+func (c *decoratedCollection) Watch(ctx context.Context, pipeline any, opts ...options.Lister[options.ChangeStreamOptions]) (
+	*mongo.ChangeStream, error) {
+	return c.Collection.Watch(ctx, pipeline, opts...)
+}
+
 func (c *decoratedCollection) logDuration(ctx context.Context, method string,
 	startTime time.Duration, err error, docs ...any) {
 	logDurationWithDocs(ctx, c.name, method, startTime, err, docs...)
@@ -546,3 +569,71 @@ func isDupKeyError(err error) bool {
 
 	return e.HasErrorCode(duplicateKeyCode)
 }
+
+// monCollection defines a MongoDB collection, used for unit test
+type monCollection interface {
+	// Aggregate executes an aggregation pipeline.
+	Aggregate(ctx context.Context, pipeline any, opts ...options.Lister[options.AggregateOptions]) (
+		*mongo.Cursor, error)
+	// BulkWrite performs a bulk write operation.
+	BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...options.Lister[options.BulkWriteOptions]) (
+		*mongo.BulkWriteResult, error)
+	// Clone creates a copy of this collection with the same settings.
+	Clone(opts ...options.Lister[options.CollectionOptions]) *mongo.Collection
+	// CountDocuments returns the number of documents in the collection that match the filter.
+	CountDocuments(ctx context.Context, filter any, opts ...options.Lister[options.CountOptions]) (int64, error)
+	// Database returns the database that this collection is a part of.
+	Database() *mongo.Database
+	// DeleteMany deletes documents from the collection that match the filter.
+	DeleteMany(ctx context.Context, filter any, opts ...options.Lister[options.DeleteManyOptions]) (
+		*mongo.DeleteResult, error)
+	// DeleteOne deletes at most one document from the collection that matches the filter.
+	DeleteOne(ctx context.Context, filter any, opts ...options.Lister[options.DeleteOneOptions]) (
+		*mongo.DeleteResult, error)
+	// Distinct returns a list of distinct values for the given key across the collection.
+	Distinct(ctx context.Context, fieldName string, filter any,
+		opts ...options.Lister[options.DistinctOptions]) *mongo.DistinctResult
+	// Drop drops this collection from database.
+	Drop(ctx context.Context, opts ...options.Lister[options.DropCollectionOptions]) error
+	// EstimatedDocumentCount returns an estimate of the count of documents in a collection
+	// using collection metadata.
+	EstimatedDocumentCount(ctx context.Context, opts ...options.Lister[options.EstimatedDocumentCountOptions]) (int64, error)
+	// Find finds the documents matching the provided filter.
+	Find(ctx context.Context, filter any, opts ...options.Lister[options.FindOptions]) (*mongo.Cursor, error)
+	// FindOne returns up to one document that matches the provided filter.
+	FindOne(ctx context.Context, filter any, opts ...options.Lister[options.FindOneOptions]) *mongo.SingleResult
+	// FindOneAndDelete returns at most one document that matches the filter. If the filter
+	// matches multiple documents, only the first document is deleted.
+	FindOneAndDelete(ctx context.Context, filter any, opts ...options.Lister[options.FindOneAndDeleteOptions]) *mongo.SingleResult
+	// FindOneAndReplace returns at most one document that matches the filter. If the filter
+	// matches multiple documents, FindOneAndReplace returns the first document in the
+	// collection that matches the filter.
+	FindOneAndReplace(ctx context.Context, filter, replacement any,
+		opts ...options.Lister[options.FindOneAndReplaceOptions]) *mongo.SingleResult
+	// FindOneAndUpdate returns at most one document that matches the filter. If the filter
+	// matches multiple documents, FindOneAndUpdate returns the first document in the
+	// collection that matches the filter.
+	FindOneAndUpdate(ctx context.Context, filter, update any,
+		opts ...options.Lister[options.FindOneAndUpdateOptions]) *mongo.SingleResult
+	// Indexes returns the index view for this collection.
+	Indexes() mongo.IndexView
+	// InsertMany inserts the provided documents.
+	InsertMany(ctx context.Context, documents interface{}, opts ...options.Lister[options.InsertManyOptions]) (*mongo.InsertManyResult, error)
+	// InsertOne inserts the provided document.
+	InsertOne(ctx context.Context, document any, opts ...options.Lister[options.InsertOneOptions]) (*mongo.InsertOneResult, error)
+	// ReplaceOne replaces at most one document that matches the filter.
+	ReplaceOne(ctx context.Context, filter, replacement any,
+		opts ...options.Lister[options.ReplaceOptions]) (*mongo.UpdateResult, error)
+	// UpdateByID updates a single document matching the provided filter.
+	UpdateByID(ctx context.Context, id, update any,
+		opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error)
+	// UpdateMany updates the provided documents.
+	UpdateMany(ctx context.Context, filter, update any,
+		opts ...options.Lister[options.UpdateManyOptions]) (*mongo.UpdateResult, error)
+	// UpdateOne updates a single document matching the provided filter.
+	UpdateOne(ctx context.Context, filter, update any,
+		opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error)
+	// Watch returns a change stream cursor used to receive notifications of changes to the collection.
+	Watch(ctx context.Context, pipeline any, opts ...options.Lister[options.ChangeStreamOptions]) (
+		*mongo.ChangeStream, error)
+}

core/stores/mon/collection_mock.go

@@ -0,0 +1,952 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: collection.go
+//
+// Generated by this command:
+//
+//	mockgen -package mon -destination collection_mock.go -source collection.go Collection,monCollection
+//
+
+// Package mon is a generated GoMock package.
+package mon
+
+import (
+	context "context"
+	reflect "reflect"
+
+	mongo "go.mongodb.org/mongo-driver/v2/mongo"
+	options "go.mongodb.org/mongo-driver/v2/mongo/options"
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MockCollection is a mock of Collection interface.
+type MockCollection struct {
+	ctrl     *gomock.Controller
+	recorder *MockCollectionMockRecorder
+	isgomock struct{}
+}
+
+// MockCollectionMockRecorder is the mock recorder for MockCollection.
+type MockCollectionMockRecorder struct {
+	mock *MockCollection
+}
+
+// NewMockCollection creates a new mock instance.
+func NewMockCollection(ctrl *gomock.Controller) *MockCollection {
+	mock := &MockCollection{ctrl: ctrl}
+	mock.recorder = &MockCollectionMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockCollection) EXPECT() *MockCollectionMockRecorder {
+	return m.recorder
+}
+
+// Aggregate mocks base method.
+func (m *MockCollection) Aggregate(ctx context.Context, pipeline any, opts ...options.Lister[options.AggregateOptions]) (*mongo.Cursor, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, pipeline}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Aggregate", varargs...)
+	ret0, _ := ret[0].(*mongo.Cursor)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Aggregate indicates an expected call of Aggregate.
+func (mr *MockCollectionMockRecorder) Aggregate(ctx, pipeline any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, pipeline}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Aggregate", reflect.TypeOf((*MockCollection)(nil).Aggregate), varargs...)
+}
+
+// BulkWrite mocks base method.
+func (m *MockCollection) BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...options.Lister[options.BulkWriteOptions]) (*mongo.BulkWriteResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, models}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "BulkWrite", varargs...)
+	ret0, _ := ret[0].(*mongo.BulkWriteResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// BulkWrite indicates an expected call of BulkWrite.
+func (mr *MockCollectionMockRecorder) BulkWrite(ctx, models any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, models}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BulkWrite", reflect.TypeOf((*MockCollection)(nil).BulkWrite), varargs...)
+}
+
+// Clone mocks base method.
+func (m *MockCollection) Clone(opts ...options.Lister[options.CollectionOptions]) *mongo.Collection {
+	m.ctrl.T.Helper()
+	varargs := []any{}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Clone", varargs...)
+	ret0, _ := ret[0].(*mongo.Collection)
+	return ret0
+}
+
+// Clone indicates an expected call of Clone.
+func (mr *MockCollectionMockRecorder) Clone(opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Clone", reflect.TypeOf((*MockCollection)(nil).Clone), opts...)
+}
+
+// CountDocuments mocks base method.
+func (m *MockCollection) CountDocuments(ctx context.Context, filter any, opts ...options.Lister[options.CountOptions]) (int64, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "CountDocuments", varargs...)
+	ret0, _ := ret[0].(int64)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// CountDocuments indicates an expected call of CountDocuments.
+func (mr *MockCollectionMockRecorder) CountDocuments(ctx, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountDocuments", reflect.TypeOf((*MockCollection)(nil).CountDocuments), varargs...)
+}
+
+// Database mocks base method.
+func (m *MockCollection) Database() *mongo.Database {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Database")
+	ret0, _ := ret[0].(*mongo.Database)
+	return ret0
+}
+
+// Database indicates an expected call of Database.
+func (mr *MockCollectionMockRecorder) Database() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Database", reflect.TypeOf((*MockCollection)(nil).Database))
+}
+
+// DeleteMany mocks base method.
+func (m *MockCollection) DeleteMany(ctx context.Context, filter any, opts ...options.Lister[options.DeleteManyOptions]) (*mongo.DeleteResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "DeleteMany", varargs...)
+	ret0, _ := ret[0].(*mongo.DeleteResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DeleteMany indicates an expected call of DeleteMany.
+func (mr *MockCollectionMockRecorder) DeleteMany(ctx, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteMany", reflect.TypeOf((*MockCollection)(nil).DeleteMany), varargs...)
+}
+
+// DeleteOne mocks base method.
+func (m *MockCollection) DeleteOne(ctx context.Context, filter any, opts ...options.Lister[options.DeleteOneOptions]) (*mongo.DeleteResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "DeleteOne", varargs...)
+	ret0, _ := ret[0].(*mongo.DeleteResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DeleteOne indicates an expected call of DeleteOne.
+func (mr *MockCollectionMockRecorder) DeleteOne(ctx, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOne", reflect.TypeOf((*MockCollection)(nil).DeleteOne), varargs...)
+}
+
+// Distinct mocks base method.
+func (m *MockCollection) Distinct(ctx context.Context, fieldName string, filter any, opts ...options.Lister[options.DistinctOptions]) (*mongo.DistinctResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, fieldName, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Distinct", varargs...)
+	ret0, _ := ret[0].(*mongo.DistinctResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Distinct indicates an expected call of Distinct.
+func (mr *MockCollectionMockRecorder) Distinct(ctx, fieldName, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, fieldName, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Distinct", reflect.TypeOf((*MockCollection)(nil).Distinct), varargs...)
+}
+
+// Drop mocks base method.
+func (m *MockCollection) Drop(ctx context.Context, opts ...options.Lister[options.DropCollectionOptions]) error {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Drop", varargs...)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Drop indicates an expected call of Drop.
+func (mr *MockCollectionMockRecorder) Drop(ctx any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Drop", reflect.TypeOf((*MockCollection)(nil).Drop), varargs...)
+}
+
+// EstimatedDocumentCount mocks base method.
+func (m *MockCollection) EstimatedDocumentCount(ctx context.Context, opts ...options.Lister[options.EstimatedDocumentCountOptions]) (int64, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "EstimatedDocumentCount", varargs...)
+	ret0, _ := ret[0].(int64)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// EstimatedDocumentCount indicates an expected call of EstimatedDocumentCount.
+func (mr *MockCollectionMockRecorder) EstimatedDocumentCount(ctx any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EstimatedDocumentCount", reflect.TypeOf((*MockCollection)(nil).EstimatedDocumentCount), varargs...)
+}
+
+// Find mocks base method.
+func (m *MockCollection) Find(ctx context.Context, filter any, opts ...options.Lister[options.FindOptions]) (*mongo.Cursor, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Find", varargs...)
+	ret0, _ := ret[0].(*mongo.Cursor)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Find indicates an expected call of Find.
+func (mr *MockCollectionMockRecorder) Find(ctx, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Find", reflect.TypeOf((*MockCollection)(nil).Find), varargs...)
+}
+
+// FindOne mocks base method.
+func (m *MockCollection) FindOne(ctx context.Context, filter any, opts ...options.Lister[options.FindOneOptions]) (*mongo.SingleResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "FindOne", varargs...)
+	ret0, _ := ret[0].(*mongo.SingleResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// FindOne indicates an expected call of FindOne.
+func (mr *MockCollectionMockRecorder) FindOne(ctx, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindOne", reflect.TypeOf((*MockCollection)(nil).FindOne), varargs...)
+}
+
+// FindOneAndDelete mocks base method.
+func (m *MockCollection) FindOneAndDelete(ctx context.Context, filter any, opts ...options.Lister[options.FindOneAndDeleteOptions]) (*mongo.SingleResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "FindOneAndDelete", varargs...)
+	ret0, _ := ret[0].(*mongo.SingleResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// FindOneAndDelete indicates an expected call of FindOneAndDelete.
+func (mr *MockCollectionMockRecorder) FindOneAndDelete(ctx, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindOneAndDelete", reflect.TypeOf((*MockCollection)(nil).FindOneAndDelete), varargs...)
+}
+
+// FindOneAndReplace mocks base method.
+func (m *MockCollection) FindOneAndReplace(ctx context.Context, filter, replacement any, opts ...options.Lister[options.FindOneAndReplaceOptions]) (*mongo.SingleResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter, replacement}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "FindOneAndReplace", varargs...)
+	ret0, _ := ret[0].(*mongo.SingleResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// FindOneAndReplace indicates an expected call of FindOneAndReplace.
+func (mr *MockCollectionMockRecorder) FindOneAndReplace(ctx, filter, replacement any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter, replacement}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindOneAndReplace", reflect.TypeOf((*MockCollection)(nil).FindOneAndReplace), varargs...)
+}
+
+// FindOneAndUpdate mocks base method.
+func (m *MockCollection) FindOneAndUpdate(ctx context.Context, filter, update any, opts ...options.Lister[options.FindOneAndUpdateOptions]) (*mongo.SingleResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter, update}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "FindOneAndUpdate", varargs...)
+	ret0, _ := ret[0].(*mongo.SingleResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// FindOneAndUpdate indicates an expected call of FindOneAndUpdate.
+func (mr *MockCollectionMockRecorder) FindOneAndUpdate(ctx, filter, update any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter, update}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindOneAndUpdate", reflect.TypeOf((*MockCollection)(nil).FindOneAndUpdate), varargs...)
+}
+
+// Indexes mocks base method.
+func (m *MockCollection) Indexes() mongo.IndexView {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Indexes")
+	ret0, _ := ret[0].(mongo.IndexView)
+	return ret0
+}
+
+// Indexes indicates an expected call of Indexes.
+func (mr *MockCollectionMockRecorder) Indexes() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Indexes", reflect.TypeOf((*MockCollection)(nil).Indexes))
+}
+
+// InsertMany mocks base method.
+func (m *MockCollection) InsertMany(ctx context.Context, documents []any, opts ...options.Lister[options.InsertManyOptions]) (*mongo.InsertManyResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, documents}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "InsertMany", varargs...)
+	ret0, _ := ret[0].(*mongo.InsertManyResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// InsertMany indicates an expected call of InsertMany.
+func (mr *MockCollectionMockRecorder) InsertMany(ctx, documents any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, documents}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertMany", reflect.TypeOf((*MockCollection)(nil).InsertMany), varargs...)
+}
+
+// InsertOne mocks base method.
+func (m *MockCollection) InsertOne(ctx context.Context, document any, opts ...options.Lister[options.InsertOneOptions]) (*mongo.InsertOneResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, document}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "InsertOne", varargs...)
+	ret0, _ := ret[0].(*mongo.InsertOneResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// InsertOne indicates an expected call of InsertOne.
+func (mr *MockCollectionMockRecorder) InsertOne(ctx, document any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, document}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertOne", reflect.TypeOf((*MockCollection)(nil).InsertOne), varargs...)
+}
+
+// ReplaceOne mocks base method.
+func (m *MockCollection) ReplaceOne(ctx context.Context, filter, replacement any, opts ...options.Lister[options.ReplaceOptions]) (*mongo.UpdateResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter, replacement}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "ReplaceOne", varargs...)
+	ret0, _ := ret[0].(*mongo.UpdateResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// ReplaceOne indicates an expected call of ReplaceOne.
+func (mr *MockCollectionMockRecorder) ReplaceOne(ctx, filter, replacement any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter, replacement}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReplaceOne", reflect.TypeOf((*MockCollection)(nil).ReplaceOne), varargs...)
+}
+
+// UpdateByID mocks base method.
+func (m *MockCollection) UpdateByID(ctx context.Context, id, update any, opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, id, update}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "UpdateByID", varargs...)
+	ret0, _ := ret[0].(*mongo.UpdateResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// UpdateByID indicates an expected call of UpdateByID.
+func (mr *MockCollectionMockRecorder) UpdateByID(ctx, id, update any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, id, update}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateByID", reflect.TypeOf((*MockCollection)(nil).UpdateByID), varargs...)
+}
+
+// UpdateMany mocks base method.
+func (m *MockCollection) UpdateMany(ctx context.Context, filter, update any, opts ...options.Lister[options.UpdateManyOptions]) (*mongo.UpdateResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter, update}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "UpdateMany", varargs...)
+	ret0, _ := ret[0].(*mongo.UpdateResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// UpdateMany indicates an expected call of UpdateMany.
+func (mr *MockCollectionMockRecorder) UpdateMany(ctx, filter, update any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter, update}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateMany", reflect.TypeOf((*MockCollection)(nil).UpdateMany), varargs...)
+}
+
+// UpdateOne mocks base method.
+func (m *MockCollection) UpdateOne(ctx context.Context, filter, update any, opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter, update}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "UpdateOne", varargs...)
+	ret0, _ := ret[0].(*mongo.UpdateResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// UpdateOne indicates an expected call of UpdateOne.
+func (mr *MockCollectionMockRecorder) UpdateOne(ctx, filter, update any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter, update}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateOne", reflect.TypeOf((*MockCollection)(nil).UpdateOne), varargs...)
+}
+
+// Watch mocks base method.
+func (m *MockCollection) Watch(ctx context.Context, pipeline any, opts ...options.Lister[options.ChangeStreamOptions]) (*mongo.ChangeStream, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, pipeline}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Watch", varargs...)
+	ret0, _ := ret[0].(*mongo.ChangeStream)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Watch indicates an expected call of Watch.
+func (mr *MockCollectionMockRecorder) Watch(ctx, pipeline any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, pipeline}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Watch", reflect.TypeOf((*MockCollection)(nil).Watch), varargs...)
+}
+
+// MockmonCollection is a mock of monCollection interface.
+type MockmonCollection struct {
+	ctrl     *gomock.Controller
+	recorder *MockmonCollectionMockRecorder
+	isgomock struct{}
+}
+
+// MockmonCollectionMockRecorder is the mock recorder for MockmonCollection.
+type MockmonCollectionMockRecorder struct {
+	mock *MockmonCollection
+}
+
+// NewMockmonCollection creates a new mock instance.
+func NewMockmonCollection(ctrl *gomock.Controller) *MockmonCollection {
+	mock := &MockmonCollection{ctrl: ctrl}
+	mock.recorder = &MockmonCollectionMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockmonCollection) EXPECT() *MockmonCollectionMockRecorder {
+	return m.recorder
+}
+
+// Aggregate mocks base method.
+func (m *MockmonCollection) Aggregate(ctx context.Context, pipeline any, opts ...options.Lister[options.AggregateOptions]) (*mongo.Cursor, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, pipeline}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Aggregate", varargs...)
+	ret0, _ := ret[0].(*mongo.Cursor)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Aggregate indicates an expected call of Aggregate.
+func (mr *MockmonCollectionMockRecorder) Aggregate(ctx, pipeline any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, pipeline}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Aggregate", reflect.TypeOf((*MockmonCollection)(nil).Aggregate), varargs...)
+}
+
+// BulkWrite mocks base method.
+func (m *MockmonCollection) BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...options.Lister[options.BulkWriteOptions]) (*mongo.BulkWriteResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, models}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "BulkWrite", varargs...)
+	ret0, _ := ret[0].(*mongo.BulkWriteResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// BulkWrite indicates an expected call of BulkWrite.
+func (mr *MockmonCollectionMockRecorder) BulkWrite(ctx, models any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, models}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BulkWrite", reflect.TypeOf((*MockmonCollection)(nil).BulkWrite), varargs...)
+}
+
+// Clone mocks base method.
+func (m *MockmonCollection) Clone(opts ...options.Lister[options.CollectionOptions]) *mongo.Collection {
+	m.ctrl.T.Helper()
+	varargs := []any{}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Clone", varargs...)
+	ret0, _ := ret[0].(*mongo.Collection)
+	return ret0
+}
+
+// Clone indicates an expected call of Clone.
+func (mr *MockmonCollectionMockRecorder) Clone(opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Clone", reflect.TypeOf((*MockmonCollection)(nil).Clone), opts...)
+}
+
+// CountDocuments mocks base method.
+func (m *MockmonCollection) CountDocuments(ctx context.Context, filter any, opts ...options.Lister[options.CountOptions]) (int64, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "CountDocuments", varargs...)
+	ret0, _ := ret[0].(int64)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// CountDocuments indicates an expected call of CountDocuments.
+func (mr *MockmonCollectionMockRecorder) CountDocuments(ctx, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountDocuments", reflect.TypeOf((*MockmonCollection)(nil).CountDocuments), varargs...)
+}
+
+// Database mocks base method.
+func (m *MockmonCollection) Database() *mongo.Database {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Database")
+	ret0, _ := ret[0].(*mongo.Database)
+	return ret0
+}
+
+// Database indicates an expected call of Database.
+func (mr *MockmonCollectionMockRecorder) Database() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Database", reflect.TypeOf((*MockmonCollection)(nil).Database))
+}
+
+// DeleteMany mocks base method.
+func (m *MockmonCollection) DeleteMany(ctx context.Context, filter any, opts ...options.Lister[options.DeleteManyOptions]) (*mongo.DeleteResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "DeleteMany", varargs...)
+	ret0, _ := ret[0].(*mongo.DeleteResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DeleteMany indicates an expected call of DeleteMany.
+func (mr *MockmonCollectionMockRecorder) DeleteMany(ctx, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteMany", reflect.TypeOf((*MockmonCollection)(nil).DeleteMany), varargs...)
+}
+
+// DeleteOne mocks base method.
+func (m *MockmonCollection) DeleteOne(ctx context.Context, filter any, opts ...options.Lister[options.DeleteOneOptions]) (*mongo.DeleteResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "DeleteOne", varargs...)
+	ret0, _ := ret[0].(*mongo.DeleteResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DeleteOne indicates an expected call of DeleteOne.
+func (mr *MockmonCollectionMockRecorder) DeleteOne(ctx, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOne", reflect.TypeOf((*MockmonCollection)(nil).DeleteOne), varargs...)
+}
+
+// Distinct mocks base method.
+func (m *MockmonCollection) Distinct(ctx context.Context, fieldName string, filter any, opts ...options.Lister[options.DistinctOptions]) *mongo.DistinctResult {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, fieldName, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Distinct", varargs...)
+	ret0, _ := ret[0].(*mongo.DistinctResult)
+	return ret0
+}
+
+// Distinct indicates an expected call of Distinct.
+func (mr *MockmonCollectionMockRecorder) Distinct(ctx, fieldName, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, fieldName, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Distinct", reflect.TypeOf((*MockmonCollection)(nil).Distinct), varargs...)
+}
+
+// Drop mocks base method.
+func (m *MockmonCollection) Drop(ctx context.Context, opts ...options.Lister[options.DropCollectionOptions]) error {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Drop", varargs...)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Drop indicates an expected call of Drop.
+func (mr *MockmonCollectionMockRecorder) Drop(ctx any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Drop", reflect.TypeOf((*MockmonCollection)(nil).Drop), varargs...)
+}
+
+// EstimatedDocumentCount mocks base method.
+func (m *MockmonCollection) EstimatedDocumentCount(ctx context.Context, opts ...options.Lister[options.EstimatedDocumentCountOptions]) (int64, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "EstimatedDocumentCount", varargs...)
+	ret0, _ := ret[0].(int64)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// EstimatedDocumentCount indicates an expected call of EstimatedDocumentCount.
+func (mr *MockmonCollectionMockRecorder) EstimatedDocumentCount(ctx any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EstimatedDocumentCount", reflect.TypeOf((*MockmonCollection)(nil).EstimatedDocumentCount), varargs...)
+}
+
+// Find mocks base method.
+func (m *MockmonCollection) Find(ctx context.Context, filter any, opts ...options.Lister[options.FindOptions]) (*mongo.Cursor, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Find", varargs...)
+	ret0, _ := ret[0].(*mongo.Cursor)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Find indicates an expected call of Find.
+func (mr *MockmonCollectionMockRecorder) Find(ctx, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Find", reflect.TypeOf((*MockmonCollection)(nil).Find), varargs...)
+}
+
+// FindOne mocks base method.
+func (m *MockmonCollection) FindOne(ctx context.Context, filter any, opts ...options.Lister[options.FindOneOptions]) *mongo.SingleResult {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "FindOne", varargs...)
+	ret0, _ := ret[0].(*mongo.SingleResult)
+	return ret0
+}
+
+// FindOne indicates an expected call of FindOne.
+func (mr *MockmonCollectionMockRecorder) FindOne(ctx, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindOne", reflect.TypeOf((*MockmonCollection)(nil).FindOne), varargs...)
+}
+
+// FindOneAndDelete mocks base method.
+func (m *MockmonCollection) FindOneAndDelete(ctx context.Context, filter any, opts ...options.Lister[options.FindOneAndDeleteOptions]) *mongo.SingleResult {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "FindOneAndDelete", varargs...)
+	ret0, _ := ret[0].(*mongo.SingleResult)
+	return ret0
+}
+
+// FindOneAndDelete indicates an expected call of FindOneAndDelete.
+func (mr *MockmonCollectionMockRecorder) FindOneAndDelete(ctx, filter any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindOneAndDelete", reflect.TypeOf((*MockmonCollection)(nil).FindOneAndDelete), varargs...)
+}
+
+// FindOneAndReplace mocks base method.
+func (m *MockmonCollection) FindOneAndReplace(ctx context.Context, filter, replacement any, opts ...options.Lister[options.FindOneAndReplaceOptions]) *mongo.SingleResult {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter, replacement}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "FindOneAndReplace", varargs...)
+	ret0, _ := ret[0].(*mongo.SingleResult)
+	return ret0
+}
+
+// FindOneAndReplace indicates an expected call of FindOneAndReplace.
+func (mr *MockmonCollectionMockRecorder) FindOneAndReplace(ctx, filter, replacement any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter, replacement}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindOneAndReplace", reflect.TypeOf((*MockmonCollection)(nil).FindOneAndReplace), varargs...)
+}
+
+// FindOneAndUpdate mocks base method.
+func (m *MockmonCollection) FindOneAndUpdate(ctx context.Context, filter, update any, opts ...options.Lister[options.FindOneAndUpdateOptions]) *mongo.SingleResult {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter, update}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "FindOneAndUpdate", varargs...)
+	ret0, _ := ret[0].(*mongo.SingleResult)
+	return ret0
+}
+
+// FindOneAndUpdate indicates an expected call of FindOneAndUpdate.
+func (mr *MockmonCollectionMockRecorder) FindOneAndUpdate(ctx, filter, update any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter, update}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindOneAndUpdate", reflect.TypeOf((*MockmonCollection)(nil).FindOneAndUpdate), varargs...)
+}
+
+// Indexes mocks base method.
+func (m *MockmonCollection) Indexes() mongo.IndexView {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Indexes")
+	ret0, _ := ret[0].(mongo.IndexView)
+	return ret0
+}
+
+// Indexes indicates an expected call of Indexes.
+func (mr *MockmonCollectionMockRecorder) Indexes() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Indexes", reflect.TypeOf((*MockmonCollection)(nil).Indexes))
+}
+
+// InsertMany mocks base method.
+func (m *MockmonCollection) InsertMany(ctx context.Context, documents any, opts ...options.Lister[options.InsertManyOptions]) (*mongo.InsertManyResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, documents}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "InsertMany", varargs...)
+	ret0, _ := ret[0].(*mongo.InsertManyResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// InsertMany indicates an expected call of InsertMany.
+func (mr *MockmonCollectionMockRecorder) InsertMany(ctx, documents any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, documents}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertMany", reflect.TypeOf((*MockmonCollection)(nil).InsertMany), varargs...)
+}
+
+// InsertOne mocks base method.
+func (m *MockmonCollection) InsertOne(ctx context.Context, document any, opts ...options.Lister[options.InsertOneOptions]) (*mongo.InsertOneResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, document}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "InsertOne", varargs...)
+	ret0, _ := ret[0].(*mongo.InsertOneResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// InsertOne indicates an expected call of InsertOne.
+func (mr *MockmonCollectionMockRecorder) InsertOne(ctx, document any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, document}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertOne", reflect.TypeOf((*MockmonCollection)(nil).InsertOne), varargs...)
+}
+
+// ReplaceOne mocks base method.
+func (m *MockmonCollection) ReplaceOne(ctx context.Context, filter, replacement any, opts ...options.Lister[options.ReplaceOptions]) (*mongo.UpdateResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter, replacement}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "ReplaceOne", varargs...)
+	ret0, _ := ret[0].(*mongo.UpdateResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// ReplaceOne indicates an expected call of ReplaceOne.
+func (mr *MockmonCollectionMockRecorder) ReplaceOne(ctx, filter, replacement any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter, replacement}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReplaceOne", reflect.TypeOf((*MockmonCollection)(nil).ReplaceOne), varargs...)
+}
+
+// UpdateByID mocks base method.
+func (m *MockmonCollection) UpdateByID(ctx context.Context, id, update any, opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, id, update}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "UpdateByID", varargs...)
+	ret0, _ := ret[0].(*mongo.UpdateResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// UpdateByID indicates an expected call of UpdateByID.
+func (mr *MockmonCollectionMockRecorder) UpdateByID(ctx, id, update any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, id, update}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateByID", reflect.TypeOf((*MockmonCollection)(nil).UpdateByID), varargs...)
+}
+
+// UpdateMany mocks base method.
+func (m *MockmonCollection) UpdateMany(ctx context.Context, filter, update any, opts ...options.Lister[options.UpdateManyOptions]) (*mongo.UpdateResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter, update}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "UpdateMany", varargs...)
+	ret0, _ := ret[0].(*mongo.UpdateResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// UpdateMany indicates an expected call of UpdateMany.
+func (mr *MockmonCollectionMockRecorder) UpdateMany(ctx, filter, update any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter, update}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateMany", reflect.TypeOf((*MockmonCollection)(nil).UpdateMany), varargs...)
+}
+
+// UpdateOne mocks base method.
+func (m *MockmonCollection) UpdateOne(ctx context.Context, filter, update any, opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, filter, update}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "UpdateOne", varargs...)
+	ret0, _ := ret[0].(*mongo.UpdateResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// UpdateOne indicates an expected call of UpdateOne.
+func (mr *MockmonCollectionMockRecorder) UpdateOne(ctx, filter, update any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, filter, update}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateOne", reflect.TypeOf((*MockmonCollection)(nil).UpdateOne), varargs...)
+}
+
+// Watch mocks base method.
+func (m *MockmonCollection) Watch(ctx context.Context, pipeline any, opts ...options.Lister[options.ChangeStreamOptions]) (*mongo.ChangeStream, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, pipeline}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Watch", varargs...)
+	ret0, _ := ret[0].(*mongo.ChangeStream)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Watch indicates an expected call of Watch.
+func (mr *MockmonCollectionMockRecorder) Watch(ctx, pipeline any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, pipeline}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Watch", reflect.TypeOf((*MockmonCollection)(nil).Watch), varargs...)
+}

core/stores/mon/collection_test.go

@@ -10,12 +10,10 @@ import (
 	"github.com/zeromicro/go-zero/core/logx/logtest"
 	"github.com/zeromicro/go-zero/core/stringx"
 	"github.com/zeromicro/go-zero/core/timex"
-	"go.mongodb.org/mongo-driver/bson"
-	"go.mongodb.org/mongo-driver/bson/primitive"
-	"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"
+	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
+	"go.uber.org/mock/gomock"
 )
 
 var errDummy = errors.New("dummy")
@@ -68,471 +66,345 @@ func TestKeepPromise_keep(t *testing.T) {
 }
 
 func TestNewCollection(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		coll := mt.Coll
-		assert.NotNil(t, coll)
-		col := newCollection(coll, breaker.GetBreaker("localhost"))
-		assert.Equal(t, t.Name()+"/test", col.(*decoratedCollection).name)
-	})
+	_ = newCollection(&mongo.Collection{}, breaker.GetBreaker("localhost"))
 }
 
 func TestCollection_Aggregate(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		coll := mt.Coll
-		assert.NotNil(t, coll)
-		col := newCollection(coll, breaker.GetBreaker("localhost"))
-		ns := mt.Coll.Database().Name() + "." + mt.Coll.Name()
-		aggRes := mtest.CreateCursorResponse(1, ns, mtest.FirstBatch)
-		mt.AddMockResponses(aggRes)
-		assert.Equal(t, t.Name()+"/test", col.(*decoratedCollection).name)
-		cursor, err := col.Aggregate(context.Background(), mongo.Pipeline{}, mopt.Aggregate())
-		assert.Nil(t, err)
-		cursor.Close(context.Background())
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().Aggregate(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.Cursor{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.Aggregate(context.Background(), []interface{}{}, options.Aggregate())
+	assert.Nil(t, err)
 }
 
 func TestCollection_BulkWrite(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
-		res, err := c.BulkWrite(context.Background(), []mongo.WriteModel{
-			mongo.NewInsertOneModel().SetDocument(bson.D{{Key: "foo", Value: 1}}),
-		})
-		assert.Nil(t, err)
-		assert.NotNil(t, res)
-		c.brk = new(dropBreaker)
-		_, err = c.BulkWrite(context.Background(), []mongo.WriteModel{
-			mongo.NewInsertOneModel().SetDocument(bson.D{{Key: "foo", Value: 1}}),
-		})
-		assert.Equal(t, errDummy, err)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().BulkWrite(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.BulkWriteResult{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.BulkWrite(context.Background(), []mongo.WriteModel{
+		mongo.NewInsertOneModel().SetDocument(bson.D{{Key: "foo", Value: 1}}),
 	})
+	assert.Nil(t, err)
+	c.brk = new(dropBreaker)
+	_, err = c.BulkWrite(context.Background(), []mongo.WriteModel{
+		mongo.NewInsertOneModel().SetDocument(bson.D{{Key: "foo", Value: 1}}),
+	})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_CountDocuments(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(mtest.CreateCursorResponse(
-			1,
-			"DBName.CollectionName",
-			mtest.FirstBatch,
-			bson.D{
-				{Key: "n", Value: 1},
-			}))
-		res, err := c.CountDocuments(context.Background(), bson.D{})
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), res)
-
-		c.brk = new(dropBreaker)
-		_, err = c.CountDocuments(context.Background(), bson.D{{Key: "foo", Value: 1}})
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().CountDocuments(gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	res, err := c.CountDocuments(context.Background(), bson.D{})
+	assert.Nil(t, err)
+	assert.Equal(t, int64(0), res)
+	c.brk = new(dropBreaker)
+	_, err = c.CountDocuments(context.Background(), bson.D{{Key: "foo", Value: 1}})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestDecoratedCollection_DeleteMany(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
-		res, err := c.DeleteMany(context.Background(), bson.D{})
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), res.DeletedCount)
-
-		c.brk = new(dropBreaker)
-		_, err = c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: 1}})
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().DeleteMany(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.DeleteResult{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.DeleteMany(context.Background(), bson.D{})
+	assert.Nil(t, err)
+	c.brk = new(dropBreaker)
+	_, err = c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: 1}})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_Distinct(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "values", Value: []int{1}}})
-		resp, err := c.Distinct(context.Background(), "foo", bson.D{})
-		assert.Nil(t, err)
-		assert.Equal(t, 1, len(resp))
-
-		c.brk = new(dropBreaker)
-		_, err = c.Distinct(context.Background(), "foo", bson.D{{Key: "foo", Value: 1}})
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().Distinct(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.DistinctResult{})
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.Distinct(context.Background(), "foo", bson.D{})
+	assert.Nil(t, err)
+	c.brk = new(dropBreaker)
+	_, err = c.Distinct(context.Background(), "foo", bson.D{{Key: "foo", Value: 1}})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_EstimatedDocumentCount(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "n", Value: 1}})
-		res, err := c.EstimatedDocumentCount(context.Background())
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), res)
-
-		c.brk = new(dropBreaker)
-		_, err = c.EstimatedDocumentCount(context.Background())
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().EstimatedDocumentCount(gomock.Any(), gomock.Any()).Return(int64(0), nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.EstimatedDocumentCount(context.Background())
+	assert.Nil(t, err)
+	c.brk = new(dropBreaker)
+	_, err = c.EstimatedDocumentCount(context.Background())
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_Find(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		find := mtest.CreateCursorResponse(
-			1,
-			"DBName.CollectionName",
-			mtest.FirstBatch,
-			bson.D{
-				{Key: "name", Value: "John"},
-			})
-		getMore := mtest.CreateCursorResponse(
-			1,
-			"DBName.CollectionName",
-			mtest.NextBatch,
-			bson.D{
-				{Key: "name", Value: "Mary"},
-			})
-		killCursors := mtest.CreateCursorResponse(
-			0,
-			"DBName.CollectionName",
-			mtest.NextBatch)
-		mt.AddMockResponses(find, getMore, killCursors)
-		filter := bson.D{{Key: "x", Value: 1}}
-		cursor, err := c.Find(context.Background(), filter, mopt.Find())
-		assert.Nil(t, err)
-		defer cursor.Close(context.Background())
-
-		var val []struct {
-			ID   primitive.ObjectID `bson:"_id"`
-			Name string             `bson:"name"`
-		}
-		assert.Nil(t, cursor.All(context.Background(), &val))
-		assert.Equal(t, 2, len(val))
-		assert.Equal(t, "John", val[0].Name)
-		assert.Equal(t, "Mary", val[1].Name)
-
-		c.brk = new(dropBreaker)
-		_, err = c.Find(context.Background(), filter, mopt.Find())
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().Find(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.Cursor{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	filter := bson.D{{Key: "x", Value: 1}}
+	_, err := c.Find(context.Background(), filter, options.Find())
+	assert.Nil(t, err)
+	c.brk = new(dropBreaker)
+	_, err = c.Find(context.Background(), filter, options.Find())
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_FindOne(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		find := mtest.CreateCursorResponse(
-			1,
-			"DBName.CollectionName",
-			mtest.FirstBatch,
-			bson.D{
-				{Key: "name", Value: "John"},
-			})
-		getMore := mtest.CreateCursorResponse(
-			1,
-			"DBName.CollectionName",
-			mtest.NextBatch,
-			bson.D{
-				{Key: "name", Value: "Mary"},
-			})
-		killCursors := mtest.CreateCursorResponse(
-			0,
-			"DBName.CollectionName",
-			mtest.NextBatch)
-		mt.AddMockResponses(find, getMore, killCursors)
-		filter := bson.D{{Key: "x", Value: 1}}
-		resp, err := c.FindOne(context.Background(), filter)
-		assert.Nil(t, err)
-		var val struct {
-			ID   primitive.ObjectID `bson:"_id"`
-			Name string             `bson:"name"`
-		}
-		assert.Nil(t, resp.Decode(&val))
-		assert.Equal(t, "John", val.Name)
-
-		c.brk = new(dropBreaker)
-		_, err = c.FindOne(context.Background(), filter)
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().FindOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.SingleResult{})
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	filter := bson.D{{Key: "x", Value: 1}}
+	_, err := c.FindOne(context.Background(), filter)
+	assert.Equal(t, mongo.ErrNoDocuments, err)
+	c.brk = new(dropBreaker)
+	_, err = c.FindOne(context.Background(), filter)
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_FindOneAndDelete(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		filter := bson.D{}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{}...))
-		_, err := c.FindOneAndDelete(context.Background(), filter, mopt.FindOneAndDelete())
-		assert.Equal(t, mongo.ErrNoDocuments, err)
-
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "name", Value: "John"}}},
-		}...))
-		resp, err := c.FindOneAndDelete(context.Background(), filter, mopt.FindOneAndDelete())
-		assert.Nil(t, err)
-		var val struct {
-			Name string `bson:"name"`
-		}
-		assert.Nil(t, resp.Decode(&val))
-		assert.Equal(t, "John", val.Name)
-
-		c.brk = new(dropBreaker)
-		_, err = c.FindOneAndDelete(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().FindOneAndDelete(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.SingleResult{})
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	filter := bson.D{}
+	mockCollection.EXPECT().FindOneAndDelete(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.SingleResult{})
+	_, err := c.FindOneAndDelete(context.Background(), filter, options.FindOneAndDelete())
+	assert.Equal(t, mongo.ErrNoDocuments, err)
+	_, err = c.FindOneAndDelete(context.Background(), filter, options.FindOneAndDelete())
+	assert.Equal(t, mongo.ErrNoDocuments, err)
+	c.brk = new(dropBreaker)
+	_, err = c.FindOneAndDelete(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_FindOneAndReplace(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{}...))
-		filter := bson.D{{Key: "x", Value: 1}}
-		replacement := bson.D{{Key: "x", Value: 2}}
-		opts := mopt.FindOneAndReplace().SetUpsert(true)
-		_, err := c.FindOneAndReplace(context.Background(), filter, replacement, opts)
-		assert.Equal(t, mongo.ErrNoDocuments, err)
-		mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "value", Value: bson.D{
-			{Key: "name", Value: "John"},
-		}}})
-		resp, err := c.FindOneAndReplace(context.Background(), filter, replacement, opts)
-		assert.Nil(t, err)
-		var val struct {
-			Name string `bson:"name"`
-		}
-		assert.Nil(t, resp.Decode(&val))
-		assert.Equal(t, "John", val.Name)
-
-		c.brk = new(dropBreaker)
-		_, err = c.FindOneAndReplace(context.Background(), filter, replacement, opts)
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().FindOneAndReplace(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.SingleResult{})
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	filter := bson.D{{Key: "x", Value: 1}}
+	replacement := bson.D{{Key: "x", Value: 2}}
+	opts := options.FindOneAndReplace().SetUpsert(true)
+	mockCollection.EXPECT().FindOneAndReplace(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.SingleResult{})
+	_, err := c.FindOneAndReplace(context.Background(), filter, replacement, opts)
+	assert.Equal(t, mongo.ErrNoDocuments, err)
+	_, err = c.FindOneAndReplace(context.Background(), filter, replacement, opts)
+	assert.Equal(t, mongo.ErrNoDocuments, err)
+	c.brk = new(dropBreaker)
+	_, err = c.FindOneAndReplace(context.Background(), filter, replacement, opts)
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_FindOneAndUpdate(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}})
-		filter := bson.D{{Key: "x", Value: 1}}
-		update := bson.D{{Key: "$x", Value: 2}}
-		opts := mopt.FindOneAndUpdate().SetUpsert(true)
-		_, err := c.FindOneAndUpdate(context.Background(), filter, update, opts)
-		assert.Equal(t, mongo.ErrNoDocuments, err)
-
-		mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "value", Value: bson.D{
-			{Key: "name", Value: "John"},
-		}}})
-		resp, err := c.FindOneAndUpdate(context.Background(), filter, update, opts)
-		assert.Nil(t, err)
-		var val struct {
-			Name string `bson:"name"`
-		}
-		assert.Nil(t, resp.Decode(&val))
-		assert.Equal(t, "John", val.Name)
-
-		c.brk = new(dropBreaker)
-		_, err = c.FindOneAndUpdate(context.Background(), filter, update, opts)
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().FindOneAndUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.SingleResult{})
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	filter := bson.D{{Key: "x", Value: 1}}
+	update := bson.D{{Key: "$x", Value: 2}}
+	mockCollection.EXPECT().FindOneAndUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.SingleResult{})
+	opts := options.FindOneAndUpdate().SetUpsert(true)
+	_, err := c.FindOneAndUpdate(context.Background(), filter, update, opts)
+	assert.Equal(t, mongo.ErrNoDocuments, err)
+	_, err = c.FindOneAndUpdate(context.Background(), filter, update, opts)
+	assert.Equal(t, mongo.ErrNoDocuments, err)
+	c.brk = new(dropBreaker)
+	_, err = c.FindOneAndUpdate(context.Background(), filter, update, opts)
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_InsertOne(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
-		res, err := c.InsertOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
-		assert.Nil(t, err)
-		assert.NotNil(t, res)
-
-		c.brk = new(dropBreaker)
-		_, err = c.InsertOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().InsertOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.InsertOneResult{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	res, err := c.InsertOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+	assert.Nil(t, err)
+	assert.NotNil(t, res)
+	c.brk = new(dropBreaker)
+	_, err = c.InsertOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_InsertMany(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
-		res, err := c.InsertMany(context.Background(), []any{
-			bson.D{{Key: "foo", Value: "bar"}},
-			bson.D{{Key: "foo", Value: "baz"}},
-		})
-		assert.Nil(t, err)
-		assert.NotNil(t, res)
-		assert.Equal(t, 2, len(res.InsertedIDs))
-
-		c.brk = new(dropBreaker)
-		_, err = c.InsertMany(context.Background(), []any{bson.D{{Key: "foo", Value: "bar"}}})
-		assert.Equal(t, errDummy, err)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().InsertMany(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.InsertManyResult{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.InsertMany(context.Background(), []any{
+		bson.D{{Key: "foo", Value: "bar"}},
+		bson.D{{Key: "foo", Value: "baz"}},
 	})
+	assert.Nil(t, err)
+	c.brk = new(dropBreaker)
+	_, err = c.InsertMany(context.Background(), []any{bson.D{{Key: "foo", Value: "bar"}}})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_DeleteOne(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
-		res, err := c.DeleteOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), res.DeletedCount)
-
-		c.brk = new(dropBreaker)
-		_, err = c.DeleteOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().DeleteOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.DeleteResult{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.DeleteOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+	assert.Nil(t, err)
+	c.brk = new(dropBreaker)
+	_, err = c.DeleteOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_DeleteMany(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
-		res, err := c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), res.DeletedCount)
-
-		c.brk = new(dropBreaker)
-		_, err = c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().DeleteMany(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.DeleteResult{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+	assert.Nil(t, err)
+	c.brk = new(dropBreaker)
+	_, err = c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_ReplaceOne(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
-		res, err := c.ReplaceOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
-			bson.D{{Key: "foo", Value: "baz"}},
-		)
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), res.MatchedCount)
-
-		c.brk = new(dropBreaker)
-		_, err = c.ReplaceOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
-			bson.D{{Key: "foo", Value: "baz"}})
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().ReplaceOne(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.ReplaceOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
+		bson.D{{Key: "foo", Value: "baz"}},
+	)
+	assert.Nil(t, err)
+	c.brk = new(dropBreaker)
+	_, err = c.ReplaceOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
+		bson.D{{Key: "foo", Value: "baz"}})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_UpdateOne(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
-		resp, err := c.UpdateOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
-			bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), resp.MatchedCount)
-
-		c.brk = new(dropBreaker)
-		_, err = c.UpdateOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
-			bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().UpdateOne(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.UpdateOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
+		bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
+	assert.Nil(t, err)
+	c.brk = new(dropBreaker)
+	_, err = c.UpdateOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
+		bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_UpdateByID(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
-		resp, err := c.UpdateByID(context.Background(), primitive.NewObjectID(),
-			bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), resp.MatchedCount)
-
-		c.brk = new(dropBreaker)
-		_, err = c.UpdateByID(context.Background(), primitive.NewObjectID(),
-			bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().UpdateByID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.UpdateByID(context.Background(), bson.NewObjectID(),
+		bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
+	assert.Nil(t, err)
+	c.brk = new(dropBreaker)
+	_, err = c.UpdateByID(context.Background(), bson.NewObjectID(),
+		bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestCollection_UpdateMany(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		c := decoratedCollection{
-			Collection: mt.Coll,
-			brk:        breaker.NewBreaker(),
-		}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
-		resp, err := c.UpdateMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
-			bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), resp.MatchedCount)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().UpdateMany(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.UpdateMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
+		bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
+	assert.Nil(t, err)
+	c.brk = new(dropBreaker)
+	_, err = c.UpdateMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
+		bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
+	assert.Equal(t, errDummy, err)
+}
 
-		c.brk = new(dropBreaker)
-		_, err = c.UpdateMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
-			bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
-		assert.Equal(t, errDummy, err)
-	})
+func TestCollection_Watch(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().Watch(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.ChangeStream{}, nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	_, err := c.Watch(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+	assert.Nil(t, err)
+}
+
+func TestCollection_Clone(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().Clone(gomock.Any()).Return(nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	cc := c.Clone()
+	assert.Nil(t, cc)
+}
+
+func TestCollection_Database(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().Database().Return(nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	db := c.Database()
+	assert.Nil(t, db)
+}
+
+func TestCollection_Drop(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	mockCollection.EXPECT().Drop(gomock.Any()).Return(nil)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	err := c.Drop(context.Background())
+	assert.Nil(t, err)
+}
+
+func TestCollection_Indexes(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	idx := mongo.IndexView{}
+	mockCollection.EXPECT().Indexes().Return(idx)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
+	index := c.Indexes()
+	assert.Equal(t, index, idx)
 }
 
 func TestDecoratedCollection_LogDuration(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	c := decoratedCollection{
-		Collection: mt.Coll,
-		brk:        breaker.NewBreaker(),
-	}
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := NewMockmonCollection(ctrl)
+	c := newTestCollection(mockCollection, breaker.GetBreaker("localhost"))
 
 	buf := logtest.NewCollector(t)
 
@@ -585,14 +457,6 @@ func TestAcceptable(t *testing.T) {
 		{"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},
 	}
@@ -623,6 +487,14 @@ func TestIsDupKeyError(t *testing.T) {
 	}
 }
 
+func newTestCollection(collection monCollection, brk breaker.Breaker) *decoratedCollection {
+	return &decoratedCollection{
+		Collection: collection,
+		name:       "test",
+		brk:        brk,
+	}
+}
+
 type mockPromise struct {
 	accepted bool
 	reason   string

core/stores/mon/collectioninserter_mock.go

@@ -0,0 +1,63 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: bulkinserter.go
+//
+// Generated by this command:
+//
+//	mockgen -package mon -destination collectioninserter_mock.go -source bulkinserter.go collectionInserter
+//
+
+// Package mon is a generated GoMock package.
+package mon
+
+import (
+	context "context"
+	reflect "reflect"
+
+	mongo "go.mongodb.org/mongo-driver/v2/mongo"
+	options "go.mongodb.org/mongo-driver/v2/mongo/options"
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MockcollectionInserter is a mock of collectionInserter interface.
+type MockcollectionInserter struct {
+	ctrl     *gomock.Controller
+	recorder *MockcollectionInserterMockRecorder
+	isgomock struct{}
+}
+
+// MockcollectionInserterMockRecorder is the mock recorder for MockcollectionInserter.
+type MockcollectionInserterMockRecorder struct {
+	mock *MockcollectionInserter
+}
+
+// NewMockcollectionInserter creates a new mock instance.
+func NewMockcollectionInserter(ctrl *gomock.Controller) *MockcollectionInserter {
+	mock := &MockcollectionInserter{ctrl: ctrl}
+	mock.recorder = &MockcollectionInserterMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockcollectionInserter) EXPECT() *MockcollectionInserterMockRecorder {
+	return m.recorder
+}
+
+// InsertMany mocks base method.
+func (m *MockcollectionInserter) InsertMany(ctx context.Context, documents any, opts ...options.Lister[options.InsertManyOptions]) (*mongo.InsertManyResult, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, documents}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "InsertMany", varargs...)
+	ret0, _ := ret[0].(*mongo.InsertManyResult)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// InsertMany indicates an expected call of InsertMany.
+func (mr *MockcollectionInserterMockRecorder) InsertMany(ctx, documents any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, documents}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertMany", reflect.TypeOf((*MockcollectionInserter)(nil).InsertMany), varargs...)
+}

core/stores/mon/migration-2.0.md

@@ -0,0 +1,19 @@
+# Migrating from 1.x to 2.0
+
+To upgrade imports of the Go Driver from v1 to v2, we recommend using [marwan-at-work/mod
+](https://github.com/marwan-at-work/mod):
+
+```
+mod upgrade --mod-name=go.mongodb.org/mongo-driver
+```
+
+# Notice
+After completing the mod upgrade, code changes are typically unnecessary in the vast majority of cases. However, if your project references packages including but not limited to those listed below, you'll need to manually replace them, as these libraries are no longer present in the v2 version.
+```go
+go.mongodb.org/mongo-driver/bson/bsonrw    => go.mongodb.org/mongo-driver/v2/bson
+go.mongodb.org/mongo-driver/bson/bsoncodec => go.mongodb.org/mongo-driver/v2/bson
+go.mongodb.org/mongo-driver/bson/primitive => go.mongodb.org/mongo-driver/v2/bson
+```
+
+See the following resources to learn more about upgrading from version 1.x to 2.0.:
+https://raw.githubusercontent.com/mongodb/mongo-go-driver/refs/heads/master/docs/migration-2.0.md

core/stores/mon/model.go

@@ -1,3 +1,4 @@
+//go:generate mockgen -package mon -destination model_mock.go -source model.go monClient monSession
 package mon
 
 import (
@@ -7,8 +8,8 @@ import (
 	"github.com/zeromicro/go-zero/core/breaker"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/timex"
-	"go.mongodb.org/mongo-driver/mongo"
-	mopt "go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
 )
 
 const (
@@ -24,15 +25,15 @@ type (
 	Model struct {
 		Collection
 		name string
-		cli  *mongo.Client
+		cli  monClient
 		brk  breaker.Breaker
 		opts []Option
 	}
 
-	wrappedSession struct {
-		mongo.Session
-		name string
-		brk  breaker.Breaker
+	Session struct {
+		session monSession
+		name    string
+		brk     breaker.Breaker
 	}
 )
 
@@ -61,14 +62,14 @@ func newModel(name string, cli *mongo.Client, coll Collection, brk breaker.Break
 	return &Model{
 		name:       name,
 		Collection: coll,
-		cli:        cli,
+		cli:        &wrappedMonClient{c: cli},
 		brk:        brk,
 		opts:       opts,
 	}
 }
 
 // StartSession starts a new session.
-func (m *Model) StartSession(opts ...*mopt.SessionOptions) (sess mongo.Session, err error) {
+func (m *Model) StartSession(opts ...options.Lister[options.SessionOptions]) (sess *Session, err error) {
 	starTime := timex.Now()
 	defer func() {
 		logDuration(context.Background(), m.name, startSession, starTime, err)
@@ -79,15 +80,16 @@ func (m *Model) StartSession(opts ...*mopt.SessionOptions) (sess mongo.Session,
 		return nil, sessionErr
 	}
 
-	return &wrappedSession{
-		Session: session,
+	return &Session{
+		session: session,
 		name:    m.name,
 		brk:     m.brk,
 	}, nil
 }
 
 // Aggregate executes an aggregation pipeline.
-func (m *Model) Aggregate(ctx context.Context, v, pipeline any, opts ...*mopt.AggregateOptions) error {
+func (m *Model) Aggregate(ctx context.Context, v, pipeline any,
+	opts ...options.Lister[options.AggregateOptions]) error {
 	cur, err := m.Collection.Aggregate(ctx, pipeline, opts...)
 	if err != nil {
 		return err
@@ -98,7 +100,8 @@ func (m *Model) Aggregate(ctx context.Context, v, pipeline any, opts ...*mopt.Ag
 }
 
 // DeleteMany deletes documents that match the filter.
-func (m *Model) DeleteMany(ctx context.Context, filter any, opts ...*mopt.DeleteOptions) (int64, error) {
+func (m *Model) DeleteMany(ctx context.Context, filter any,
+	opts ...options.Lister[options.DeleteManyOptions]) (int64, error) {
 	res, err := m.Collection.DeleteMany(ctx, filter, opts...)
 	if err != nil {
 		return 0, err
@@ -108,7 +111,8 @@ func (m *Model) DeleteMany(ctx context.Context, filter any, opts ...*mopt.Delete
 }
 
 // DeleteOne deletes the first document that matches the filter.
-func (m *Model) DeleteOne(ctx context.Context, filter any, opts ...*mopt.DeleteOptions) (int64, error) {
+func (m *Model) DeleteOne(ctx context.Context, filter any,
+	opts ...options.Lister[options.DeleteOneOptions]) (int64, error) {
 	res, err := m.Collection.DeleteOne(ctx, filter, opts...)
 	if err != nil {
 		return 0, err
@@ -118,7 +122,8 @@ func (m *Model) DeleteOne(ctx context.Context, filter any, opts ...*mopt.DeleteO
 }
 
 // Find finds documents that match the filter.
-func (m *Model) Find(ctx context.Context, v, filter any, opts ...*mopt.FindOptions) error {
+func (m *Model) Find(ctx context.Context, v, filter any,
+	opts ...options.Lister[options.FindOptions]) error {
 	cur, err := m.Collection.Find(ctx, filter, opts...)
 	if err != nil {
 		return err
@@ -129,7 +134,8 @@ func (m *Model) Find(ctx context.Context, v, filter any, opts ...*mopt.FindOptio
 }
 
 // FindOne finds the first document that matches the filter.
-func (m *Model) FindOne(ctx context.Context, v, filter any, opts ...*mopt.FindOneOptions) error {
+func (m *Model) FindOne(ctx context.Context, v, filter any,
+	opts ...options.Lister[options.FindOneOptions]) error {
 	res, err := m.Collection.FindOne(ctx, filter, opts...)
 	if err != nil {
 		return err
@@ -140,7 +146,7 @@ func (m *Model) FindOne(ctx context.Context, v, filter any, opts ...*mopt.FindOn
 
 // FindOneAndDelete finds a single document and deletes it.
 func (m *Model) FindOneAndDelete(ctx context.Context, v, filter any,
-	opts ...*mopt.FindOneAndDeleteOptions) error {
+	opts ...options.Lister[options.FindOneAndDeleteOptions]) error {
 	res, err := m.Collection.FindOneAndDelete(ctx, filter, opts...)
 	if err != nil {
 		return err
@@ -151,7 +157,7 @@ func (m *Model) FindOneAndDelete(ctx context.Context, v, filter any,
 
 // FindOneAndReplace finds a single document and replaces it.
 func (m *Model) FindOneAndReplace(ctx context.Context, v, filter, replacement any,
-	opts ...*mopt.FindOneAndReplaceOptions) error {
+	opts ...options.Lister[options.FindOneAndReplaceOptions]) error {
 	res, err := m.Collection.FindOneAndReplace(ctx, filter, replacement, opts...)
 	if err != nil {
 		return err
@@ -162,7 +168,7 @@ func (m *Model) FindOneAndReplace(ctx context.Context, v, filter, replacement an
 
 // FindOneAndUpdate finds a single document and updates it.
 func (m *Model) FindOneAndUpdate(ctx context.Context, v, filter, update any,
-	opts ...*mopt.FindOneAndUpdateOptions) error {
+	opts ...options.Lister[options.FindOneAndUpdateOptions]) error {
 	res, err := m.Collection.FindOneAndUpdate(ctx, filter, update, opts...)
 	if err != nil {
 		return err
@@ -171,8 +177,8 @@ func (m *Model) FindOneAndUpdate(ctx context.Context, v, filter, update any,
 	return res.Decode(v)
 }
 
-// AbortTransaction implements the mongo.Session interface.
-func (w *wrappedSession) AbortTransaction(ctx context.Context) (err error) {
+// AbortTransaction implements the mongo.session interface.
+func (w *Session) AbortTransaction(ctx context.Context) (err error) {
 	ctx, span := startSpan(ctx, abortTransaction)
 	defer func() {
 		endSpan(span, err)
@@ -184,12 +190,12 @@ func (w *wrappedSession) AbortTransaction(ctx context.Context) (err error) {
 			logDuration(ctx, w.name, abortTransaction, starTime, err)
 		}()
 
-		return w.Session.AbortTransaction(ctx)
+		return w.session.AbortTransaction(ctx)
 	}, acceptable)
 }
 
-// CommitTransaction implements the mongo.Session interface.
-func (w *wrappedSession) CommitTransaction(ctx context.Context) (err error) {
+// CommitTransaction implements the mongo.session interface.
+func (w *Session) CommitTransaction(ctx context.Context) (err error) {
 	ctx, span := startSpan(ctx, commitTransaction)
 	defer func() {
 		endSpan(span, err)
@@ -201,15 +207,15 @@ func (w *wrappedSession) CommitTransaction(ctx context.Context) (err error) {
 			logDuration(ctx, w.name, commitTransaction, starTime, err)
 		}()
 
-		return w.Session.CommitTransaction(ctx)
+		return w.session.CommitTransaction(ctx)
 	}, acceptable)
 }
 
-// WithTransaction implements the mongo.Session interface.
-func (w *wrappedSession) WithTransaction(
+// WithTransaction implements the mongo.session interface.
+func (w *Session) WithTransaction(
 	ctx context.Context,
-	fn func(sessCtx mongo.SessionContext) (any, error),
-	opts ...*mopt.TransactionOptions,
+	fn func(sessCtx context.Context) (any, error),
+	opts ...options.Lister[options.TransactionOptions],
 ) (res any, err error) {
 	ctx, span := startSpan(ctx, withTransaction)
 	defer func() {
@@ -222,15 +228,15 @@ func (w *wrappedSession) WithTransaction(
 			logDuration(ctx, w.name, withTransaction, starTime, err)
 		}()
 
-		res, err = w.Session.WithTransaction(ctx, fn, opts...)
+		res, err = w.session.WithTransaction(ctx, fn, opts...)
 		return err
 	}, acceptable)
 
 	return
 }
 
-// EndSession implements the mongo.Session interface.
-func (w *wrappedSession) EndSession(ctx context.Context) {
+// EndSession implements the mongo.session interface.
+func (w *Session) EndSession(ctx context.Context) {
 	var err error
 	ctx, span := startSpan(ctx, endSession)
 	defer func() {
@@ -243,7 +249,34 @@ func (w *wrappedSession) EndSession(ctx context.Context) {
 			logDuration(ctx, w.name, endSession, starTime, err)
 		}()
 
-		w.Session.EndSession(ctx)
+		w.session.EndSession(ctx)
 		return nil
 	}, acceptable)
 }
+
+type (
+	// for unit test
+	monClient interface {
+		StartSession(opts ...options.Lister[options.SessionOptions]) (monSession, error)
+	}
+
+	monSession interface {
+		AbortTransaction(ctx context.Context) error
+		CommitTransaction(ctx context.Context) error
+		EndSession(ctx context.Context)
+		WithTransaction(ctx context.Context, fn func(sessCtx context.Context) (any, error),
+			opts ...options.Lister[options.TransactionOptions]) (any, error)
+	}
+)
+
+type wrappedMonClient struct {
+	c *mongo.Client
+}
+
+// StartSession starts a new session using the underlying *mongo.Client.
+// It implements the monClient interface.
+// This is used to allow mocking in unit tests.
+func (m *wrappedMonClient) StartSession(opts ...options.Lister[options.SessionOptions]) (
+	monSession, error) {
+	return m.c.StartSession(opts...)
+}

core/stores/mon/model_mock.go

@@ -0,0 +1,145 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: model.go
+//
+// Generated by this command:
+//
+//	mockgen -package mon -destination model_mock.go -source model.go monClient monSession
+//
+
+// Package mon is a generated GoMock package.
+package mon
+
+import (
+	context "context"
+	reflect "reflect"
+
+	options "go.mongodb.org/mongo-driver/v2/mongo/options"
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MockmonClient is a mock of monClient interface.
+type MockmonClient struct {
+	ctrl     *gomock.Controller
+	recorder *MockmonClientMockRecorder
+	isgomock struct{}
+}
+
+// MockmonClientMockRecorder is the mock recorder for MockmonClient.
+type MockmonClientMockRecorder struct {
+	mock *MockmonClient
+}
+
+// NewMockmonClient creates a new mock instance.
+func NewMockmonClient(ctrl *gomock.Controller) *MockmonClient {
+	mock := &MockmonClient{ctrl: ctrl}
+	mock.recorder = &MockmonClientMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockmonClient) EXPECT() *MockmonClientMockRecorder {
+	return m.recorder
+}
+
+// StartSession mocks base method.
+func (m *MockmonClient) StartSession(opts ...options.Lister[options.SessionOptions]) (monSession, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "StartSession", varargs...)
+	ret0, _ := ret[0].(monSession)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// StartSession indicates an expected call of StartSession.
+func (mr *MockmonClientMockRecorder) StartSession(opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StartSession", reflect.TypeOf((*MockmonClient)(nil).StartSession), opts...)
+}
+
+// MockmonSession is a mock of monSession interface.
+type MockmonSession struct {
+	ctrl     *gomock.Controller
+	recorder *MockmonSessionMockRecorder
+	isgomock struct{}
+}
+
+// MockmonSessionMockRecorder is the mock recorder for MockmonSession.
+type MockmonSessionMockRecorder struct {
+	mock *MockmonSession
+}
+
+// NewMockmonSession creates a new mock instance.
+func NewMockmonSession(ctrl *gomock.Controller) *MockmonSession {
+	mock := &MockmonSession{ctrl: ctrl}
+	mock.recorder = &MockmonSessionMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockmonSession) EXPECT() *MockmonSessionMockRecorder {
+	return m.recorder
+}
+
+// AbortTransaction mocks base method.
+func (m *MockmonSession) AbortTransaction(ctx context.Context) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "AbortTransaction", ctx)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// AbortTransaction indicates an expected call of AbortTransaction.
+func (mr *MockmonSessionMockRecorder) AbortTransaction(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AbortTransaction", reflect.TypeOf((*MockmonSession)(nil).AbortTransaction), ctx)
+}
+
+// CommitTransaction mocks base method.
+func (m *MockmonSession) CommitTransaction(ctx context.Context) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CommitTransaction", ctx)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// CommitTransaction indicates an expected call of CommitTransaction.
+func (mr *MockmonSessionMockRecorder) CommitTransaction(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CommitTransaction", reflect.TypeOf((*MockmonSession)(nil).CommitTransaction), ctx)
+}
+
+// EndSession mocks base method.
+func (m *MockmonSession) EndSession(ctx context.Context) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "EndSession", ctx)
+}
+
+// EndSession indicates an expected call of EndSession.
+func (mr *MockmonSessionMockRecorder) EndSession(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EndSession", reflect.TypeOf((*MockmonSession)(nil).EndSession), ctx)
+}
+
+// WithTransaction mocks base method.
+func (m *MockmonSession) WithTransaction(ctx context.Context, fn func(context.Context) (any, error), opts ...options.Lister[options.TransactionOptions]) (any, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, fn}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "WithTransaction", varargs...)
+	ret0, _ := ret[0].(any)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// WithTransaction indicates an expected call of WithTransaction.
+func (mr *MockmonSessionMockRecorder) WithTransaction(ctx, fn any, opts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, fn}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WithTransaction", reflect.TypeOf((*MockmonSession)(nil).WithTransaction), varargs...)
+}

core/stores/mon/model_test.go

@@ -2,224 +2,242 @@ package mon
 
 import (
 	"context"
+	"errors"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"go.mongodb.org/mongo-driver/bson"
-	"go.mongodb.org/mongo-driver/mongo"
-	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
+	"github.com/zeromicro/go-zero/core/breaker"
+	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
+	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/drivertest"
+	"go.uber.org/mock/gomock"
 )
 
 func TestModel_StartSession(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		m := createModel(mt)
-		sess, err := m.StartSession()
-		assert.Nil(t, err)
-		defer sess.EndSession(context.Background())
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockMonCollection := NewMockmonCollection(ctrl)
+	mockedMonClient := NewMockmonClient(ctrl)
+	mockMonSession := NewMockmonSession(ctrl)
+	warpSession := &Session{
+		session: mockMonSession,
+		name:    "",
+		brk:     breaker.GetBreaker("localhost"),
+	}
 
-		_, err = sess.WithTransaction(context.Background(), func(sessCtx mongo.SessionContext) (any, error) {
-			_ = sessCtx.StartTransaction()
-			sessCtx.Client().Database("1")
-			sessCtx.EndSession(context.Background())
-			return nil, nil
-		})
-		assert.Nil(t, err)
-		assert.NoError(t, sess.CommitTransaction(context.Background()))
-		assert.Error(t, sess.AbortTransaction(context.Background()))
+	m := newTestModel("foo", mockedMonClient, mockMonCollection, breaker.GetBreaker("test"))
+	mockedMonClient.EXPECT().StartSession(gomock.Any()).Return(warpSession, errors.New("error"))
+	_, err := m.StartSession()
+	assert.NotNil(t, err)
+	mockedMonClient.EXPECT().StartSession(gomock.Any()).Return(warpSession, nil)
+	sess, err := m.StartSession()
+	assert.Nil(t, err)
+	defer sess.EndSession(context.Background())
+	mockMonSession.EXPECT().WithTransaction(gomock.Any(), gomock.Any()).Return(nil, nil)
+	mockMonSession.EXPECT().CommitTransaction(gomock.Any()).Return(nil)
+	mockMonSession.EXPECT().AbortTransaction(gomock.Any()).Return(nil)
+	mockMonSession.EXPECT().EndSession(gomock.Any())
+	_, err = sess.WithTransaction(context.Background(), func(sessCtx context.Context) (any, error) {
+		// _ = sessCtx.StartTransaction()
+		// sessCtx.Client().Database("1")
+		// sessCtx.EndSession(context.Background())
+		return nil, nil
 	})
+	assert.Nil(t, err)
+	assert.NoError(t, sess.CommitTransaction(context.Background()))
+	assert.NoError(t, sess.AbortTransaction(context.Background()))
 }
 
 func TestModel_Aggregate(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		m := createModel(mt)
-		find := mtest.CreateCursorResponse(
-			1,
-			"DBName.CollectionName",
-			mtest.FirstBatch,
-			bson.D{
-				{Key: "name", Value: "John"},
-			})
-		getMore := mtest.CreateCursorResponse(
-			1,
-			"DBName.CollectionName",
-			mtest.NextBatch,
-			bson.D{
-				{Key: "name", Value: "Mary"},
-			})
-		killCursors := mtest.CreateCursorResponse(
-			0,
-			"DBName.CollectionName",
-			mtest.NextBatch)
-		mt.AddMockResponses(find, getMore, killCursors)
-		var result []any
-		err := m.Aggregate(context.Background(), &result, mongo.Pipeline{})
-		assert.Nil(t, err)
-		assert.Equal(t, 2, len(result))
-		assert.Equal(t, "John", result[0].(bson.D).Map()["name"])
-		assert.Equal(t, "Mary", result[1].(bson.D).Map()["name"])
-
-		triggerBreaker(m)
-		assert.Equal(t, errDummy, m.Aggregate(context.Background(), &result, mongo.Pipeline{}))
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockMonCollection := NewMockmonCollection(ctrl)
+	mockedMonClient := NewMockmonClient(ctrl)
+	cursor, err := mongo.NewCursorFromDocuments([]any{
+		bson.M{
+			"name": "John",
+		},
+		bson.M{
+			"name": "Mary",
+		},
+	}, nil, nil)
+	assert.NoError(t, err)
+	mockMonCollection.EXPECT().Aggregate(gomock.Any(), gomock.Any(), gomock.Any()).Return(cursor, nil)
+	m := newTestModel("foo", mockedMonClient, mockMonCollection, breaker.GetBreaker("test"))
+	var result []bson.M
+	err = m.Aggregate(context.Background(), &result, bson.D{})
+	assert.Nil(t, err)
+	assert.Equal(t, 2, len(result))
+	assert.Equal(t, "John", result[0]["name"])
+	assert.Equal(t, "Mary", result[1]["name"])
+	triggerBreaker(m)
+	assert.Equal(t, errDummy, m.Aggregate(context.Background(), &result, bson.D{}))
 }
 
 func TestModel_DeleteMany(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		m := createModel(mt)
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
-		val, err := m.DeleteMany(context.Background(), bson.D{})
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), val)
-
-		triggerBreaker(m)
-		_, err = m.DeleteMany(context.Background(), bson.D{})
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockMonCollection := NewMockmonCollection(ctrl)
+	mockedMonClient := NewMockmonClient(ctrl)
+	mockMonCollection.EXPECT().DeleteMany(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.DeleteResult{}, nil)
+	m := newTestModel("foo", mockedMonClient, mockMonCollection, breaker.GetBreaker("test"))
+	_, err := m.DeleteMany(context.Background(), bson.D{})
+	assert.Nil(t, err)
+	triggerBreaker(m)
+	_, err = m.DeleteMany(context.Background(), bson.D{})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestModel_DeleteOne(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		m := createModel(mt)
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
-		val, err := m.DeleteOne(context.Background(), bson.D{})
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), val)
-
-		triggerBreaker(m)
-		_, err = m.DeleteOne(context.Background(), bson.D{})
-		assert.Equal(t, errDummy, err)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockMonCollection := NewMockmonCollection(ctrl)
+	mockedMonClient := NewMockmonClient(ctrl)
+	mockMonCollection.EXPECT().DeleteOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.DeleteResult{}, nil)
+	m := newTestModel("foo", mockedMonClient, mockMonCollection, breaker.GetBreaker("test"))
+	_, err := m.DeleteOne(context.Background(), bson.D{})
+	assert.Nil(t, err)
+	triggerBreaker(m)
+	_, err = m.DeleteOne(context.Background(), bson.D{})
+	assert.Equal(t, errDummy, err)
 }
 
 func TestModel_Find(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		m := createModel(mt)
-		find := mtest.CreateCursorResponse(
-			1,
-			"DBName.CollectionName",
-			mtest.FirstBatch,
-			bson.D{
-				{Key: "name", Value: "John"},
-			})
-		getMore := mtest.CreateCursorResponse(
-			1,
-			"DBName.CollectionName",
-			mtest.NextBatch,
-			bson.D{
-				{Key: "name", Value: "Mary"},
-			})
-		killCursors := mtest.CreateCursorResponse(
-			0,
-			"DBName.CollectionName",
-			mtest.NextBatch)
-		mt.AddMockResponses(find, getMore, killCursors)
-		var result []any
-		err := m.Find(context.Background(), &result, bson.D{})
-		assert.Nil(t, err)
-		assert.Equal(t, 2, len(result))
-		assert.Equal(t, "John", result[0].(bson.D).Map()["name"])
-		assert.Equal(t, "Mary", result[1].(bson.D).Map()["name"])
-
-		triggerBreaker(m)
-		assert.Equal(t, errDummy, m.Find(context.Background(), &result, bson.D{}))
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockMonCollection := NewMockmonCollection(ctrl)
+	mockedMonClient := NewMockmonClient(ctrl)
+	cursor, err := mongo.NewCursorFromDocuments([]any{
+		bson.M{
+			"name": "John",
+		},
+		bson.M{
+			"name": "Mary",
+		},
+	}, nil, nil)
+	assert.NoError(t, err)
+	mockMonCollection.EXPECT().Find(gomock.Any(), gomock.Any(), gomock.Any()).Return(cursor, nil)
+	m := newTestModel("foo", mockedMonClient, mockMonCollection, breaker.GetBreaker("test"))
+	var result []bson.M
+	err = m.Find(context.Background(), &result, bson.D{})
+	assert.Nil(t, err)
+	assert.Equal(t, 2, len(result))
+	assert.Equal(t, "John", result[0]["name"])
+	assert.Equal(t, "Mary", result[1]["name"])
+	triggerBreaker(m)
+	assert.Equal(t, errDummy, m.Find(context.Background(), &result, bson.D{}))
 }
 
 func TestModel_FindOne(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		m := createModel(mt)
-		find := mtest.CreateCursorResponse(
-			1,
-			"DBName.CollectionName",
-			mtest.FirstBatch,
-			bson.D{
-				{Key: "name", Value: "John"},
-			})
-		killCursors := mtest.CreateCursorResponse(
-			0,
-			"DBName.CollectionName",
-			mtest.NextBatch)
-		mt.AddMockResponses(find, killCursors)
-		var result bson.D
-		err := m.FindOne(context.Background(), &result, bson.D{})
-		assert.Nil(t, err)
-		assert.Equal(t, "John", result.Map()["name"])
-
-		triggerBreaker(m)
-		assert.Equal(t, errDummy, m.FindOne(context.Background(), &result, bson.D{}))
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockMonCollection := NewMockmonCollection(ctrl)
+	mockedMonClient := NewMockmonClient(ctrl)
+	mockMonCollection.EXPECT().FindOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(mongo.NewSingleResultFromDocument(bson.M{"name": "John"}, nil, nil))
+	m := newTestModel("foo", mockedMonClient, mockMonCollection, breaker.GetBreaker("test"))
+	var result bson.M
+	err := m.FindOne(context.Background(), &result, bson.D{})
+	assert.Nil(t, err)
+	assert.Equal(t, "John", result["name"])
+	triggerBreaker(m)
+	assert.Equal(t, errDummy, m.FindOne(context.Background(), &result, bson.D{}))
 }
 
 func TestModel_FindOneAndDelete(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		m := createModel(mt)
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "name", Value: "John"}}},
-		}...))
-		var result bson.D
-		err := m.FindOneAndDelete(context.Background(), &result, bson.D{})
-		assert.Nil(t, err)
-		assert.Equal(t, "John", result.Map()["name"])
-
-		triggerBreaker(m)
-		assert.Equal(t, errDummy, m.FindOneAndDelete(context.Background(), &result, bson.D{}))
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockMonCollection := NewMockmonCollection(ctrl)
+	mockedMonClient := NewMockmonClient(ctrl)
+	mockMonCollection.EXPECT().FindOneAndDelete(gomock.Any(), gomock.Any(), gomock.Any()).Return(mongo.NewSingleResultFromDocument(bson.M{"name": "John"}, nil, nil))
+	m := newTestModel("foo", mockedMonClient, mockMonCollection, breaker.GetBreaker("test"))
+	var result bson.M
+	err := m.FindOneAndDelete(context.Background(), &result, bson.M{})
+	assert.Nil(t, err)
+	assert.Equal(t, "John", result["name"])
+	triggerBreaker(m)
+	assert.Equal(t, errDummy, m.FindOneAndDelete(context.Background(), &result, bson.D{}))
 }
 
 func TestModel_FindOneAndReplace(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		m := createModel(mt)
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "name", Value: "John"}}},
-		}...))
-		var result bson.D
-		err := m.FindOneAndReplace(context.Background(), &result, bson.D{}, bson.D{
-			{Key: "name", Value: "Mary"},
-		})
-		assert.Nil(t, err)
-		assert.Equal(t, "John", result.Map()["name"])
-
-		triggerBreaker(m)
-		assert.Equal(t, errDummy, m.FindOneAndReplace(context.Background(), &result, bson.D{}, bson.D{
-			{Key: "name", Value: "Mary"},
-		}))
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockMonCollection := NewMockmonCollection(ctrl)
+	mockedMonClient := NewMockmonClient(ctrl)
+	mockMonCollection.EXPECT().FindOneAndReplace(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(mongo.NewSingleResultFromDocument(bson.M{"name": "John"}, nil, nil))
+	m := newTestModel("foo", mockedMonClient, mockMonCollection, breaker.GetBreaker("test"))
+	var result bson.M
+	err := m.FindOneAndReplace(context.Background(), &result, bson.D{}, bson.D{
+		{Key: "name", Value: "Mary"},
 	})
+	assert.Nil(t, err)
+	assert.Equal(t, "John", result["name"])
+	triggerBreaker(m)
+	assert.Equal(t, errDummy, m.FindOneAndReplace(context.Background(), &result, bson.D{}, bson.D{
+		{Key: "name", Value: "Mary"},
+	}))
 }
 
 func TestModel_FindOneAndUpdate(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		m := createModel(mt)
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "name", Value: "John"}}},
-		}...))
-		var result bson.D
-		err := m.FindOneAndUpdate(context.Background(), &result, bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
-		})
-		assert.Nil(t, err)
-		assert.Equal(t, "John", result.Map()["name"])
-
-		triggerBreaker(m)
-		assert.Equal(t, errDummy, m.FindOneAndUpdate(context.Background(), &result, bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
-		}))
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockMonCollection := NewMockmonCollection(ctrl)
+	mockedMonClient := NewMockmonClient(ctrl)
+	mockMonCollection.EXPECT().FindOneAndUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(mongo.NewSingleResultFromDocument(bson.M{"name": "John"}, nil, nil))
+	m := newTestModel("foo", mockedMonClient, mockMonCollection, breaker.GetBreaker("test"))
+	var result bson.M
+	err := m.FindOneAndUpdate(context.Background(), &result, bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
 	})
-}
+	assert.Nil(t, err)
+	assert.Equal(t, "John", result["name"])
+	triggerBreaker(m)
+	assert.Equal(t, errDummy, m.FindOneAndUpdate(context.Background(), &result, bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
+	}))
 
-func createModel(mt *mtest.T) *Model {
-	Inject(mt.Name(), mt.Client)
-	return MustNewModel(mt.Name(), mt.DB.Name(), mt.Coll.Name())
 }
 
 func triggerBreaker(m *Model) {
 	m.Collection.(*decoratedCollection).brk = new(dropBreaker)
 }
+
+func TestMustNewModel(t *testing.T) {
+	Inject("mongodb://localhost:27017", &mongo.Client{})
+	MustNewModel("mongodb://localhost:27017", "test", "test")
+}
+
+func TestNewModel(t *testing.T) {
+	NewModel("mongo://localhost:27018", "test", "test")
+	Inject("mongodb://localhost:27018", &mongo.Client{})
+	NewModel("mongodb://localhost:27018", "test", "test")
+}
+
+func Test_newModel(t *testing.T) {
+	Inject("mongodb://localhost:27019", &mongo.Client{})
+	newModel("mongodb://localhost:27019", nil, nil, nil)
+}
+
+func Test_mockMonClient_StartSession(t *testing.T) {
+	md := drivertest.NewMockDeployment()
+	opts := options.Client()
+	opts.Deployment = md
+	client, err := mongo.Connect(opts)
+	assert.Nil(t, err)
+	m := wrappedMonClient{
+		c: client,
+	}
+	_, err = m.StartSession()
+	assert.Nil(t, err)
+}
+
+func newTestModel(name string, cli monClient, coll monCollection, brk breaker.Breaker,
+	opts ...Option) *Model {
+	return &Model{
+		name:       name,
+		Collection: newTestCollection(coll, breaker.GetBreaker("localhost")),
+		cli:        cli,
+		brk:        brk,
+		opts:       opts,
+	}
+}

core/stores/mon/options.go

@@ -5,9 +5,8 @@ import (
 	"time"
 
 	"github.com/zeromicro/go-zero/core/syncx"
-	"go.mongodb.org/mongo-driver/bson"
-	"go.mongodb.org/mongo-driver/bson/bsoncodec"
-	mopt "go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
 )
 
 const defaultTimeout = time.Second * 3
@@ -20,16 +19,16 @@ var (
 
 type (
 	// Option defines the method to customize a mongo model.
-	Option func(opts *options)
+	Option func(opts *clientOptions)
 
 	// TypeCodec is a struct that stores specific type Encoder/Decoder.
 	TypeCodec struct {
 		ValueType reflect.Type
-		Encoder   bsoncodec.ValueEncoder
-		Decoder   bsoncodec.ValueDecoder
+		Encoder   bson.ValueEncoder
+		Decoder   bson.ValueDecoder
 	}
 
-	options = mopt.ClientOptions
+	clientOptions = options.ClientOptions
 )
 
 // DisableLog disables logging of mongo commands, includes info and slow logs.
@@ -50,14 +49,14 @@ func SetSlowThreshold(threshold time.Duration) {
 
 // WithTimeout set the mon client operation timeout.
 func WithTimeout(timeout time.Duration) Option {
-	return func(opts *options) {
+	return func(opts *clientOptions) {
 		opts.SetTimeout(timeout)
 	}
 }
 
 // WithTypeCodec registers TypeCodecs to convert custom types.
 func WithTypeCodec(typeCodecs ...TypeCodec) Option {
-	return func(opts *options) {
+	return func(opts *clientOptions) {
 		registry := bson.NewRegistry()
 		for _, v := range typeCodecs {
 			registry.RegisterTypeEncoder(v.ValueType, v.Encoder)
@@ -68,7 +67,7 @@ func WithTypeCodec(typeCodecs ...TypeCodec) Option {
 }
 
 func defaultTimeoutOption() Option {
-	return func(opts *options) {
+	return func(opts *clientOptions) {
 		opts.SetTimeout(defaultTimeout)
 	}
 }

core/stores/mon/options_test.go

@@ -7,9 +7,8 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"go.mongodb.org/mongo-driver/bson/bsoncodec"
-	"go.mongodb.org/mongo-driver/bson/bsonrw"
-	mopt "go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
 )
 
 func TestSetSlowThreshold(t *testing.T) {
@@ -19,13 +18,13 @@ func TestSetSlowThreshold(t *testing.T) {
 }
 
 func Test_defaultTimeoutOption(t *testing.T) {
-	opts := mopt.Client()
+	opts := options.Client()
 	defaultTimeoutOption()(opts)
 	assert.Equal(t, defaultTimeout, *opts.Timeout)
 }
 
 func TestWithTimeout(t *testing.T) {
-	opts := mopt.Client()
+	opts := options.Client()
 	WithTimeout(time.Second)(opts)
 	assert.Equal(t, time.Second, *opts.Timeout)
 }
@@ -57,10 +56,11 @@ func TestDisableInfoLog(t *testing.T) {
 }
 
 func TestWithRegistryForTimestampRegisterType(t *testing.T) {
-	opts := mopt.Client()
+	opts := options.Client()
 
 	// mongoDateTimeEncoder allow user convert time.Time to primitive.DateTime.
-	var mongoDateTimeEncoder bsoncodec.ValueEncoderFunc = func(ect bsoncodec.EncodeContext, w bsonrw.ValueWriter, value reflect.Value) error {
+	var mongoDateTimeEncoder bson.ValueEncoderFunc = func(ect bson.EncodeContext,
+		w bson.ValueWriter, value reflect.Value) error {
 		// Use reflect, determine if it can be converted to time.Time.
 		dec, ok := value.Interface().(time.Time)
 		if !ok {
@@ -70,7 +70,8 @@ func TestWithRegistryForTimestampRegisterType(t *testing.T) {
 	}
 
 	// mongoDateTimeEncoder allow user convert primitive.DateTime to time.Time.
-	var mongoDateTimeDecoder bsoncodec.ValueDecoderFunc = func(ect bsoncodec.DecodeContext, r bsonrw.ValueReader, value reflect.Value) error {
+	var mongoDateTimeDecoder bson.ValueDecoderFunc = func(ect bson.DecodeContext,
+		r bson.ValueReader, value reflect.Value) error {
 		primTime, err := r.ReadDateTime()
 		if err != nil {
 			return fmt.Errorf("error reading primitive.DateTime from ValueReader: %v", err)

core/stores/mon/trace.go

@@ -5,7 +5,7 @@ import (
 
 	"github.com/zeromicro/go-zero/core/errorx"
 	"github.com/zeromicro/go-zero/core/trace"
-	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
 	oteltrace "go.opentelemetry.io/otel/trace"

core/stores/monc/cachedmodel.go

@@ -8,8 +8,8 @@ import (
 	"github.com/zeromicro/go-zero/core/stores/mon"
 	"github.com/zeromicro/go-zero/core/stores/redis"
 	"github.com/zeromicro/go-zero/core/syncx"
-	"go.mongodb.org/mongo-driver/mongo"
-	mopt "go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
 )
 
 var (
@@ -78,7 +78,7 @@ func (mm *Model) DelCache(ctx context.Context, keys ...string) error {
 
 // DeleteOne deletes the document with given filter, and remove it from cache.
 func (mm *Model) DeleteOne(ctx context.Context, key string, filter any,
-	opts ...*mopt.DeleteOptions) (int64, error) {
+	opts ...options.Lister[options.DeleteOneOptions]) (int64, error) {
 	val, err := mm.Model.DeleteOne(ctx, filter, opts...)
 	if err != nil {
 		return 0, err
@@ -93,13 +93,13 @@ func (mm *Model) DeleteOne(ctx context.Context, key string, filter any,
 
 // DeleteOneNoCache deletes the document with given filter.
 func (mm *Model) DeleteOneNoCache(ctx context.Context, filter any,
-	opts ...*mopt.DeleteOptions) (int64, error) {
+	opts ...options.Lister[options.DeleteOneOptions]) (int64, error) {
 	return mm.Model.DeleteOne(ctx, filter, opts...)
 }
 
 // FindOne unmarshals a record into v with given key and query.
 func (mm *Model) FindOne(ctx context.Context, key string, v, filter any,
-	opts ...*mopt.FindOneOptions) error {
+	opts ...options.Lister[options.FindOneOptions]) error {
 	return mm.cache.TakeCtx(ctx, v, key, func(v any) error {
 		return mm.Model.FindOne(ctx, v, filter, opts...)
 	})
@@ -107,13 +107,13 @@ func (mm *Model) FindOne(ctx context.Context, key string, v, filter any,
 
 // FindOneNoCache unmarshals a record into v with query, without cache.
 func (mm *Model) FindOneNoCache(ctx context.Context, v, filter any,
-	opts ...*mopt.FindOneOptions) error {
+	opts ...options.Lister[options.FindOneOptions]) error {
 	return mm.Model.FindOne(ctx, v, filter, opts...)
 }
 
 // FindOneAndDelete deletes the document with given filter, and unmarshals it into v.
 func (mm *Model) FindOneAndDelete(ctx context.Context, key string, v, filter any,
-	opts ...*mopt.FindOneAndDeleteOptions) error {
+	opts ...options.Lister[options.FindOneAndDeleteOptions]) error {
 	if err := mm.Model.FindOneAndDelete(ctx, v, filter, opts...); err != nil {
 		return err
 	}
@@ -123,13 +123,13 @@ func (mm *Model) FindOneAndDelete(ctx context.Context, key string, v, filter any
 
 // FindOneAndDeleteNoCache deletes the document with given filter, and unmarshals it into v.
 func (mm *Model) FindOneAndDeleteNoCache(ctx context.Context, v, filter any,
-	opts ...*mopt.FindOneAndDeleteOptions) error {
+	opts ...options.Lister[options.FindOneAndDeleteOptions]) error {
 	return mm.Model.FindOneAndDelete(ctx, v, filter, opts...)
 }
 
 // FindOneAndReplace replaces the document with given filter with replacement, and unmarshals it into v.
 func (mm *Model) FindOneAndReplace(ctx context.Context, key string, v, filter any,
-	replacement any, opts ...*mopt.FindOneAndReplaceOptions) error {
+	replacement any, opts ...options.Lister[options.FindOneAndReplaceOptions]) error {
 	if err := mm.Model.FindOneAndReplace(ctx, v, filter, replacement, opts...); err != nil {
 		return err
 	}
@@ -139,13 +139,13 @@ func (mm *Model) FindOneAndReplace(ctx context.Context, key string, v, filter an
 
 // FindOneAndReplaceNoCache replaces the document with given filter with replacement, and unmarshals it into v.
 func (mm *Model) FindOneAndReplaceNoCache(ctx context.Context, v, filter any,
-	replacement any, opts ...*mopt.FindOneAndReplaceOptions) error {
+	replacement any, opts ...options.Lister[options.FindOneAndReplaceOptions]) error {
 	return mm.Model.FindOneAndReplace(ctx, v, filter, replacement, opts...)
 }
 
 // FindOneAndUpdate updates the document with given filter with update, and unmarshals it into v.
 func (mm *Model) FindOneAndUpdate(ctx context.Context, key string, v, filter any,
-	update any, opts ...*mopt.FindOneAndUpdateOptions) error {
+	update any, opts ...options.Lister[options.FindOneAndUpdateOptions]) error {
 	if err := mm.Model.FindOneAndUpdate(ctx, v, filter, update, opts...); err != nil {
 		return err
 	}
@@ -155,7 +155,7 @@ func (mm *Model) FindOneAndUpdate(ctx context.Context, key string, v, filter any
 
 // FindOneAndUpdateNoCache updates the document with given filter with update, and unmarshals it into v.
 func (mm *Model) FindOneAndUpdateNoCache(ctx context.Context, v, filter any,
-	update any, opts ...*mopt.FindOneAndUpdateOptions) error {
+	update any, opts ...options.Lister[options.FindOneAndUpdateOptions]) error {
 	return mm.Model.FindOneAndUpdate(ctx, v, filter, update, opts...)
 }
 
@@ -166,7 +166,7 @@ func (mm *Model) GetCache(key string, v any) error {
 
 // InsertOne inserts a single document into the collection, and remove the cache placeholder.
 func (mm *Model) InsertOne(ctx context.Context, key string, document any,
-	opts ...*mopt.InsertOneOptions) (*mongo.InsertOneResult, error) {
+	opts ...options.Lister[options.InsertOneOptions]) (*mongo.InsertOneResult, error) {
 	res, err := mm.Model.InsertOne(ctx, document, opts...)
 	if err != nil {
 		return nil, err
@@ -181,13 +181,13 @@ func (mm *Model) InsertOne(ctx context.Context, key string, document any,
 
 // InsertOneNoCache inserts a single document into the collection.
 func (mm *Model) InsertOneNoCache(ctx context.Context, document any,
-	opts ...*mopt.InsertOneOptions) (*mongo.InsertOneResult, error) {
+	opts ...options.Lister[options.InsertOneOptions]) (*mongo.InsertOneResult, error) {
 	return mm.Model.InsertOne(ctx, document, opts...)
 }
 
 // ReplaceOne replaces a single document in the collection, and remove the cache.
 func (mm *Model) ReplaceOne(ctx context.Context, key string, filter, replacement any,
-	opts ...*mopt.ReplaceOptions) (*mongo.UpdateResult, error) {
+	opts ...options.Lister[options.ReplaceOptions]) (*mongo.UpdateResult, error) {
 	res, err := mm.Model.ReplaceOne(ctx, filter, replacement, opts...)
 	if err != nil {
 		return nil, err
@@ -202,7 +202,7 @@ func (mm *Model) ReplaceOne(ctx context.Context, key string, filter, replacement
 
 // ReplaceOneNoCache replaces a single document in the collection.
 func (mm *Model) ReplaceOneNoCache(ctx context.Context, filter, replacement any,
-	opts ...*mopt.ReplaceOptions) (*mongo.UpdateResult, error) {
+	opts ...options.Lister[options.ReplaceOptions]) (*mongo.UpdateResult, error) {
 	return mm.Model.ReplaceOne(ctx, filter, replacement, opts...)
 }
 
@@ -213,7 +213,7 @@ func (mm *Model) SetCache(key string, v any) error {
 
 // UpdateByID updates the document with given id with update, and remove the cache.
 func (mm *Model) UpdateByID(ctx context.Context, key string, id, update any,
-	opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error) {
+	opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error) {
 	res, err := mm.Model.UpdateByID(ctx, id, update, opts...)
 	if err != nil {
 		return nil, err
@@ -228,13 +228,13 @@ func (mm *Model) UpdateByID(ctx context.Context, key string, id, update any,
 
 // UpdateByIDNoCache updates the document with given id with update.
 func (mm *Model) UpdateByIDNoCache(ctx context.Context, id, update any,
-	opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error) {
+	opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error) {
 	return mm.Model.UpdateByID(ctx, id, update, opts...)
 }
 
 // UpdateMany updates the documents that match filter with update, and remove the cache.
 func (mm *Model) UpdateMany(ctx context.Context, keys []string, filter, update any,
-	opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error) {
+	opts ...options.Lister[options.UpdateManyOptions]) (*mongo.UpdateResult, error) {
 	res, err := mm.Model.UpdateMany(ctx, filter, update, opts...)
 	if err != nil {
 		return nil, err
@@ -249,13 +249,13 @@ func (mm *Model) UpdateMany(ctx context.Context, keys []string, filter, update a
 
 // UpdateManyNoCache updates the documents that match filter with update.
 func (mm *Model) UpdateManyNoCache(ctx context.Context, filter, update any,
-	opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error) {
+	opts ...options.Lister[options.UpdateManyOptions]) (*mongo.UpdateResult, error) {
 	return mm.Model.UpdateMany(ctx, filter, update, opts...)
 }
 
 // UpdateOne updates the first document that matches filter with update, and remove the cache.
 func (mm *Model) UpdateOne(ctx context.Context, key string, filter, update any,
-	opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error) {
+	opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error) {
 	res, err := mm.Model.UpdateOne(ctx, filter, update, opts...)
 	if err != nil {
 		return nil, err
@@ -270,6 +270,6 @@ func (mm *Model) UpdateOne(ctx context.Context, key string, filter, update any,
 
 // UpdateOneNoCache updates the first document that matches filter with update.
 func (mm *Model) UpdateOneNoCache(ctx context.Context, filter, update any,
-	opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error) {
+	opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error) {
 	return mm.Model.UpdateOne(ctx, filter, update, opts...)
 }

core/stores/monc/cachedmodel_test.go

@@ -8,506 +8,519 @@ import (
 
 	"github.com/alicebob/miniredis/v2"
 	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stores/cache"
 	"github.com/zeromicro/go-zero/core/stores/mon"
 	"github.com/zeromicro/go-zero/core/stores/redis"
-	"go.mongodb.org/mongo-driver/bson"
-	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
+	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.uber.org/mock/gomock"
 )
 
-func TestNewModel(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		_, err := newModel("foo", mt.DB.Name(), mt.Coll.Name(), nil)
-		assert.NotNil(mt, err)
+func TestMustNewModel(t *testing.T) {
+	s, err := miniredis.Run()
+	assert.Nil(t, err)
+	original := logx.ExitOnFatal.True()
+	logx.ExitOnFatal.Set(false)
+	defer logx.ExitOnFatal.Set(original)
+
+	assert.Panics(t, func() {
+		MustNewModel("foo", "db", "collectino", cache.CacheConf{
+			cache.NodeConf{
+				RedisConf: redis.RedisConf{
+					Host: s.Addr(),
+					Type: redis.NodeType,
+				},
+				Weight: 100,
+			}})
 	})
 }
 
+func TestMustNewNodeModel(t *testing.T) {
+	s, err := miniredis.Run()
+	assert.Nil(t, err)
+	original := logx.ExitOnFatal.True()
+	logx.ExitOnFatal.Set(false)
+	defer logx.ExitOnFatal.Set(original)
+
+	assert.Panics(t, func() {
+		MustNewNodeModel("foo", "db", "collectino", redis.New(s.Addr()))
+	})
+}
+
+func TestNewModel(t *testing.T) {
+	s, err := miniredis.Run()
+	assert.Nil(t, err)
+	_, err = NewModel("foo", "db", "coll", cache.CacheConf{
+		cache.NodeConf{
+			RedisConf: redis.RedisConf{
+				Host: s.Addr(),
+				Type: redis.NodeType,
+			},
+			Weight: 100,
+		},
+	})
+	assert.Error(t, err)
+}
+
+func TestNewNodeModel(t *testing.T) {
+	_, err := NewNodeModel("foo", "db", "coll", nil)
+	assert.NotNil(t, err)
+}
+
+func TestNewModelWithCache(t *testing.T) {
+	_, err := NewModelWithCache("foo", "db", "coll", nil)
+	assert.NotNil(t, err)
+}
+
+func Test_newModel(t *testing.T) {
+	mon.Inject("mongodb://localhost:27018", &mongo.Client{})
+	model, err := newModel("mongodb://localhost:27018", "db", "collection", nil)
+	assert.Nil(t, err)
+	assert.NotNil(t, model)
+}
+
 func TestModel_DelCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		m := createModel(t, mt)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-		assert.Nil(t, m.cache.Set("bar", "baz"))
-		assert.Nil(t, m.DelCache(context.Background(), "foo", "bar"))
-		var v string
-		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
-		assert.True(t, m.cache.IsNotFound(m.cache.Get("bar", &v)))
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
+	assert.Nil(t, m.cache.Set("bar", "baz"))
+	assert.Nil(t, m.DelCache(context.Background(), "foo", "bar"))
+	var v string
+	assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+	assert.True(t, m.cache.IsNotFound(m.cache.Get("bar", &v)))
 }
 
 func TestModel_DeleteOne(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
-		m := createModel(t, mt)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-		val, err := m.DeleteOne(context.Background(), "foo", bson.D{{Key: "foo", Value: "bar"}})
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), val)
-		var v string
-		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
-		_, err = m.DeleteOne(context.Background(), "foo", bson.D{{Key: "foo", Value: "bar"}})
-		assert.NotNil(t, err)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	mockCollection.EXPECT().DeleteOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.DeleteResult{DeletedCount: 1}, nil)
+	m := createModel(t, mockCollection)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
+	val, err := m.DeleteOne(context.Background(), "foo", bson.D{{Key: "foo", Value: "bar"}})
+	assert.Nil(t, err)
+	assert.Equal(t, int64(1), val)
+	var v string
+	assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+	mockCollection.EXPECT().DeleteOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.DeleteResult{}, errMocked)
+	_, err = m.DeleteOne(context.Background(), "foo", bson.D{{Key: "foo", Value: "bar"}})
+	assert.NotNil(t, err)
 
-		m.cache = mockedCache{m.cache}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
-		_, err = m.DeleteOne(context.Background(), "foo", bson.D{{Key: "foo", Value: "bar"}})
-		assert.Equal(t, errMocked, err)
-	})
+	m.cache = mockedCache{m.cache}
+	mockCollection.EXPECT().DeleteOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.DeleteResult{}, nil)
+	_, err = m.DeleteOne(context.Background(), "foo", bson.D{{Key: "foo", Value: "bar"}})
+	assert.Equal(t, errMocked, err)
 }
 
 func TestModel_DeleteOneNoCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
-		m := createModel(t, mt)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-		val, err := m.DeleteOneNoCache(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
-		assert.Nil(t, err)
-		assert.Equal(t, int64(1), val)
-		var v string
-		assert.Nil(t, m.cache.Get("foo", &v))
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	mockCollection.EXPECT().DeleteOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.DeleteResult{DeletedCount: 1}, nil)
+	m := createModel(t, mockCollection)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
+	val, err := m.DeleteOneNoCache(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+	assert.Nil(t, err)
+	assert.Equal(t, int64(1), val)
+	var v string
+	assert.Nil(t, m.cache.Get("foo", &v))
 }
 
 func TestModel_FindOne(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		resp := mtest.CreateCursorResponse(
-			1,
-			"DBName.CollectionName",
-			mtest.FirstBatch,
-			bson.D{
-				{Key: "foo", Value: "bar"},
-			})
-		mt.AddMockResponses(resp)
-		m := createModel(t, mt)
-		var v struct {
-			Foo string `bson:"foo"`
-		}
-		assert.Nil(t, m.FindOne(context.Background(), "foo", &v, bson.D{}))
-		assert.Equal(t, "bar", v.Foo)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	mockCollection.EXPECT().FindOne(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, nil), nil)
+	m := createModel(t, mockCollection)
+	var v struct {
+		Foo string `bson:"foo"`
+	}
+	assert.Nil(t, m.FindOne(context.Background(), "foo", &v, bson.D{}))
+	assert.Equal(t, "bar", v.Foo)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
 }
 
 func TestModel_FindOneNoCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		resp := mtest.CreateCursorResponse(
-			1,
-			"DBName.CollectionName",
-			mtest.FirstBatch,
-			bson.D{
-				{Key: "foo", Value: "bar"},
-			})
-		mt.AddMockResponses(resp)
-		m := createModel(t, mt)
-		v := struct {
-			Foo string `bson:"foo"`
-		}{}
-		assert.Nil(t, m.FindOneNoCache(context.Background(), &v, bson.D{}))
-		assert.Equal(t, "bar", v.Foo)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	mockCollection.EXPECT().FindOne(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, nil), nil)
+	m := createModel(t, mockCollection)
+	v := struct {
+		Foo string `bson:"foo"`
+	}{}
+	assert.Nil(t, m.FindOneNoCache(context.Background(), &v, bson.D{}))
+	assert.Equal(t, "bar", v.Foo)
 }
 
 func TestModel_FindOneAndDelete(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-		v := struct {
-			Foo string `bson:"foo"`
-		}{}
-		assert.Nil(t, m.FindOneAndDelete(context.Background(), "foo", &v, bson.D{}))
-		assert.Equal(t, "bar", v.Foo)
-		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
-		assert.NotNil(t, m.FindOneAndDelete(context.Background(), "foo", &v, bson.D{}))
-
-		m.cache = mockedCache{m.cache}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		assert.Equal(t, errMocked, m.FindOneAndDelete(context.Background(), "foo", &v, bson.D{}))
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	mockCollection.EXPECT().FindOneAndDelete(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, bson.NewRegistry()), nil)
+	m := createModel(t, mockCollection)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
+	v := struct {
+		Foo string `bson:"foo"`
+	}{}
+	assert.Nil(t, m.FindOneAndDelete(context.Background(), "foo", &v, bson.D{}))
+	assert.Equal(t, "bar", v.Foo)
+	assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+	mockCollection.EXPECT().FindOneAndDelete(gomock.Any(), gomock.Any(), gomock.Any()).Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, bson.NewRegistry()), errMocked)
+	assert.NotNil(t, m.FindOneAndDelete(context.Background(), "foo", &v, bson.D{}))
+	m.cache = mockedCache{m.cache}
+	mockCollection.EXPECT().FindOneAndDelete(gomock.Any(), gomock.Any(), gomock.Any()).Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, bson.NewRegistry()), nil)
+	assert.Equal(t, errMocked, m.FindOneAndDelete(context.Background(), "foo", &v, bson.D{}))
 }
 
 func TestModel_FindOneAndDeleteNoCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		v := struct {
-			Foo string `bson:"foo"`
-		}{}
-		assert.Nil(t, m.FindOneAndDeleteNoCache(context.Background(), &v, bson.D{}))
-		assert.Equal(t, "bar", v.Foo)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	mockCollection.EXPECT().FindOneAndDelete(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, nil), nil)
+	m := createModel(t, mockCollection)
+	v := struct {
+		Foo string `bson:"foo"`
+	}{}
+	assert.Nil(t, m.FindOneAndDeleteNoCache(context.Background(), &v, bson.D{}))
+	assert.Equal(t, "bar", v.Foo)
 }
 
 func TestModel_FindOneAndReplace(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-		v := struct {
-			Foo string `bson:"foo"`
-		}{}
-		assert.Nil(t, m.FindOneAndReplace(context.Background(), "foo", &v, bson.D{}, bson.D{
-			{Key: "name", Value: "Mary"},
-		}))
-		assert.Equal(t, "bar", v.Foo)
-		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
-		assert.NotNil(t, m.FindOneAndReplace(context.Background(), "foo", &v, bson.D{}, bson.D{
-			{Key: "name", Value: "Mary"},
-		}))
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
+	v := struct {
+		Foo string `bson:"foo"`
+	}{}
+	mockCollection.EXPECT().FindOneAndReplace(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, nil), nil)
+	assert.Nil(t, m.FindOneAndReplace(context.Background(), "foo", &v, bson.D{}, bson.D{
+		{Key: "name", Value: "Mary"},
+	}))
+	assert.Equal(t, "bar", v.Foo)
+	assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+	mockCollection.EXPECT().FindOneAndReplace(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(mongo.NewSingleResultFromDocument(bson.M{"name": "Mary"}, nil, nil), errMocked)
+	assert.NotNil(t, m.FindOneAndReplace(context.Background(), "foo", &v, bson.D{}, bson.D{
+		{Key: "name", Value: "Mary"},
+	}))
 
-		m.cache = mockedCache{m.cache}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		assert.Equal(t, errMocked, m.FindOneAndReplace(context.Background(), "foo", &v, bson.D{}, bson.D{
-			{Key: "name", Value: "Mary"},
-		}))
-	})
+	m.cache = mockedCache{m.cache}
+	mockCollection.EXPECT().FindOneAndReplace(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, nil), nil)
+	assert.Equal(t, errMocked, m.FindOneAndReplace(context.Background(), "foo", &v, bson.D{}, bson.D{
+		{Key: "name", Value: "Mary"},
+	}))
 }
 
 func TestModel_FindOneAndReplaceNoCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		v := struct {
-			Foo string `bson:"foo"`
-		}{}
-		assert.Nil(t, m.FindOneAndReplaceNoCache(context.Background(), &v, bson.D{}, bson.D{
-			{Key: "name", Value: "Mary"},
-		}))
-		assert.Equal(t, "bar", v.Foo)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	v := struct {
+		Foo string `bson:"foo"`
+	}{}
+	mockCollection.EXPECT().FindOneAndReplace(gomock.Any(), gomock.Any(), gomock.Any()).Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, nil), nil)
+	assert.Nil(t, m.FindOneAndReplaceNoCache(context.Background(), &v, bson.D{}, bson.D{
+		{Key: "name", Value: "Mary"},
+	}))
+	assert.Equal(t, "bar", v.Foo)
 }
 
 func TestModel_FindOneAndUpdate(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-		v := struct {
-			Foo string `bson:"foo"`
-		}{}
-		assert.Nil(t, m.FindOneAndUpdate(context.Background(), "foo", &v, bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
-		}))
-		assert.Equal(t, "bar", v.Foo)
-		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
-		assert.NotNil(t, m.FindOneAndUpdate(context.Background(), "foo", &v, bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
-		}))
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
+	v := struct {
+		Foo string `bson:"foo"`
+	}{}
+	mockCollection.EXPECT().FindOneAndUpdate(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, nil), nil)
+	assert.Nil(t, m.FindOneAndUpdate(context.Background(), "foo", &v, bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
+	}))
+	assert.Equal(t, "bar", v.Foo)
+	assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+	mockCollection.EXPECT().FindOneAndUpdate(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, nil), errMocked)
+	assert.NotNil(t, m.FindOneAndUpdate(context.Background(), "foo", &v, bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
+	}))
 
-		m.cache = mockedCache{m.cache}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		assert.Equal(t, errMocked, m.FindOneAndUpdate(context.Background(), "foo", &v, bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
-		}))
-	})
+	m.cache = mockedCache{m.cache}
+	mockCollection.EXPECT().FindOneAndUpdate(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, nil), nil)
+	assert.Equal(t, errMocked, m.FindOneAndUpdate(context.Background(), "foo", &v, bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
+	}))
 }
 
 func TestModel_FindOneAndUpdateNoCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		v := struct {
-			Foo string `bson:"foo"`
-		}{}
-		assert.Nil(t, m.FindOneAndUpdateNoCache(context.Background(), &v, bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
-		}))
-		assert.Equal(t, "bar", v.Foo)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	v := struct {
+		Foo string `bson:"foo"`
+	}{}
+	mockCollection.EXPECT().FindOneAndUpdate(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(mongo.NewSingleResultFromDocument(bson.M{"foo": "bar"}, nil, nil), nil)
+	assert.Nil(t, m.FindOneAndUpdateNoCache(context.Background(), &v, bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
+	}))
+	assert.Equal(t, "bar", v.Foo)
 }
 
 func TestModel_GetCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		m := createModel(t, mt)
-		assert.NotNil(t, m.cache)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-		var s string
-		assert.Nil(t, m.cache.Get("foo", &s))
-		assert.Equal(t, "bar", s)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	assert.NotNil(t, m.cache)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
+	var s string
+	assert.Nil(t, m.cache.Get("foo", &s))
+	assert.Equal(t, "bar", s)
 }
 
 func TestModel_InsertOne(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-		resp, err := m.InsertOne(context.Background(), "foo", bson.D{
-			{Key: "name", Value: "Mary"},
-		})
-		assert.Nil(t, err)
-		assert.NotNil(t, resp)
-		var v string
-		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
-		_, err = m.InsertOne(context.Background(), "foo", bson.D{
-			{Key: "name", Value: "Mary"},
-		})
-		assert.NotNil(t, err)
-
-		m.cache = mockedCache{m.cache}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		_, err = m.InsertOne(context.Background(), "foo", bson.D{
-			{Key: "name", Value: "Mary"},
-		})
-		assert.Equal(t, errMocked, err)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
+	mockCollection.EXPECT().InsertOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.InsertOneResult{}, nil)
+	resp, err := m.InsertOne(context.Background(), "foo", bson.D{
+		{Key: "name", Value: "Mary"},
 	})
+	assert.Nil(t, err)
+	assert.NotNil(t, resp)
+	var v string
+	assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+	mockCollection.EXPECT().InsertOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.InsertOneResult{}, errMocked)
+	_, err = m.InsertOne(context.Background(), "foo", bson.D{
+		{Key: "name", Value: "Mary"},
+	})
+	assert.NotNil(t, err)
+
+	m.cache = mockedCache{m.cache}
+	mockCollection.EXPECT().InsertOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.InsertOneResult{}, nil)
+	_, err = m.InsertOne(context.Background(), "foo", bson.D{
+		{Key: "name", Value: "Mary"},
+	})
+	assert.Equal(t, errMocked, err)
 }
 
 func TestModel_InsertOneNoCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		resp, err := m.InsertOneNoCache(context.Background(), bson.D{
-			{Key: "name", Value: "Mary"},
-		})
-		assert.Nil(t, err)
-		assert.NotNil(t, resp)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	mockCollection.EXPECT().InsertOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.InsertOneResult{}, nil)
+	resp, err := m.InsertOneNoCache(context.Background(), bson.D{
+		{Key: "name", Value: "Mary"},
 	})
+	assert.Nil(t, err)
+	assert.NotNil(t, resp)
 }
 
 func TestModel_ReplaceOne(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-		resp, err := m.ReplaceOne(context.Background(), "foo", bson.D{}, bson.D{
-			{Key: "foo", Value: "baz"},
-		})
-		assert.Nil(t, err)
-		assert.NotNil(t, resp)
-		var v string
-		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
-		_, err = m.ReplaceOne(context.Background(), "foo", bson.D{}, bson.D{
-			{Key: "foo", Value: "baz"},
-		})
-		assert.NotNil(t, err)
-
-		m.cache = mockedCache{m.cache}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		_, err = m.ReplaceOne(context.Background(), "foo", bson.D{}, bson.D{
-			{Key: "foo", Value: "baz"},
-		})
-		assert.Equal(t, errMocked, err)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
+	mockCollection.EXPECT().ReplaceOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	resp, err := m.ReplaceOne(context.Background(), "foo", bson.D{}, bson.D{
+		{Key: "foo", Value: "baz"},
 	})
+	assert.Nil(t, err)
+	assert.NotNil(t, resp)
+	var v string
+	assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+	mockCollection.EXPECT().ReplaceOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, errMocked)
+	_, err = m.ReplaceOne(context.Background(), "foo", bson.D{}, bson.D{
+		{Key: "foo", Value: "baz"},
+	})
+	assert.NotNil(t, err)
+
+	m.cache = mockedCache{m.cache}
+	mockCollection.EXPECT().ReplaceOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	_, err = m.ReplaceOne(context.Background(), "foo", bson.D{}, bson.D{
+		{Key: "foo", Value: "baz"},
+	})
+	assert.Equal(t, errMocked, err)
 }
 
 func TestModel_ReplaceOneNoCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		resp, err := m.ReplaceOneNoCache(context.Background(), bson.D{}, bson.D{
-			{Key: "foo", Value: "baz"},
-		})
-		assert.Nil(t, err)
-		assert.NotNil(t, resp)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	mockCollection.EXPECT().ReplaceOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	resp, err := m.ReplaceOneNoCache(context.Background(), bson.D{}, bson.D{
+		{Key: "foo", Value: "baz"},
 	})
+	assert.Nil(t, err)
+	assert.NotNil(t, resp)
 }
 
 func TestModel_SetCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		m := createModel(t, mt)
-		assert.Nil(t, m.SetCache("foo", "bar"))
-		var v string
-		assert.Nil(t, m.GetCache("foo", &v))
-		assert.Equal(t, "bar", v)
-	})
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	assert.Nil(t, m.SetCache("foo", "bar"))
+	var v string
+	assert.Nil(t, m.GetCache("foo", &v))
+	assert.Equal(t, "bar", v)
 }
 
 func TestModel_UpdateByID(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-		resp, err := m.UpdateByID(context.Background(), "foo", bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
-		})
-		assert.Nil(t, err)
-		assert.NotNil(t, resp)
-		var v string
-		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
-		_, err = m.UpdateByID(context.Background(), "foo", bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
-		})
-		assert.NotNil(t, err)
-
-		m.cache = mockedCache{m.cache}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		_, err = m.UpdateByID(context.Background(), "foo", bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
-		})
-		assert.Equal(t, errMocked, err)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
+	mockCollection.EXPECT().UpdateByID(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	resp, err := m.UpdateByID(context.Background(), "foo", bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
 	})
+	assert.Nil(t, err)
+	assert.NotNil(t, resp)
+	var v string
+	assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+	mockCollection.EXPECT().UpdateByID(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, errMocked)
+	_, err = m.UpdateByID(context.Background(), "foo", bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+	})
+	assert.NotNil(t, err)
+
+	m.cache = mockedCache{m.cache}
+	mockCollection.EXPECT().UpdateByID(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	_, err = m.UpdateByID(context.Background(), "foo", bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+	})
+	assert.Equal(t, errMocked, err)
 }
 
 func TestModel_UpdateByIDNoCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		resp, err := m.UpdateByIDNoCache(context.Background(), bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
-		})
-		assert.Nil(t, err)
-		assert.NotNil(t, resp)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	mockCollection.EXPECT().UpdateByID(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	resp, err := m.UpdateByIDNoCache(context.Background(), bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
 	})
+	assert.Nil(t, err)
+	assert.NotNil(t, resp)
 }
 
 func TestModel_UpdateMany(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-		assert.Nil(t, m.cache.Set("bar", "baz"))
-		resp, err := m.UpdateMany(context.Background(), []string{"foo", "bar"}, bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
-		})
-		assert.Nil(t, err)
-		assert.NotNil(t, resp)
-		var v string
-		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
-		assert.True(t, m.cache.IsNotFound(m.cache.Get("bar", &v)))
-		_, err = m.UpdateMany(context.Background(), []string{"foo", "bar"}, bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
-		})
-		assert.NotNil(t, err)
-
-		m.cache = mockedCache{m.cache}
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		_, err = m.UpdateMany(context.Background(), []string{"foo", "bar"}, bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
-		})
-		assert.Equal(t, errMocked, err)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
+	assert.Nil(t, m.cache.Set("bar", "baz"))
+	mockCollection.EXPECT().UpdateMany(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	resp, err := m.UpdateMany(context.Background(), []string{"foo", "bar"}, bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
 	})
+	assert.Nil(t, err)
+	assert.NotNil(t, resp)
+	var v string
+	assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+	assert.True(t, m.cache.IsNotFound(m.cache.Get("bar", &v)))
+	mockCollection.EXPECT().UpdateMany(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, errMocked)
+	_, err = m.UpdateMany(context.Background(), []string{"foo", "bar"}, bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+	})
+	assert.NotNil(t, err)
+
+	m.cache = mockedCache{m.cache}
+	mockCollection.EXPECT().UpdateMany(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	_, err = m.UpdateMany(context.Background(), []string{"foo", "bar"}, bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+	})
+	assert.Equal(t, errMocked, err)
 }
 
 func TestModel_UpdateManyNoCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		resp, err := m.UpdateManyNoCache(context.Background(), bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
-		})
-		assert.Nil(t, err)
-		assert.NotNil(t, resp)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	mockCollection.EXPECT().UpdateMany(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	resp, err := m.UpdateManyNoCache(context.Background(), bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
 	})
+	assert.Nil(t, err)
+	assert.NotNil(t, resp)
 }
 
 func TestModel_UpdateOne(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		assert.Nil(t, m.cache.Set("foo", "bar"))
-		resp, err := m.UpdateOne(context.Background(), "foo", bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
-		})
-		assert.Nil(t, err)
-		assert.NotNil(t, resp)
-		var v string
-		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
-		_, err = m.UpdateOne(context.Background(), "foo", bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
-		})
-		assert.NotNil(t, err)
-
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m.cache = mockedCache{m.cache}
-		_, err = m.UpdateOne(context.Background(), "foo", bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
-		})
-		assert.Equal(t, errMocked, err)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	assert.Nil(t, m.cache.Set("foo", "bar"))
+	mockCollection.EXPECT().UpdateOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	resp, err := m.UpdateOne(context.Background(), "foo", bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
 	})
+	assert.Nil(t, err)
+	assert.NotNil(t, resp)
+	var v string
+	mockCollection.EXPECT().UpdateOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, errMocked)
+	assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+	_, err = m.UpdateOne(context.Background(), "foo", bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+	})
+	assert.NotNil(t, err)
+
+	mockCollection.EXPECT().UpdateOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	m.cache = mockedCache{m.cache}
+	_, err = m.UpdateOne(context.Background(), "foo", bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+	})
+	assert.Equal(t, errMocked, err)
 }
 
 func TestModel_UpdateOneNoCache(t *testing.T) {
-	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
-	mt.Run("test", func(mt *mtest.T) {
-		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
-			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
-		}...))
-		m := createModel(t, mt)
-		resp, err := m.UpdateOneNoCache(context.Background(), bson.D{}, bson.D{
-			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
-		})
-		assert.Nil(t, err)
-		assert.NotNil(t, resp)
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockCollection := mon.NewMockCollection(ctrl)
+	m := createModel(t, mockCollection)
+	mockCollection.EXPECT().UpdateOne(gomock.Any(), gomock.Any(), gomock.Any()).Return(&mongo.UpdateResult{}, nil)
+	resp, err := m.UpdateOneNoCache(context.Background(), bson.D{}, bson.D{
+		{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
 	})
+	assert.Nil(t, err)
+	assert.NotNil(t, resp)
 }
 
-func createModel(t *testing.T, mt *mtest.T) *Model {
+func createModel(t *testing.T, coll mon.Collection) *Model {
 	s, err := miniredis.Run()
 	assert.Nil(t, err)
-	mon.Inject(mt.Name(), mt.Client)
 	if atomic.AddInt32(&index, 1)%2 == 0 {
-		return MustNewNodeModel(mt.Name(), mt.DB.Name(), mt.Coll.Name(), redis.New(s.Addr()))
+		return mustNewTestNodeModel(coll, redis.New(s.Addr()))
 	} else {
-		return MustNewModel(mt.Name(), mt.DB.Name(), mt.Coll.Name(), cache.CacheConf{
+		return mustNewTestModel(coll, cache.CacheConf{
 			cache.NodeConf{
 				RedisConf: redis.RedisConf{
 					Host: s.Addr(),
@@ -519,6 +532,27 @@ func createModel(t *testing.T, mt *mtest.T) *Model {
 	}
 }
 
+// mustNewTestModel returns a test Model with the given cache.
+func mustNewTestModel(collection mon.Collection, c cache.CacheConf, opts ...cache.Option) *Model {
+	return &Model{
+		Model: &mon.Model{
+			Collection: collection,
+		},
+		cache: cache.New(c, singleFlight, stats, mongo.ErrNoDocuments, opts...),
+	}
+}
+
+// NewNodeModel returns a test Model with a cache node.
+func mustNewTestNodeModel(collection mon.Collection, rds *redis.Redis, opts ...cache.Option) *Model {
+	c := cache.NewNode(rds, singleFlight, stats, mongo.ErrNoDocuments, opts...)
+	return &Model{
+		Model: &mon.Model{
+			Collection: collection,
+		},
+		cache: c,
+	}
+}
+
 var (
 	errMocked = errors.New("mocked error")
 	index     int32

core/stores/monc/migration-2.0.md

@@ -0,0 +1,19 @@
+# Migrating from 1.x to 2.0
+
+To upgrade imports of the Go Driver from v1 to v2, we recommend using [marwan-at-work/mod
+](https://github.com/marwan-at-work/mod):
+
+```
+mod upgrade --mod-name=go.mongodb.org/mongo-driver
+```
+
+# Notice
+After completing the mod upgrade, code changes are typically unnecessary in the vast majority of cases. However, if your project references packages including but not limited to those listed below, you'll need to manually replace them, as these libraries are no longer present in the v2 version.
+```go
+go.mongodb.org/mongo-driver/bson/bsonrw    => go.mongodb.org/mongo-driver/v2/bson
+go.mongodb.org/mongo-driver/bson/bsoncodec => go.mongodb.org/mongo-driver/v2/bson
+go.mongodb.org/mongo-driver/bson/primitive => go.mongodb.org/mongo-driver/v2/bson
+```   
+
+See the following resources to learn more about upgrading from version 1.x to 2.0.:
+https://raw.githubusercontent.com/mongodb/mongo-go-driver/refs/heads/master/docs/migration-2.0.md

core/stores/redis/redis.go

@@ -65,7 +65,6 @@ type (
 	// RedisNode interface represents a redis node.
 	RedisNode interface {
 		red.Cmdable
-		red.BitMapCmdable
 	}
 
 	// GeoLocation is used with GeoAdd to add geospatial location.
@@ -260,12 +259,34 @@ func (s *Redis) BitPosCtx(ctx context.Context, key string, bit, start, end int64
 }
 
 // Blpop uses passed in redis connection to execute blocking queries.
+//
+// For blocking operations, you must create a dedicated RedisNode using CreateBlockingNode to avoid
+// exhausting the connection pool. Blocking commands hold connections for extended periods and should
+// not share the regular connection pool.
+//
+// Example usage:
+//
+//	node, err := redis.CreateBlockingNode(rds)
+//	if err != nil {
+//	    // handle error
+//	}
+//	defer node.Close()
+//
+//	value, err := rds.Blpop(node, "mylist")
+//	if err != nil {
+//	    // handle error
+//	}
+//
 // Doesn't benefit from pooling redis connections of blocking queries
 func (s *Redis) Blpop(node RedisNode, key string) (string, error) {
 	return s.BlpopCtx(context.Background(), node, key)
 }
 
 // BlpopCtx uses passed in redis connection to execute blocking queries.
+//
+// For blocking operations, you must create a dedicated RedisNode using CreateBlockingNode.
+// See Blpop for usage examples.
+//
 // Doesn't benefit from pooling redis connections of blocking queries
 func (s *Redis) BlpopCtx(ctx context.Context, node RedisNode, key string) (string, error) {
 	return s.BlpopWithTimeoutCtx(ctx, node, blockingQueryTimeout, key)
@@ -273,12 +294,18 @@ func (s *Redis) BlpopCtx(ctx context.Context, node RedisNode, key string) (strin
 
 // BlpopEx uses passed in redis connection to execute blpop command.
 // The difference against Blpop is that this method returns a bool to indicate success.
+//
+// For blocking operations, you must create a dedicated RedisNode using CreateBlockingNode.
+// See Blpop for usage examples.
 func (s *Redis) BlpopEx(node RedisNode, key string) (string, bool, error) {
 	return s.BlpopExCtx(context.Background(), node, key)
 }
 
 // BlpopExCtx uses passed in redis connection to execute blpop command.
 // The difference against Blpop is that this method returns a bool to indicate success.
+//
+// For blocking operations, you must create a dedicated RedisNode using CreateBlockingNode.
+// See Blpop for usage examples.
 func (s *Redis) BlpopExCtx(ctx context.Context, node RedisNode, key string) (string, bool, error) {
 	if node == nil {
 		return "", false, ErrNilNode
@@ -298,12 +325,18 @@ func (s *Redis) BlpopExCtx(ctx context.Context, node RedisNode, key string) (str
 
 // BlpopWithTimeout uses passed in redis connection to execute blpop command.
 // Control blocking query timeout
+//
+// For blocking operations, you must create a dedicated RedisNode using CreateBlockingNode.
+// See Blpop for usage examples.
 func (s *Redis) BlpopWithTimeout(node RedisNode, timeout time.Duration, key string) (string, error) {
 	return s.BlpopWithTimeoutCtx(context.Background(), node, timeout, key)
 }
 
 // BlpopWithTimeoutCtx uses passed in redis connection to execute blpop command.
 // Control blocking query timeout
+//
+// For blocking operations, you must create a dedicated RedisNode using CreateBlockingNode.
+// See Blpop for usage examples.
 func (s *Redis) BlpopWithTimeoutCtx(ctx context.Context, node RedisNode, timeout time.Duration,
 	key string) (string, error) {
 	if node == nil {
@@ -631,6 +664,28 @@ func (s *Redis) GetDelCtx(ctx context.Context, key string) (string, error) {
 	return val, err
 }
 
+// GetEx is the implementation of redis getex command.
+// Available since: redis version 6.2.0
+func (s *Redis) GetEx(key string, seconds int) (string, error) {
+	return s.GetExCtx(context.Background(), key, seconds)
+}
+
+// GetExCtx is the implementation of redis getex command.
+// Available since: redis version 6.2.0
+func (s *Redis) GetExCtx(ctx context.Context, key string, seconds int) (string, error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return "", err
+	}
+
+	val, err := conn.GetEx(ctx, key, time.Duration(seconds)*time.Second).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)
@@ -1285,10 +1340,12 @@ func (s *Redis) RpushCtx(ctx context.Context, key string, values ...any) (int, e
 	return int(v), nil
 }
 
+// RPopLPush atomically removes the last element from source list and prepends it to destination list.
 func (s *Redis) RPopLPush(source string, destination string) (string, error) {
 	return s.RPopLPushCtx(context.Background(), source, destination)
 }
 
+// RPopLPushCtx is the context-aware version of RPopLPush.
 func (s *Redis) RPopLPushCtx(ctx context.Context, source string, destination string) (string, error) {
 	conn, err := getRedis(s)
 	if err != nil {
@@ -1695,14 +1752,17 @@ func (s *Redis) TtlCtx(ctx context.Context, key string) (int, error) {
 	return int(duration), nil
 }
 
+// TxPipeline returns a Redis transaction pipeline for executing multiple commands atomically.
 func (s *Redis) TxPipeline() (pipe Pipeliner, err error) {
 	conn, err := getRedis(s)
 	if err != nil {
 		return nil, err
 	}
+
 	return conn.TxPipeline(), nil
 }
 
+// Unlink is similar to Del but removes keys asynchronously in a separate thread.
 func (s *Redis) Unlink(keys ...string) (int64, error) {
 	return s.UnlinkCtx(context.Background(), keys...)
 }
@@ -1712,9 +1772,181 @@ func (s *Redis) UnlinkCtx(ctx context.Context, keys ...string) (int64, error) {
 	if err != nil {
 		return 0, err
 	}
+
 	return conn.Unlink(ctx, keys...).Result()
 }
 
+// XAck acknowledges one or more messages in a Redis stream consumer group.
+// It marks the specified messages as successfully processed.
+func (s *Redis) XAck(stream string, group string, ids ...string) (int64, error) {
+	return s.XAckCtx(context.Background(), stream, group, ids...)
+}
+
+// XAckCtx is the context-aware version of XAck.
+func (s *Redis) XAckCtx(ctx context.Context, stream string, group string, ids ...string) (int64, error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return 0, err
+	}
+
+	return conn.XAck(ctx, stream, group, ids...).Result()
+}
+
+// XAdd adds a new entry to a Redis stream with the specified ID and field-value pairs.
+// If noMkStream is true, the command will fail if the stream doesn't exist.
+func (s *Redis) XAdd(stream string, noMkStream bool, id string, values any) (string, error) {
+	return s.XAddCtx(context.Background(), stream, noMkStream, id, values)
+}
+
+// XAddCtx is the context-aware version of XAdd.
+func (s *Redis) XAddCtx(ctx context.Context, stream string, noMkStream bool, id string, values any) (
+	string, error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return "", err
+	}
+
+	return conn.XAdd(ctx, &red.XAddArgs{
+		Stream:     stream,
+		ID:         id,
+		Values:     values,
+		NoMkStream: noMkStream,
+	}).Result()
+}
+
+// XGroupCreateMkStream creates a consumer group for a Redis stream.
+// If the stream doesn't exist, it will be created automatically.
+func (s *Redis) XGroupCreateMkStream(stream string, group string, start string) (string, error) {
+	return s.XGroupCreateMkStreamCtx(context.Background(), stream, group, start)
+}
+
+// XGroupCreateMkStreamCtx is the context-aware version of XGroupCreateMkStream.
+func (s *Redis) XGroupCreateMkStreamCtx(ctx context.Context, stream string, group string,
+	start string) (string, error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return "", err
+	}
+
+	return conn.XGroupCreateMkStream(ctx, stream, group, start).Result()
+}
+
+// XGroupCreate creates a consumer group for a Redis stream.
+// The stream must already exist, otherwise the command will fail.
+func (s *Redis) XGroupCreate(stream string, group string, start string) (string, error) {
+	return s.XGroupCreateCtx(context.Background(), stream, group, start)
+}
+
+// XGroupCreateCtx is the context-aware version of XGroupCreate.
+func (s *Redis) XGroupCreateCtx(ctx context.Context, stream string, group string, start string) (
+	string, error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return "", err
+	}
+
+	return conn.XGroupCreate(ctx, stream, group, start).Result()
+}
+
+// XInfoConsumers returns information about consumers in a Redis stream consumer group.
+func (s *Redis) XInfoConsumers(stream string, group string) ([]red.XInfoConsumer, error) {
+	return s.XInfoConsumersCtx(context.Background(), stream, group)
+}
+
+// XInfoConsumersCtx is the context-aware version of XInfoConsumers.
+func (s *Redis) XInfoConsumersCtx(ctx context.Context, stream string, group string) (
+	[]red.XInfoConsumer, error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return nil, err
+	}
+
+	return conn.XInfoConsumers(ctx, stream, group).Result()
+}
+
+// XInfoGroups returns information about consumer groups for a Redis stream.
+func (s *Redis) XInfoGroups(stream string) ([]red.XInfoGroup, error) {
+	return s.XInfoGroupsCtx(context.Background(), stream)
+}
+
+// XInfoGroupsCtx is the context-aware version of XInfoGroups.
+func (s *Redis) XInfoGroupsCtx(ctx context.Context, stream string) ([]red.XInfoGroup, error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return nil, err
+	}
+
+	return conn.XInfoGroups(ctx, stream).Result()
+}
+
+// XInfoStream returns general information about a Redis stream.
+func (s *Redis) XInfoStream(stream string) (*red.XInfoStream, error) {
+	return s.XInfoStreamCtx(context.Background(), stream)
+}
+
+// XInfoStreamCtx is the context-aware version of XInfoStream.
+func (s *Redis) XInfoStreamCtx(ctx context.Context, stream string) (*red.XInfoStream, error) {
+	conn, err := getRedis(s)
+	if err != nil {
+		return nil, err
+	}
+
+	return conn.XInfoStream(ctx, stream).Result()
+}
+
+// XReadGroup reads messages from Redis streams as part of a consumer group.
+// It allows for distributed processing of stream messages with automatic message delivery semantics.
+//
+// For blocking operations, you must create a dedicated RedisNode using CreateBlockingNode to avoid
+// exhausting the connection pool. Blocking commands hold connections for extended periods and should
+// not share the regular connection pool.
+//
+// Example usage:
+//
+//	node, err := redis.CreateBlockingNode(rds)
+//	if err != nil {
+//	    // handle error
+//	}
+//	defer node.Close()
+//
+//	streams, err := rds.XReadGroup(
+//	    node,              // RedisNode created with CreateBlockingNode
+//	    "mygroup",         // consumer group name
+//	    "consumer1",       // consumer ID
+//	    10,                // max number of messages to read
+//	    5*time.Second,     // block duration
+//	    false,             // noAck flag
+//	    "mystream",        // stream name
+//	)
+//
+// Doesn't benefit from pooling redis connections of blocking queries.
+func (s *Redis) XReadGroup(node RedisNode, group string, consumerId string, count int64,
+	block time.Duration, noAck bool, streams ...string) ([]red.XStream, error) {
+	return s.XReadGroupCtx(context.Background(), node, group, consumerId, count, block, noAck, streams...)
+}
+
+// XReadGroupCtx is the context-aware version of XReadGroup.
+//
+// For blocking operations, you must create a dedicated RedisNode using CreateBlockingNode to avoid
+// exhausting the connection pool. See XReadGroup for usage examples.
+//
+// Doesn't benefit from pooling redis connections of blocking queries.
+func (s *Redis) XReadGroupCtx(ctx context.Context, node RedisNode, group string, consumerId string,
+	count int64, block time.Duration, noAck bool, streams ...string) ([]red.XStream, error) {
+	if node == nil {
+		return nil, ErrNilNode
+	}
+
+	return node.XReadGroup(ctx, &red.XReadGroupArgs{
+		Group:    group,
+		Consumer: consumerId,
+		Count:    count,
+		Block:    block,
+		NoAck:    noAck,
+		Streams:  streams,
+	}).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)
@@ -1795,7 +2027,7 @@ func (s *Redis) ZaddsCtx(ctx context.Context, key string, ps ...Pair) (int64, er
 		return 0, err
 	}
 
-	var zs []red.Z
+	zs := make([]red.Z, 0, len(ps))
 	for _, p := range ps {
 		z := red.Z{Score: float64(p.Score), Member: p.Key}
 		zs = append(zs, z)

core/stores/redis/redis_test.go

@@ -916,6 +916,11 @@ func TestRedis_Ping(t *testing.T) {
 			ok := client.Ping()
 			assert.True(t, ok)
 		})
+
+		runOnRedisWithError(t, func(client *Redis) {
+			ok := client.Ping()
+			assert.False(t, ok)
+		})
 	})
 }
 
@@ -1099,6 +1104,45 @@ func TestRedis_GetDel(t *testing.T) {
 	})
 }
 
+func TestRedis_GetEx(t *testing.T) {
+	t.Run("get_ex", func(t *testing.T) {
+		runOnRedis(t, func(client *Redis) {
+			val, err := client.GetEx("getex_key", 10)
+			assert.Equal(t, "", val)
+			assert.Nil(t, err)
+
+			err = client.Set("getex_key", "getex_value")
+			assert.Nil(t, err)
+
+			val, err = client.GetEx("getex_key", 10)
+			assert.Nil(t, err)
+			assert.Equal(t, "getex_value", val)
+			val, err = client.Get("getex_key")
+			assert.Nil(t, err)
+			assert.Equal(t, "getex_value", val)
+
+			ttl, err := client.Ttl("getex_key")
+			assert.Nil(t, err)
+			assert.True(t, ttl > 0 && ttl <= 10)
+
+			val, err = client.GetEx("getex_key", 5)
+			assert.Nil(t, err)
+			assert.Equal(t, "getex_value", val)
+
+			ttl, err = client.Ttl("getex_key")
+			assert.Nil(t, err)
+			assert.True(t, ttl > 0 && ttl <= 5)
+		})
+	})
+
+	t.Run("get_ex_with_error", func(t *testing.T) {
+		runOnRedisWithError(t, func(client *Redis) {
+			_, err := newRedis(client.Addr, badType()).GetEx("hello", 10)
+			assert.Error(t, err)
+		})
+	})
+}
+
 func TestRedis_GetSet(t *testing.T) {
 	t.Run("set_get", func(t *testing.T) {
 		runOnRedis(t, func(client *Redis) {
@@ -2029,6 +2073,16 @@ func TestRedis_WithUserPass(t *testing.T) {
 		err := newRedis(client.Addr, WithUser("any"), WithPass("any")).Ping()
 		assert.NotNil(t, err)
 	})
+
+	runOnRedisWithAccount(t, "foo", "bar", func(client *Redis) {
+		err := client.Set("key1", "value1")
+		assert.Nil(t, err)
+		_, err = newRedis(client.Addr, badType()).Keys("*")
+		assert.NotNil(t, err)
+		keys, err := client.Keys("*")
+		assert.Nil(t, err)
+		assert.ElementsMatch(t, []string{"key1"}, keys)
+	})
 }
 
 func TestRedis_checkConnection(t *testing.T) {
@@ -2057,6 +2111,19 @@ func runOnRedis(t *testing.T, fn func(client *Redis)) {
 	}))
 }
 
+func runOnRedisWithAccount(t *testing.T, user, pass string, fn func(client *Redis)) {
+	logx.Disable()
+
+	s := miniredis.RunT(t)
+	s.RequireUserAuth(user, pass)
+	fn(MustNewRedis(RedisConf{
+		Host: s.Addr(),
+		Type: NodeType,
+		User: user,
+		Pass: pass,
+	}))
+}
+
 func runOnRedisWithError(t *testing.T, fn func(client *Redis)) {
 	logx.Disable()
 
@@ -2175,3 +2242,115 @@ func TestRedisTxPipeline(t *testing.T) {
 		assert.Equal(t, hashValue, value)
 	})
 }
+
+func TestRedisXGroupCreate(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		_, err := newRedis(client.Addr, badType()).XGroupCreate("Source", "Destination", "0")
+		assert.NotNil(t, err)
+
+		redisCli := newRedis(client.Addr)
+
+		_, err = redisCli.XGroupCreate("aa", "bb", "0")
+		assert.NotNil(t, err)
+
+		_, err = newRedis(client.Addr, badType()).XGroupCreateMkStream("Source", "Destination", "0")
+		assert.NotNil(t, err)
+
+		_, err = redisCli.XGroupCreateMkStream("aa", "bb", "0")
+		assert.Nil(t, err)
+
+		_, err = redisCli.XGroupCreate("aa", "cc", "0")
+		assert.Nil(t, err)
+	})
+}
+
+func TestRedisXInfo(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		_, err := newRedis(client.Addr, badType()).XInfoStream("Source")
+		assert.NotNil(t, err)
+		_, err = newRedis(client.Addr, badType()).XInfoGroups("Source")
+		assert.NotNil(t, err)
+
+		redisCli := newRedis(client.Addr)
+
+		stream := "aa"
+		group := "bb"
+
+		_, err = redisCli.XGroupCreateMkStream(stream, group, "$")
+		assert.Nil(t, err)
+
+		_, err = redisCli.XAdd(stream, true, "*", []string{"key1", "value1", "key2", "value2"})
+		assert.Nil(t, err)
+
+		infoStream, err := redisCli.XInfoStream(stream)
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), infoStream.Length)
+
+		infoGroups, err := redisCli.XInfoGroups(stream)
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), infoGroups[0].Lag)
+		assert.Equal(t, group, infoGroups[0].Name)
+
+		node, err := getRedis(redisCli)
+		assert.NoError(t, err)
+		redisCli.XAdd(stream, true, "*", []string{"key1", "value1", "key2", "value2"})
+		streamRes, err := redisCli.XReadGroup(node, group, "consumer", 1, 2000, false, stream, ">")
+		assert.Nil(t, err)
+		assert.Equal(t, 1, len(streamRes))
+		assert.Equal(t, "value1", streamRes[0].Messages[0].Values["key1"])
+
+		infoConsumers, err := redisCli.XInfoConsumers(stream, group)
+		assert.Nil(t, err)
+		assert.Equal(t, 1, len(infoConsumers))
+
+		_, err = newRedis(client.Addr, badType()).XInfoConsumers(stream, group)
+		assert.NotNil(t, err)
+	})
+}
+
+func TestRedisXReadGroup(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		_, err := newRedis(client.Addr, badType()).XAdd("bb", true, "*", []string{"key1", "value1", "key2", "value2"})
+		assert.NotNil(t, err)
+		_, err = newRedis(client.Addr, badType()).XAck("bb", "aa", "123")
+		assert.NotNil(t, err)
+
+		redisCli := newRedis(client.Addr)
+
+		stream := "aa"
+		group := "bb"
+
+		_, err = redisCli.XGroupCreateMkStream(stream, group, "$")
+		assert.Nil(t, err)
+
+		_, err = redisCli.XAdd(stream, true, "*", []string{"key1", "value1", "key2", "value2"})
+		assert.Nil(t, err)
+
+		node, err := getRedis(redisCli)
+		assert.NoError(t, err)
+		redisCli.XAdd(stream, true, "*", []string{"key1", "value1", "key2", "value2"})
+		_, err = redisCli.XReadGroup(nil, group, "consumer", 1, 2000, false, stream, ">")
+		assert.Error(t, err)
+		streamRes, err := redisCli.XReadGroup(node, group, "consumer", 1, 2000, false, stream, ">")
+		assert.Nil(t, err)
+		assert.Equal(t, 1, len(streamRes))
+		assert.Equal(t, "value1", streamRes[0].Messages[0].Values["key1"])
+
+		_, err = redisCli.XReadGroup(nil, group, "consumer", 1, 2000, false, stream, "0")
+		assert.Error(t, err)
+		streamRes1, err := redisCli.XReadGroup(node, group, "consumer", 1, 2000, false, stream, "0")
+		assert.Nil(t, err)
+		assert.Equal(t, 1, len(streamRes1))
+		assert.Equal(t, "value1", streamRes1[0].Messages[0].Values["key1"])
+
+		_, err = redisCli.XAck(stream, group, streamRes[0].Messages[0].ID)
+		assert.Nil(t, err)
+
+		_, err = redisCli.XReadGroup(nil, group, "consumer", 1, 2000, false, stream, "0")
+		assert.Error(t, err)
+		streamRes2, err := redisCli.XReadGroup(node, group, "consumer", 1, 2000, false, stream, "0")
+		assert.Nil(t, err)
+		assert.Greater(t, len(streamRes2), 0, "streamRes2 is empty")
+		assert.Equal(t, 0, len(streamRes2[0].Messages))
+	})
+}

core/stores/redis/redisblockingnode.go

@@ -13,7 +13,37 @@ type ClosableNode interface {
 	Close()
 }
 
-// CreateBlockingNode returns a ClosableNode.
+// CreateBlockingNode creates a dedicated RedisNode for blocking operations.
+//
+// Blocking Redis commands (like BLPOP, BRPOP, XREADGROUP with block parameter) hold connections
+// for extended periods while waiting for data. Using them with the regular Redis connection pool
+// can exhaust all available connections, causing other operations to fail or timeout.
+//
+// CreateBlockingNode creates a separate Redis client with a minimal connection pool (size 1) that
+// is dedicated to blocking operations. This ensures blocking commands don't interfere with regular
+// Redis operations.
+//
+// Example usage:
+//
+//	rds := redis.MustNewRedis(redis.RedisConf{
+//	    Host: "localhost:6379",
+//	    Type: redis.NodeType,
+//	})
+//
+//	// Create a dedicated node for blocking operations
+//	node, err := redis.CreateBlockingNode(rds)
+//	if err != nil {
+//	    // handle error
+//	}
+//	defer node.Close() // Important: close the node when done
+//
+//	// Use the node for blocking operations
+//	value, err := rds.Blpop(node, "mylist")
+//	if err != nil {
+//	    // handle error
+//	}
+//
+// The returned ClosableNode must be closed when no longer needed to release resources.
 func CreateBlockingNode(r *Redis) (ClosableNode, error) {
 	timeout := readWriteTimeout + blockingQueryTimeout
 

core/stores/sqlx/bulkinserter.go

@@ -25,8 +25,8 @@ type (
 	ResultHandler func(sql.Result, error)
 
 	// A BulkInserter is used to batch insert records.
-	// Postgresql is not supported yet, because of the sql is formated with symbol `$`.
-	// Oracle is not supported yet, because of the sql is formated with symbol `:`.
+	// Postgresql is not supported yet, because of the sql is formatted with symbol `$`.
+	// Oracle is not supported yet, because of the sql is formatted with symbol `:`.
 	BulkInserter struct {
 		executor *executors.PeriodicalExecutor
 		inserter *dbInserter

core/stores/sqlx/orm.go

@@ -9,7 +9,10 @@ import (
 	"github.com/zeromicro/go-zero/core/mapping"
 )
 
-const tagName = "db"
+const (
+	tagIgnore = "-"
+	tagName   = "db"
+)
 
 var (
 	// ErrNotMatchDestination is an error that indicates not matching destination to scan.
@@ -67,25 +70,16 @@ func getTaggedFieldValueMap(v reflect.Value) (map[string]any, error) {
 }
 
 func getValueInterface(value reflect.Value) (any, error) {
-	switch value.Kind() {
-	case reflect.Ptr:
-		if !value.CanInterface() {
-			return nil, ErrNotReadableValue
-		}
-
-		if value.IsNil() {
-			baseValueType := mapping.Deref(value.Type())
-			value.Set(reflect.New(baseValueType))
-		}
-
-		return value.Interface(), nil
-	default:
-		if !value.CanAddr() || !value.Addr().CanInterface() {
-			return nil, ErrNotReadableValue
-		}
-
-		return value.Addr().Interface(), nil
+	if !value.CanAddr() || !value.Addr().CanInterface() {
+		return nil, ErrNotReadableValue
 	}
+
+	if value.Kind() == reflect.Pointer && value.IsNil() {
+		baseValueType := mapping.Deref(value.Type())
+		value.Set(reflect.New(baseValueType))
+	}
+
+	return value.Addr().Interface(), nil
 }
 
 func isScanFailed(err error) bool {
@@ -269,13 +263,17 @@ func unwrapFields(v reflect.Value) []reflect.Value {
 			continue
 		}
 
+		childType := indirect.Type().Field(i)
+		if parseTagName(childType) == tagIgnore {
+			continue
+		}
+
 		if child.Kind() == reflect.Ptr && child.IsNil() {
 			baseValueType := mapping.Deref(child.Type())
 			child.Set(reflect.New(baseValueType))
 		}
 
 		child = reflect.Indirect(child)
-		childType := indirect.Type().Field(i)
 		if child.Kind() == reflect.Struct && childType.Anonymous {
 			fields = append(fields, unwrapFields(child)...)
 		} else {

core/syncx/immutableresource.go

@@ -47,30 +47,31 @@ func (ir *ImmutableResource) Get() (any, error) {
 		return resource, nil
 	}
 
-	ir.maybeRefresh(func() {
-		res, err := ir.fetch()
-		ir.lock.Lock()
-		if err != nil {
-			ir.err = err
-		} else {
-			ir.resource, ir.err = res, nil
-		}
-		ir.lock.Unlock()
-	})
+	ir.lock.Lock()
+	defer ir.lock.Unlock()
 
-	ir.lock.RLock()
-	resource, err := ir.resource, ir.err
-	ir.lock.RUnlock()
-	return resource, err
+	// double check
+	if ir.resource != nil {
+		return ir.resource, nil
+	}
+	if ir.err != nil && !ir.shouldRefresh() {
+		return ir.resource, ir.err
+	}
+
+	res, err := ir.fetch()
+	ir.lastTime.Set(timex.Now())
+	if err != nil {
+		ir.err = err
+		return nil, err
+	}
+
+	ir.resource, ir.err = res, nil
+	return res, nil
 }
 
-func (ir *ImmutableResource) maybeRefresh(execute func()) {
-	now := timex.Now()
+func (ir *ImmutableResource) shouldRefresh() bool {
 	lastTime := ir.lastTime.Load()
-	if lastTime == 0 || lastTime+ir.refreshInterval < now {
-		ir.lastTime.Set(now)
-		execute()
-	}
+	return lastTime == 0 || lastTime+ir.refreshInterval < timex.Now()
 }
 
 // WithRefreshIntervalOnFailure sets refresh interval on failure.

core/syncx/immutableresource_test.go

@@ -2,6 +2,8 @@ package syncx
 
 import (
 	"errors"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -56,6 +58,50 @@ func TestImmutableResourceError(t *testing.T) {
 	assert.Equal(t, 2, count)
 }
 
+// It's hard to test more than one goroutine fetching the resource at the same time,
+// because it's difficult to make more than one goroutine to pass the first read lock
+// and wait another to pass the read lock before it gets the write lock.
+func TestImmutableResourceConcurrent(t *testing.T) {
+	const message = "hello"
+	var count int32
+	ready := make(chan struct{})
+	r := NewImmutableResource(func() (any, error) {
+		atomic.AddInt32(&count, 1)
+		close(ready)                      // signal that fetch started
+		time.Sleep(10 * time.Millisecond) // simulate slow fetch
+		return message, nil
+	})
+
+	const goroutines = 10
+	var wg sync.WaitGroup
+	results := make([]any, goroutines)
+	errs := make([]error, goroutines)
+
+	wg.Add(goroutines)
+	for i := 0; i < goroutines; i++ {
+		go func(idx int) {
+			defer wg.Done()
+			res, err := r.Get()
+			results[idx] = res
+			errs[idx] = err
+		}(i)
+	}
+
+	// wait for fetch to start
+	<-ready
+
+	wg.Wait()
+
+	// fetch should only be called once despite concurrent access
+	assert.Equal(t, int32(1), atomic.LoadInt32(&count))
+
+	// all goroutines should eventually get the same result
+	for i := 0; i < goroutines; i++ {
+		assert.Nil(t, errs[i])
+		assert.Equal(t, message, results[i])
+	}
+}
+
 func TestImmutableResourceErrorRefreshAlways(t *testing.T) {
 	var count int
 	r := NewImmutableResource(func() (any, error) {

core/threading/routinegroup_test.go

@@ -1,13 +1,12 @@
 package threading
 
 import (
-	"io"
-	"log"
 	"sync"
 	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/logx/logtest"
 )
 
 func TestRoutineGroupRun(t *testing.T) {
@@ -25,7 +24,7 @@ func TestRoutineGroupRun(t *testing.T) {
 }
 
 func TestRoutingGroupRunSafe(t *testing.T) {
-	log.SetOutput(io.Discard)
+	logtest.Discard(t)
 
 	var count int32
 	group := NewRoutineGroup()

core/threading/routines_test.go

@@ -3,13 +3,12 @@ package threading
 import (
 	"bytes"
 	"context"
-	"io"
-	"log"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/lang"
 	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logx/logtest"
 )
 
 func TestRoutineId(t *testing.T) {
@@ -17,7 +16,7 @@ func TestRoutineId(t *testing.T) {
 }
 
 func TestRunSafe(t *testing.T) {
-	log.SetOutput(io.Discard)
+	logtest.Discard(t)
 
 	i := 0
 

core/trace/agent.go

@@ -7,7 +7,6 @@ import (
 	"os"
 	"sync"
 
-	"github.com/zeromicro/go-zero/core/lang"
 	"github.com/zeromicro/go-zero/core/logx"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/exporters/jaeger"
@@ -30,42 +29,36 @@ const (
 )
 
 var (
-	agents = make(map[string]lang.PlaceholderType)
-	lock   sync.Mutex
-	tp     *sdktrace.TracerProvider
+	once           sync.Once
+	tp             *sdktrace.TracerProvider
+	shutdownOnceFn = sync.OnceFunc(func() {
+		if tp != nil {
+			_ = tp.Shutdown(context.Background())
+		}
+	})
 )
 
 // StartAgent starts an opentelemetry agent.
+// It uses sync.Once to ensure the agent is initialized only once,
+// similar to prometheus.StartAgent and logx.SetUp.
+// This prevents multiple ServiceConf.SetUp() calls from reinitializing
+// the global tracer provider when running multiple servers (e.g., REST + RPC)
+// in the same process.
 func StartAgent(c Config) {
 	if c.Disabled {
 		return
 	}
 
-	lock.Lock()
-	defer lock.Unlock()
-
-	_, ok := agents[c.Endpoint]
-	if ok {
-		return
-	}
-
-	// if error happens, let later calls run.
-	if err := startAgent(c); err != nil {
-		return
-	}
-
-	agents[c.Endpoint] = lang.Placeholder
+	once.Do(func() {
+		if err := startAgent(c); err != nil {
+			logx.Error(err)
+		}
+	})
 }
 
 // StopAgent shuts down the span processors in the order they were registered.
 func StopAgent() {
-	lock.Lock()
-	defer lock.Unlock()
-
-	if tp != nil {
-		_ = tp.Shutdown(context.Background())
-		tp = nil
-	}
+	shutdownOnceFn()
 }
 
 func createExporter(c Config) (sdktrace.SpanExporter, error) {

core/trace/agent_test.go

@@ -1,10 +1,13 @@
 package trace
 
 import (
+	"context"
+	"errors"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/logx"
+	"go.opentelemetry.io/otel"
 )
 
 func TestStartAgent(t *testing.T) {
@@ -89,23 +92,305 @@ func TestStartAgent(t *testing.T) {
 	StartAgent(c10)
 	defer StopAgent()
 
-	lock.Lock()
-	defer lock.Unlock()
-
-	// because remotehost cannot be resolved
-	assert.Equal(t, 6, len(agents))
-	_, ok := agents[""]
-	assert.True(t, ok)
-	_, ok = agents[endpoint1]
-	assert.True(t, ok)
-	_, ok = agents[endpoint2]
-	assert.False(t, ok)
-	_, ok = agents[endpoint5]
-	assert.True(t, ok)
-	_, ok = agents[endpoint6]
-	assert.False(t, ok)
-	_, ok = agents[endpoint71]
-	assert.True(t, ok)
-	_, ok = agents[endpoint72]
-	assert.False(t, ok)
+	// With sync.Once, only the first non-disabled config (c1) takes effect.
+	// Subsequent calls are ignored, which is the desired behavior to prevent
+	// multiple servers (REST + RPC) from reinitializing the global tracer.
+	assert.NotNil(t, tp)
+}
+
+func TestCreateExporter_InvalidFilePath(t *testing.T) {
+	logx.Disable()
+
+	c := Config{
+		Name:     "test-invalid-file",
+		Endpoint: "/non-existent-directory/trace.log",
+		Batcher:  kindFile,
+	}
+
+	_, err := createExporter(c)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "file exporter endpoint error")
+}
+
+func TestCreateExporter_UnknownBatcher(t *testing.T) {
+	logx.Disable()
+
+	c := Config{
+		Name:     "test-unknown",
+		Endpoint: "localhost:1234",
+		Batcher:  "unknown-batcher-type",
+	}
+
+	_, err := createExporter(c)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "unknown exporter")
+}
+
+func TestCreateExporter_ValidExporters(t *testing.T) {
+	logx.Disable()
+
+	tests := []struct {
+		name    string
+		config  Config
+		wantErr bool
+		errMsg  string
+	}{
+		{
+			name: "valid file exporter",
+			config: Config{
+				Name:     "file-test",
+				Endpoint: "/tmp/trace-test.log",
+				Batcher:  kindFile,
+			},
+			wantErr: false,
+		},
+		{
+			name: "invalid file path",
+			config: Config{
+				Name:     "file-test-invalid",
+				Endpoint: "/invalid-path/that/does/not/exist/trace.log",
+				Batcher:  kindFile,
+			},
+			wantErr: true,
+			errMsg:  "file exporter endpoint error",
+		},
+		{
+			name: "unknown batcher",
+			config: Config{
+				Name:     "unknown-test",
+				Endpoint: "localhost:1234",
+				Batcher:  "invalid-batcher",
+			},
+			wantErr: true,
+			errMsg:  "unknown exporter",
+		},
+		{
+			name: "jaeger http",
+			config: Config{
+				Name:     "jaeger-http",
+				Endpoint: "http://localhost:14268/api/traces",
+				Batcher:  kindJaeger,
+			},
+			wantErr: false,
+		},
+		{
+			name: "jaeger udp",
+			config: Config{
+				Name:     "jaeger-udp",
+				Endpoint: "udp://localhost:6831",
+				Batcher:  kindJaeger,
+			},
+			wantErr: false,
+		},
+		{
+			name: "zipkin",
+			config: Config{
+				Name:     "zipkin",
+				Endpoint: "http://localhost:9411/api/v2/spans",
+				Batcher:  kindZipkin,
+			},
+			wantErr: false,
+		},
+		{
+			name: "otlpgrpc",
+			config: Config{
+				Name:     "otlpgrpc",
+				Endpoint: "localhost:4317",
+				Batcher:  kindOtlpGrpc,
+			},
+			wantErr: false,
+		},
+		{
+			name: "otlpgrpc with headers",
+			config: Config{
+				Name:     "otlpgrpc-headers",
+				Endpoint: "localhost:4317",
+				Batcher:  kindOtlpGrpc,
+				OtlpHeaders: map[string]string{
+					"authorization": "Bearer token123",
+					"x-custom-key":  "custom-value",
+				},
+			},
+			wantErr: false,
+		},
+		{
+			name: "otlphttp",
+			config: Config{
+				Name:     "otlphttp",
+				Endpoint: "localhost:4318",
+				Batcher:  kindOtlpHttp,
+			},
+			wantErr: false,
+		},
+		{
+			name: "otlphttp with headers",
+			config: Config{
+				Name:     "otlphttp-headers",
+				Endpoint: "localhost:4318",
+				Batcher:  kindOtlpHttp,
+				OtlpHeaders: map[string]string{
+					"authorization": "Bearer token456",
+					"x-api-key":     "api-key-value",
+				},
+			},
+			wantErr: false,
+		},
+		{
+			name: "otlphttp with headers and path",
+			config: Config{
+				Name:         "otlphttp-headers-path",
+				Endpoint:     "localhost:4318",
+				Batcher:      kindOtlpHttp,
+				OtlpHttpPath: "/v1/traces",
+				OtlpHeaders: map[string]string{
+					"authorization":  "Bearer token789",
+					"x-custom-trace": "trace-id",
+				},
+			},
+			wantErr: false,
+		},
+		{
+			name: "otlphttp with secure connection",
+			config: Config{
+				Name:           "otlphttp-secure",
+				Endpoint:       "localhost:4318",
+				Batcher:        kindOtlpHttp,
+				OtlpHttpSecure: true,
+				OtlpHeaders: map[string]string{
+					"authorization": "Bearer secure-token",
+				},
+			},
+			wantErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			exporter, err := createExporter(tt.config)
+			if tt.wantErr {
+				assert.Error(t, err)
+				if tt.errMsg != "" {
+					assert.Contains(t, err.Error(), tt.errMsg)
+				}
+				assert.Nil(t, exporter)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, exporter)
+				// Clean up the exporter
+				if exporter != nil {
+					_ = exporter.Shutdown(context.Background())
+				}
+			}
+		})
+	}
+}
+
+func TestStopAgent(t *testing.T) {
+	logx.Disable()
+
+	// StopAgent should be idempotent and safe to call multiple times
+	assert.NotPanics(t, func() {
+		StopAgent()
+		StopAgent()
+		StopAgent()
+	})
+}
+
+func TestStartAgent_WithEndpoint(t *testing.T) {
+	logx.Disable()
+
+	tests := []struct {
+		name    string
+		config  Config
+		wantErr bool
+	}{
+		{
+			name: "empty endpoint - no exporter created",
+			config: Config{
+				Name:    "test-no-endpoint",
+				Sampler: 1.0,
+			},
+			wantErr: false,
+		},
+		{
+			name: "valid endpoint with file exporter",
+			config: Config{
+				Name:     "test-with-endpoint",
+				Endpoint: "/tmp/test-trace.log",
+				Batcher:  kindFile,
+				Sampler:  1.0,
+			},
+			wantErr: false,
+		},
+		{
+			name: "endpoint with invalid exporter type",
+			config: Config{
+				Name:     "test-invalid-batcher",
+				Endpoint: "localhost:1234",
+				Batcher:  "invalid-type",
+				Sampler:  1.0,
+			},
+			wantErr: true,
+		},
+		{
+			name: "endpoint with invalid file path",
+			config: Config{
+				Name:     "test-invalid-path",
+				Endpoint: "/non/existent/path/trace.log",
+				Batcher:  kindFile,
+				Sampler:  1.0,
+			},
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Reset tp for each test
+			originalTp := tp
+			tp = nil
+			defer func() {
+				if tp != nil {
+					_ = tp.Shutdown(context.Background())
+				}
+				tp = originalTp
+			}()
+
+			err := startAgent(tt.config)
+			if tt.wantErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, tp, "TracerProvider should be created")
+			}
+		})
+	}
+}
+
+func TestStartAgent_ErrorHandler(t *testing.T) {
+	// Setup a tracer provider to test error handler
+	originalTp := tp
+	tp = nil
+	defer func() {
+		if tp != nil {
+			_ = tp.Shutdown(context.Background())
+		}
+		tp = originalTp
+	}()
+
+	// Call startAgent to set up the error handler
+	config := Config{
+		Name:    "test-error-handler",
+		Sampler: 1.0,
+	}
+	err := startAgent(config)
+	assert.NoError(t, err)
+	assert.NotNil(t, tp)
+
+	// Verify the error handler was set and can be called without panicking
+	// We test this by calling otel.Handle which will invoke the registered error handler
+	testErr := errors.New("test otel error")
+	assert.NotPanics(t, func() {
+		otel.Handle(testErr)
+	}, "Error handler should handle errors without panicking")
 }

core/trace/tracer_test.go

@@ -82,7 +82,7 @@ func TestExtractValidTraceContext(t *testing.T) {
 			}),
 		},
 		{
-			name:        "invalid tracestate perserves traceparent",
+			name:        "invalid tracestate preserves traceparent",
 			traceparent: "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00",
 			tracestate:  "invalid$@#=invalid",
 			sc: trace.NewSpanContext(trace.SpanContextConfig{

gateway/internal/eventhandler.go

@@ -2,6 +2,7 @@ package internal
 
 import (
 	"io"
+	"net/http"
 
 	"github.com/golang/protobuf/jsonpb"
 	"github.com/golang/protobuf/proto"
@@ -11,6 +12,16 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+const (
+	// MetadataHeaderPrefix is the http prefix that represents custom metadata
+	// parameters to or from a gRPC call.
+	MetadataHeaderPrefix = "Grpc-Metadata-"
+
+	// MetadataTrailerPrefix is prepended to gRPC metadata as it is converted to
+	// HTTP headers in a response handled by go-zero gateway
+	MetadataTrailerPrefix = "Grpc-Trailer-"
+)
+
 type EventHandler struct {
 	Status    *status.Status
 	writer    io.Writer
@@ -27,13 +38,35 @@ func NewEventHandler(writer io.Writer, resolver jsonpb.AnyResolver) *EventHandle
 	}
 }
 
+func (h *EventHandler) OnReceiveHeaders(md metadata.MD) {
+	w, ok := h.writer.(http.ResponseWriter)
+	if ok {
+		for k, vs := range md {
+			header := defaultOutgoingHeaderMatcher(k)
+			for _, v := range vs {
+				w.Header().Add(header, v)
+			}
+		}
+	}
+}
+
 func (h *EventHandler) OnReceiveResponse(message proto.Message) {
 	if err := h.marshaler.Marshal(h.writer, message); err != nil {
 		logx.Error(err)
 	}
 }
 
-func (h *EventHandler) OnReceiveTrailers(status *status.Status, _ metadata.MD) {
+func (h *EventHandler) OnReceiveTrailers(status *status.Status, md metadata.MD) {
+	w, ok := h.writer.(http.ResponseWriter)
+	if ok {
+		for k, vs := range md {
+			header := defaultOutgoingTrailerMatcher(k)
+			for _, v := range vs {
+				w.Header().Add(header, v)
+			}
+		}
+	}
+
 	h.Status = status
 }
 
@@ -43,5 +76,10 @@ func (h *EventHandler) OnResolveMethod(_ *desc.MethodDescriptor) {
 func (h *EventHandler) OnSendHeaders(_ metadata.MD) {
 }
 
-func (h *EventHandler) OnReceiveHeaders(_ metadata.MD) {
+func defaultOutgoingHeaderMatcher(key string) string {
+	return MetadataHeaderPrefix + key
+}
+
+func defaultOutgoingTrailerMatcher(key string) string {
+	return MetadataTrailerPrefix + key
 }

gateway/internal/eventhandler_test.go

@@ -2,10 +2,12 @@ package internal
 
 import (
 	"io"
+	"net/http/httptest"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 )
 
@@ -18,3 +20,219 @@ func TestEventHandler(t *testing.T) {
 	assert.Equal(t, codes.OK, h.Status.Code())
 	h.OnReceiveResponse(nil)
 }
+
+func TestEventHandler_OnReceiveTrailers(t *testing.T) {
+	tests := []struct {
+		name           string
+		writer         io.Writer
+		status         *status.Status
+		metadata       metadata.MD
+		expectedStatus codes.Code
+		expectedHeader map[string][]string
+	}{
+		{
+			name:   "with http.ResponseWriter and metadata",
+			writer: httptest.NewRecorder(),
+			status: status.New(codes.OK, "success"),
+			metadata: metadata.MD{
+				"x-custom-header":  []string{"value1", "value2"},
+				"x-another-header": []string{"single-value"},
+			},
+			expectedStatus: codes.OK,
+			expectedHeader: map[string][]string{
+				"Grpc-Trailer-X-Custom-Header":  {"value1", "value2"},
+				"Grpc-Trailer-X-Another-Header": {"single-value"},
+			},
+		},
+		{
+			name:           "with http.ResponseWriter and nil metadata",
+			writer:         httptest.NewRecorder(),
+			status:         status.New(codes.Internal, "error"),
+			metadata:       nil,
+			expectedStatus: codes.Internal,
+			expectedHeader: map[string][]string{},
+		},
+		{
+			name:           "with non-http.ResponseWriter",
+			writer:         io.Discard,
+			status:         status.New(codes.OK, "success"),
+			metadata:       metadata.MD{"x-header": []string{"value"}},
+			expectedStatus: codes.OK,
+			expectedHeader: nil, // headers should not be set
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			h := NewEventHandler(tt.writer, nil)
+			
+			h.OnReceiveTrailers(tt.status, tt.metadata)
+			
+			// Check status is set correctly
+			assert.Equal(t, tt.expectedStatus, h.Status.Code())
+			
+			// Check headers are set correctly if writer is http.ResponseWriter
+			if recorder, ok := tt.writer.(*httptest.ResponseRecorder); ok {
+				if tt.expectedHeader != nil {
+					for key, expectedValues := range tt.expectedHeader {
+						actualValues := recorder.Header()[key]
+						assert.Equal(t, expectedValues, actualValues, "Header %s should match", key)
+					}
+				}
+			}
+		})
+	}
+}
+
+func TestEventHandler_OnReceiveHeaders(t *testing.T) {
+	tests := []struct {
+		name           string
+		writer         io.Writer
+		metadata       metadata.MD
+		expectedHeader map[string][]string
+	}{
+		{
+			name:   "with http.ResponseWriter and metadata",
+			writer: httptest.NewRecorder(),
+			metadata: metadata.MD{
+				"content-type":     []string{"application/json"},
+				"x-custom-header":  []string{"value1", "value2"},
+				"x-another-header": []string{"single-value"},
+			},
+			expectedHeader: map[string][]string{
+				"Grpc-Metadata-Content-Type":     {"application/json"},
+				"Grpc-Metadata-X-Custom-Header":  {"value1", "value2"},
+				"Grpc-Metadata-X-Another-Header": {"single-value"},
+			},
+		},
+		{
+			name:           "with http.ResponseWriter and nil metadata",
+			writer:         httptest.NewRecorder(),
+			metadata:       nil,
+			expectedHeader: map[string][]string{},
+		},
+		{
+			name:           "with http.ResponseWriter and empty metadata",
+			writer:         httptest.NewRecorder(),
+			metadata:       metadata.MD{},
+			expectedHeader: map[string][]string{},
+		},
+		{
+			name:           "with non-http.ResponseWriter",
+			writer:         io.Discard,
+			metadata:       metadata.MD{"x-header": []string{"value"}},
+			expectedHeader: nil, // headers should not be set
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			h := NewEventHandler(tt.writer, nil)
+			
+			h.OnReceiveHeaders(tt.metadata)
+			
+			// Check headers are set correctly if writer is http.ResponseWriter
+			if recorder, ok := tt.writer.(*httptest.ResponseRecorder); ok {
+				if tt.expectedHeader != nil {
+					for key, expectedValues := range tt.expectedHeader {
+						actualValues := recorder.Header()[key]
+						assert.Equal(t, expectedValues, actualValues, "Header %s should match", key)
+					}
+				}
+			}
+		})
+	}
+}
+
+func TestEventHandler_OnReceiveHeaders_MultipleValues(t *testing.T) {
+	recorder := httptest.NewRecorder()
+	h := NewEventHandler(recorder, nil)
+	
+	// Test that multiple calls to OnReceiveHeaders accumulate headers
+	h.OnReceiveHeaders(metadata.MD{
+		"x-header-1": []string{"value1"},
+	})
+	
+	h.OnReceiveHeaders(metadata.MD{
+		"x-header-1": []string{"value2"}, // Should add to existing header
+		"x-header-2": []string{"value3"},
+	})
+	
+	// Check that headers are accumulated (not overwritten) with proper prefix
+	assert.Equal(t, []string{"value1", "value2"}, recorder.Header()["Grpc-Metadata-X-Header-1"])
+	assert.Equal(t, []string{"value3"}, recorder.Header()["Grpc-Metadata-X-Header-2"])
+}
+
+func TestEventHandler_OnReceiveHeaders_MetadataPrefix(t *testing.T) {
+	tests := []struct {
+		name           string
+		metadata       metadata.MD
+		expectedHeader map[string][]string
+	}{
+		{
+			name: "all metadata headers should be prefixed with Grpc-Metadata-",
+			metadata: metadata.MD{
+				"content-type":    []string{"application/grpc"},
+				"x-custom-header": []string{"value1"},
+				"authorization":   []string{"Bearer token"},
+			},
+			expectedHeader: map[string][]string{
+				"Grpc-Metadata-Content-Type":    {"application/grpc"},
+				"Grpc-Metadata-X-Custom-Header": {"value1"},
+				"Grpc-Metadata-Authorization":   {"Bearer token"},
+			},
+		},
+		{
+			name: "mixed case headers should be prefixed",
+			metadata: metadata.MD{
+				"Content-Type":    []string{"APPLICATION/JSON"},
+				"X-Custom-Header": []string{"value1"},
+			},
+			expectedHeader: map[string][]string{
+				"Grpc-Metadata-Content-Type":    {"APPLICATION/JSON"},
+				"Grpc-Metadata-X-Custom-Header": {"value1"},
+			},
+		},
+		{
+			name: "multiple values for same header",
+			metadata: metadata.MD{
+				"x-multi-header": []string{"value1", "value2", "value3"},
+			},
+			expectedHeader: map[string][]string{
+				"Grpc-Metadata-X-Multi-Header": {"value1", "value2", "value3"},
+			},
+		},
+		{
+			name: "empty metadata",
+			metadata: metadata.MD{},
+			expectedHeader: map[string][]string{},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			recorder := httptest.NewRecorder()
+			h := NewEventHandler(recorder, nil)
+			
+			h.OnReceiveHeaders(tt.metadata)
+			
+			// Check that headers are set correctly
+			for key, expectedValues := range tt.expectedHeader {
+				actualValues := recorder.Header()[key]
+				assert.Equal(t, expectedValues, actualValues, "Header %s should match", key)
+			}
+			
+			// Ensure no unexpected headers are set
+			for actualKey := range recorder.Header() {
+				found := false
+				for expectedKey := range tt.expectedHeader {
+					if actualKey == expectedKey {
+						found = true
+						break
+					}
+				}
+				assert.True(t, found, "Unexpected header found: %s", actualKey)
+			}
+		})
+	}
+}

gateway/internal/headerprocessor.go

@@ -11,16 +11,40 @@ const (
 	metadataPrefix       = "gateway-"
 )
 
+// OpenTelemetry trace propagation headers that need to be forwarded to gRPC metadata.
+// These headers are used by the W3C Trace Context standard for distributed tracing.
+var traceHeaders = map[string]bool{
+	"traceparent": true,
+	"tracestate":  true,
+	"baggage":     true,
+}
+
 // ProcessHeaders builds the headers for the gateway from HTTP headers.
+// It forwards both custom metadata headers (with Grpc-Metadata- prefix)
+// and OpenTelemetry trace propagation headers (traceparent, tracestate, baggage)
+// to ensure distributed tracing works correctly across the gateway.
 func ProcessHeaders(header http.Header) []string {
 	var headers []string
 
 	for k, v := range header {
+		// Forward OpenTelemetry trace propagation headers
+		// These must be lowercase per gRPC metadata conventions
+		if lowerKey := strings.ToLower(k); traceHeaders[lowerKey] {
+			for _, vv := range v {
+				headers = append(headers, lowerKey+":"+vv)
+			}
+			continue
+		}
+
+		// Forward custom metadata headers with Grpc-Metadata- prefix
 		if !strings.HasPrefix(k, metadataHeaderPrefix) {
 			continue
 		}
 
-		key := fmt.Sprintf("%s%s", metadataPrefix, strings.TrimPrefix(k, metadataHeaderPrefix))
+		// gRPC metadata keys are case-insensitive and stored as lowercase,
+		// so we lowercase the key to match gRPC conventions
+		trimmedKey := strings.TrimPrefix(k, metadataHeaderPrefix)
+		key := strings.ToLower(fmt.Sprintf("%s%s", metadataPrefix, trimmedKey))
 		for _, vv := range v {
 			headers = append(headers, key+":"+vv)
 		}

gateway/internal/headerprocessor_test.go

@@ -18,5 +18,93 @@ func TestBuildHeadersWithValues(t *testing.T) {
 	req := httptest.NewRequest("GET", "/", http.NoBody)
 	req.Header.Add("grpc-metadata-a", "b")
 	req.Header.Add("grpc-metadata-b", "b")
-	assert.ElementsMatch(t, []string{"gateway-A:b", "gateway-B:b"}, ProcessHeaders(req.Header))
+	assert.ElementsMatch(t, []string{"gateway-a:b", "gateway-b:b"}, ProcessHeaders(req.Header))
+}
+
+func TestProcessHeadersWithTraceContext(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", http.NoBody)
+	req.Header.Set("traceparent", "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01")
+	req.Header.Set("tracestate", "key1=value1,key2=value2")
+	req.Header.Set("baggage", "userId=alice,serverNode=DF:28")
+
+	headers := ProcessHeaders(req.Header)
+
+	assert.Len(t, headers, 3)
+	assert.Contains(t, headers, "traceparent:00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01")
+	assert.Contains(t, headers, "tracestate:key1=value1,key2=value2")
+	assert.Contains(t, headers, "baggage:userId=alice,serverNode=DF:28")
+}
+
+func TestProcessHeadersWithMixedHeaders(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", http.NoBody)
+	req.Header.Set("traceparent", "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01")
+	req.Header.Set("grpc-metadata-custom", "value1")
+	req.Header.Set("content-type", "application/json")
+	req.Header.Set("tracestate", "key1=value1")
+
+	headers := ProcessHeaders(req.Header)
+
+	// Should include trace headers and grpc-metadata headers, but not regular headers
+	assert.Len(t, headers, 3)
+	assert.Contains(t, headers, "traceparent:00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01")
+	assert.Contains(t, headers, "tracestate:key1=value1")
+	assert.Contains(t, headers, "gateway-custom:value1")
+}
+
+func TestProcessHeadersTraceparentCaseInsensitive(t *testing.T) {
+	tests := []struct {
+		name        string
+		headerKey   string
+		headerVal   string
+		expectedKey string
+	}{
+		{
+			name:        "lowercase traceparent",
+			headerKey:   "traceparent",
+			headerVal:   "00-trace-span-01",
+			expectedKey: "traceparent",
+		},
+		{
+			name:        "uppercase Traceparent",
+			headerKey:   "Traceparent",
+			headerVal:   "00-trace-span-01",
+			expectedKey: "traceparent",
+		},
+		{
+			name:        "mixed case TraceParent",
+			headerKey:   "TraceParent",
+			headerVal:   "00-trace-span-01",
+			expectedKey: "traceparent",
+		},
+		{
+			name:        "lowercase tracestate",
+			headerKey:   "tracestate",
+			headerVal:   "key=value",
+			expectedKey: "tracestate",
+		},
+		{
+			name:        "mixed case TraceState",
+			headerKey:   "TraceState",
+			headerVal:   "key=value",
+			expectedKey: "tracestate",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req := httptest.NewRequest("GET", "/", http.NoBody)
+			req.Header.Set(tt.headerKey, tt.headerVal)
+
+			headers := ProcessHeaders(req.Header)
+
+			assert.Len(t, headers, 1)
+			assert.Contains(t, headers, tt.expectedKey+":"+tt.headerVal)
+		})
+	}
+}
+
+func TestProcessHeadersEmptyHeaders(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", http.NoBody)
+	headers := ProcessHeaders(req.Header)
+	assert.Empty(t, headers)
 }

gateway/internal/requestparser.go

@@ -26,11 +26,11 @@ func NewRequestParser(r *http.Request, resolver jsonpb.AnyResolver) (grpcurl.Req
 
 	body, ok := getBody(r)
 	if !ok {
-		return buildJsonRequestParser(params, resolver)
+		return buildJsonRequestParserFromMap(params, resolver)
 	}
 
 	if len(params) == 0 {
-		return grpcurl.NewJSONRequestParser(body, resolver), nil
+		return buildJsonRequestParserFromReader(body, resolver)
 	}
 
 	m := make(map[string]any)
@@ -42,17 +42,28 @@ func NewRequestParser(r *http.Request, resolver jsonpb.AnyResolver) (grpcurl.Req
 		m[k] = v
 	}
 
-	return buildJsonRequestParser(m, resolver)
+	return buildJsonRequestParserFromMap(m, resolver)
 }
 
-func buildJsonRequestParser(m map[string]any, resolver jsonpb.AnyResolver) (
+func buildJsonRequestParserFromMap(data map[string]any, resolver jsonpb.AnyResolver) (
 	grpcurl.RequestParser, error) {
 	var buf bytes.Buffer
-	if err := json.NewEncoder(&buf).Encode(m); err != nil {
+	if err := json.NewEncoder(&buf).Encode(data); err != nil {
 		return nil, err
 	}
 
-	return grpcurl.NewJSONRequestParser(&buf, resolver), nil
+	return buildJsonRequestParserFromReader(&buf, resolver)
+}
+
+// buildJsonRequestParserFromReader creates a JSON request parser with ignoring unknown fields.
+func buildJsonRequestParserFromReader(data io.Reader, resolver jsonpb.AnyResolver) (
+	grpcurl.RequestParser, error) {
+	unmarshaler := jsonpb.Unmarshaler{
+		AllowUnknownFields: true,
+		AnyResolver:        resolver,
+	}
+
+	return grpcurl.NewJSONRequestParserWithUnmarshaler(data, unmarshaler), nil
 }
 
 func getBody(r *http.Request) (io.Reader, bool) {

gateway/internal/requestparser_test.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/golang/protobuf/proto"
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/rest/pathvar"
 )
@@ -87,12 +88,123 @@ func TestNewRequestParserWithBadForm(t *testing.T) {
 	assert.Nil(t, parser)
 }
 
-func TestRequestParser_buildJsonRequestParser(t *testing.T) {
-	parser, err := buildJsonRequestParser(map[string]any{"a": make(chan int)}, nil)
+func TestRequestParser_buildJsonRequestParserFromMap(t *testing.T) {
+	parser, err := buildJsonRequestParserFromMap(map[string]any{"a": make(chan int)}, nil)
 	assert.NotNil(t, err)
 	assert.Nil(t, parser)
 }
 
+// mockAnyResolver is a simple implementation of jsonpb.AnyResolver for testing
+type mockAnyResolver struct{}
+
+func (m *mockAnyResolver) Resolve(typeUrl string) (proto.Message, error) {
+	return nil, nil
+}
+
+func TestNewRequestParserWithIgnoreUnknownFields(t *testing.T) {
+	// Create a concrete resolver for testing
+	resolver := &mockAnyResolver{}
+
+	// Test case 1: No body, no vars - should work with both true and false
+	req1 := httptest.NewRequest("GET", "/", http.NoBody)
+	parser1, err1 := NewRequestParser(req1, resolver)
+	assert.Nil(t, err1)
+	assert.NotNil(t, parser1)
+
+	req2 := httptest.NewRequest("GET", "/", http.NoBody)
+	parser2, err2 := NewRequestParser(req2, resolver)
+	assert.Nil(t, err2)
+	assert.NotNil(t, parser2)
+
+	// Test case 2: With JSON body - tests the body parsing path
+	req3 := httptest.NewRequest("POST", "/", strings.NewReader(`{"field": "value"}`))
+	parser3, err3 := NewRequestParser(req3, resolver)
+	assert.Nil(t, err3)
+	assert.NotNil(t, parser3)
+
+	req4 := httptest.NewRequest("POST", "/", strings.NewReader(`{"field": "value"}`))
+	parser4, err4 := NewRequestParser(req4, resolver)
+	assert.Nil(t, err4)
+	assert.NotNil(t, parser4)
+}
+
+func TestNewRequestParserWithVarsAndIgnoreUnknownFields(t *testing.T) {
+	resolver := &mockAnyResolver{}
+
+	// Test with path variables and ignoreUnknownFields = true
+	req := httptest.NewRequest("GET", "/", http.NoBody)
+	req = pathvar.WithVars(req, map[string]string{"a": "b"})
+	parser, err := NewRequestParser(req, resolver)
+	assert.Nil(t, err)
+	assert.NotNil(t, parser)
+
+	// Test with path variables and ignoreUnknownFields = false
+	req2 := httptest.NewRequest("GET", "/", http.NoBody)
+	req2 = pathvar.WithVars(req2, map[string]string{"c": "d"})
+	parser2, err2 := NewRequestParser(req2, resolver)
+	assert.Nil(t, err2)
+	assert.NotNil(t, parser2)
+}
+
+func TestNewRequestParserWithBodyAndIgnoreUnknownFields(t *testing.T) {
+	resolver := &mockAnyResolver{}
+
+	// Test with body and ignoreUnknownFields = true
+	req := httptest.NewRequest("POST", "/", strings.NewReader(`{"a": "b"}`))
+	parser, err := NewRequestParser(req, resolver)
+	assert.Nil(t, err)
+	assert.NotNil(t, parser)
+
+	// Test with body and ignoreUnknownFields = false
+	req2 := httptest.NewRequest("POST", "/", strings.NewReader(`{"c": "d"}`))
+	parser2, err2 := NewRequestParser(req2, resolver)
+	assert.Nil(t, err2)
+	assert.NotNil(t, parser2)
+}
+
+func TestNewRequestParserWithVarsBodyAndIgnoreUnknownFields(t *testing.T) {
+	resolver := &mockAnyResolver{}
+
+	// Test with both path variables and body, ignoreUnknownFields = true
+	req := httptest.NewRequest("POST", "/", strings.NewReader(`{"a": "b"}`))
+	req = pathvar.WithVars(req, map[string]string{"c": "d"})
+	parser, err := NewRequestParser(req, resolver)
+	assert.Nil(t, err)
+	assert.NotNil(t, parser)
+
+	// Test with both path variables and body, ignoreUnknownFields = false
+	req2 := httptest.NewRequest("POST", "/", strings.NewReader(`{"e": "f"}`))
+	req2 = pathvar.WithVars(req2, map[string]string{"g": "h"})
+	parser2, err2 := NewRequestParser(req2, resolver)
+	assert.Nil(t, err2)
+	assert.NotNil(t, parser2)
+}
+
+func TestBuildJsonRequestParserFromMapWithIgnoreUnknownFields(t *testing.T) {
+	resolver := &mockAnyResolver{}
+
+	// Test buildJsonRequestParserFromMap with ignoreUnknownFields = true
+	data := map[string]any{"key": "value"}
+	parser, err := buildJsonRequestParserFromMap(data, resolver)
+	assert.Nil(t, err)
+	assert.NotNil(t, parser)
+
+	// Test buildJsonRequestParserFromMap with ignoreUnknownFields = false
+	parser2, err2 := buildJsonRequestParserFromMap(data, resolver)
+	assert.Nil(t, err2)
+	assert.NotNil(t, parser2)
+}
+
+func TestBuildJsonRequestParserWithUnknownFields(t *testing.T) {
+	resolver := &mockAnyResolver{}
+
+	// Test buildJsonRequestParserWithUnknownFields
+	data := strings.NewReader(`{"test": "value"}`)
+	parser, err := buildJsonRequestParserFromReader(data, resolver)
+	assert.Nil(t, err)
+	assert.NotNil(t, parser)
+}
+
 type badBody struct{}
 
 func (badBody) Read([]byte) (int, error) { return 0, errors.New("something bad") }

gateway/server.go

@@ -34,6 +34,7 @@ type (
 		conns         []zrpc.Client
 		processHeader func(http.Header) []string
 		dialer        func(conf zrpc.RpcClientConf) zrpc.Client
+		middlewares   []rest.Middleware
 	}
 
 	// Option defines the method to customize Server.
@@ -103,9 +104,16 @@ func (s *Server) build() error {
 	})
 }
 
+func (s *Server) buildChainHandler(handler http.HandlerFunc) http.HandlerFunc {
+	for i := len(s.middlewares) - 1; i >= 0; i-- {
+		handler = s.middlewares[i](handler)
+	}
+	return handler
+}
+
 func (s *Server) buildGrpcHandler(source grpcurl.DescriptorSource, resolver jsonpb.AnyResolver,
 	cli zrpc.Client, rpcPath string) func(http.ResponseWriter, *http.Request) {
-	return func(w http.ResponseWriter, r *http.Request) {
+	handler := func(w http.ResponseWriter, r *http.Request) {
 		parser, err := internal.NewRequestParser(r, resolver)
 		if err != nil {
 			httpx.ErrorCtx(r.Context(), w, err)
@@ -124,6 +132,8 @@ func (s *Server) buildGrpcHandler(source grpcurl.DescriptorSource, resolver json
 			httpx.ErrorCtx(r.Context(), w, st.Err())
 		}
 	}
+
+	return s.buildChainHandler(handler)
 }
 
 func (s *Server) buildGrpcRoute(up Upstream, writer mr.Writer[rest.Route], cancel func(error)) {
@@ -177,7 +187,7 @@ func (s *Server) buildGrpcRoute(up Upstream, writer mr.Writer[rest.Route], cance
 }
 
 func (s *Server) buildHttpHandler(target *HttpClientConf) http.HandlerFunc {
-	return func(w http.ResponseWriter, r *http.Request) {
+	handler := func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set(httpx.ContentType, httpx.JsonContentType)
 		req, err := buildRequestWithNewTarget(r, target)
 		if err != nil {
@@ -213,6 +223,7 @@ func (s *Server) buildHttpHandler(target *HttpClientConf) http.HandlerFunc {
 			logc.Error(r.Context(), err)
 		}
 	}
+	return s.buildChainHandler(handler)
 }
 
 func (s *Server) buildHttpRoute(up Upstream, writer mr.Writer[rest.Route]) {
@@ -263,6 +274,14 @@ func WithHeaderProcessor(processHeader func(http.Header) []string) func(*Server)
 	}
 }
 
+// WithMiddleware adds one or more middleware functions to process HTTP requests.
+// Multiple middlewares will be executed in the order they were passed (like an onion model).
+func WithMiddleware(middlewares ...rest.Middleware) func(*Server) {
+	return func(s *Server) {
+		s.middlewares = append(s.middlewares, middlewares...)
+	}
+}
+
 func buildRequestWithNewTarget(r *http.Request, target *HttpClientConf) (*http.Request, error) {
 	u := *r.URL
 	u.Host = target.Target

gateway/server_test.go

@@ -325,3 +325,50 @@ type badResponseWriter struct {
 func (w *badResponseWriter) Write([]byte) (int, error) {
 	return 0, errors.New("bad writer")
 }
+
+func TestWithMiddleware(t *testing.T) {
+	var callOrder []string
+
+	firstMiddleware := func(next http.HandlerFunc) http.HandlerFunc {
+		return func(w http.ResponseWriter, r *http.Request) {
+			callOrder = append(callOrder, "first-start")
+			w.Header().Set("X-First-Middleware", "called")
+			next(w, r)
+			callOrder = append(callOrder, "first-end")
+		}
+	}
+
+	secondMiddleware := func(next http.HandlerFunc) http.HandlerFunc {
+		return func(w http.ResponseWriter, r *http.Request) {
+			callOrder = append(callOrder, "second-start")
+			w.Header().Set("X-Second-Middleware", "called")
+			next(w, r)
+			callOrder = append(callOrder, "second-end")
+		}
+	}
+
+	var c GatewayConf
+	err := conf.FillDefault(&c)
+	assert.Nil(t, err)
+	// Test multiple middlewares in one call
+	server1 := MustNewServer(c, WithMiddleware(firstMiddleware, secondMiddleware))
+	assert.Len(t, server1.middlewares, 2, "Should have 2 middlewares from one call")
+	// Test multiple middleware calls
+	server2 := MustNewServer(c, WithMiddleware(firstMiddleware), WithMiddleware(secondMiddleware))
+	assert.Len(t, server2.middlewares, 2, "Should have 2 middlewares from separate calls")
+	// Test execution order (onion model)
+	finalHandler := func(w http.ResponseWriter, r *http.Request) {
+		callOrder = append(callOrder, "handler")
+		w.WriteHeader(http.StatusOK)
+	}
+
+	testHandler := server1.buildChainHandler(finalHandler)
+	w := httptest.NewRecorder()
+	r := httptest.NewRequest("GET", "/test", nil)
+	testHandler(w, r)
+
+	expectedOrder := []string{"first-start", "second-start", "handler", "second-end", "first-end"}
+	assert.Equal(t, expectedOrder, callOrder, "Middleware execution should follow onion model")
+	assert.Equal(t, "called", w.Header().Get("X-First-Middleware"))
+	assert.Equal(t, "called", w.Header().Get("X-Second-Middleware"))
+}

go.mod

@@ -9,20 +9,19 @@ require (
 	github.com/fullstorydev/grpcurl v1.9.3
 	github.com/go-sql-driver/mysql v1.9.0
 	github.com/golang-jwt/jwt/v4 v4.5.2
-	github.com/golang/mock v1.6.0
 	github.com/golang/protobuf v1.5.4
 	github.com/google/uuid v1.6.0
-	github.com/grafana/pyroscope-go v1.2.2
+	github.com/grafana/pyroscope-go v1.2.7
 	github.com/jackc/pgx/v5 v5.7.4
 	github.com/jhump/protoreflect v1.17.0
 	github.com/pelletier/go-toml/v2 v2.2.2
 	github.com/prometheus/client_golang v1.21.1
-	github.com/redis/go-redis/v9 v9.11.0
+	github.com/redis/go-redis/v9 v9.17.2
 	github.com/spaolacci/murmur3 v1.1.0
-	github.com/stretchr/testify v1.10.0
+	github.com/stretchr/testify v1.11.1
 	go.etcd.io/etcd/api/v3 v3.5.15
 	go.etcd.io/etcd/client/v3 v3.5.15
-	go.mongodb.org/mongo-driver v1.17.4
+	go.mongodb.org/mongo-driver/v2 v2.4.1
 	go.opentelemetry.io/otel v1.24.0
 	go.opentelemetry.io/otel/exporters/jaeger v1.17.0
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.24.0
@@ -33,6 +32,7 @@ require (
 	go.opentelemetry.io/otel/trace v1.24.0
 	go.uber.org/automaxprocs v1.6.0
 	go.uber.org/goleak v1.3.0
+	go.uber.org/mock v0.4.0
 	golang.org/x/net v0.35.0
 	golang.org/x/sys v0.30.0
 	golang.org/x/time v0.10.0
@@ -68,11 +68,11 @@ require (
 	github.com/go-openapi/jsonreference v0.20.2 // indirect
 	github.com/go-openapi/swag v0.22.4 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
-	github.com/golang/snappy v0.0.4 // indirect
+	github.com/golang/snappy v1.0.0 // indirect
 	github.com/google/gnostic-models v0.6.8 // indirect
 	github.com/google/go-cmp v0.6.0 // indirect
 	github.com/google/gofuzz v1.2.0 // indirect
-	github.com/grafana/pyroscope-go/godeltaprof v0.1.8 // indirect
+	github.com/grafana/pyroscope-go/godeltaprof v0.1.9 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0 // indirect
 	github.com/h2non/parth v0.0.0-20190131123155-b4df798d6542 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
@@ -88,7 +88,6 @@ require (
 	github.com/mattn/go-runewidth v0.0.15 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
-	github.com/montanaflynn/stats v0.7.1 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/openzipkin/zipkin-go v0.4.3 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect

go.sum

@@ -62,12 +62,10 @@ github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
 github.com/golang-jwt/jwt/v4 v4.5.2 h1:YtQM7lnr8iZ+j5q71MGKkNw9Mn7AjHM68uc9g5fXeUI=
 github.com/golang-jwt/jwt/v4 v4.5.2/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
-github.com/golang/mock v1.6.0 h1:ErTB+efbowRARo13NNdxyJji2egdxLGQhRaY+DUumQc=
-github.com/golang/mock v1.6.0/go.mod h1:p6yTPP+5HYm5mzsMV8JkE6ZKdX+/wYM6Hr+LicevLPs=
 github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
 github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
-github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM=
-github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
+github.com/golang/snappy v1.0.0 h1:Oy607GVXHs7RtbggtPBnr2RmDArIsAefDwvrdWvRhGs=
+github.com/golang/snappy v1.0.0/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/google/gnostic-models v0.6.8 h1:yo/ABAfM5IMRsS1VnXjTBvUb61tFIHozhlYvRgGre9I=
 github.com/google/gnostic-models v0.6.8/go.mod h1:5n7qKqH0f5wFt+aWF8CW6pZLLNOfYuF5OpfBSENuI8U=
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
@@ -80,10 +78,10 @@ github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1 h1:K6RDEckDVWvDI9JAJY
 github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/grafana/pyroscope-go v1.2.2 h1:uvKCyZMD724RkaCEMrSTC38Yn7AnFe8S2wiAIYdDPCE=
-github.com/grafana/pyroscope-go v1.2.2/go.mod h1:zzT9QXQAp2Iz2ZdS216UiV8y9uXJYQiGE1q8v1FyhqU=
-github.com/grafana/pyroscope-go/godeltaprof v0.1.8 h1:iwOtYXeeVSAeYefJNaxDytgjKtUuKQbJqgAIjlnicKg=
-github.com/grafana/pyroscope-go/godeltaprof v0.1.8/go.mod h1:2+l7K7twW49Ct4wFluZD3tZ6e0SjanjcUUBPVD/UuGU=
+github.com/grafana/pyroscope-go v1.2.7 h1:VWBBlqxjyR0Cwk2W6UrE8CdcdD80GOFNutj0Kb1T8ac=
+github.com/grafana/pyroscope-go v1.2.7/go.mod h1:o/bpSLiJYYP6HQtvcoVKiE9s5RiNgjYTj1DhiddP2Pc=
+github.com/grafana/pyroscope-go/godeltaprof v0.1.9 h1:c1Us8i6eSmkW+Ez05d3co8kasnuOY813tbMN8i/a3Og=
+github.com/grafana/pyroscope-go/godeltaprof v0.1.9/go.mod h1:2+l7K7twW49Ct4wFluZD3tZ6e0SjanjcUUBPVD/UuGU=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0 h1:bkypFPDjIYGfCYD5mRBvpqxfYX1YCS1PXdKYWi8FsN0=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0/go.mod h1:P+Lt/0by1T8bfcF3z737NnSbmxQAppXMRziHUxPOC8k=
 github.com/h2non/parth v0.0.0-20190131123155-b4df798d6542 h1:2VTzZjLZBgl62/EtslCrtky5vbi9dd7HrQPQIx6wqiw=
@@ -130,8 +128,6 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w
 github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
 github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
 github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
-github.com/montanaflynn/stats v0.7.1 h1:etflOAAHORrCC44V+aR6Ftzort912ZU+YLiSTuV8eaE=
-github.com/montanaflynn/stats v0.7.1/go.mod h1:etXPPgVO6n31NxCd9KQUMvCM+ve0ruNzt6R8Bnaayow=
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
 github.com/nbio/st v0.0.0-20140626010706-e9e8d9816f32 h1:W6apQkHrMkS0Muv8G/TipAy/FJl/rCYT0+EuS8+Z0z4=
@@ -158,8 +154,8 @@ github.com/prometheus/common v0.62.0 h1:xasJaQlnWAeyHdUBeGjXmutelfJHWMRr+Fg4QszZ
 github.com/prometheus/common v0.62.0/go.mod h1:vyBcEuLSvWos9B1+CyL7JZ2up+uFzXhkqml0W5zIY1I=
 github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0learggepc=
 github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk=
-github.com/redis/go-redis/v9 v9.11.0 h1:E3S08Gl/nJNn5vkxd2i78wZxWAPNZgUNTp8WIJUAiIs=
-github.com/redis/go-redis/v9 v9.11.0/go.mod h1:huWgSWd8mW6+m0VPhJjSSQ+d6Nh1VICQ6Q5lHuCH/Iw=
+github.com/redis/go-redis/v9 v9.17.2 h1:P2EGsA4qVIM3Pp+aPocCJ7DguDHhqrXNhVcEp4ViluI=
+github.com/redis/go-redis/v9 v9.17.2/go.mod h1:u410H11HMLoB+TP67dz8rL9s6QW2j76l0//kSOd3370=
 github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
 github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ=
@@ -180,8 +176,8 @@ github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
-github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
+github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
 github.com/xdg-go/pbkdf2 v1.0.0 h1:Su7DPu48wXMwC3bs7MCNG+z4FhcyEuz5dlvchbq0B0c=
 github.com/xdg-go/pbkdf2 v1.0.0/go.mod h1:jrpuAogTd400dnrH08LKmI/xc1MbPOebTwRqcT5RDeI=
 github.com/xdg-go/scram v1.1.2 h1:FHX5I5B4i4hKRVRBCFRxq1iQRej7WO3hhBuJf+UUySY=
@@ -192,7 +188,6 @@ github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 h1:ilQV1hzziu+LLM3zU
 github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78/go.mod h1:aL8wCCfTfSfmXjznFBSZNN13rSJjlIOI1fUNAtF7rmI=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
-github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
 github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M=
 github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw=
@@ -202,8 +197,8 @@ go.etcd.io/etcd/client/pkg/v3 v3.5.15 h1:fo0HpWz/KlHGMCC+YejpiCmyWDEuIpnTDzpJLB5
 go.etcd.io/etcd/client/pkg/v3 v3.5.15/go.mod h1:mXDI4NAOwEiszrHCb0aqfAYNCrZP4e9hRca3d1YK8EU=
 go.etcd.io/etcd/client/v3 v3.5.15 h1:23M0eY4Fd/inNv1ZfU3AxrbbOdW79r9V9Rl62Nm6ip4=
 go.etcd.io/etcd/client/v3 v3.5.15/go.mod h1:CLSJxrYjvLtHsrPKsy7LmZEE+DK2ktfd2bN4RhBMwlU=
-go.mongodb.org/mongo-driver v1.17.4 h1:jUorfmVzljjr0FLzYQsGP8cgN/qzzxlY9Vh0C9KFXVw=
-go.mongodb.org/mongo-driver v1.17.4/go.mod h1:Hy04i7O2kC4RS06ZrhPRqj/u4DTYkFDAAccj+rVKqgQ=
+go.mongodb.org/mongo-driver/v2 v2.4.1 h1:hGDMngUao03OVQ6sgV5csk+RWOIkF+CuLsTPobNMGNI=
+go.mongodb.org/mongo-driver/v2 v2.4.1/go.mod h1:jHeEDJHJq7tm6ZF45Issun9dbogjfnPySb1vXA7EeAI=
 go.opentelemetry.io/otel v1.24.0 h1:0LAOdjNmQeSTzGBzduGe/rU4tZhMwL5rWgtp9Ku5Jfo=
 go.opentelemetry.io/otel v1.24.0/go.mod h1:W7b9Ozg4nkF5tWI5zsXkaKKDjdVjpD4oAt9Qi/MArHo=
 go.opentelemetry.io/otel/exporters/jaeger v1.17.0 h1:D7UpUy2Xc2wsi1Ras6V40q806WM07rqoCWzXu7Sqy+4=
@@ -232,6 +227,8 @@ go.uber.org/automaxprocs v1.6.0 h1:O3y2/QNTOdbF+e/dpXNNW7Rx2hZ4sTIPyybbxyNqTUs=
 go.uber.org/automaxprocs v1.6.0/go.mod h1:ifeIMSnPZuznNm6jmdzmU3/bfk01Fe2fotchwEFJ8r8=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
+go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=
+go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc=
 go.uber.org/multierr v1.9.0 h1:7fIwc/ZtS0q++VgcfqFDxSBZVv/Xo49/SYnDFupUwlI=
 go.uber.org/multierr v1.9.0/go.mod h1:X2jQV1h+kxSjClGpnseKVIxpmcjrj7MNnI0bnlfKTVQ=
 go.uber.org/zap v1.24.0 h1:FiJd5l1UOLj0wCgbSE0rwwXHzEdAZS6hiiSnxJN/D60=
@@ -244,14 +241,12 @@ golang.org/x/crypto v0.33.0 h1:IOBPskki6Lysi0lo9qQvbxiQ+FvsCC/YWOecCHAixus=
 golang.org/x/crypto v0.33.0/go.mod h1:bVdXmD7IV/4GdElGPozy6U7lWdRXA4qyRVGJV57uQ5M=
 golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
-golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
-golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
 golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
 golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8=
 golang.org/x/net v0.35.0/go.mod h1:EglIi67kWsHKlRzzVMUD93VMSWGFOMSZgxFjparz1Qk=
@@ -260,7 +255,6 @@ golang.org/x/oauth2 v0.24.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbht
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.11.0 h1:GGz8+XQP4FvTTrjZPzNKTMFtSXH80RAzG+5ghFPgK9w=
 golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
@@ -268,8 +262,6 @@ golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
@@ -293,7 +285,6 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
 golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
-golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
 golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d h1:vU5i/LfpvrRCpgM/VPfJLg5KjxD3E+hfT1SH+d9zLwg=
 golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=

internal/health/health.go

@@ -43,7 +43,7 @@ func AddProbe(probe Probe) {
 	defaultHealthManager.addProbe(probe)
 }
 
-// CreateHttpHandler create health http handler base on given probe.
+// CreateHttpHandler creates a health http handler based on the given probe.
 func CreateHttpHandler(healthResponse string) http.HandlerFunc {
 	return func(w http.ResponseWriter, _ *http.Request) {
 		if defaultHealthManager.IsReady() {

mcp/server.go

@@ -392,8 +392,8 @@ func (s *sseMcpServer) processListTools(ctx context.Context, client *mcpClient,
 		}
 	}
 
-	var toolsList []Tool
 	s.toolsLock.Lock()
+	toolsList := make([]Tool, 0, len(s.tools))
 	for _, tool := range s.tools {
 		if len(tool.InputSchema.Type) == 0 {
 			tool.InputSchema.Type = ContentTypeObject
@@ -437,8 +437,8 @@ func (s *sseMcpServer) processListPrompts(ctx context.Context, client *mcpClient
 	}
 
 	// Prepare prompt list
-	var promptsList []Prompt
 	s.promptsLock.Lock()
+	promptsList := make([]Prompt, 0, len(s.prompts))
 	for _, prompt := range s.prompts {
 		promptsList = append(promptsList, prompt)
 	}
@@ -475,8 +475,8 @@ func (s *sseMcpServer) processListResources(ctx context.Context, client *mcpClie
 		}
 	}
 
-	var resourcesList []Resource
 	s.resourcesLock.Lock()
+	resourcesList := make([]Resource, 0, len(s.resources))
 	for _, resource := range s.resources {
 		// Create a copy without the handler function which shouldn't be sent to clients
 		resourceCopy := Resource{

mcp/server_test.go

@@ -1741,7 +1741,7 @@ func TestHandleSSEPing(t *testing.T) {
 	mock.server.clientsLock.Unlock()
 
 	assert.Equal(t, 1, clientCount, "One client should be added during the test")
-	assert.Equal(t, 0, finalClientCount, "Client should be cleaned up after context cancelation")
+	assert.Equal(t, 0, finalClientCount, "Client should be cleaned up after context cancellation")
 }
 
 // TestHandleSSEPing tests the automatic ping functionality in the SSE handler

readme-cn.md

@@ -17,7 +17,7 @@
 <a href="https://trendshift.io/repositories/3263" target="_blank"><img src="https://trendshift.io/api/badge/repositories/3263" alt="zeromicro%2Fgo-zero | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
 <a href="https://www.producthunt.com/posts/go-zero?utm_source=badge-featured&utm_medium=badge&utm_souce=badge-go&#0045;zero" target="_blank"><img src="https://api.producthunt.com/widgets/embed-image/v1/featured.svg?post_id=334030&theme=light" alt="go&#0045;zero - A&#0032;web&#0032;&#0038;&#0032;rpc&#0032;framework&#0032;written&#0032;in&#0032;Go&#0046; | Product Hunt" style="width: 250px; height: 54px;" width="250" height="54" /></a>
 
-## 0. go-zero 介绍
+## go-zero 介绍
 
 go-zero（收录于 CNCF 云原生技术全景图：[https://landscape.cncf.io/?selected=go-zero](https://landscape.cncf.io/?selected=go-zero)）是一个集成了各种工程实践的 web 和 rpc 框架。通过弹性设计保障了大并发服务端的稳定性，经受了充分的实战检验。
 
@@ -25,72 +25,50 @@ go-zero 包含极简的 API 定义和生成工具 goctl，可以根据定义的
 
 使用 go-zero 的好处：
 
-* 轻松获得支撑千万日活服务的稳定性
-* 内建级联超时控制、限流、自适应熔断、自适应降载等微服务治理能力，无需配置和额外代码
-* 微服务治理中间件可无缝集成到其它现有框架使用
-* 极简的 API 描述，一键生成各端代码
-* 自动校验客户端请求参数合法性
-* 大量微服务治理和并发工具包
+* 经过千万日活服务验证的稳定性
+* 内建弹性保护:级联超时、限流、熔断、降载(无需配置)
+* 极简 API 语法生成多端代码
+* 自动参数校验和丰富的微服务工具包
 
 ![架构图](https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/architecture.png)
 
-## 1. go-zero 框架背景
+## go-zero 框架背景
 
-18 年初，我们决定从 `Java+MongoDB` 的单体架构迁移到微服务架构，经过仔细思考和对比，我们决定：
+18 年初，我们决定从 `Java+MongoDB` 的单体架构迁移到微服务架构，选择:
 
-* 基于 Go 语言
-  * 高效的性能
-  * 简洁的语法
-  * 广泛验证的工程效率
-  * 极致的部署体验
-  * 极低的服务端资源成本
-* 自研微服务框架
-  * 有过很多微服务框架自研经验
-  * 需要有更快速的问题定位能力
-  * 更便捷的增加新特性
+* **基于 Go 语言** - 高效性能、简洁语法、极致部署体验、极低资源成本
+* **自研微服务框架** - 更快速的问题定位、更便捷的新特性增加
 
-## 2. go-zero 框架设计思考
+## go-zero 框架设计思考
 
-对于微服务框架的设计，我们期望保障微服务稳定性的同时，也要特别注重研发效率。所以设计之初，我们就有如下一些准则：
+go-zero 遵循以下核心设计准则:
 
-* 保持简单，第一原则
-* 弹性设计，面向故障编程
-* 工具大于约定和文档
-* 高可用、高并发、易扩展
-* 对业务开发友好，封装复杂度
-* 约束做一件事只有一种方式
+* **保持简单** - 简单是第一原则
+* **高可用** - 高并发、易扩展
+* **弹性设计** - 面向故障编程
+* **工具驱动** - 工具大于约定和文档
+* **业务友好** - 封装复杂度、一事一法
 
-我们经历不到半年时间，彻底完成了从 `Java+MongoDB` 到 `Golang+MySQL` 为主的微服务体系迁移，并于 18 年 8 月底完全上线，稳定保障了业务后续迅速增长，确保了整个服务的高可用。
+## go-zero 项目实现和特点
 
-## 3. go-zero 项目实现和特点
+go-zero 集成各种工程实践，主要特点:
 
-go-zero 是一个集成了各种工程实践的包含 web 和 rpc 框架，有如下主要特点：
-
-* 强大的工具支持，尽可能少的代码编写
-* 极简的接口
-* 完全兼容 net/http
-* 支持中间件，方便扩展
-* 高性能
-* 面向故障编程，弹性设计
-* 内建服务发现、负载均衡
-* 内建限流、熔断、降载，且自动触发，自动恢复
-* API 参数自动校验
-* 超时级联控制
-* 自动缓存控制
-* 链路跟踪、统计报警等
-* 高并发支撑，稳定保障了疫情期间每天的流量洪峰
-
-如下图，我们从多个层面保障了整体服务的高可用：
+* **强大工具支持** - 尽可能少的代码编写
+* **极简接口** - 完全兼容 net/http
+* **高性能** - 优化的速度和效率
+* **弹性设计** - 内建限流、熔断、降载，自动触发、自动恢复
+* **服务治理** - 内建服务发现、负载均衡、链路跟踪
+* **开发工具** - API 参数自动校验、超时级联控制、自动缓存控制
 
 ![弹性设计](https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/resilience.jpg)
 
-## 4. 我们使用 go-zero 的基本架构图
+## 我们使用 go-zero 的基本架构图
 
 <img width="1067" alt="image" src="https://user-images.githubusercontent.com/1918356/171880582-11a86658-41c3-466c-95e7-7b1220eecc52.png">
 
 觉得不错的话，别忘 **star** 👏
 
-## 5. Installation
+## Installation
 
 在项目目录下通过如下命令安装：
 
@@ -98,7 +76,57 @@ go-zero 是一个集成了各种工程实践的包含 web 和 rpc 框架，有
 GO111MODULE=on GOPROXY=https://goproxy.cn/,direct go get -u github.com/zeromicro/go-zero
 ```
 
-## 6. Quick Start
+## AI 原生开发
+
+go-zero 团队构建了完整的 AI 工具生态，让 Claude、GitHub Copilot、Cursor 生成符合 go-zero 规范的代码。
+
+### 三大核心项目
+
+**[ai-context](https://github.com/zeromicro/ai-context)** - AI 的工作流程指南
+
+**[zero-skills](https://github.com/zeromicro/zero-skills)** - 模式库和示例
+
+**[mcp-zero](https://github.com/zeromicro/mcp-zero)** - 基于 MCP 的代码生成工具
+
+### 快速配置
+
+#### GitHub Copilot
+```bash
+git submodule add https://github.com/zeromicro/ai-context.git .github/ai-context
+ln -s ai-context/00-instructions.md .github/copilot-instructions.md  # macOS/Linux
+# Windows: mklink .github\copilot-instructions.md .github\ai-context\00-instructions.md
+git submodule update --remote .github/ai-context  # 更新
+```
+
+#### Cursor
+```bash
+git submodule add https://github.com/zeromicro/ai-context.git .cursorrules
+git submodule update --remote .cursorrules  # 更新
+```
+
+#### Windsurf
+```bash
+git submodule add https://github.com/zeromicro/ai-context.git .windsurfrules
+git submodule update --remote .windsurfrules  # 更新
+```
+
+#### Claude Desktop
+```bash
+git clone https://github.com/zeromicro/mcp-zero.git && cd mcp-zero && go build
+# 配置: ~/Library/Application Support/Claude/claude_desktop_config.json
+# 或: claude mcp add --transport stdio mcp-zero --env GOCTL_PATH=/path/to/goctl -- /path/to/mcp-zero
+```
+
+### 协同工作原理
+
+AI 助手通过三个工具协同配合:
+1. **ai-context** - 工作流程指导
+2. **zero-skills** - 实现模式
+3. **mcp-zero** - 实时代码生成
+
+**示例**:创建新的 REST API → AI 读取 **ai-context** 了解工作流 → 调用 **mcp-zero** 生成代码 → 参考 **zero-skills** 实现模式 → 生成符合规范的代码 ✅
+
+## Quick Start
 
 0. 完整示例请查看
 
@@ -108,23 +136,22 @@ GO111MODULE=on GOPROXY=https://goproxy.cn/,direct go get -u github.com/zeromicro
 
 1. 安装 goctl 工具
 
-    `goctl` 读作 `go control`，不要读成 `go C-T-L`。`goctl` 的意思是不要被代码控制，而是要去控制它。其中的 `go` 不是指 `golang`。在设计 `goctl` 之初，我就希望通过 `工具` 来解放我们的双手👈
 
     ```shell
     # Go
     GOPROXY=https://goproxy.cn/,direct go install github.com/zeromicro/go-zero/tools/goctl@latest
-    
+
     # For Mac
     brew install goctl
-    
+
     # docker for all platforms
     docker pull kevinwan/goctl
     # run goctl
     docker run --rm -it -v `pwd`:/app kevinwan/goctl --help
     ```
-    
-    确保 goctl 可执行，并且在 $PATH 环境变量里。
-    
+
+    确保 goctl 可执行并在 $PATH 环境变量里。
+
 2. 快速生成 api 服务
 
     ```shell
@@ -157,7 +184,7 @@ GO111MODULE=on GOPROXY=https://goproxy.cn/,direct go get -u github.com/zeromicro
       * 可以在 `servicecontext.go` 里面传递依赖给 logic，比如 mysql, redis 等
       * 在 api 定义的 `get/post/put/delete` 等请求对应的 logic 里增加业务处理逻辑
 
-3. 可以根据 api 文件生成前端需要的 Java, TypeScript, Dart, JavaScript 代码
+3. 生成多语言客户端代码
 
     ```shell
     goctl api java -api greet.api -dir greet
@@ -165,17 +192,17 @@ GO111MODULE=on GOPROXY=https://goproxy.cn/,direct go get -u github.com/zeromicro
     ...
     ```
 
-## 7. Benchmark
+## Benchmark
 
 ![benchmark](https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/benchmark.png)
 
 [测试代码见这里](https://github.com/smallnest/go-web-framework-benchmark)
 
-## 8. 文档
+## 文档
 
 * API 文档
 
-  [https://go-zero.dev/cn/](https://go-zero.dev/cn/)
+  [https://go-zero.dev](https://go-zero.dev)
 
 * awesome 系列（更多文章见『微服务实践』公众号）
 
@@ -192,9 +219,9 @@ GO111MODULE=on GOPROXY=https://goproxy.cn/,direct go get -u github.com/zeromicro
   | [goctl-android](https://github.com/zeromicro/goctl-android) | 生成 `java (android)` 端 `http client` 请求代码 |
   | [goctl-go-compact](https://github.com/zeromicro/goctl-go-compact) | 合并 `api` 里同一个 `group` 里的 `handler` 到一个 `go` 文件 |
 
-## 9. go-zero 用户
+## go-zero 用户
 
-go-zero 已被许多公司用于生产部署，接入场景如在线教育、电商业务、游戏、区块链等，目前为止，已使用 go-zero 的公司包括但不限于：
+go-zero 已被众多公司用于生产部署，场景涵盖在线教育、电商、游戏、区块链等。目前使用 go-zero 的公司包括但不限于:
 
 >1. 好未来
 >2. 上海晓信信息科技有限公司（晓黑板）
@@ -304,10 +331,14 @@ go-zero 已被许多公司用于生产部署，接入场景如在线教育、电
 >106. 无锡盛算信息技术有限公司
 >107. 深圳市聚货通信息科技有限公司
 >108. 浙江银盾云科技有限公司
+>109. 南京造世网络科技有限公司
+>110. 温州飞儿云信息技术有限公司
+>111. 统信软件
+>112. 深圳坐标软件集团有限公司
 
 如果贵公司也已使用 go-zero，欢迎在 [登记地址](https://github.com/zeromicro/go-zero/issues/602) 登记，仅仅为了推广，不做其它用途。
 
-## 10. CNCF 云原生技术全景图
+## CNCF 云原生技术全景图
 
 <p float="left">
 <img src="https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/cncf-logo.svg" width="200"/>&nbsp;&nbsp;&nbsp;
@@ -316,13 +347,13 @@ go-zero 已被许多公司用于生产部署，接入场景如在线教育、电
 
 go-zero 收录在 [CNCF Cloud Native 云原生技术全景图](https://landscape.cncf.io/?selected=go-zero)。
 
-## 11. 微信公众号
+## 微信公众号
 
 `go-zero` 相关文章和视频都会在 `微服务实践` 公众号整理呈现，欢迎扫码关注 👏
 
 <img src="https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/zeromicro.jpg" alt="wechat" width="600" />
 
-## 12. 微信交流群
+## 微信交流群
 
 如果文档中未能覆盖的任何疑问，欢迎您在群里提出，我们会尽快答复。
 
@@ -332,10 +363,4 @@ go-zero 收录在 [CNCF Cloud Native 云原生技术全景图](https://landscape
 
 加群之前有劳点一下 ***star***，一个小小的 ***star*** 是作者们回答海量问题的动力！🤝
 
-<img src="https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/wechat.jpg" alt="wechat" width="300" />
-
-## 13. 知识星球
-
-官方团队运营的知识星球
-
-<img src="https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/zsxq.jpg" alt="知识星球" width="300" />
+<img src="https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/wechat.jpg" alt="wechat" width="300" />

readme.md

@@ -42,61 +42,39 @@ go-zero contains simple API description syntax and code generation tool called `
 
 ## Backgrounds of go-zero
 
-In early 2018, we embarked on a transformative journey to redesign our system, transitioning from a monolithic architecture built with Java and MongoDB to a microservices architecture. After careful research and comparison, we made a deliberate choice to:
+In early 2018, we transitioned from a Java+MongoDB monolithic architecture to microservices, choosing:
 
-* Go Beyond with Golang
-  * Great performance
-  * Simple syntax
-  * Proven engineering efficiency
-  * Extreme deployment experience
-  * Less server resource consumption
-
-* Self-Design Our Microservice Architecture
-  * Microservice architecture facilitates the creation of scalable, flexible, and maintainable software systems with independent, reusable components.
-  * Easy to locate the problems within microservices.
-  * Easy to extend the features by adding or modifying specific microservices without impacting the entire system.
+* **Golang** - High performance, simple syntax, excellent deployment experience, and low resource consumption
+* **Self-designed microservice framework** - Better problem isolation, easier feature extension, and faster issue resolution
 
 ## Design considerations on go-zero
 
-By designing the microservice architecture, we expected to ensure stability, as well as productivity. And from just the beginning, we have the following design principles:
+go-zero follows these core design principles:
 
-* Keep it simple
-* High availability
-* Stable on high concurrency
-* Easy to extend
-* Resilience design, failure-oriented programming
-* Try best to be friendly to the business logic development, encapsulate the complexity
-* One thing, one way
-
-After almost half a year, we finished the transfer from a monolithic system to microservice system and deployed on August 2018. The new system guaranteed business growth and system stability.
+* **Simplicity** - Keep it simple, first principle
+* **High availability** - Stable under high concurrency
+* **Resilience** - Failure-oriented programming with adaptive protection
+* **Developer friendly** - Encapsulate complexity, one way to do one thing
+* **Easy to extend** - Flexible architecture for growth
 
 ## The implementation and features of go-zero
 
-go-zero is a web and rpc framework that integrates lots of engineering practices. The features are mainly listed below:
+go-zero integrates engineering best practices:
 
-* Powerful tool included, less code to write
-* Simple interfaces
-* Fully compatible with net/http
-* Middlewares are supported, easy to extend
-* High performance
-* Failure-oriented programming, resilience design
-* Builtin service discovery, load balancing
-* Builtin concurrency control, adaptive circuit breaker, adaptive load shedding, auto-trigger, auto recover
-* Auto validation of API request parameters
-* Chained timeout control
-* Auto management of data caching
-* Call tracing, metrics, and monitoring
-* High concurrency protected
-
-As below, go-zero protects the system with a couple of layers and mechanisms:
+* **Code generation** - Powerful tools to minimize boilerplate
+* **Simple API** - Clean interfaces, fully compatible with net/http
+* **High performance** - Optimized for speed and efficiency
+* **Resilience** - Built-in circuit breaker, rate limiting, load shedding, timeout control
+* **Service mesh** - Service discovery, load balancing, call tracing
+* **Developer tools** - Auto parameter validation, cache management, metrics and monitoring
 
 ![Resilience](https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/resilience-en.png)
 
-##  The simplified architecture that we use with go-zero
+## Architecture with go-zero
 
 <img width="1067" alt="image" src="https://user-images.githubusercontent.com/1918356/171880372-5010d846-e8b1-4942-8fe2-e2bbb584f762.png">
 
-##  Installation
+## Installation
 
 Run the following command under your project:
 
@@ -104,9 +82,59 @@ Run the following command under your project:
 go get -u github.com/zeromicro/go-zero
 ```
 
-##  Quick Start
+## AI-Native Development
 
-1. Full examples can be checked out from below:
+The go-zero team provides AI tooling for Claude, GitHub Copilot, Cursor to generate framework-compliant code.
+
+### Three Core Projects
+
+**[ai-context](https://github.com/zeromicro/ai-context)** - Workflow guide for AI assistants
+
+**[zero-skills](https://github.com/zeromicro/zero-skills)** - Pattern library with examples
+
+**[mcp-zero](https://github.com/zeromicro/mcp-zero)** - Code generation tools via Model Context Protocol
+
+### Quick Setup
+
+#### GitHub Copilot
+```bash
+git submodule add https://github.com/zeromicro/ai-context.git .github/ai-context
+ln -s ai-context/00-instructions.md .github/copilot-instructions.md  # macOS/Linux
+# Windows: mklink .github\copilot-instructions.md .github\ai-context\00-instructions.md
+git submodule update --remote .github/ai-context  # Update
+```
+
+#### Cursor
+```bash
+git submodule add https://github.com/zeromicro/ai-context.git .cursorrules
+git submodule update --remote .cursorrules  # Update
+```
+
+#### Windsurf
+```bash
+git submodule add https://github.com/zeromicro/ai-context.git .windsurfrules
+git submodule update --remote .windsurfrules  # Update
+```
+
+#### Claude Desktop
+```bash
+git clone https://github.com/zeromicro/mcp-zero.git && cd mcp-zero && go build
+# Configure: ~/Library/Application Support/Claude/claude_desktop_config.json
+# Or: claude mcp add --transport stdio mcp-zero --env GOCTL_PATH=/path/to/goctl -- /path/to/mcp-zero
+```
+
+### How It Works
+
+AI assistants use these tools together:
+1. **ai-context** - workflow guidance
+2. **zero-skills** - implementation patterns
+3. **mcp-zero** - real-time code generation
+
+**Example**: Creating a REST API → AI reads **ai-context** for workflow → calls **mcp-zero** to generate code → references **zero-skills** for patterns → produces production-ready code ✅
+
+## Quick Start
+
+1. Full examples:
 
      [Rapid development of microservice systems](https://github.com/zeromicro/zero-doc/blob/main/doc/shorturl-en.md)
 
@@ -114,24 +142,22 @@ go get -u github.com/zeromicro/go-zero
 
 2. Install goctl
 
-   `goctl`can be read as `go control`. `goctl` means not to be controlled by code, instead, we control it. The inside `go` is not `golang`. At the very beginning, I was expecting it to help us improve productivity, and make our lives easier.
-
    ```shell
    # for Go
    go install github.com/zeromicro/go-zero/tools/goctl@latest
-   
+
    # For Mac
    brew install goctl
-   
+
    # docker for all platforms
    docker pull kevinwan/goctl
    # run goctl
    docker run --rm -it -v `pwd`:/app kevinwan/goctl --help
    ```
-   
-   make sure goctl is executable and in your $PATH.
-   
-3. Create the API file, like greet.api, you can install the plugin of goctl in vs code, api syntax is supported.
+
+   Ensure goctl is executable and in your $PATH.
+
+3. Create the API file (greet.api):
 
    ```go
    type (
@@ -150,19 +176,19 @@ go get -u github.com/zeromicro/go-zero
    }
    ```
 
-   the .api files also can be generated by goctl, like below:
+   Generate .api template:
 
    ```shell
    goctl api -o greet.api
    ```
 
-4. Generate the go server-side code
+4. Generate Go server code
 
    ```shell
    goctl api go -api greet.api -dir greet
    ```
 
-   the generated files look like:
+   Generated structure:
 
    ```Plain Text
    ├── greet
@@ -184,7 +210,7 @@ go get -u github.com/zeromicro/go-zero
    └── greet.api                     // api description file
    ```
 
-   the generated code can be run directly:
+   Run the service:
 
    ```shell
    cd greet
@@ -192,15 +218,15 @@ go get -u github.com/zeromicro/go-zero
    go run greet.go -f etc/greet-api.yaml
    ```
 
-   by default, it’s listening on port 8888, while it can be changed in the configuration file.
+   Default port: 8888 (configurable in etc/greet-api.yaml)
 
-   you can check it by curl:
+   Test with curl:
 
    ```shell
    curl -i http://localhost:8888/greet/from/you
    ```
 
-   the response looks like below:
+   Response:
 
    ```http
    HTTP/1.1 200 OK
@@ -208,12 +234,12 @@ go get -u github.com/zeromicro/go-zero
    Content-Length: 0
    ```
 
-5. Write the business logic code
+5. Write business logic
 
-    * the dependencies can be passed into the logic within servicecontext.go, like mysql, redis, etc.
-    * add the logic code in a logic package according to .api file
+    * Pass dependencies (mysql, redis, etc.) via servicecontext.go
+    * Add logic code in the logic package per .api definition
 
-6. Generate code like Java, TypeScript, Dart, JavaScript, etc. just from the api file
+6. Generate client code for multiple languages
 
    ```shell
    goctl api java -api greet.api -dir greet
@@ -234,11 +260,11 @@ go get -u github.com/zeromicro/go-zero
 * [Rapid development of microservice systems - multiple RPCs](https://github.com/zeromicro/zero-doc/blob/main/docs/zero/bookstore-en.md)
 * [Examples](https://github.com/zeromicro/zero-examples)
 
-##  Chat group
+## Chat group
 
 Join the chat via https://discord.gg/4JQvC5A4Fe
 
-##  Cloud Native Landscape
+## Cloud Native Landscape
 
 <p float="left">
 <img src="https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/cncf-logo.svg" width="200"/>&nbsp;&nbsp;&nbsp;