chore(deps): bump github.com/go-sql-driver/mysql in /tools/goctl

Bumps [github.com/go-sql-driver/mysql](https://github.com/go-sql-driver/mysql) from 1.10.0 to 1.10.1.
- [Release notes](https://github.com/go-sql-driver/mysql/releases)
- [Changelog](https://github.com/go-sql-driver/mysql/blob/master/CHANGELOG.md)
- [Commits](https://github.com/go-sql-driver/mysql/compare/v1.10.0...v1.10.1)

---
updated-dependencies:
- dependency-name: github.com/go-sql-driver/mysql
  dependency-version: 1.10.1
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>

.codecov.yml

@@ -1,4 +0,0 @@
-ignore:
-  - "doc"
-  - "example"
-  - "tools"

.dockerignore

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

.github/FUNDING.yml

@@ -0,0 +1,3 @@
+# These are supported funding model platforms
+
+github: [zeromicro]

.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,40 @@
+---
+name: Bug report
+about: Create a report to help us improve
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**Describe the bug**
+A clear and concise description of what the bug is.
+
+**To Reproduce**
+Steps to reproduce the behavior, if applicable:
+
+1. The code is
+
+   ```go
+   
+   ```
+
+2. The error is
+
+   ```
+   
+   ```
+
+**Expected behavior**
+A clear and concise description of what you expected to happen.
+
+**Screenshots**
+If applicable, add screenshots to help explain your problem.
+
+**Environments (please complete the following information):**
+ - OS: [e.g. Linux]
+ - go-zero version [e.g. 1.2.1]
+ - goctl version [e.g. 1.2.1, optional]
+
+**More description**
+Add any other context about the problem here.

.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context or screenshots about the feature request here.

.github/ISSUE_TEMPLATE/question.md

@@ -0,0 +1,10 @@
+---
+name: Question
+about: Ask a question on using go-zero or goctl
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+

.github/copilot-instructions.md

@@ -0,0 +1,344 @@
+# 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
+
+## GitHub Issue Management
+
+### Understanding and Categorizing Issues
+
+When analyzing GitHub issues, consider these common categories:
+
+1. **Bug Reports**: Stack traces, version info, reproduction steps
+2. **Feature Requests**: Use case, proposed solution, alternatives
+3. **Questions**: Usage, configuration, or architecture
+4. **Documentation Issues**: Missing, unclear, or incorrect docs
+5. **Performance Issues**: Benchmarks, profiling data, resource usage
+
+### Issue Analysis Checklist
+
+- Identify affected component (REST, RPC, Gateway, MCP, Core utilities, goctl)
+- Check versions (go-zero, Go)
+- Look for reproduction steps or code examples
+- Review code snippets, logs, or stack traces
+- Check if related to resilience features (breaker, load shedding, rate limiting)
+- Determine production impact
+
+### Responding to Issues
+
+Be helpful and professional. Ask clarifying questions when needed. Reference relevant documentation and code files. Provide code examples following project conventions. Suggest workarounds when applicable.
+
+### Chinese to English Translation
+
+go-zero has an international user base. When encountering issues or comments written in Chinese, translate them to English to ensure all contributors can participate in discussions.
+
+#### Translation Guidelines
+
+1. **Update issue titles**: Edit the issue title to include English translation only
+2. **Translate comments in place**: Add a comment with the English translation, followed by the original Chinese text
+3. **Keep original Chinese**: After translating, include the original Chinese text in a blockquote for verification
+4. **Encourage English communication**: Politely suggest users write in English for better collaboration
+5. **Maintain technical accuracy**: Preserve technical terms, component names, and code exactly
+6. **Translate naturally**: Avoid literal word-by-word translation; use idiomatic English
+7. **Preserve formatting**: Keep markdown formatting, code blocks, and links intact
+8. **Keep URLs unchanged**: Don't translate URLs or file paths
+
+#### Common Technical Terms (Chinese → English)
+
+- 框架 → **Framework** | 中间件 → **Middleware** | 负载均衡 → **Load Balancing**
+- 熔断器 → **Circuit Breaker** | 限流 → **Rate Limiting** | 降载/过载保护 → **Load Shedding**
+- 服务发现 → **Service Discovery** | 配置 → **Configuration** | 弹性/容错 → **Resilience** | 微服务 → **Microservices**
+
+#### Translation Example
+
+**Original Chinese Title:** `goctl 执行环境问题`
+**Updated Title:** `goctl Execution Environment Issue`
+
+**Original Chinese Comment:** `我在项目中遇到熔断器配置问题`
+**Translation in Comment:**
+```markdown
+I encountered a circuit breaker configuration issue in my project.
+
+> Original (原文): 我在项目中遇到熔断器配置问题
+```
+
+### Common Issue Patterns and Solutions
+
+#### Configuration Issues
+- Check `service.ServiceConf` embedding and struct tags
+- Verify YAML syntax, defaults, and validation rules
+- Reference: [rest/config.go](rest/config.go), [zrpc/config.go](zrpc/config.go)
+
+#### Code Generation (goctl) Issues
+- Verify `.api` or `.proto` file syntax and goctl version
+- Reference: `tools/goctl/` directory
+
+#### RPC Connection Issues
+- Check etcd configuration, service discovery, and endpoints
+- Verify load balancing settings (p2c_ewma)
+
+#### Database/Cache Issues
+- Verify `sqlx.SqlConn` usage with context
+- Check cache key generation, invalidation, and connection pools
+- Use test helpers (`redistest`, `mongtest`)
+
+#### Performance Issues
+- Check if load shedding is enabled (mode: `pre`/`pro`)
+- Review circuit breaker thresholds, rate limiting, and context timeouts
+
+### Referencing Codebase
+
+When explaining issues, reference specific files and patterns:
+- REST API: `rest/`, `rest/handler/`, `rest/httpx/`
+- RPC: `zrpc/`, `zrpc/internal/`
+- Core utilities: `core/breaker/`, `core/limit/`, `core/load/`, etc.
+- Gateway: `gateway/`
+- MCP: `mcp/`
+- Code generation: `tools/goctl/`
+- Examples: `adhoc/` directory contains various examples
+
+### Encouraging Best Practices
+
+When responding to issues, gently guide users toward:
+- Proper error handling with context
+- Using resilience features (breakers, rate limiters)
+- Following testing patterns with table-driven tests
+- Implementing proper resource cleanup
+- Reading existing documentation in `docs/` and `readme.md`
+
+## 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/dependabot.yml

@@ -0,0 +1,23 @@
+# To get started with Dependabot version updates, you'll need to specify which
+# package ecosystems to update and where the package manifests are located.
+# Please see the documentation for all configuration options:
+# https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates
+
+version: 2
+updates:
+  - package-ecosystem: "docker" # Update image tags in Dockerfile
+    directory: "/"
+    schedule:
+      interval: "weekly"
+  - package-ecosystem: "github-actions" # Update GitHub Actions
+    directory: "/"
+    schedule:
+      interval: "weekly"
+  - package-ecosystem: "gomod" # See documentation for possible values
+    directory: "/" # Location of package manifests
+    schedule:
+      interval: "daily"
+  - package-ecosystem: "gomod" # See documentation for possible values
+    directory: "/tools/goctl" # Location of package manifests
+    schedule:
+      interval: "daily"

.github/workflows/codeql-analysis.yml

@@ -0,0 +1,67 @@
+# For most projects, this workflow file will not need changing; you simply need
+# to commit it to your repository.
+#
+# You may wish to alter this file to override the set of languages analyzed,
+# or to provide custom queries or build logic.
+#
+# ******** NOTE ********
+# We have attempted to detect the languages in your repository. Please check
+# the `language` matrix defined below to confirm you have the correct set of
+# supported CodeQL languages.
+#
+name: "CodeQL"
+
+on:
+  push:
+    branches: [ master ]
+  pull_request:
+    # The branches below must be a subset of the branches above
+    branches: [ master ]
+  schedule:
+    - cron: '18 19 * * 6'
+
+jobs:
+  analyze:
+    name: Analyze
+    runs-on: ubuntu-latest
+
+    strategy:
+      fail-fast: false
+      matrix:
+        language: [ 'go' ]
+        # CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python' ]
+        # Learn more:
+        # https://docs.github.com/en/free-pro-team@latest/github/finding-security-vulnerabilities-and-errors-in-your-code/configuring-code-scanning#changing-the-languages-that-are-analyzed
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v7
+
+    # Initializes the CodeQL tools for scanning.
+    - name: Initialize CodeQL
+      uses: github/codeql-action/init@v4.37.9
+      with:
+        languages: ${{ matrix.language }}
+        # If you wish to specify custom queries, you can do so here or in a config file.
+        # By default, queries listed here will override any specified in a config file.
+        # Prefix the list here with "+" to use these queries and those in the config file.
+        # queries: ./path/to/local/query, your-org/your-repo/queries@main
+
+    # 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@v4.37.9
+
+    # ℹ️ Command-line programs to run using the OS shell.
+    # 📚 https://git.io/JvXDl
+
+    # ✏️ If the Autobuild fails above, remove it and uncomment the following three lines
+    #    and modify them (or add more) to build your code if your project
+    #    uses a compiled language
+
+    #- run: |
+    #   make bootstrap
+    #   make release
+
+    - name: Perform CodeQL Analysis
+      uses: github/codeql-action/analyze@v4.37.9

.github/workflows/go.yml

@@ -7,28 +7,64 @@ on:
     branches: [ master ]
 
 jobs:
-  build:
-    name: Build
+  test-linux:
+    name: Linux
     runs-on: ubuntu-latest
     steps:
+      - name: Check out code into the Go module directory
+        uses: actions/checkout@v7
 
-    - name: Set up Go 1.x
-      uses: actions/setup-go@v2
-      with:
-        go-version: ^1.13
-      id: go
+      - name: Set up Go 1.x
+        uses: actions/setup-go@v7
+        with:
+          go-version-file: go.mod
+          check-latest: true
+          cache: true
+        id: go
 
-    - name: Check out code into the Go module directory
-      uses: actions/checkout@v2
+      - name: Get dependencies
+        run: |
+          go get -v -t -d ./...
 
-    - name: Get dependencies
-      run: |
-        go get -v -t -d ./...
+      - name: Lint
+        run: |
+          go vet -stdmethods=false $(go list ./...)
 
-    - name: Test
-      run: go test -race -coverprofile=coverage.txt -covermode=atomic ./...
+          go mod tidy
+          if ! test -z "$(git status --porcelain)"; then
+            echo "Please run 'go mod tidy'"
+            exit 1
+          fi
 
-    - name: Codecov
-      uses: codecov/codecov-action@v1.0.6
-      with:
-        token: ${{secrets.CODECOV_TOKEN}}
+      - name: Test
+        run: go test -race -coverprofile=coverage.txt -covermode=atomic ./...
+
+      - name: Codecov
+        uses: codecov/codecov-action@v7
+        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@v7
+
+      - name: Set up Go 1.x
+        uses: actions/setup-go@v7
+        with:
+          # make sure Go version compatible with go-zero
+          go-version-file: go.mod
+          check-latest: true
+          cache: true
+
+      - name: Test
+        run: |
+          go mod verify
+          go mod download
+          go test ./...
+          cd tools/goctl && go build -v goctl.go

.github/workflows/issues.yml

@@ -0,0 +1,19 @@
+name: Close inactive issues
+on:
+  schedule:
+    - cron: "30 1 * * *"
+
+jobs:
+  close-issues:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/stale@v11
+        with:
+          days-before-issue-stale: 365
+          days-before-issue-close: 90
+          stale-issue-label: "stale"
+          stale-issue-message: "This issue is stale because it has been open for 30 days with no activity."
+          close-issue-message: "This issue was closed because it has been inactive for 14 days since being marked as stale."
+          days-before-pr-stale: -1
+          days-before-pr-close: -1
+          repo-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/release.yaml

@@ -0,0 +1,28 @@
+on:
+  push:
+    tags:
+      - "tools/goctl/*"
+jobs:
+  releases-matrix:
+    name: Release goctl binary
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        # build and publish in parallel: linux/386, linux/amd64, linux/arm64,
+        # windows/386, windows/amd64, windows/arm64, darwin/amd64, darwin/arm64
+        goos: [ linux, windows, darwin ]
+        goarch: [ "386", amd64, arm64 ]
+        exclude:
+          - goarch: "386"
+            goos: darwin
+    steps:
+      - uses: actions/checkout@v7
+      - uses: zeromicro/go-zero-release-action@master
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+          goos: ${{ matrix.goos }}
+          goarch: ${{ matrix.goarch }}
+          goversion: "https://dl.google.com/go/go1.25.12.linux-amd64.tar.gz"
+          project_path: "tools/goctl"
+          binary_name: "goctl"
+          extra_files: tools/goctl/readme.md tools/goctl/readme-cn.md

.github/workflows/reviewdog.yml

@@ -0,0 +1,24 @@
+name: reviewdog
+on: [pull_request]
+jobs:
+  staticcheck:
+    name: runner / staticcheck
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: actions/setup-go@v7
+        with:
+          go-version-file: go.mod
+          check-latest: true
+          cache: true
+      - uses: reviewdog/action-staticcheck@v1
+        with:
+          github_token: ${{ secrets.github_token }}
+          # Change reviewdog reporter if you need [github-pr-check,github-check,github-pr-review].
+          reporter: github-pr-review
+          # Report all results.
+          filter_mode: nofilter
+          # Exit with 1 when it find at least one finding.
+          fail_level: any
+          # Set staticcheck flags
+          staticcheck_flags: -checks=inherit,-SA1019,-SA1029,-SA5008

.github/workflows/version-check.yml

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

.gitignore

@@ -4,18 +4,28 @@
 # Unignore all with extensions
 !*.*
 !**/Dockerfile
+!**/Makefile
 
 # Unignore all dirs
 !*/
 !api
 
-.idea
+# ignore
+**/.idea
+**/.vscode
 **/.DS_Store
 **/logs
-!Makefile
+**/adhoc
+**/coverage.txt
+**/WARP.md
+
+# for test purpose
+go.work
+go.work.sum
 
 # gitlab ci
 .cache
+.golangci.yml
 
 # vim auto backup file
 *~

CONTRIBUTING.md

@@ -0,0 +1,76 @@
+# 🚀 Contributing to go-zero
+
+Welcome to the go-zero community! We're thrilled to have you here. Contributing to our project is a fantastic way to be a part of the go-zero journey. Let's make this guide exciting and fun!
+
+## 📜 Before You Dive In
+
+### 🤝 Code of Conduct
+
+Let's start on the right foot. Please take a moment to read and embrace our [Code of Conduct](/code-of-conduct.md). We're all about creating a welcoming and respectful environment.
+
+### 🌟 Community Expectations
+
+At go-zero, we're like a close-knit family, and we believe in creating a healthy, friendly, and productive atmosphere. It's all about sharing knowledge and building amazing things together.
+
+## 🚀 Getting Started
+
+Get your adventure rolling! Here's how to begin:
+
+1. 🍴 **Fork the Repository**: Head over to the GitHub repository and fork it to your own space.
+
+2. 🛠️ **Make Your Magic**: Work your magic in your forked repository. Create new features, squash bugs, or improve documentation - it's your world to conquer!
+
+3. 🚀 **Submit a PR (Pull Request)**: When you're ready to unveil your creation, submit a Pull Request. We can't wait to see your awesome work!
+
+## 🌟 Your First Contribution
+
+We're here to guide you on your quest to become a go-zero contributor. Whether you want to file issues, develop features, or tame some critical bugs, we've got you covered.
+
+If you have questions or need guidance at any stage, don't hesitate to [open an issue](https://github.com/zeromicro/go-zero/issues/new/choose).
+
+## 🔍 Find Something to Work On
+
+Ready to dive into the action? There are several ways to contribute:
+
+### 💼 Find a Good First Topic
+
+Discover easy-entry issues labeled as [help wanted](https://github.com/zeromicro/go-zero/issues?q=is%3Aopen+is%3Aissue+label%3A%22help+wanted%22) or [good first issue](https://github.com/zeromicro/go-zero/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22). These issues are perfect for newcomers and don't require deep knowledge of the system. We're here to assist you with these tasks.
+
+### 🪄 Work on an Issue
+
+Once you've picked an issue that excites you, let us know by commenting on it. Our maintainers will assign it to you, and you can embark on your mission!
+
+### 📢 File an Issue
+
+Reporting an issue is just as valuable as code contributions. If you discover a problem, don't hesitate to [open an issue](https://github.com/zeromicro/go-zero/issues/new/choose). Be sure to follow our guidelines when submitting an issue.
+
+## 🎯 Contributor Workflow
+
+Here's a rough guide to your contributor journey:
+
+1. 🌱 Create a New Branch: Start by creating a topic branch, usually based on the 'master' branch. This is where your contribution will grow.
+
+2. 💡 Make Commits: Commit your work in logical units. Each commit should tell a story.
+
+3. 🚀 Push Changes: Push the changes in your topic branch to your personal fork of the repository.
+
+4. 📦 Submit a Pull Request: When your creation is complete, submit a Pull Request to the [go-zero repository](https://github.com/zeromicro/go-zero).
+
+## 🌠 Creating Pull Requests
+
+Pull Requests (PRs) are your way of making a grand entrance with your contribution. Here's how to do it:
+
+- 💼 Format Your Code: Ensure your code is beautifully formatted with `gofmt`.
+- 🏃 Run Tests: Verify that your changes pass all the tests, including data race tests. Run `go test -race ./...` for the ultimate validation.
+
+## 👁️‍🗨️ Code Review
+
+Getting your PR reviewed is the final step before your contribution becomes part of go-zero's magical world. To make the process smooth, keep these things in mind:
+
+- 🧙‍♀️ Follow Good Coding Practices: Stick to [good coding guidelines](https://github.com/golang/go/wiki/CodeReviewComments).
+- 📝 Write Awesome Commit Messages: Craft [impressive commit messages](https://chris.beams.io/posts/git-commit/) - they're like spells in the wizard's book!
+- 🔍 Break It Down: For larger changes, consider breaking them into a series of smaller, logical patches. Each patch should make an understandable and meaningful improvement.
+
+Congratulations on your contribution journey! We're thrilled to have you as part of our go-zero community. Let's make amazing things together! 🌟
+
+Now, go out there and start your adventure! If you have any more magical ideas to enhance this guide, please share them. 🔥

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2020 xiaoheiban_server_go
+Copyright (c) 2022 zeromicro
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

SECURITY.md

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

code-of-conduct.md

@@ -0,0 +1,127 @@
+
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, caste, color, religion, or sexual
+identity and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the overall
+  community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or advances of
+  any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email address,
+  without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at
+[INSERT CONTACT METHOD].
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series of
+actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or permanent
+ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior, harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within the
+community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.1, available at
+[https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1].
+
+Community Impact Guidelines were inspired by
+[Mozilla's code of conduct enforcement ladder][Mozilla CoC].
+
+For answers to common questions about this code of conduct, see the FAQ at
+[https://www.contributor-covenant.org/faq][FAQ]. Translations are available at
+[https://www.contributor-covenant.org/translations][translations].

core/bloom/bloom.go

@@ -1,82 +1,86 @@
 package bloom
 
 import (
+	"context"
+	_ "embed"
 	"errors"
 	"strconv"
 
-	"github.com/tal-tech/go-zero/core/hash"
-	"github.com/tal-tech/go-zero/core/stores/redis"
+	"github.com/zeromicro/go-zero/core/hash"
+	"github.com/zeromicro/go-zero/core/stores/redis"
 )
 
-const (
-	// for detailed error rate table, see http://pages.cs.wisc.edu/~cao/papers/summary-cache/node8.html
-	// maps as k in the error rate table
-	maps      = 14
-	setScript = `
-for _, offset in ipairs(ARGV) do
-	redis.call("setbit", KEYS[1], offset, 1)
-end
-`
-	testScript = `
-for _, offset in ipairs(ARGV) do
-	if tonumber(redis.call("getbit", KEYS[1], offset)) == 0 then
-		return false
-	end
-end
-return true
-`
-)
+// for detailed error rate table, see http://pages.cs.wisc.edu/~cao/papers/summary-cache/node8.html
+// maps as k in the error rate table
+const maps = 14
 
-var ErrTooLargeOffset = errors.New("too large offset")
+var (
+	// ErrTooLargeOffset indicates the offset is too large in bitset.
+	ErrTooLargeOffset = errors.New("too large offset")
+
+	//go:embed setscript.lua
+	setLuaScript string
+	setScript    = redis.NewScript(setLuaScript)
+
+	//go:embed testscript.lua
+	testLuaScript string
+	testScript    = redis.NewScript(testLuaScript)
+)
 
 type (
-	BitSetProvider interface {
-		check([]uint) (bool, error)
-		set([]uint) error
+	// A Filter is a bloom filter.
+	Filter struct {
+		bits   uint
+		bitSet bitSetProvider
 	}
 
-	BloomFilter struct {
-		bits   uint
-		bitSet BitSetProvider
+	bitSetProvider interface {
+		check(ctx context.Context, offsets []uint) (bool, error)
+		set(ctx context.Context, offsets []uint) error
 	}
 )
 
-// New create a BloomFilter, 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
 // when maps = 14, formula: 0.7*(bits/maps), bits = 20*elements, the error rate is 0.000067 < 1e-4
 // for detailed error rate table, see http://pages.cs.wisc.edu/~cao/papers/summary-cache/node8.html
-func New(store *redis.Redis, key string, bits uint) *BloomFilter {
-	return &BloomFilter{
+func New(store *redis.Redis, key string, bits uint) *Filter {
+	return &Filter{
 		bits:   bits,
 		bitSet: newRedisBitSet(store, key, bits),
 	}
 }
 
-func (f *BloomFilter) Add(data []byte) error {
-	locations := f.getLocations(data)
-	err := f.bitSet.set(locations)
-	if err != nil {
-		return err
-	}
-	return nil
+// Add adds data into f.
+func (f *Filter) Add(data []byte) error {
+	return f.AddCtx(context.Background(), data)
 }
 
-func (f *BloomFilter) Exists(data []byte) (bool, error) {
+// AddCtx adds data into f with context.
+func (f *Filter) AddCtx(ctx context.Context, data []byte) error {
 	locations := f.getLocations(data)
-	isSet, err := f.bitSet.check(locations)
+	return f.bitSet.set(ctx, locations)
+}
+
+// Exists checks if data is in f.
+func (f *Filter) Exists(data []byte) (bool, error) {
+	return f.ExistsCtx(context.Background(), data)
+}
+
+// ExistsCtx checks if data is in f with context.
+func (f *Filter) ExistsCtx(ctx context.Context, data []byte) (bool, error) {
+	locations := f.getLocations(data)
+	isSet, err := f.bitSet.check(ctx, locations)
 	if err != nil {
 		return false, err
 	}
-	if !isSet {
-		return false, nil
-	}
 
-	return true, nil
+	return isSet, nil
 }
 
-func (f *BloomFilter) getLocations(data []byte) []uint {
+func (f *Filter) getLocations(data []byte) []uint {
 	locations := make([]uint, maps)
 	for i := uint(0); i < maps; i++ {
 		hashValue := hash.Hash(append(data, byte(i)))
@@ -101,7 +105,7 @@ func newRedisBitSet(store *redis.Redis, key string, bits uint) *redisBitSet {
 }
 
 func (r *redisBitSet) buildOffsetArgs(offsets []uint) ([]string, error) {
-	var args []string
+	args := make([]string, 0, len(offsets))
 
 	for _, offset := range offsets {
 		if offset >= r.bits {
@@ -114,45 +118,48 @@ func (r *redisBitSet) buildOffsetArgs(offsets []uint) ([]string, error) {
 	return args, nil
 }
 
-func (r *redisBitSet) check(offsets []uint) (bool, error) {
+func (r *redisBitSet) check(ctx context.Context, offsets []uint) (bool, error) {
 	args, err := r.buildOffsetArgs(offsets)
 	if err != nil {
 		return false, err
 	}
 
-	resp, err := r.store.Eval(testScript, []string{r.key}, args)
-	if err == redis.Nil {
+	resp, err := r.store.ScriptRunCtx(ctx, testScript, []string{r.key}, args)
+	if errors.Is(err, redis.Nil) {
 		return false, nil
 	} else if err != nil {
 		return false, err
 	}
 
-	if exists, ok := resp.(int64); !ok {
+	exists, ok := resp.(int64)
+	if !ok {
 		return false, nil
-	} else {
-		return exists == 1, nil
 	}
+
+	return exists == 1, nil
 }
 
+// del only use for testing.
 func (r *redisBitSet) del() error {
 	_, err := r.store.Del(r.key)
 	return err
 }
 
+// expire only use for testing.
 func (r *redisBitSet) expire(seconds int) error {
 	return r.store.Expire(r.key, seconds)
 }
 
-func (r *redisBitSet) set(offsets []uint) error {
+func (r *redisBitSet) set(ctx context.Context, offsets []uint) error {
 	args, err := r.buildOffsetArgs(offsets)
 	if err != nil {
 		return err
 	}
 
-	_, err = r.store.Eval(setScript, []string{r.key}, args)
-	if err == redis.Nil {
+	_, err = r.store.ScriptRunCtx(ctx, setScript, []string{r.key}, args)
+	if errors.Is(err, redis.Nil) {
 		return nil
-	} else {
-		return err
 	}
+
+	return err
 }

core/bloom/bloom_test.go

@@ -1,30 +1,31 @@
 package bloom
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/stores/redis/redistest"
+	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/stores/redis/redistest"
 )
 
 func TestRedisBitSet_New_Set_Test(t *testing.T) {
-	store, clean, err := redistest.CreateRedis()
-	assert.Nil(t, err)
-	defer clean()
+	store := redistest.CreateRedis(t)
+	ctx := context.Background()
 
 	bitSet := newRedisBitSet(store, "test_key", 1024)
-	isSetBefore, err := bitSet.check([]uint{0})
+	isSetBefore, err := bitSet.check(ctx, []uint{0})
 	if err != nil {
 		t.Fatal(err)
 	}
 	if isSetBefore {
 		t.Fatal("Bit should not be set")
 	}
-	err = bitSet.set([]uint{512})
+	err = bitSet.set(ctx, []uint{512})
 	if err != nil {
 		t.Fatal(err)
 	}
-	isSetAfter, err := bitSet.check([]uint{512})
+	isSetAfter, err := bitSet.check(ctx, []uint{512})
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -42,9 +43,7 @@ func TestRedisBitSet_New_Set_Test(t *testing.T) {
 }
 
 func TestRedisBitSet_Add(t *testing.T) {
-	store, clean, err := redistest.CreateRedis()
-	assert.Nil(t, err)
-	defer clean()
+	store := redistest.CreateRedis(t)
 
 	filter := New(store, "test_key", 64)
 	assert.Nil(t, filter.Add([]byte("hello")))
@@ -53,3 +52,51 @@ func TestRedisBitSet_Add(t *testing.T) {
 	assert.Nil(t, err)
 	assert.True(t, ok)
 }
+
+func TestFilter_Exists(t *testing.T) {
+	store, clean := redistest.CreateRedisWithClean(t)
+
+	rbs := New(store, "test", 64)
+	_, err := rbs.Exists([]byte{0, 1, 2})
+	assert.NoError(t, err)
+
+	clean()
+	rbs = New(store, "test", 64)
+	_, err = rbs.Exists([]byte{0, 1, 2})
+	assert.Error(t, err)
+}
+
+func TestRedisBitSet_check(t *testing.T) {
+	store, clean := redistest.CreateRedisWithClean(t)
+	ctx := context.Background()
+
+	rbs := newRedisBitSet(store, "test", 0)
+	assert.Error(t, rbs.set(ctx, []uint{0, 1, 2}))
+	_, err := rbs.check(ctx, []uint{0, 1, 2})
+	assert.Error(t, err)
+
+	rbs = newRedisBitSet(store, "test", 64)
+	_, err = rbs.check(ctx, []uint{0, 1, 2})
+	assert.NoError(t, err)
+
+	clean()
+	rbs = newRedisBitSet(store, "test", 64)
+	_, err = rbs.check(ctx, []uint{0, 1, 2})
+	assert.Error(t, err)
+}
+
+func TestRedisBitSet_set(t *testing.T) {
+	logx.Disable()
+	store, clean := redistest.CreateRedisWithClean(t)
+	ctx := context.Background()
+
+	rbs := newRedisBitSet(store, "test", 0)
+	assert.Error(t, rbs.set(ctx, []uint{0, 1, 2}))
+
+	rbs = newRedisBitSet(store, "test", 64)
+	assert.NoError(t, rbs.set(ctx, []uint{0, 1, 2}))
+
+	clean()
+	rbs = newRedisBitSet(store, "test", 64)
+	assert.Error(t, rbs.set(ctx, []uint{0, 1, 2}))
+}

core/bloom/setscript.lua

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

core/bloom/testscript.lua

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

core/breaker/breaker.go

@@ -1,70 +1,88 @@
 package breaker
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strings"
 	"sync"
+	"time"
 
-	"github.com/tal-tech/go-zero/core/mathx"
-	"github.com/tal-tech/go-zero/core/proc"
-	"github.com/tal-tech/go-zero/core/stat"
-	"github.com/tal-tech/go-zero/core/stringx"
-	"github.com/tal-tech/go-zero/core/timex"
+	"github.com/zeromicro/go-zero/core/proc"
+	"github.com/zeromicro/go-zero/core/stat"
+	"github.com/zeromicro/go-zero/core/stringx"
 )
 
-const (
-	numHistoryReasons = 5
-	timeFormat        = "15:04:05"
-)
+const numHistoryReasons = 5
 
-// ErrServiceUnavailable is returned when the CB state is open
+// ErrServiceUnavailable is returned when the Breaker state is open.
 var ErrServiceUnavailable = errors.New("circuit breaker is open")
 
 type (
+	// Acceptable is the func to check if the error can be accepted.
 	Acceptable func(err error) bool
 
+	// A Breaker represents a circuit breaker.
 	Breaker interface {
 		// Name returns the name of the Breaker.
 		Name() string
 
 		// Allow checks if the request is allowed.
-		// If allowed, a promise will be returned, the caller needs to call promise.Accept()
-		// on success, or call promise.Reject() on failure.
-		// If not allow, ErrServiceUnavailable will be returned.
+		// If allowed, a promise will be returned,
+		// otherwise ErrServiceUnavailable will be returned as the error.
+		// The caller needs to call promise.Accept() on success,
+		// or call promise.Reject() on failure.
 		Allow() (Promise, error)
+		// AllowCtx checks if the request is allowed when ctx isn't done.
+		AllowCtx(ctx context.Context) (Promise, error)
 
 		// Do runs the given request if the Breaker accepts it.
 		// Do returns an error instantly if the Breaker rejects the request.
 		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
 		Do(req func() error) error
+		// DoCtx runs the given request if the Breaker accepts it when ctx isn't done.
+		DoCtx(ctx context.Context, req func() error) error
 
 		// DoWithAcceptable runs the given request if the Breaker accepts it.
 		// DoWithAcceptable returns an error instantly if the Breaker rejects the request.
 		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
-		// acceptable checks if it's a successful call, even if the err is not nil.
+		// acceptable checks if it's a successful call, even if the error is not nil.
 		DoWithAcceptable(req func() error, acceptable Acceptable) error
+		// DoWithAcceptableCtx runs the given request if the Breaker accepts it when ctx isn't done.
+		DoWithAcceptableCtx(ctx context.Context, req func() error, acceptable Acceptable) error
 
 		// DoWithFallback runs the given request if the Breaker accepts it.
 		// DoWithFallback runs the fallback if the Breaker rejects the request.
 		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
-		DoWithFallback(req func() error, fallback func(err error) error) error
+		DoWithFallback(req func() error, fallback Fallback) error
+		// DoWithFallbackCtx runs the given request if the Breaker accepts it when ctx isn't done.
+		DoWithFallbackCtx(ctx context.Context, req func() error, fallback Fallback) error
 
 		// DoWithFallbackAcceptable runs the given request if the Breaker accepts it.
 		// DoWithFallbackAcceptable runs the fallback if the Breaker rejects the request.
 		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
-		// acceptable checks if it's a successful call, even if the err is not nil.
-		DoWithFallbackAcceptable(req func() error, fallback func(err error) error, acceptable Acceptable) error
+		// acceptable checks if it's a successful call, even if the error is not nil.
+		DoWithFallbackAcceptable(req func() error, fallback Fallback, acceptable Acceptable) error
+		// DoWithFallbackAcceptableCtx runs the given request if the Breaker accepts it when ctx isn't done.
+		DoWithFallbackAcceptableCtx(ctx context.Context, req func() error, fallback Fallback,
+			acceptable Acceptable) error
 	}
 
+	// Fallback is the func to be called if the request is rejected.
+	Fallback func(err error) error
+
+	// Option defines the method to customize a Breaker.
 	Option func(breaker *circuitBreaker)
 
+	// Promise interface defines the callbacks that returned by Breaker.Allow.
 	Promise interface {
+		// Accept tells the Breaker that the call is successful.
 		Accept()
+		// Reject tells the Breaker that the call is failed.
 		Reject(reason string)
 	}
 
@@ -80,15 +98,17 @@ type (
 
 	internalThrottle interface {
 		allow() (internalPromise, error)
-		doReq(req func() error, fallback func(err error) error, acceptable Acceptable) error
+		doReq(req func() error, fallback Fallback, acceptable Acceptable) error
 	}
 
 	throttle interface {
 		allow() (Promise, error)
-		doReq(req func() error, fallback func(err error) error, acceptable Acceptable) error
+		doReq(req func() error, fallback Fallback, acceptable Acceptable) error
 	}
 )
 
+// NewBreaker returns a Breaker object.
+// opts can be used to customize the Breaker.
 func NewBreaker(opts ...Option) Breaker {
 	var b circuitBreaker
 	for _, opt := range opts {
@@ -106,27 +126,76 @@ func (cb *circuitBreaker) Allow() (Promise, error) {
 	return cb.throttle.allow()
 }
 
+func (cb *circuitBreaker) AllowCtx(ctx context.Context) (Promise, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+		return cb.Allow()
+	}
+}
+
 func (cb *circuitBreaker) Do(req func() error) error {
 	return cb.throttle.doReq(req, nil, defaultAcceptable)
 }
 
+func (cb *circuitBreaker) DoCtx(ctx context.Context, req func() error) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return cb.Do(req)
+	}
+}
+
 func (cb *circuitBreaker) DoWithAcceptable(req func() error, acceptable Acceptable) error {
 	return cb.throttle.doReq(req, nil, acceptable)
 }
 
-func (cb *circuitBreaker) DoWithFallback(req func() error, fallback func(err error) error) error {
+func (cb *circuitBreaker) DoWithAcceptableCtx(ctx context.Context, req func() error,
+	acceptable Acceptable) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return cb.DoWithAcceptable(req, acceptable)
+	}
+}
+
+func (cb *circuitBreaker) DoWithFallback(req func() error, fallback Fallback) error {
 	return cb.throttle.doReq(req, fallback, defaultAcceptable)
 }
 
-func (cb *circuitBreaker) DoWithFallbackAcceptable(req func() error, fallback func(err error) error,
+func (cb *circuitBreaker) DoWithFallbackCtx(ctx context.Context, req func() error,
+	fallback Fallback) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return cb.DoWithFallback(req, fallback)
+	}
+}
+
+func (cb *circuitBreaker) DoWithFallbackAcceptable(req func() error, fallback Fallback,
 	acceptable Acceptable) error {
 	return cb.throttle.doReq(req, fallback, acceptable)
 }
 
+func (cb *circuitBreaker) DoWithFallbackAcceptableCtx(ctx context.Context, req func() error,
+	fallback Fallback, acceptable Acceptable) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return cb.DoWithFallbackAcceptable(req, fallback, acceptable)
+	}
+}
+
 func (cb *circuitBreaker) Name() string {
 	return cb.name
 }
 
+// WithName returns a function to set the name of a Breaker.
 func WithName(name string) Option {
 	return func(b *circuitBreaker) {
 		b.name = name
@@ -159,10 +228,10 @@ func (lt loggedThrottle) allow() (Promise, error) {
 	}, lt.logError(err)
 }
 
-func (lt loggedThrottle) doReq(req func() error, fallback func(err error) error, acceptable Acceptable) error {
+func (lt loggedThrottle) doReq(req func() error, fallback Fallback, acceptable Acceptable) error {
 	return lt.logError(lt.internalThrottle.doReq(req, fallback, func(err error) bool {
 		accept := acceptable(err)
-		if !accept {
+		if !accept && err != nil {
 			lt.errWin.add(err.Error())
 		}
 		return accept
@@ -170,7 +239,7 @@ func (lt loggedThrottle) doReq(req func() error, fallback func(err error) error,
 }
 
 func (lt loggedThrottle) logError(err error) error {
-	if err == ErrServiceUnavailable {
+	if errors.Is(err, ErrServiceUnavailable) {
 		// if circuit open, not possible to have empty error window
 		stat.Report(fmt.Sprintf(
 			"proc(%s/%d), callee: %s, breaker is open and requests dropped\nlast errors:\n%s",
@@ -189,14 +258,14 @@ type errorWindow struct {
 
 func (ew *errorWindow) add(reason string) {
 	ew.lock.Lock()
-	ew.reasons[ew.index] = fmt.Sprintf("%s %s", timex.Time().Format(timeFormat), reason)
+	ew.reasons[ew.index] = fmt.Sprintf("%s %s", time.Now().Format(time.TimeOnly), reason)
 	ew.index = (ew.index + 1) % numHistoryReasons
-	ew.count = mathx.MinInt(ew.count+1, numHistoryReasons)
+	ew.count = min(ew.count+1, numHistoryReasons)
 	ew.lock.Unlock()
 }
 
 func (ew *errorWindow) String() string {
-	var reasons []string
+	reasons := make([]string, 0, ew.count)
 
 	ew.lock.Lock()
 	// reverse order

core/breaker/breaker_test.go

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

core/breaker/breakers.go

@@ -1,37 +1,75 @@
 package breaker
 
-import "sync"
+import (
+	"context"
+	"sync"
+)
 
 var (
 	lock     sync.RWMutex
 	breakers = make(map[string]Breaker)
 )
 
+// Do calls Breaker.Do on the Breaker with given name.
 func Do(name string, req func() error) error {
 	return do(name, func(b Breaker) error {
 		return b.Do(req)
 	})
 }
 
+// DoCtx calls Breaker.DoCtx on the Breaker with given name.
+func DoCtx(ctx context.Context, name string, req func() error) error {
+	return do(name, func(b Breaker) error {
+		return b.DoCtx(ctx, req)
+	})
+}
+
+// DoWithAcceptable calls Breaker.DoWithAcceptable on the Breaker with given name.
 func DoWithAcceptable(name string, req func() error, acceptable Acceptable) error {
 	return do(name, func(b Breaker) error {
 		return b.DoWithAcceptable(req, acceptable)
 	})
 }
 
-func DoWithFallback(name string, req func() error, fallback func(err error) error) error {
+// DoWithAcceptableCtx calls Breaker.DoWithAcceptableCtx on the Breaker with given name.
+func DoWithAcceptableCtx(ctx context.Context, name string, req func() error,
+	acceptable Acceptable) error {
+	return do(name, func(b Breaker) error {
+		return b.DoWithAcceptableCtx(ctx, req, acceptable)
+	})
+}
+
+// DoWithFallback calls Breaker.DoWithFallback on the Breaker with given name.
+func DoWithFallback(name string, req func() error, fallback Fallback) error {
 	return do(name, func(b Breaker) error {
 		return b.DoWithFallback(req, fallback)
 	})
 }
 
-func DoWithFallbackAcceptable(name string, req func() error, fallback func(err error) error,
+// DoWithFallbackCtx calls Breaker.DoWithFallbackCtx on the Breaker with given name.
+func DoWithFallbackCtx(ctx context.Context, name string, req func() error, fallback Fallback) error {
+	return do(name, func(b Breaker) error {
+		return b.DoWithFallbackCtx(ctx, req, fallback)
+	})
+}
+
+// DoWithFallbackAcceptable calls Breaker.DoWithFallbackAcceptable on the Breaker with given name.
+func DoWithFallbackAcceptable(name string, req func() error, fallback Fallback,
 	acceptable Acceptable) error {
 	return do(name, func(b Breaker) error {
 		return b.DoWithFallbackAcceptable(req, fallback, acceptable)
 	})
 }
 
+// DoWithFallbackAcceptableCtx calls Breaker.DoWithFallbackAcceptableCtx on the Breaker with given name.
+func DoWithFallbackAcceptableCtx(ctx context.Context, name string, req func() error,
+	fallback Fallback, acceptable Acceptable) error {
+	return do(name, func(b Breaker) error {
+		return b.DoWithFallbackAcceptableCtx(ctx, req, fallback, acceptable)
+	})
+}
+
+// GetBreaker returns the Breaker with the given name.
 func GetBreaker(name string) Breaker {
 	lock.RLock()
 	b, ok := breakers[name]
@@ -51,9 +89,10 @@ func GetBreaker(name string) Breaker {
 	return b
 }
 
-func NoBreakFor(name string) {
+// NoBreakerFor disables the circuit breaker for the given name.
+func NoBreakerFor(name string) {
 	lock.Lock()
-	breakers[name] = newNoOpBreaker()
+	breakers[name] = NopBreaker()
 	lock.Unlock()
 }
 

core/breaker/breakers_test.go

@@ -1,12 +1,13 @@
 package breaker
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/stat"
+	"github.com/zeromicro/go-zero/core/stat"
 )
 
 func init() {
@@ -22,6 +23,9 @@ func TestBreakersDo(t *testing.T) {
 	assert.Equal(t, errDummy, Do("any", func() error {
 		return errDummy
 	}))
+	assert.Equal(t, errDummy, DoCtx(context.Background(), "any", func() error {
+		return errDummy
+	}))
 }
 
 func TestBreakersDoWithAcceptable(t *testing.T) {
@@ -30,7 +34,7 @@ func TestBreakersDoWithAcceptable(t *testing.T) {
 		assert.Equal(t, errDummy, GetBreaker("anyone").DoWithAcceptable(func() error {
 			return errDummy
 		}, func(err error) bool {
-			return err == nil || err == errDummy
+			return err == nil || errors.Is(err, errDummy)
 		}))
 	}
 	verify(t, func() bool {
@@ -38,6 +42,13 @@ func TestBreakersDoWithAcceptable(t *testing.T) {
 			return nil
 		}) == nil
 	})
+	verify(t, func() bool {
+		return DoWithAcceptableCtx(context.Background(), "anyone", func() error {
+			return nil
+		}, func(err error) bool {
+			return true
+		}) == nil
+	})
 
 	for i := 0; i < 10000; i++ {
 		err := DoWithAcceptable("another", func() error {
@@ -45,17 +56,17 @@ func TestBreakersDoWithAcceptable(t *testing.T) {
 		}, func(err error) bool {
 			return err == nil
 		})
-		assert.True(t, err == errDummy || err == ErrServiceUnavailable)
+		assert.True(t, errors.Is(err, errDummy) || errors.Is(err, ErrServiceUnavailable))
 	}
 	verify(t, func() bool {
-		return ErrServiceUnavailable == Do("another", func() error {
+		return errors.Is(Do("another", func() error {
 			return nil
-		})
+		}), ErrServiceUnavailable)
 	})
 }
 
 func TestBreakersNoBreakerFor(t *testing.T) {
-	NoBreakFor("any")
+	NoBreakerFor("any")
 	errDummy := errors.New("any")
 	for i := 0; i < 10000; i++ {
 		assert.Equal(t, errDummy, GetBreaker("any").Do(func() error {
@@ -75,18 +86,24 @@ func TestBreakersFallback(t *testing.T) {
 		}, func(err error) error {
 			return nil
 		})
-		assert.True(t, err == nil || err == errDummy)
-	}
-	verify(t, func() bool {
-		return ErrServiceUnavailable == Do("fallback", func() error {
+		assert.True(t, err == nil || errors.Is(err, errDummy))
+		err = DoWithFallbackCtx(context.Background(), "fallback", func() error {
+			return errDummy
+		}, func(err error) error {
 			return nil
 		})
+		assert.True(t, err == nil || errors.Is(err, errDummy))
+	}
+	verify(t, func() bool {
+		return errors.Is(Do("fallback", func() error {
+			return nil
+		}), ErrServiceUnavailable)
 	})
 }
 
 func TestBreakersAcceptableFallback(t *testing.T) {
 	errDummy := errors.New("any")
-	for i := 0; i < 10000; i++ {
+	for i := 0; i < 5000; i++ {
 		err := DoWithFallbackAcceptable("acceptablefallback", func() error {
 			return errDummy
 		}, func(err error) error {
@@ -94,12 +111,20 @@ func TestBreakersAcceptableFallback(t *testing.T) {
 		}, func(err error) bool {
 			return err == nil
 		})
-		assert.True(t, err == nil || err == errDummy)
+		assert.True(t, err == nil || errors.Is(err, errDummy))
+		err = DoWithFallbackAcceptableCtx(context.Background(), "acceptablefallback", func() error {
+			return errDummy
+		}, func(err error) error {
+			return nil
+		}, func(err error) bool {
+			return err == nil
+		})
+		assert.True(t, err == nil || errors.Is(err, errDummy))
 	}
 	verify(t, func() bool {
-		return ErrServiceUnavailable == Do("acceptablefallback", func() error {
+		return errors.Is(Do("acceptablefallback", func() error {
 			return nil
-		})
+		}), ErrServiceUnavailable)
 	})
 }
 
@@ -110,5 +135,5 @@ func verify(t *testing.T, fn func() bool) {
 			count++
 		}
 	}
-	assert.True(t, count >= 80, fmt.Sprintf("should be greater than 80, actual %d", count))
+	assert.True(t, count >= 75, fmt.Sprintf("should be greater than 75, actual %d", count))
 }

core/breaker/bucket.go

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

core/breaker/bucket_test.go

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

core/breaker/googlebreaker.go

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

core/breaker/googlebreaker_test.go

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

core/breaker/nopbreaker.go

@@ -1,35 +1,58 @@
 package breaker
 
-const noOpBreakerName = "nopBreaker"
+import "context"
 
-type noOpBreaker struct{}
+const nopBreakerName = "nopBreaker"
 
-func newNoOpBreaker() Breaker {
-	return noOpBreaker{}
+type nopBreaker struct{}
+
+// NopBreaker returns a breaker that never trigger breaker circuit.
+func NopBreaker() Breaker {
+	return nopBreaker{}
 }
 
-func (b noOpBreaker) Name() string {
-	return noOpBreakerName
+func (b nopBreaker) Name() string {
+	return nopBreakerName
 }
 
-func (b noOpBreaker) Allow() (Promise, error) {
+func (b nopBreaker) Allow() (Promise, error) {
 	return nopPromise{}, nil
 }
 
-func (b noOpBreaker) Do(req func() error) error {
+func (b nopBreaker) AllowCtx(_ context.Context) (Promise, error) {
+	return nopPromise{}, nil
+}
+
+func (b nopBreaker) Do(req func() error) error {
 	return req()
 }
 
-func (b noOpBreaker) DoWithAcceptable(req func() error, acceptable Acceptable) error {
+func (b nopBreaker) DoCtx(_ context.Context, req func() error) error {
 	return req()
 }
 
-func (b noOpBreaker) DoWithFallback(req func() error, fallback func(err error) error) error {
+func (b nopBreaker) DoWithAcceptable(req func() error, _ Acceptable) error {
 	return req()
 }
 
-func (b noOpBreaker) DoWithFallbackAcceptable(req func() error, fallback func(err error) error,
-	acceptable Acceptable) error {
+func (b nopBreaker) DoWithAcceptableCtx(_ context.Context, req func() error, _ Acceptable) error {
+	return req()
+}
+
+func (b nopBreaker) DoWithFallback(req func() error, _ Fallback) error {
+	return req()
+}
+
+func (b nopBreaker) DoWithFallbackCtx(_ context.Context, req func() error, _ Fallback) error {
+	return req()
+}
+
+func (b nopBreaker) DoWithFallbackAcceptable(req func() error, _ Fallback, _ Acceptable) error {
+	return req()
+}
+
+func (b nopBreaker) DoWithFallbackAcceptableCtx(_ context.Context, req func() error,
+	_ Fallback, _ Acceptable) error {
 	return req()
 }
 
@@ -38,5 +61,5 @@ type nopPromise struct{}
 func (p nopPromise) Accept() {
 }
 
-func (p nopPromise) Reject(reason string) {
+func (p nopPromise) Reject(_ string) {
 }

core/breaker/nopbreaker_test.go

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

core/cmdline/input.go

@@ -7,11 +7,13 @@ import (
 	"strings"
 )
 
+// EnterToContinue let stdin waiting for an enter key to continue.
 func EnterToContinue() {
 	fmt.Print("Press 'Enter' to continue...")
 	bufio.NewReader(os.Stdin).ReadBytes('\n')
 }
 
+// ReadLine shows prompt to stdout and read a line from stdin.
 func ReadLine(prompt string) string {
 	fmt.Print(prompt)
 	input, _ := bufio.NewReader(os.Stdin).ReadString('\n')

core/cmdline/input_test.go

@@ -8,8 +8,8 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/iox"
-	"github.com/tal-tech/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/iox"
+	"github.com/zeromicro/go-zero/core/lang"
 )
 
 func TestEnterToContinue(t *testing.T) {

core/codec/aesecb.go

@@ -6,10 +6,9 @@ import (
 	"crypto/cipher"
 	"encoding/base64"
 	"errors"
-
-	"github.com/tal-tech/go-zero/core/logx"
 )
 
+// ErrPaddingSize indicates bad padding size.
 var ErrPaddingSize = errors.New("padding size error")
 
 type ecb struct {
@@ -26,21 +25,23 @@ func newECB(b cipher.Block) *ecb {
 
 type ecbEncrypter ecb
 
+// Deprecated: NewECBEncrypter returns an ECB encrypter.
+// ECB mode is insecure for multi-block data. Use AES-GCM instead.
 func NewECBEncrypter(b cipher.Block) cipher.BlockMode {
 	return (*ecbEncrypter)(newECB(b))
 }
 
+// BlockSize returns the mode's block size.
 func (x *ecbEncrypter) BlockSize() int { return x.blockSize }
 
-// why we don't return error is because cipher.BlockMode doesn't allow this
+// CryptBlocks encrypts a number of blocks. The length of src must be a multiple of
+// the block size. Dst and src must overlap entirely or not at all.
 func (x *ecbEncrypter) CryptBlocks(dst, src []byte) {
 	if len(src)%x.blockSize != 0 {
-		logx.Error("crypto/cipher: input not full blocks")
-		return
+		panic("crypto/cipher: input not full blocks")
 	}
 	if len(dst) < len(src) {
-		logx.Error("crypto/cipher: output smaller than input")
-		return
+		panic("crypto/cipher: output smaller than input")
 	}
 
 	for len(src) > 0 {
@@ -52,23 +53,25 @@ func (x *ecbEncrypter) CryptBlocks(dst, src []byte) {
 
 type ecbDecrypter ecb
 
+// Deprecated: NewECBDecrypter returns an ECB decrypter.
+// ECB mode is insecure for multi-block data. Use AES-GCM instead.
 func NewECBDecrypter(b cipher.Block) cipher.BlockMode {
 	return (*ecbDecrypter)(newECB(b))
 }
 
+// BlockSize returns the mode's block size.
 func (x *ecbDecrypter) BlockSize() int {
 	return x.blockSize
 }
 
-// why we don't return error is because cipher.BlockMode doesn't allow this
+// CryptBlocks decrypts a number of blocks. The length of src must be a multiple of
+// the block size. Dst and src must overlap entirely or not at all.
 func (x *ecbDecrypter) CryptBlocks(dst, src []byte) {
 	if len(src)%x.blockSize != 0 {
-		logx.Error("crypto/cipher: input not full blocks")
-		return
+		panic("crypto/cipher: input not full blocks")
 	}
 	if len(dst) < len(src) {
-		logx.Error("crypto/cipher: output smaller than input")
-		return
+		panic("crypto/cipher: output smaller than input")
 	}
 
 	for len(src) > 0 {
@@ -78,13 +81,18 @@ func (x *ecbDecrypter) CryptBlocks(dst, src []byte) {
 	}
 }
 
+// Deprecated: EcbDecrypt decrypts src with the given key.
+// ECB mode is insecure for multi-block data. Use AES-GCM instead.
 func EcbDecrypt(key, src []byte) ([]byte, error) {
 	block, err := aes.NewCipher(key)
 	if err != nil {
-		logx.Errorf("Decrypt key error: % x", key)
 		return nil, err
 	}
 
+	if len(src)%block.BlockSize() != 0 {
+		return nil, ErrPaddingSize
+	}
+
 	decrypter := NewECBDecrypter(block)
 	decrypted := make([]byte, len(src))
 	decrypter.CryptBlocks(decrypted, src)
@@ -92,6 +100,9 @@ func EcbDecrypt(key, src []byte) ([]byte, error) {
 	return pkcs5Unpadding(decrypted, decrypter.BlockSize())
 }
 
+// Deprecated: EcbDecryptBase64 decrypts base64 encoded src with the given base64 encoded key.
+// The returned string is also base64 encoded.
+// ECB mode is insecure for multi-block data. Use AES-GCM instead.
 func EcbDecryptBase64(key, src string) (string, error) {
 	keyBytes, err := getKeyBytes(key)
 	if err != nil {
@@ -111,10 +122,11 @@ func EcbDecryptBase64(key, src string) (string, error) {
 	return base64.StdEncoding.EncodeToString(decryptedBytes), nil
 }
 
+// Deprecated: EcbEncrypt encrypts src with the given key.
+// ECB mode is insecure for multi-block data. Use AES-GCM instead.
 func EcbEncrypt(key, src []byte) ([]byte, error) {
 	block, err := aes.NewCipher(key)
 	if err != nil {
-		logx.Errorf("Encrypt key error: % x", key)
 		return nil, err
 	}
 
@@ -126,6 +138,9 @@ func EcbEncrypt(key, src []byte) ([]byte, error) {
 	return crypted, nil
 }
 
+// Deprecated: EcbEncryptBase64 encrypts base64 encoded src with the given base64 encoded key.
+// The returned string is also base64 encoded.
+// ECB mode is insecure for multi-block data. Use AES-GCM instead.
 func EcbEncryptBase64(key, src string) (string, error) {
 	keyBytes, err := getKeyBytes(key)
 	if err != nil {
@@ -146,15 +161,16 @@ func EcbEncryptBase64(key, src string) (string, error) {
 }
 
 func getKeyBytes(key string) ([]byte, error) {
-	if len(key) > 32 {
-		if keyBytes, err := base64.StdEncoding.DecodeString(key); err != nil {
-			return nil, err
-		} else {
-			return keyBytes, nil
-		}
+	if len(key) <= 32 {
+		return []byte(key), nil
 	}
 
-	return []byte(key), nil
+	keyBytes, err := base64.StdEncoding.DecodeString(key)
+	if err != nil {
+		return nil, err
+	}
+
+	return keyBytes, nil
 }
 
 func pkcs5Padding(ciphertext []byte, blockSize int) []byte {
@@ -165,10 +181,20 @@ func pkcs5Padding(ciphertext []byte, blockSize int) []byte {
 
 func pkcs5Unpadding(src []byte, blockSize int) ([]byte, error) {
 	length := len(src)
-	unpadding := int(src[length-1])
-	if unpadding >= length || unpadding > blockSize {
+	if length == 0 {
 		return nil, ErrPaddingSize
 	}
 
+	unpadding := int(src[length-1])
+	if unpadding < 1 || unpadding > blockSize || unpadding > length {
+		return nil, ErrPaddingSize
+	}
+
+	for _, b := range src[length-unpadding:] {
+		if int(b) != unpadding {
+			return nil, ErrPaddingSize
+		}
+	}
+
 	return src[:length-unpadding], nil
 }

core/codec/aesecb_test.go

@@ -0,0 +1,136 @@
+package codec
+
+import (
+	"crypto/aes"
+	"encoding/base64"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAesEcb(t *testing.T) {
+	var (
+		key     = []byte("q4t7w!z%C*F-JaNdRgUjXn2r5u8x/A?D")
+		val     = []byte("helloworld")
+		valLong = []byte("helloworldlong..")
+		badKey1 = []byte("aaaaaaaaa")
+		// more than 32 chars
+		badKey2 = []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")
+	)
+	_, err := EcbEncrypt(badKey1, val)
+	assert.NotNil(t, err)
+	_, err = EcbEncrypt(badKey2, val)
+	assert.NotNil(t, err)
+	dst, err := EcbEncrypt(key, val)
+	assert.Nil(t, err)
+	_, err = EcbDecrypt(badKey1, dst)
+	assert.NotNil(t, err)
+	_, err = EcbDecrypt(badKey2, dst)
+	assert.NotNil(t, err)
+	_, err = EcbDecrypt(key, val)
+	// not a multiple of block size
+	assert.NotNil(t, err)
+	src, err := EcbDecrypt(key, dst)
+	assert.Nil(t, err)
+	assert.Equal(t, val, src)
+	block, err := aes.NewCipher(key)
+	assert.NoError(t, err)
+	encrypter := NewECBEncrypter(block)
+	assert.Equal(t, 16, encrypter.BlockSize())
+	decrypter := NewECBDecrypter(block)
+	assert.Equal(t, 16, decrypter.BlockSize())
+
+	dst = make([]byte, 8)
+	assert.Panics(t, func() {
+		encrypter.CryptBlocks(dst, val)
+	})
+
+	dst = make([]byte, 8)
+	assert.Panics(t, func() {
+		encrypter.CryptBlocks(dst, valLong)
+	})
+
+	dst = make([]byte, 8)
+	assert.Panics(t, func() {
+		decrypter.CryptBlocks(dst, val)
+	})
+
+	dst = make([]byte, 8)
+	assert.Panics(t, func() {
+		decrypter.CryptBlocks(dst, valLong)
+	})
+
+	_, err = EcbEncryptBase64("cTR0N3dDKkYtSmFOZFJnVWpYbjJyNXU4eC9BP0QK", "aGVsbG93b3JsZGxvbmcuLgo=")
+	assert.Error(t, err)
+}
+func TestAesEcbBase64(t *testing.T) {
+	const (
+		val     = "hello"
+		badKey1 = "aaaaaaaaa"
+		// more than 32 chars
+		badKey2 = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+	)
+	key := []byte("q4t7w!z%C*F-JaNdRgUjXn2r5u8x/A?D")
+	b64Key := base64.StdEncoding.EncodeToString(key)
+	b64Val := base64.StdEncoding.EncodeToString([]byte(val))
+	_, err := EcbEncryptBase64(badKey1, val)
+	assert.NotNil(t, err)
+	_, err = EcbEncryptBase64(badKey2, val)
+	assert.NotNil(t, err)
+	_, err = EcbEncryptBase64(b64Key, val)
+	assert.NotNil(t, err)
+	dst, err := EcbEncryptBase64(b64Key, b64Val)
+	assert.Nil(t, err)
+	_, err = EcbDecryptBase64(badKey1, dst)
+	assert.NotNil(t, err)
+	_, err = EcbDecryptBase64(badKey2, dst)
+	assert.NotNil(t, err)
+	_, err = EcbDecryptBase64(b64Key, val)
+	assert.NotNil(t, err)
+	src, err := EcbDecryptBase64(b64Key, dst)
+	assert.Nil(t, err)
+	b, err := base64.StdEncoding.DecodeString(src)
+	assert.Nil(t, err)
+	assert.Equal(t, val, string(b))
+}
+
+func TestPkcs5UnpaddingEmptyInput(t *testing.T) {
+	_, err := pkcs5Unpadding([]byte{}, 16)
+	assert.Equal(t, ErrPaddingSize, err)
+}
+
+func TestPkcs5UnpaddingMalformedPadding(t *testing.T) {
+	// Valid PKCS5 padding of 3: last 3 bytes should all be 0x03
+	// Here we corrupt one padding byte
+	malformed := []byte{0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+		0x41, 0x41, 0x41, 0x41, 0x41, 0x02, 0x03, 0x03}
+	_, err := pkcs5Unpadding(malformed, 16)
+	assert.Equal(t, ErrPaddingSize, err)
+
+	// All padding bytes correct
+	valid := []byte{0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+		0x41, 0x41, 0x41, 0x41, 0x41, 0x03, 0x03, 0x03}
+	result, err := pkcs5Unpadding(valid, 16)
+	assert.NoError(t, err)
+	assert.Equal(t, valid[:13], result)
+}
+
+func TestPkcs5UnpaddingInvalidPaddingValue(t *testing.T) {
+	// padding value = 0 (< 1)
+	_, err := pkcs5Unpadding([]byte{0x41, 0x00}, 16)
+	assert.Equal(t, ErrPaddingSize, err)
+
+	// padding value > blockSize
+	_, err = pkcs5Unpadding([]byte{0x41, 0x41, 0x41, 0x41, 17}, 4)
+	assert.Equal(t, ErrPaddingSize, err)
+
+	// padding value > length
+	_, err = pkcs5Unpadding([]byte{0x41, 0x03}, 16)
+	assert.Equal(t, ErrPaddingSize, err)
+}
+
+func TestEcbDecryptEmptyInput(t *testing.T) {
+	key := []byte("q4t7w!z%C*F-JaNdRgUjXn2r5u8x/A?D")
+	_, err := EcbDecrypt(key, []byte{})
+	assert.Equal(t, ErrPaddingSize, err)
+}

core/codec/dh.go

@@ -11,8 +11,11 @@ import (
 // 2048-bit MODP Group
 
 var (
-	ErrInvalidPriKey    = errors.New("invalid private key")
-	ErrInvalidPubKey    = errors.New("invalid public key")
+	// ErrInvalidPriKey indicates the invalid private key.
+	ErrInvalidPriKey = errors.New("invalid private key")
+	// ErrInvalidPubKey indicates the invalid public key.
+	ErrInvalidPubKey = errors.New("invalid public key")
+	// ErrPubKeyOutOfBound indicates the public key is out of bound.
 	ErrPubKeyOutOfBound = errors.New("public key out of bound")
 
 	p, _ = new(big.Int).SetString("FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129024E088A67CC74020BBEA63B139B22514A08798E3404DDEF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7EDEE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3DC2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F83655D23DCA3AD961C62F356208552BB9ED529077096966D670C354E4ABC9804F1746C08CA18217C32905E462E36CE3BE39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9DE2BCBF6955817183995497CEA956AE515D2261898FA051015728E5A8AACAA68FFFFFFFFFFFFFFFF", 16)
@@ -20,17 +23,19 @@ var (
 	zero = big.NewInt(0)
 )
 
+// DhKey defines the Diffie-Hellman key.
 type DhKey struct {
 	PriKey *big.Int
 	PubKey *big.Int
 }
 
+// ComputeKey returns a key from public key and private key.
 func ComputeKey(pubKey, priKey *big.Int) (*big.Int, error) {
 	if pubKey == nil {
 		return nil, ErrInvalidPubKey
 	}
 
-	if pubKey.Sign() <= 0 && p.Cmp(pubKey) <= 0 {
+	if pubKey.Sign() <= 0 || p.Cmp(pubKey) <= 0 {
 		return nil, ErrPubKeyOutOfBound
 	}
 
@@ -41,6 +46,7 @@ func ComputeKey(pubKey, priKey *big.Int) (*big.Int, error) {
 	return new(big.Int).Exp(pubKey, priKey, p), nil
 }
 
+// GenerateKey returns a Diffie-Hellman key.
 func GenerateKey() (*DhKey, error) {
 	var err error
 	var x *big.Int
@@ -63,10 +69,12 @@ func GenerateKey() (*DhKey, error) {
 	return key, nil
 }
 
+// NewPublicKey returns a public key from the given bytes.
 func NewPublicKey(bs []byte) *big.Int {
 	return new(big.Int).SetBytes(bs)
 }
 
+// Bytes returns public key bytes.
 func (k *DhKey) Bytes() []byte {
 	if k.PubKey == nil {
 		return nil

core/codec/dh_test.go

@@ -80,3 +80,46 @@ func TestKeyBytes(t *testing.T) {
 	assert.Nil(t, err)
 	assert.True(t, len(key.Bytes()) > 0)
 }
+
+func TestDHOnErrors(t *testing.T) {
+	key, err := GenerateKey()
+	assert.Nil(t, err)
+	assert.NotEmpty(t, key.Bytes())
+	_, err = ComputeKey(key.PubKey, key.PriKey)
+	assert.NoError(t, err)
+	_, err = ComputeKey(nil, key.PriKey)
+	assert.Error(t, err)
+	_, err = ComputeKey(key.PubKey, nil)
+	assert.Error(t, err)
+
+	assert.NotNil(t, NewPublicKey([]byte("")))
+}
+
+func TestDHPubKeyBoundary(t *testing.T) {
+	key, err := GenerateKey()
+	assert.Nil(t, err)
+
+	// pubKey = 0 should be rejected
+	_, err = ComputeKey(big.NewInt(0), key.PriKey)
+	assert.ErrorIs(t, err, ErrPubKeyOutOfBound)
+
+	// pubKey = -1 should be rejected
+	_, err = ComputeKey(big.NewInt(-1), key.PriKey)
+	assert.ErrorIs(t, err, ErrPubKeyOutOfBound)
+
+	// pubKey = p should be rejected
+	_, err = ComputeKey(new(big.Int).Set(p), key.PriKey)
+	assert.ErrorIs(t, err, ErrPubKeyOutOfBound)
+
+	// pubKey = p+1 should be rejected
+	_, err = ComputeKey(new(big.Int).Add(p, big.NewInt(1)), key.PriKey)
+	assert.ErrorIs(t, err, ErrPubKeyOutOfBound)
+
+	// pubKey = 1 should be accepted
+	_, err = ComputeKey(big.NewInt(1), key.PriKey)
+	assert.NoError(t, err)
+
+	// pubKey = p-1 should be accepted
+	_, err = ComputeKey(new(big.Int).Sub(p, big.NewInt(1)), key.PriKey)
+	assert.NoError(t, err)
+}

core/codec/gzip.go

@@ -8,6 +8,7 @@ import (
 
 const unzipLimit = 100 * 1024 * 1024 // 100MB
 
+// Gzip compresses bs.
 func Gzip(bs []byte) []byte {
 	var b bytes.Buffer
 
@@ -18,6 +19,7 @@ func Gzip(bs []byte) []byte {
 	return b.Bytes()
 }
 
+// Gunzip uncompresses bs.
 func Gunzip(bs []byte) ([]byte, error) {
 	r, err := gzip.NewReader(bytes.NewBuffer(bs))
 	if err != nil {

core/codec/gzip_test.go

@@ -2,6 +2,8 @@ package codec
 
 import (
 	"bytes"
+	"compress/gzip"
+	"errors"
 	"fmt"
 	"testing"
 
@@ -21,3 +23,45 @@ func TestGzip(t *testing.T) {
 	assert.True(t, len(bs) < buf.Len())
 	assert.Equal(t, buf.Bytes(), actual)
 }
+
+func TestGunzip(t *testing.T) {
+	tests := []struct {
+		name        string
+		input       []byte
+		expected    []byte
+		expectedErr error
+	}{
+		{
+			name: "valid input",
+			input: func() []byte {
+				var buf bytes.Buffer
+				gz := gzip.NewWriter(&buf)
+				gz.Write([]byte("hello"))
+				gz.Close()
+				return buf.Bytes()
+			}(),
+			expected:    []byte("hello"),
+			expectedErr: nil,
+		},
+		{
+			name:        "invalid input",
+			input:       []byte("invalid input"),
+			expected:    nil,
+			expectedErr: gzip.ErrHeader,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			result, err := Gunzip(test.input)
+
+			if !bytes.Equal(result, test.expected) {
+				t.Errorf("unexpected result: %v", result)
+			}
+
+			if !errors.Is(err, test.expectedErr) {
+				t.Errorf("unexpected error: %v", err)
+			}
+		})
+	}
+}

core/codec/hmac.go

@@ -7,12 +7,14 @@ import (
 	"io"
 )
 
+// Hmac returns HMAC bytes for body with the given key.
 func Hmac(key []byte, body string) []byte {
 	h := hmac.New(sha256.New, key)
 	io.WriteString(h, body)
 	return h.Sum(nil)
 }
 
+// HmacBase64 returns the base64 encoded string of HMAC for body with the given key.
 func HmacBase64(key []byte, body string) string {
 	return base64.StdEncoding.EncodeToString(Hmac(key, body))
 }

core/codec/rsa.go

@@ -3,25 +3,31 @@ package codec
 import (
 	"crypto/rand"
 	"crypto/rsa"
+	"crypto/sha256"
 	"crypto/x509"
 	"encoding/base64"
 	"encoding/pem"
 	"errors"
-	"io/ioutil"
+	"os"
 )
 
 var (
+	// ErrPrivateKey indicates the invalid private key.
 	ErrPrivateKey = errors.New("private key error")
-	ErrPublicKey  = errors.New("failed to parse PEM block containing the public key")
-	ErrNotRsaKey  = errors.New("key type is not RSA")
+	// ErrPublicKey indicates the invalid public key.
+	ErrPublicKey = errors.New("failed to parse PEM block containing the public key")
+	// ErrNotRsaKey indicates the invalid RSA key.
+	ErrNotRsaKey = errors.New("key type is not RSA")
 )
 
 type (
+	// RsaDecrypter represents a RSA decrypter.
 	RsaDecrypter interface {
 		Decrypt(input []byte) ([]byte, error)
 		DecryptBase64(input string) ([]byte, error)
 	}
 
+	// RsaEncrypter represents a RSA encrypter.
 	RsaEncrypter interface {
 		Encrypt(input []byte) ([]byte, error)
 	}
@@ -41,8 +47,11 @@ type (
 	}
 )
 
+// Deprecated: NewRsaDecrypter returns a RsaDecrypter with the given file.
+// PKCS#1 v1.5 padding is vulnerable to padding oracle attacks.
+// Use NewRsaOAEPDecrypter instead.
 func NewRsaDecrypter(file string) (RsaDecrypter, error) {
-	content, err := ioutil.ReadFile(file)
+	content, err := os.ReadFile(file)
 	if err != nil {
 		return nil, err
 	}
@@ -84,6 +93,9 @@ func (r *rsaDecrypter) DecryptBase64(input string) ([]byte, error) {
 	return r.Decrypt(base64Decoded)
 }
 
+// Deprecated: NewRsaEncrypter returns a RsaEncrypter with the given key.
+// PKCS#1 v1.5 padding is vulnerable to padding oracle attacks.
+// Use NewRsaOAEPEncrypter instead.
 func NewRsaEncrypter(key []byte) (RsaEncrypter, error) {
 	block, _ := pem.Decode(key)
 	if block == nil {
@@ -147,3 +159,90 @@ func rsaDecryptBlock(privateKey *rsa.PrivateKey, block []byte) ([]byte, error) {
 func rsaEncryptBlock(publicKey *rsa.PublicKey, msg []byte) ([]byte, error) {
 	return rsa.EncryptPKCS1v15(rand.Reader, publicKey, msg)
 }
+
+// NewRsaOAEPDecrypter returns a RsaDecrypter using OAEP with SHA-256.
+func NewRsaOAEPDecrypter(file string) (RsaDecrypter, error) {
+	content, err := os.ReadFile(file)
+	if err != nil {
+		return nil, err
+	}
+
+	block, _ := pem.Decode(content)
+	if block == nil {
+		return nil, ErrPrivateKey
+	}
+
+	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
+	if err != nil {
+		return nil, err
+	}
+
+	return &rsaOAEPDecrypter{
+		rsaBase: rsaBase{
+			bytesLimit: privateKey.N.BitLen() >> 3,
+		},
+		privateKey: privateKey,
+	}, nil
+}
+
+// NewRsaOAEPEncrypter returns a RsaEncrypter using OAEP with SHA-256.
+func NewRsaOAEPEncrypter(key []byte) (RsaEncrypter, error) {
+	block, _ := pem.Decode(key)
+	if block == nil {
+		return nil, ErrPublicKey
+	}
+
+	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
+	if err != nil {
+		return nil, err
+	}
+
+	switch pubKey := pub.(type) {
+	case *rsa.PublicKey:
+		// OAEP overhead: 2*hash_size + 2
+		hashSize := sha256.New().Size()
+		return &rsaOAEPEncrypter{
+			rsaBase: rsaBase{
+				bytesLimit: (pubKey.N.BitLen() >> 3) - 2*hashSize - 2,
+			},
+			publicKey: pubKey,
+		}, nil
+	default:
+		return nil, ErrNotRsaKey
+	}
+}
+
+type rsaOAEPDecrypter struct {
+	rsaBase
+	privateKey *rsa.PrivateKey
+}
+
+func (r *rsaOAEPDecrypter) Decrypt(input []byte) ([]byte, error) {
+	return r.crypt(input, func(block []byte) ([]byte, error) {
+		return rsa.DecryptOAEP(sha256.New(), rand.Reader, r.privateKey, block, nil)
+	})
+}
+
+func (r *rsaOAEPDecrypter) DecryptBase64(input string) ([]byte, error) {
+	if len(input) == 0 {
+		return nil, nil
+	}
+
+	base64Decoded, err := base64.StdEncoding.DecodeString(input)
+	if err != nil {
+		return nil, err
+	}
+
+	return r.Decrypt(base64Decoded)
+}
+
+type rsaOAEPEncrypter struct {
+	rsaBase
+	publicKey *rsa.PublicKey
+}
+
+func (r *rsaOAEPEncrypter) Encrypt(input []byte) ([]byte, error) {
+	return r.crypt(input, func(block []byte) ([]byte, error) {
+		return rsa.EncryptOAEP(sha256.New(), rand.Reader, r.publicKey, block, nil)
+	})
+}

core/codec/rsa_test.go

@@ -1,11 +1,17 @@
 package codec
 
 import (
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/x509"
 	"encoding/base64"
+	"encoding/pem"
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/fs"
+	"github.com/zeromicro/go-zero/core/fs"
 )
 
 const (
@@ -41,6 +47,7 @@ func TestCryption(t *testing.T) {
 
 	file, err := fs.TempFilenameWithText(priKey)
 	assert.Nil(t, err)
+	defer os.Remove(file)
 	dec, err := NewRsaDecrypter(file)
 	assert.Nil(t, err)
 	actual, err := dec.Decrypt(ret)
@@ -56,3 +63,78 @@ func TestBadPubKey(t *testing.T) {
 	_, err := NewRsaEncrypter([]byte("foo"))
 	assert.Equal(t, ErrPublicKey, err)
 }
+
+func TestOAEPCryption(t *testing.T) {
+	enc, err := NewRsaOAEPEncrypter([]byte(pubKey))
+	assert.Nil(t, err)
+	ret, err := enc.Encrypt([]byte(testBody))
+	assert.Nil(t, err)
+
+	file, err := fs.TempFilenameWithText(priKey)
+	assert.Nil(t, err)
+	defer os.Remove(file)
+	dec, err := NewRsaOAEPDecrypter(file)
+	assert.Nil(t, err)
+	actual, err := dec.Decrypt(ret)
+	assert.Nil(t, err)
+	assert.Equal(t, testBody, string(actual))
+
+	actual, err = dec.DecryptBase64(base64.StdEncoding.EncodeToString(ret))
+	assert.Nil(t, err)
+	assert.Equal(t, testBody, string(actual))
+
+	// empty input
+	actual, err = dec.DecryptBase64("")
+	assert.Nil(t, err)
+	assert.Nil(t, actual)
+}
+
+func TestOAEPBadKeys(t *testing.T) {
+	_, err := NewRsaOAEPEncrypter([]byte("bad"))
+	assert.Equal(t, ErrPublicKey, err)
+
+	_, err = NewRsaOAEPDecrypter("nonexistent")
+	assert.Error(t, err)
+
+	// valid PEM but invalid private key content
+	badPem, err := fs.TempFilenameWithText("-----BEGIN RSA PRIVATE KEY-----\nYmFk\n-----END RSA PRIVATE KEY-----")
+	assert.Nil(t, err)
+	defer os.Remove(badPem)
+	_, err = NewRsaOAEPDecrypter(badPem)
+	assert.Error(t, err)
+
+	// not PEM content at all
+	notPem, err := fs.TempFilenameWithText("not a pem file")
+	assert.Nil(t, err)
+	defer os.Remove(notPem)
+	_, err = NewRsaOAEPDecrypter(notPem)
+	assert.Equal(t, ErrPrivateKey, err)
+}
+
+func TestOAEPEncrypterParseError(t *testing.T) {
+	// valid PEM block but invalid public key content
+	badPub := []byte("-----BEGIN PUBLIC KEY-----\nYmFk\n-----END PUBLIC KEY-----")
+	_, err := NewRsaOAEPEncrypter(badPub)
+	assert.Error(t, err)
+}
+
+func TestOAEPEncrypterNonRsaKey(t *testing.T) {
+	ecKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	assert.Nil(t, err)
+	derBytes, err := x509.MarshalPKIXPublicKey(&ecKey.PublicKey)
+	assert.Nil(t, err)
+	ecPem := pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: derBytes})
+	_, err = NewRsaOAEPEncrypter(ecPem)
+	assert.Equal(t, ErrNotRsaKey, err)
+}
+
+func TestOAEPDecryptBase64Error(t *testing.T) {
+	file, err := fs.TempFilenameWithText(priKey)
+	assert.Nil(t, err)
+	defer os.Remove(file)
+	dec, err := NewRsaOAEPDecrypter(file)
+	assert.Nil(t, err)
+
+	_, err = dec.DecryptBase64("not-valid-base64!!!")
+	assert.Error(t, err)
+}

core/collection/cache.go

@@ -6,9 +6,9 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/tal-tech/go-zero/core/logx"
-	"github.com/tal-tech/go-zero/core/mathx"
-	"github.com/tal-tech/go-zero/core/syncx"
+	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/mathx"
+	"github.com/zeromicro/go-zero/core/syncx"
 )
 
 const (
@@ -23,27 +23,30 @@ const (
 var emptyLruCache = emptyLru{}
 
 type (
+	// CacheOption defines the method to customize a Cache.
 	CacheOption func(cache *Cache)
 
+	// A Cache object is an in-memory cache.
 	Cache struct {
 		name           string
 		lock           sync.Mutex
-		data           map[string]interface{}
+		data           map[string]any
 		expire         time.Duration
 		timingWheel    *TimingWheel
 		lruCache       lru
-		barrier        syncx.SharedCalls
+		barrier        syncx.SingleFlight
 		unstableExpiry mathx.Unstable
 		stats          *cacheStat
 	}
 )
 
+// NewCache returns a Cache with given expire.
 func NewCache(expire time.Duration, opts ...CacheOption) (*Cache, error) {
 	cache := &Cache{
-		data:           make(map[string]interface{}),
+		data:           make(map[string]any),
 		expire:         expire,
 		lruCache:       emptyLruCache,
-		barrier:        syncx.NewSharedCalls(),
+		barrier:        syncx.NewSingleFlight(),
 		unstableExpiry: mathx.NewUnstable(expiryDeviation),
 	}
 
@@ -56,7 +59,7 @@ func NewCache(expire time.Duration, opts ...CacheOption) (*Cache, error) {
 	}
 	cache.stats = newCacheStat(cache.name, cache.size)
 
-	timingWheel, err := NewTimingWheel(time.Second, slots, func(k, v interface{}) {
+	timingWheel, err := NewTimingWheel(time.Second, slots, func(k, v any) {
 		key, ok := k.(string)
 		if !ok {
 			return
@@ -72,15 +75,21 @@ func NewCache(expire time.Duration, opts ...CacheOption) (*Cache, error) {
 	return cache, nil
 }
 
+// Del deletes the item with the given key from c.
 func (c *Cache) Del(key string) {
 	c.lock.Lock()
 	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)
 }
 
-func (c *Cache) Get(key string) (interface{}, bool) {
+// Get returns the item with the given key from c.
+func (c *Cache) Get(key string) (any, bool) {
 	value, ok := c.doGet(key)
 	if ok {
 		c.stats.IncrementHit()
@@ -91,14 +100,20 @@ func (c *Cache) Get(key string) (interface{}, bool) {
 	return value, ok
 }
 
-func (c *Cache) Set(key string, value interface{}) {
+// Set sets value into c with key.
+func (c *Cache) Set(key string, value any) {
+	c.SetWithExpire(key, value, c.expire)
+}
+
+// SetWithExpire sets value into c with key and expire with the given value.
+func (c *Cache) SetWithExpire(key string, value any, expire time.Duration) {
 	c.lock.Lock()
 	_, ok := c.data[key]
 	c.data[key] = value
 	c.lruCache.add(key)
 	c.lock.Unlock()
 
-	expiry := c.unstableExpiry.AroundDuration(c.expire)
+	expiry := c.unstableExpiry.AroundDuration(expire)
 	if ok {
 		c.timingWheel.MoveTimer(key, expiry)
 	} else {
@@ -106,16 +121,19 @@ func (c *Cache) Set(key string, value interface{}) {
 	}
 }
 
-func (c *Cache) Take(key string, fetch func() (interface{}, error)) (interface{}, error) {
+// Take returns the item with the given key.
+// If the item is in c, return it directly.
+// If not, use fetch method to get the item, set into c and return it.
+func (c *Cache) Take(key string, fetch func() (any, error)) (any, error) {
 	if val, ok := c.doGet(key); ok {
 		c.stats.IncrementHit()
 		return val, nil
 	}
 
 	var fresh bool
-	val, err := c.barrier.Do(key, func() (interface{}, error) {
-		// because O(1) on map search in memory, and fetch is an IO query
-		// so we do double check, cache might be taken by another call
+	val, err := c.barrier.Do(key, func() (any, error) {
+		// because O(1) on map search in memory, and fetch is an IO query,
+		// so we do double-check, cache might be taken by another call
 		if val, ok := c.doGet(key); ok {
 			return val, nil
 		}
@@ -136,15 +154,14 @@ func (c *Cache) Take(key string, fetch func() (interface{}, error)) (interface{}
 	if fresh {
 		c.stats.IncrementMiss()
 		return val, nil
-	} else {
-		// got the result from previous ongoing query
-		c.stats.IncrementHit()
 	}
 
+	// got the result from previous ongoing query
+	c.stats.IncrementHit()
 	return val, nil
 }
 
-func (c *Cache) doGet(key string) (interface{}, bool) {
+func (c *Cache) doGet(key string) (any, bool) {
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
@@ -168,6 +185,7 @@ func (c *Cache) size() int {
 	return len(c.data)
 }
 
+// WithLimit customizes a Cache with items up to limit.
 func WithLimit(limit int) CacheOption {
 	return func(cache *Cache) {
 		if limit > 0 {
@@ -176,6 +194,7 @@ func WithLimit(limit int) CacheOption {
 	}
 }
 
+// WithName customizes a Cache with the given name.
 func WithName(name string) CacheOption {
 	return func(cache *Cache) {
 		cache.name = name

core/collection/cache_test.go

@@ -18,7 +18,7 @@ func TestCacheSet(t *testing.T) {
 	assert.Nil(t, err)
 
 	cache.Set("first", "first element")
-	cache.Set("second", "second element")
+	cache.SetWithExpire("second", "second element", time.Second*3)
 
 	value, ok := cache.Get("first")
 	assert.True(t, ok)
@@ -52,7 +52,7 @@ func TestCacheTake(t *testing.T) {
 	for i := 0; i < 100; i++ {
 		wg.Add(1)
 		go func() {
-			cache.Take("first", func() (interface{}, error) {
+			cache.Take("first", func() (any, error) {
 				atomic.AddInt32(&count, 1)
 				time.Sleep(time.Millisecond * 100)
 				return "first element", nil
@@ -76,7 +76,7 @@ func TestCacheTakeExists(t *testing.T) {
 		wg.Add(1)
 		go func() {
 			cache.Set("first", "first element")
-			cache.Take("first", func() (interface{}, error) {
+			cache.Take("first", func() (any, error) {
 				atomic.AddInt32(&count, 1)
 				time.Sleep(time.Millisecond * 100)
 				return "first element", nil
@@ -99,7 +99,7 @@ func TestCacheTakeError(t *testing.T) {
 	for i := 0; i < 100; i++ {
 		wg.Add(1)
 		go func() {
-			_, err := cache.Take("first", func() (interface{}, error) {
+			_, err := cache.Take("first", func() (any, error) {
 				atomic.AddInt32(&count, 1)
 				time.Sleep(time.Millisecond * 100)
 				return "", errDummy

core/collection/fifo.go

@@ -2,22 +2,29 @@ package collection
 
 import "sync"
 
+const queueGrowThreshold = 256
+
+// A Queue is a FIFO queue.
 type Queue struct {
 	lock     sync.Mutex
-	elements []interface{}
-	size     int
+	elements []any
 	head     int
 	tail     int
 	count    int
 }
 
+// NewQueue returns a Queue object.
 func NewQueue(size int) *Queue {
+	if size < 1 {
+		panic("size must be greater than 0")
+	}
+
 	return &Queue{
-		elements: make([]interface{}, size),
-		size:     size,
+		elements: make([]any, size),
 	}
 }
 
+// Empty checks if q is empty.
 func (q *Queue) Empty() bool {
 	q.lock.Lock()
 	empty := q.count == 0
@@ -26,16 +33,17 @@ func (q *Queue) Empty() bool {
 	return empty
 }
 
-func (q *Queue) Put(element interface{}) {
+// Put puts element into q at the last position.
+func (q *Queue) Put(element any) {
 	q.lock.Lock()
 	defer q.lock.Unlock()
 
-	if q.head == q.tail && q.count > 0 {
-		nodes := make([]interface{}, len(q.elements)+q.size)
-		copy(nodes, q.elements[q.head:])
-		copy(nodes[len(q.elements)-q.head:], q.elements[:q.head])
+	if q.count == len(q.elements) {
+		nodes := make([]any, nextQueueCapacity(len(q.elements)))
+		n := copy(nodes, q.elements[q.head:])
+		copy(nodes[n:], q.elements[:q.head])
 		q.head = 0
-		q.tail = len(q.elements)
+		q.tail = q.count
 		q.elements = nodes
 	}
 
@@ -44,7 +52,8 @@ func (q *Queue) Put(element interface{}) {
 	q.count++
 }
 
-func (q *Queue) Take() (interface{}, bool) {
+// Take takes the first element out of q if not empty.
+func (q *Queue) Take() (any, bool) {
 	q.lock.Lock()
 	defer q.lock.Unlock()
 
@@ -53,8 +62,19 @@ func (q *Queue) Take() (interface{}, bool) {
 	}
 
 	element := q.elements[q.head]
+	q.elements[q.head] = nil
 	q.head = (q.head + 1) % len(q.elements)
 	q.count--
 
 	return element, true
 }
+
+func nextQueueCapacity(capacity int) int {
+	if capacity < queueGrowThreshold {
+		return capacity << 1
+	}
+
+	// Use a growth curve similar to Go slices: double small queues, then
+	// transition smoothly toward 1.25x growth for larger queues.
+	return capacity + ((capacity + 3*queueGrowThreshold) >> 2)
+}

core/collection/fifo_test.go

@@ -6,58 +6,263 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestFifo(t *testing.T) {
-	elements := [][]byte{
-		[]byte("hello"),
-		[]byte("world"),
-		[]byte("again"),
-	}
-	queue := NewQueue(8)
-	for i := range elements {
-		queue.Put(elements[i])
+func TestQueueOrder(t *testing.T) {
+	tests := []struct {
+		name         string
+		size         int
+		initial      []int
+		takeBefore   int
+		additional   []int
+		wantCapacity int
+	}{
+		{
+			name:         "within initial capacity",
+			size:         8,
+			initial:      []int{1, 2, 3},
+			wantCapacity: 8,
+		},
+		{
+			name:         "grow from beginning",
+			size:         2,
+			initial:      []int{1, 2, 3},
+			wantCapacity: 4,
+		},
+		{
+			name:         "grow after wrapping",
+			size:         4,
+			initial:      []int{1, 2, 3, 4},
+			takeBefore:   1,
+			additional:   []int{5, 6},
+			wantCapacity: 8,
+		},
+		{
+			name:         "grow repeatedly",
+			size:         1,
+			initial:      sequence(20),
+			wantCapacity: 32,
+		},
+		{
+			name:         "grow above threshold",
+			size:         queueGrowThreshold,
+			initial:      sequence(queueGrowThreshold + 1),
+			wantCapacity: queueGrowThreshold * 2,
+		},
+		{
+			name:         "grow well above threshold",
+			size:         1024,
+			initial:      sequence(1025),
+			wantCapacity: 1472,
+		},
 	}
 
-	for _, element := range elements {
-		body, ok := queue.Take()
-		assert.True(t, ok)
-		assert.Equal(t, string(element), string(body.([]byte)))
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			queue := NewQueue(test.size)
+			assert.True(t, queue.Empty())
+
+			for _, value := range test.initial {
+				queue.Put(value)
+			}
+			for i := 0; i < test.takeBefore; i++ {
+				value, ok := queue.Take()
+				assert.True(t, ok)
+				assert.Equal(t, test.initial[i], value)
+			}
+			for _, value := range test.additional {
+				queue.Put(value)
+			}
+
+			want := append([]int(nil), test.initial[test.takeBefore:]...)
+			want = append(want, test.additional...)
+			for _, expected := range want {
+				actual, ok := queue.Take()
+				assert.True(t, ok)
+				assert.Equal(t, expected, actual)
+			}
+
+			assert.Equal(t, test.wantCapacity, len(queue.elements))
+			assert.True(t, queue.Empty())
+			_, ok := queue.Take()
+			assert.False(t, ok)
+		})
 	}
 }
 
-func TestTakeTooMany(t *testing.T) {
-	elements := [][]byte{
-		[]byte("hello"),
-		[]byte("world"),
-		[]byte("again"),
-	}
-	queue := NewQueue(8)
-	for i := range elements {
-		queue.Put(elements[i])
+func TestQueueTakeClearsElement(t *testing.T) {
+	tests := []struct {
+		name       string
+		size       int
+		operations string
+	}{
+		{
+			name:       "take from beginning",
+			size:       2,
+			operations: "ppt",
+		},
+		{
+			name:       "take after wrapping",
+			size:       2,
+			operations: "pptptt",
+		},
+		{
+			name:       "take after growing",
+			size:       2,
+			operations: "pppt",
+		},
 	}
 
-	for range elements {
-		queue.Take()
-	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			queue := NewQueue(test.size)
+			value := 0
 
-	assert.True(t, queue.Empty())
-	_, ok := queue.Take()
-	assert.False(t, ok)
+			for _, operation := range test.operations {
+				switch operation {
+				case 'p':
+					value++
+					element := new(int)
+					*element = value
+					queue.Put(element)
+				case 't':
+					index := queue.head
+					_, ok := queue.Take()
+					assert.True(t, ok)
+					assert.Nil(t, queue.elements[index])
+				default:
+					t.Fatalf("unknown operation: %q", operation)
+				}
+			}
+		})
+	}
 }
 
-func TestPutMore(t *testing.T) {
-	elements := [][]byte{
-		[]byte("hello"),
-		[]byte("world"),
-		[]byte("again"),
-	}
-	queue := NewQueue(2)
-	for i := range elements {
-		queue.Put(elements[i])
+func TestNewQueueWithInvalidSize(t *testing.T) {
+	tests := []struct {
+		name string
+		size int
+	}{
+		{
+			name: "zero",
+		},
+		{
+			name: "negative",
+			size: -1,
+		},
 	}
 
-	for _, element := range elements {
-		body, ok := queue.Take()
-		assert.True(t, ok)
-		assert.Equal(t, string(element), string(body.([]byte)))
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			assert.Panics(t, func() {
+				NewQueue(test.size)
+			})
+		})
 	}
 }
+
+func BenchmarkQueueGrowth(b *testing.B) {
+	tests := []struct {
+		name  string
+		size  int
+		count int
+	}{
+		{
+			name:  "initial_1_count_256",
+			size:  1,
+			count: 256,
+		},
+		{
+			name:  "initial_1_count_4096",
+			size:  1,
+			count: 4096,
+		},
+		{
+			name:  "initial_1_count_65536",
+			size:  1,
+			count: 65536,
+		},
+		{
+			name:  "initial_8_count_4096",
+			size:  8,
+			count: 4096,
+		},
+		{
+			name:  "initial_256_count_4096",
+			size:  256,
+			count: 4096,
+		},
+	}
+
+	for _, test := range tests {
+		b.Run(test.name, func(b *testing.B) {
+			elements := make([]any, test.count)
+			for i := range elements {
+				elements[i] = i
+			}
+
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				queue := NewQueue(test.size)
+				for _, element := range elements {
+					queue.Put(element)
+				}
+			}
+		})
+	}
+}
+
+func BenchmarkQueueWrappedGrowth(b *testing.B) {
+	tests := []struct {
+		name string
+		size int
+	}{
+		{
+			name: "capacity_8",
+			size: 8,
+		},
+		{
+			name: "capacity_256",
+			size: 256,
+		},
+		{
+			name: "capacity_4096",
+			size: 4096,
+		},
+	}
+
+	for _, test := range tests {
+		b.Run(test.name, func(b *testing.B) {
+			elements := make([]any, test.size)
+			for i := range elements {
+				elements[i] = i
+			}
+
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				queue := NewQueue(test.size)
+				for _, element := range elements {
+					queue.Put(element)
+				}
+				for j := 0; j < test.size/2; j++ {
+					queue.Take()
+				}
+				for j := 0; j < test.size/2; j++ {
+					queue.Put(elements[j])
+				}
+
+				// The queue is full and wrapped. One more Put triggers growth
+				// and copies both sides of the ring into FIFO order.
+				queue.Put(elements[0])
+			}
+		})
+	}
+}
+
+func sequence(size int) []int {
+	values := make([]int, size)
+	for i := range values {
+		values[i] = i
+	}
+	return values
+}

core/collection/ring.go

@@ -1,38 +1,59 @@
 package collection
 
+import "sync"
+
+// A Ring can be used as fixed size ring.
 type Ring struct {
-	elements []interface{}
+	elements []any
 	index    int
+	lock     sync.RWMutex
 }
 
+// NewRing returns a Ring object with the given size n.
 func NewRing(n int) *Ring {
 	if n < 1 {
 		panic("n should be greater than 0")
 	}
 
 	return &Ring{
-		elements: make([]interface{}, n),
+		elements: make([]any, n),
 	}
 }
 
-func (r *Ring) Add(v interface{}) {
-	r.elements[r.index%len(r.elements)] = v
+// Add adds v into r.
+func (r *Ring) Add(v any) {
+	r.lock.Lock()
+	defer r.lock.Unlock()
+
+	rlen := len(r.elements)
+	r.elements[r.index%rlen] = v
 	r.index++
+
+	// prevent ring index overflow
+	if r.index >= rlen<<1 {
+		r.index -= rlen
+	}
 }
 
-func (r *Ring) Take() []interface{} {
+// Take takes all items from r.
+func (r *Ring) Take() []any {
+	r.lock.RLock()
+	defer r.lock.RUnlock()
+
 	var size int
 	var start int
-	if r.index > len(r.elements) {
-		size = len(r.elements)
-		start = r.index % len(r.elements)
+	rlen := len(r.elements)
+
+	if r.index > rlen {
+		size = rlen
+		start = r.index % rlen
 	} else {
 		size = r.index
 	}
 
-	elements := make([]interface{}, size)
+	elements := make([]any, size)
 	for i := 0; i < size; i++ {
-		elements[i] = r.elements[(start+i)%len(r.elements)]
+		elements[i] = r.elements[(start+i)%rlen]
 	}
 
 	return elements

core/collection/ring_test.go

@@ -1,6 +1,7 @@
 package collection
 
 import (
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -18,7 +19,7 @@ func TestRingLess(t *testing.T) {
 		ring.Add(i)
 	}
 	elements := ring.Take()
-	assert.ElementsMatch(t, []interface{}{0, 1, 2}, elements)
+	assert.ElementsMatch(t, []any{0, 1, 2}, elements)
 }
 
 func TestRingMore(t *testing.T) {
@@ -27,5 +28,32 @@ func TestRingMore(t *testing.T) {
 		ring.Add(i)
 	}
 	elements := ring.Take()
-	assert.ElementsMatch(t, []interface{}{6, 7, 8, 9, 10}, elements)
+	assert.ElementsMatch(t, []any{6, 7, 8, 9, 10}, elements)
+}
+
+func TestRingAdd(t *testing.T) {
+	ring := NewRing(5051)
+	wg := sync.WaitGroup{}
+	for i := 1; i <= 100; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			for j := 1; j <= i; j++ {
+				ring.Add(i)
+			}
+		}(i)
+	}
+	wg.Wait()
+	assert.Equal(t, 5050, len(ring.Take()))
+}
+
+func BenchmarkRingAdd(b *testing.B) {
+	ring := NewRing(500)
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			for i := 0; i < b.N; i++ {
+				ring.Add(i)
+			}
+		}
+	})
 }

core/collection/rollingwindow.go

@@ -4,31 +4,46 @@ import (
 	"sync"
 	"time"
 
-	"github.com/tal-tech/go-zero/core/timex"
+	"github.com/zeromicro/go-zero/core/mathx"
+	"github.com/zeromicro/go-zero/core/timex"
 )
 
 type (
-	RollingWindowOption func(rollingWindow *RollingWindow)
+	// BucketInterface is the interface that defines the buckets.
+	BucketInterface[T Numerical] interface {
+		Add(v T)
+		Reset()
+	}
 
-	RollingWindow struct {
+	// Numerical is the interface that restricts the numerical type.
+	Numerical = mathx.Numerical
+
+	// RollingWindowOption let callers customize the RollingWindow.
+	RollingWindowOption[T Numerical, B BucketInterface[T]] func(rollingWindow *RollingWindow[T, B])
+
+	// RollingWindow defines a rolling window to calculate the events in buckets with the time interval.
+	RollingWindow[T Numerical, B BucketInterface[T]] struct {
 		lock          sync.RWMutex
 		size          int
-		win           *window
+		win           *window[T, B]
 		interval      time.Duration
 		offset        int
 		ignoreCurrent bool
-		lastTime      time.Duration
+		lastTime      time.Duration // start time of the last bucket
 	}
 )
 
-func NewRollingWindow(size int, interval time.Duration, opts ...RollingWindowOption) *RollingWindow {
+// NewRollingWindow returns a RollingWindow that with size buckets and time interval,
+// use opts to customize the RollingWindow.
+func NewRollingWindow[T Numerical, B BucketInterface[T]](newBucket func() B, size int,
+	interval time.Duration, opts ...RollingWindowOption[T, B]) *RollingWindow[T, B] {
 	if size < 1 {
 		panic("size must be greater than 0")
 	}
 
-	w := &RollingWindow{
+	w := &RollingWindow[T, B]{
 		size:     size,
-		win:      newWindow(size),
+		win:      newWindow[T, B](newBucket, size),
 		interval: interval,
 		lastTime: timex.Now(),
 	}
@@ -38,20 +53,22 @@ func NewRollingWindow(size int, interval time.Duration, opts ...RollingWindowOpt
 	return w
 }
 
-func (rw *RollingWindow) Add(v float64) {
+// Add adds value to current bucket.
+func (rw *RollingWindow[T, B]) Add(v T) {
 	rw.lock.Lock()
 	defer rw.lock.Unlock()
 	rw.updateOffset()
 	rw.win.add(rw.offset, v)
 }
 
-func (rw *RollingWindow) Reduce(fn func(b *Bucket)) {
+// Reduce runs fn on all buckets, ignore current bucket if ignoreCurrent was set.
+func (rw *RollingWindow[T, B]) Reduce(fn func(b B)) {
 	rw.lock.RLock()
 	defer rw.lock.RUnlock()
 
 	var diff int
 	span := rw.span()
-	// ignore current bucket, because of partial data
+	// ignore the current bucket, because of partial data
 	if span == 0 && rw.ignoreCurrent {
 		diff = rw.size - 1
 	} else {
@@ -63,87 +80,82 @@ func (rw *RollingWindow) Reduce(fn func(b *Bucket)) {
 	}
 }
 
-func (rw *RollingWindow) span() int {
+func (rw *RollingWindow[T, B]) span() int {
 	offset := int(timex.Since(rw.lastTime) / rw.interval)
 	if 0 <= offset && offset < rw.size {
 		return offset
-	} else {
-		return rw.size
 	}
+
+	return rw.size
 }
 
-func (rw *RollingWindow) updateOffset() {
+func (rw *RollingWindow[T, B]) updateOffset() {
 	span := rw.span()
-	if span > 0 {
-		offset := rw.offset
-		// reset expired buckets
-		start := offset + 1
-		steps := start + span
-		var remainder int
-		if steps > rw.size {
-			remainder = steps - rw.size
-			steps = rw.size
-		}
-		for i := start; i < steps; i++ {
-			rw.win.resetBucket(i)
-			offset = i
-		}
-		for i := 0; i < remainder; i++ {
-			rw.win.resetBucket(i)
-			offset = i
-		}
-		rw.offset = offset
-		rw.lastTime = timex.Now()
+	if span <= 0 {
+		return
 	}
+
+	offset := rw.offset
+	// reset expired buckets
+	for i := 0; i < span; i++ {
+		rw.win.resetBucket((offset + i + 1) % rw.size)
+	}
+
+	rw.offset = (offset + span) % rw.size
+	now := timex.Now()
+	// align to interval time boundary
+	rw.lastTime = now - (now-rw.lastTime)%rw.interval
 }
 
-type Bucket struct {
-	Sum   float64
+// Bucket defines the bucket that holds sum and num of additions.
+type Bucket[T Numerical] struct {
+	Sum   T
 	Count int64
 }
 
-func (b *Bucket) add(v float64) {
+func (b *Bucket[T]) Add(v T) {
 	b.Sum += v
 	b.Count++
 }
 
-func (b *Bucket) reset() {
+func (b *Bucket[T]) Reset() {
 	b.Sum = 0
 	b.Count = 0
 }
 
-type window struct {
-	buckets []*Bucket
+type window[T Numerical, B BucketInterface[T]] struct {
+	buckets []B
 	size    int
 }
 
-func newWindow(size int) *window {
-	var buckets []*Bucket
+func newWindow[T Numerical, B BucketInterface[T]](newBucket func() B, size int) *window[T, B] {
+	buckets := make([]B, size)
 	for i := 0; i < size; i++ {
-		buckets = append(buckets, new(Bucket))
+		buckets[i] = newBucket()
 	}
-	return &window{
+	return &window[T, B]{
 		buckets: buckets,
 		size:    size,
 	}
 }
 
-func (w *window) add(offset int, v float64) {
-	w.buckets[offset%w.size].add(v)
+func (w *window[T, B]) add(offset int, v T) {
+	w.buckets[offset%w.size].Add(v)
 }
 
-func (w *window) reduce(start, count int, fn func(b *Bucket)) {
+func (w *window[T, B]) reduce(start, count int, fn func(b B)) {
 	for i := 0; i < count; i++ {
-		fn(w.buckets[(start+i)%len(w.buckets)])
+		fn(w.buckets[(start+i)%w.size])
 	}
 }
 
-func (w *window) resetBucket(offset int) {
-	w.buckets[offset].reset()
+func (w *window[T, B]) resetBucket(offset int) {
+	w.buckets[offset%w.size].Reset()
 }
 
-func IgnoreCurrentBucket() RollingWindowOption {
-	return func(w *RollingWindow) {
+// IgnoreCurrentBucket lets the Reduce call ignore current bucket.
+func IgnoreCurrentBucket[T Numerical, B BucketInterface[T]]() RollingWindowOption[T, B] {
+	return func(w *RollingWindow[T, B]) {
 		w.ignoreCurrent = true
 	}
 }

core/collection/rollingwindow_test.go

@@ -6,24 +6,30 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/stringx"
+	"github.com/zeromicro/go-zero/core/stringx"
 )
 
 const duration = time.Millisecond * 50
 
 func TestNewRollingWindow(t *testing.T) {
-	assert.NotNil(t, NewRollingWindow(10, time.Second))
+	assert.NotNil(t, NewRollingWindow[int64, *Bucket[int64]](func() *Bucket[int64] {
+		return new(Bucket[int64])
+	}, 10, time.Second))
 	assert.Panics(t, func() {
-		NewRollingWindow(0, time.Second)
+		NewRollingWindow[int64, *Bucket[int64]](func() *Bucket[int64] {
+			return new(Bucket[int64])
+		}, 0, time.Second)
 	})
 }
 
 func TestRollingWindowAdd(t *testing.T) {
 	const size = 3
-	r := NewRollingWindow(size, duration)
+	r := NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+		return new(Bucket[float64])
+	}, size, duration)
 	listBuckets := func() []float64 {
 		var buckets []float64
-		r.Reduce(func(b *Bucket) {
+		r.Reduce(func(b *Bucket[float64]) {
 			buckets = append(buckets, b.Sum)
 		})
 		return buckets
@@ -47,10 +53,12 @@ func TestRollingWindowAdd(t *testing.T) {
 
 func TestRollingWindowReset(t *testing.T) {
 	const size = 3
-	r := NewRollingWindow(size, duration, IgnoreCurrentBucket())
+	r := NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+		return new(Bucket[float64])
+	}, size, duration, IgnoreCurrentBucket[float64, *Bucket[float64]]())
 	listBuckets := func() []float64 {
 		var buckets []float64
-		r.Reduce(func(b *Bucket) {
+		r.Reduce(func(b *Bucket[float64]) {
 			buckets = append(buckets, b.Sum)
 		})
 		return buckets
@@ -72,15 +80,19 @@ func TestRollingWindowReset(t *testing.T) {
 func TestRollingWindowReduce(t *testing.T) {
 	const size = 4
 	tests := []struct {
-		win    *RollingWindow
+		win    *RollingWindow[float64, *Bucket[float64]]
 		expect float64
 	}{
 		{
-			win:    NewRollingWindow(size, duration),
+			win: NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+				return new(Bucket[float64])
+			}, size, duration),
 			expect: 10,
 		},
 		{
-			win:    NewRollingWindow(size, duration, IgnoreCurrentBucket()),
+			win: NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+				return new(Bucket[float64])
+			}, size, duration, IgnoreCurrentBucket[float64, *Bucket[float64]]()),
 			expect: 4,
 		},
 	}
@@ -97,7 +109,7 @@ func TestRollingWindowReduce(t *testing.T) {
 				}
 			}
 			var result float64
-			r.Reduce(func(b *Bucket) {
+			r.Reduce(func(b *Bucket[float64]) {
 				result += b.Sum
 			})
 			assert.Equal(t, test.expect, result)
@@ -105,10 +117,45 @@ func TestRollingWindowReduce(t *testing.T) {
 	}
 }
 
+func TestRollingWindowBucketTimeBoundary(t *testing.T) {
+	const size = 3
+	interval := time.Millisecond * 30
+	r := NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+		return new(Bucket[float64])
+	}, size, interval)
+	listBuckets := func() []float64 {
+		var buckets []float64
+		r.Reduce(func(b *Bucket[float64]) {
+			buckets = append(buckets, b.Sum)
+		})
+		return buckets
+	}
+	assert.Equal(t, []float64{0, 0, 0}, listBuckets())
+	r.Add(1)
+	assert.Equal(t, []float64{0, 0, 1}, listBuckets())
+	time.Sleep(time.Millisecond * 45)
+	r.Add(2)
+	r.Add(3)
+	assert.Equal(t, []float64{0, 1, 5}, listBuckets())
+	// sleep time should be less than interval, and make the bucket change happen
+	time.Sleep(time.Millisecond * 20)
+	r.Add(4)
+	r.Add(5)
+	r.Add(6)
+	assert.Equal(t, []float64{1, 5, 15}, listBuckets())
+	time.Sleep(time.Millisecond * 100)
+	r.Add(7)
+	r.Add(8)
+	r.Add(9)
+	assert.Equal(t, []float64{0, 0, 24}, listBuckets())
+}
+
 func TestRollingWindowDataRace(t *testing.T) {
 	const size = 3
-	r := NewRollingWindow(size, duration)
-	var stop = make(chan bool)
+	r := NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+		return new(Bucket[float64])
+	}, size, duration)
+	stop := make(chan bool)
 	go func() {
 		for {
 			select {
@@ -126,7 +173,7 @@ func TestRollingWindowDataRace(t *testing.T) {
 			case <-stop:
 				return
 			default:
-				r.Reduce(func(b *Bucket) {})
+				r.Reduce(func(b *Bucket[float64]) {})
 			}
 		}
 	}()

core/collection/safemap.go

@@ -1,6 +1,9 @@
 package collection
 
-import "sync"
+import (
+	"maps"
+	"sync"
+)
 
 const (
 	copyThreshold = 1000
@@ -14,19 +17,23 @@ type SafeMap struct {
 	lock        sync.RWMutex
 	deletionOld int
 	deletionNew int
-	dirtyOld    map[interface{}]interface{}
-	dirtyNew    map[interface{}]interface{}
+	dirtyOld    map[any]any
+	dirtyNew    map[any]any
 }
 
+// NewSafeMap returns a SafeMap.
 func NewSafeMap() *SafeMap {
 	return &SafeMap{
-		dirtyOld: make(map[interface{}]interface{}),
-		dirtyNew: make(map[interface{}]interface{}),
+		dirtyOld: make(map[any]any),
+		dirtyNew: make(map[any]any),
 	}
 }
 
-func (m *SafeMap) Del(key interface{}) {
+// Del deletes the value with the given key from m.
+func (m *SafeMap) Del(key any) {
 	m.lock.Lock()
+	defer m.lock.Unlock()
+
 	if _, ok := m.dirtyOld[key]; ok {
 		delete(m.dirtyOld, key)
 		m.deletionOld++
@@ -35,38 +42,55 @@ func (m *SafeMap) Del(key interface{}) {
 		m.deletionNew++
 	}
 	if m.deletionOld >= maxDeletion && len(m.dirtyOld) < copyThreshold {
-		for k, v := range m.dirtyOld {
-			m.dirtyNew[k] = v
-		}
+		maps.Copy(m.dirtyNew, m.dirtyOld)
 		m.dirtyOld = m.dirtyNew
 		m.deletionOld = m.deletionNew
-		m.dirtyNew = make(map[interface{}]interface{})
+		m.dirtyNew = make(map[any]any)
 		m.deletionNew = 0
 	}
 	if m.deletionNew >= maxDeletion && len(m.dirtyNew) < copyThreshold {
-		for k, v := range m.dirtyNew {
-			m.dirtyOld[k] = v
-		}
-		m.dirtyNew = make(map[interface{}]interface{})
+		maps.Copy(m.dirtyOld, m.dirtyNew)
+		m.dirtyNew = make(map[any]any)
 		m.deletionNew = 0
 	}
-	m.lock.Unlock()
 }
 
-func (m *SafeMap) Get(key interface{}) (interface{}, bool) {
+// Get gets the value with the given key from m.
+func (m *SafeMap) Get(key any) (any, bool) {
 	m.lock.RLock()
 	defer m.lock.RUnlock()
 
 	if val, ok := m.dirtyOld[key]; ok {
 		return val, true
-	} else {
-		val, ok := m.dirtyNew[key]
-		return val, ok
+	}
+
+	val, ok := m.dirtyNew[key]
+	return val, ok
+}
+
+// Range calls f sequentially for each key and value present in the map.
+// If f returns false, range stops the iteration.
+func (m *SafeMap) Range(f func(key, val any) bool) {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
+
+	for k, v := range m.dirtyOld {
+		if !f(k, v) {
+			return
+		}
+	}
+	for k, v := range m.dirtyNew {
+		if !f(k, v) {
+			return
+		}
 	}
 }
 
-func (m *SafeMap) Set(key, value interface{}) {
+// Set sets the value into m with the given key.
+func (m *SafeMap) Set(key, value any) {
 	m.lock.Lock()
+	defer m.lock.Unlock()
+
 	if m.deletionOld <= maxDeletion {
 		if _, ok := m.dirtyNew[key]; ok {
 			delete(m.dirtyNew, key)
@@ -80,9 +104,9 @@ func (m *SafeMap) Set(key, value interface{}) {
 		}
 		m.dirtyNew[key] = value
 	}
-	m.lock.Unlock()
 }
 
+// Size returns the size of m.
 func (m *SafeMap) Size() int {
 	m.lock.RLock()
 	size := len(m.dirtyOld) + len(m.dirtyNew)

core/collection/safemap_test.go

@@ -1,10 +1,11 @@
 package collection
 
 import (
+	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/stringx"
+	"github.com/zeromicro/go-zero/core/stringx"
 )
 
 func TestSafeMap(t *testing.T) {
@@ -107,3 +108,104 @@ func testSafeMapWithParameters(t *testing.T, size, exception int) {
 		}
 	}
 }
+
+func TestSafeMap_Range(t *testing.T) {
+	const (
+		size       = 100000
+		exception1 = 5
+		exception2 = 500
+	)
+
+	m := NewSafeMap()
+	newMap := NewSafeMap()
+
+	for i := 0; i < size; i++ {
+		m.Set(i, i)
+	}
+	for i := 0; i < size; i++ {
+		if i%exception1 == 0 {
+			m.Del(i)
+		}
+	}
+
+	for i := size; i < size<<1; i++ {
+		m.Set(i, i)
+	}
+	for i := size; i < size<<1; i++ {
+		if i%exception2 != 0 {
+			m.Del(i)
+		}
+	}
+
+	var count int32
+	m.Range(func(k, v any) bool {
+		atomic.AddInt32(&count, 1)
+		newMap.Set(k, v)
+		return true
+	})
+	assert.Equal(t, int(atomic.LoadInt32(&count)), m.Size())
+	assert.Equal(t, m.dirtyNew, newMap.dirtyNew)
+	assert.Equal(t, m.dirtyOld, newMap.dirtyOld)
+}
+
+func TestSetManyTimes(t *testing.T) {
+	const iteration = maxDeletion * 2
+	m := NewSafeMap()
+	for i := 0; i < iteration; i++ {
+		m.Set(i, i)
+		if i%3 == 0 {
+			m.Del(i / 2)
+		}
+	}
+	var count int
+	m.Range(func(k, v any) bool {
+		count++
+		return count < maxDeletion/2
+	})
+	assert.Equal(t, maxDeletion/2, count)
+	for i := 0; i < iteration; i++ {
+		m.Set(i, i)
+		if i%3 == 0 {
+			m.Del(i / 2)
+		}
+	}
+	for i := 0; i < iteration; i++ {
+		m.Set(i, i)
+		if i%3 == 0 {
+			m.Del(i / 2)
+		}
+	}
+	for i := 0; i < iteration; i++ {
+		m.Set(i, i)
+		if i%3 == 0 {
+			m.Del(i / 2)
+		}
+	}
+
+	count = 0
+	m.Range(func(k, v any) bool {
+		count++
+		return count < maxDeletion
+	})
+	assert.Equal(t, maxDeletion, count)
+}
+
+func TestSetManyTimesNew(t *testing.T) {
+	m := NewSafeMap()
+	for i := 0; i < maxDeletion*3; i++ {
+		m.Set(i, i)
+	}
+	for i := 0; i < maxDeletion*2; i++ {
+		m.Del(i)
+	}
+	for i := 0; i < maxDeletion*3; i++ {
+		m.Set(i+maxDeletion*3, i+maxDeletion*3)
+	}
+	for i := 0; i < maxDeletion*2; i++ {
+		m.Del(i + maxDeletion*2)
+	}
+	for i := 0; i < maxDeletion-copyThreshold+1; i++ {
+		m.Del(i + maxDeletion*2)
+	}
+	assert.Equal(t, 0, len(m.dirtyNew))
+}

core/collection/set.go

@@ -1,228 +1,53 @@
 package collection
 
-import (
-	"github.com/tal-tech/go-zero/core/lang"
-	"github.com/tal-tech/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[interface{}]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
 }
 
-func NewSet() *Set {
-	return &Set{
-		data: make(map[interface{}]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),
 	}
 }
 
-func NewUnmanagedSet() *Set {
-	return &Set{
-		data: make(map[interface{}]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
 	}
 }
 
-func (s *Set) Add(i ...interface{}) {
-	for _, each := range i {
-		s.add(each)
-	}
+// Clear removes all items from the set.
+func (s *Set[T]) Clear() {
+	clear(s.data)
 }
 
-func (s *Set) AddInt(ii ...int) {
-	for _, each := range ii {
-		s.add(each)
-	}
-}
-
-func (s *Set) AddInt64(ii ...int64) {
-	for _, each := range ii {
-		s.add(each)
-	}
-}
-
-func (s *Set) AddUint(ii ...uint) {
-	for _, each := range ii {
-		s.add(each)
-	}
-}
-
-func (s *Set) AddUint64(ii ...uint64) {
-	for _, each := range ii {
-		s.add(each)
-	}
-}
-
-func (s *Set) AddStr(ss ...string) {
-	for _, each := range ss {
-		s.add(each)
-	}
-}
-
-func (s *Set) Contains(i interface{}) 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
 }
 
-func (s *Set) Keys() []interface{} {
-	var keys []interface{}
-
-	for key := range s.data {
-		keys = append(keys, key)
-	}
-
-	return keys
-}
-
-func (s *Set) KeysInt() []int {
-	var keys []int
-
-	for key := range s.data {
-		if intKey, ok := key.(int); !ok {
-			continue
-		} else {
-			keys = append(keys, intKey)
-		}
-	}
-
-	return keys
-}
-
-func (s *Set) KeysInt64() []int64 {
-	var keys []int64
-
-	for key := range s.data {
-		if intKey, ok := key.(int64); !ok {
-			continue
-		} else {
-			keys = append(keys, intKey)
-		}
-	}
-
-	return keys
-}
-
-func (s *Set) KeysUint() []uint {
-	var keys []uint
-
-	for key := range s.data {
-		if intKey, ok := key.(uint); !ok {
-			continue
-		} else {
-			keys = append(keys, intKey)
-		}
-	}
-
-	return keys
-}
-
-func (s *Set) KeysUint64() []uint64 {
-	var keys []uint64
-
-	for key := range s.data {
-		if intKey, ok := key.(uint64); !ok {
-			continue
-		} else {
-			keys = append(keys, intKey)
-		}
-	}
-
-	return keys
-}
-
-func (s *Set) KeysStr() []string {
-	var keys []string
-
-	for key := range s.data {
-		if strKey, ok := key.(string); !ok {
-			continue
-		} else {
-			keys = append(keys, strKey)
-		}
-	}
-
-	return keys
-}
-
-func (s *Set) Remove(i interface{}) {
-	s.validate(i)
-	delete(s.data, i)
-}
-
-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 interface{}) {
-	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 interface{}) {
-	// 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 interface{}) {
-	if s.tp == unmanaged {
-		return
-	}
-
-	switch i.(type) {
-	case int:
-		if s.tp != intType {
-			logx.Errorf("Error: element is int, but set contains elements with type %d", s.tp)
-		}
-	case int64:
-		if s.tp != int64Type {
-			logx.Errorf("Error: element is int64, but set contains elements with type %d", s.tp)
-		}
-	case uint:
-		if s.tp != uintType {
-			logx.Errorf("Error: element is uint, but set contains elements with type %d", s.tp)
-		}
-	case uint64:
-		if s.tp != uint64Type {
-			logx.Errorf("Error: element is uint64, but set contains elements with type %d", s.tp)
-		}
-	case string:
-		if s.tp != stringType {
-			logx.Errorf("Error: 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

@@ -5,41 +5,124 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logx"
 )
 
 func init() {
 	logx.Disable()
 }
 
-func BenchmarkRawSet(b *testing.B) {
-	m := make(map[interface{}]struct{})
-	for i := 0; i < b.N; i++ {
-		m[i] = struct{}{}
-		_ = m[i]
-	}
+// 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 BenchmarkUnmanagedSet(b *testing.B) {
-	s := NewUnmanagedSet()
+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 BenchmarkSet(b *testing.B) {
-	s := NewSet()
+func BenchmarkTypedStringSet(b *testing.B) {
+	s := NewSet[string]()
 	for i := 0; i < b.N; i++ {
-		s.AddInt(i)
-		_ = s.Contains(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++ {
+		m[i] = struct{}{}
+		_ = m[i]
 	}
 }
 
 func TestAdd(t *testing.T) {
 	// given
-	set := NewUnmanagedSet()
-	values := []interface{}{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([]interface{}{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([]interface{}{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, []interface{}{1, "2"}, vals)
-}

core/collection/timingwheel.go

@@ -2,19 +2,27 @@ package collection
 
 import (
 	"container/list"
+	"errors"
 	"fmt"
 	"time"
 
-	"github.com/tal-tech/go-zero/core/lang"
-	"github.com/tal-tech/go-zero/core/threading"
-	"github.com/tal-tech/go-zero/core/timex"
+	"github.com/zeromicro/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/threading"
+	"github.com/zeromicro/go-zero/core/timex"
 )
 
 const drainWorkers = 8
 
-type (
-	Execute func(key, value interface{})
+var (
+	ErrClosed   = errors.New("TimingWheel is closed already")
+	ErrArgument = errors.New("incorrect task argument")
+)
 
+type (
+	// Execute defines the method to execute the task.
+	Execute func(key, value any)
+
+	// A TimingWheel is a timing wheel object to schedule tasks.
 	TimingWheel struct {
 		interval      time.Duration
 		ticker        timex.Ticker
@@ -25,14 +33,14 @@ type (
 		execute       Execute
 		setChannel    chan timingEntry
 		moveChannel   chan baseEntry
-		removeChannel chan interface{}
-		drainChannel  chan func(key, value interface{})
+		removeChannel chan any
+		drainChannel  chan func(key, value any)
 		stopChannel   chan lang.PlaceholderType
 	}
 
 	timingEntry struct {
 		baseEntry
-		value   interface{}
+		value   any
 		circle  int
 		diff    int
 		removed bool
@@ -40,7 +48,7 @@ type (
 
 	baseEntry struct {
 		delay time.Duration
-		key   interface{}
+		key   any
 	}
 
 	positionEntry struct {
@@ -49,21 +57,24 @@ type (
 	}
 
 	timingTask struct {
-		key   interface{}
-		value interface{}
+		key   any
+		value any
 	}
 )
 
+// NewTimingWheel returns a TimingWheel.
 func NewTimingWheel(interval time.Duration, numSlots int, execute Execute) (*TimingWheel, error) {
 	if interval <= 0 || numSlots <= 0 || execute == nil {
-		return nil, fmt.Errorf("interval: %v, slots: %d, execute: %p", interval, numSlots, execute)
+		return nil, fmt.Errorf("interval: %v, slots: %d, execute: %p",
+			interval, numSlots, execute)
 	}
 
-	return newTimingWheelWithClock(interval, numSlots, execute, timex.NewTicker(interval))
+	return NewTimingWheelWithTicker(interval, numSlots, execute, timex.NewTicker(interval))
 }
 
-func newTimingWheelWithClock(interval time.Duration, numSlots int, execute Execute, ticker timex.Ticker) (
-	*TimingWheel, error) {
+// NewTimingWheelWithTicker returns a TimingWheel with the given ticker.
+func NewTimingWheelWithTicker(interval time.Duration, numSlots int, execute Execute,
+	ticker timex.Ticker) (*TimingWheel, error) {
 	tw := &TimingWheel{
 		interval:      interval,
 		ticker:        ticker,
@@ -74,8 +85,8 @@ func newTimingWheelWithClock(interval time.Duration, numSlots int, execute Execu
 		numSlots:      numSlots,
 		setChannel:    make(chan timingEntry),
 		moveChannel:   make(chan baseEntry),
-		removeChannel: make(chan interface{}),
-		drainChannel:  make(chan func(key, value interface{})),
+		removeChannel: make(chan any),
+		drainChannel:  make(chan func(key, value any)),
 		stopChannel:   make(chan lang.PlaceholderType),
 	}
 
@@ -85,54 +96,86 @@ func newTimingWheelWithClock(interval time.Duration, numSlots int, execute Execu
 	return tw, nil
 }
 
-func (tw *TimingWheel) Drain(fn func(key, value interface{})) {
-	tw.drainChannel <- fn
+// Drain drains all items and executes them.
+func (tw *TimingWheel) Drain(fn func(key, value any)) error {
+	select {
+	case tw.drainChannel <- fn:
+		return nil
+	case <-tw.stopChannel:
+		return ErrClosed
+	}
 }
 
-func (tw *TimingWheel) MoveTimer(key interface{}, delay time.Duration) {
+// MoveTimer moves the task with the given key to the given delay.
+func (tw *TimingWheel) MoveTimer(key any, delay time.Duration) error {
 	if delay <= 0 || key == nil {
-		return
+		return ErrArgument
 	}
 
-	tw.moveChannel <- baseEntry{
+	select {
+	case tw.moveChannel <- baseEntry{
 		delay: delay,
 		key:   key,
+	}:
+		return nil
+	case <-tw.stopChannel:
+		return ErrClosed
 	}
 }
 
-func (tw *TimingWheel) RemoveTimer(key interface{}) {
+// RemoveTimer removes the task with the given key.
+func (tw *TimingWheel) RemoveTimer(key any) error {
 	if key == nil {
-		return
+		return ErrArgument
 	}
 
-	tw.removeChannel <- key
+	select {
+	case tw.removeChannel <- key:
+		return nil
+	case <-tw.stopChannel:
+		return ErrClosed
+	}
 }
 
-func (tw *TimingWheel) SetTimer(key, value interface{}, delay time.Duration) {
+// SetTimer sets the task value with the given key to the delay.
+func (tw *TimingWheel) SetTimer(key, value any, delay time.Duration) error {
 	if delay <= 0 || key == nil {
-		return
+		return ErrArgument
 	}
 
-	tw.setChannel <- timingEntry{
+	select {
+	case tw.setChannel <- timingEntry{
 		baseEntry: baseEntry{
 			delay: delay,
 			key:   key,
 		},
 		value: value,
+	}:
+		return nil
+	case <-tw.stopChannel:
+		return ErrClosed
 	}
 }
 
+// Stop stops tw. No more actions after stopping a TimingWheel.
 func (tw *TimingWheel) Stop() {
 	close(tw.stopChannel)
 }
 
-func (tw *TimingWheel) drainAll(fn func(key, value interface{})) {
+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)
 			next := e.Next()
 			slot.Remove(e)
+			if val, ok := tw.timers.Get(task.key); ok {
+				timer := val.(*positionEntry)
+				if timer.item == task {
+					tw.timers.Del(task.key)
+				}
+			}
 			e = next
 			if !task.removed {
 				runner.Schedule(func() {
@@ -141,9 +184,11 @@ func (tw *TimingWheel) drainAll(fn func(key, value interface{})) {
 			}
 		}
 	}
+
+	runner.Wait()
 }
 
-func (tw *TimingWheel) getPositionAndCircle(d time.Duration) (pos int, circle int) {
+func (tw *TimingWheel) getPositionAndCircle(d time.Duration) (pos, circle int) {
 	steps := int(d / tw.interval)
 	pos = (tw.tickedPos + steps) % tw.numSlots
 	circle = (steps - 1) / tw.numSlots
@@ -196,7 +241,7 @@ func (tw *TimingWheel) onTick() {
 	tw.scanAndRunTasks(l)
 }
 
-func (tw *TimingWheel) removeTask(key interface{}) {
+func (tw *TimingWheel) removeTask(key any) {
 	val, ok := tw.timers.Get(key)
 	if !ok {
 		return
@@ -204,6 +249,7 @@ func (tw *TimingWheel) removeTask(key interface{}) {
 
 	timer := val.(*positionEntry)
 	timer.item.removed = true
+	tw.timers.Del(key)
 }
 
 func (tw *TimingWheel) run() {
@@ -248,7 +294,6 @@ func (tw *TimingWheel) scanAndRunTasks(l *list.List) {
 		if task.removed {
 			next := e.Next()
 			l.Remove(e)
-			tw.timers.Del(task.key)
 			e = next
 			continue
 		} else if task.circle > 0 {
@@ -301,6 +346,7 @@ func (tw *TimingWheel) setTask(task *timingEntry) {
 func (tw *TimingWheel) setTimerPosition(pos int, task *timingEntry) {
 	if val, ok := tw.timers.Get(task.key); ok {
 		timer := val.(*positionEntry)
+		timer.item = task
 		timer.pos = pos
 	} else {
 		tw.timers.Set(task.key, &positionEntry{

core/collection/timingwheel_test.go

@@ -8,10 +8,10 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/lang"
-	"github.com/tal-tech/go-zero/core/stringx"
-	"github.com/tal-tech/go-zero/core/syncx"
-	"github.com/tal-tech/go-zero/core/timex"
+	"github.com/zeromicro/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/stringx"
+	"github.com/zeromicro/go-zero/core/syncx"
+	"github.com/zeromicro/go-zero/core/timex"
 )
 
 const (
@@ -20,15 +20,14 @@ const (
 )
 
 func TestNewTimingWheel(t *testing.T) {
-	_, err := NewTimingWheel(0, 10, func(key, value interface{}) {})
+	_, err := NewTimingWheel(0, 10, func(key, value any) {})
 	assert.NotNil(t, err)
 }
 
 func TestTimingWheel_Drain(t *testing.T) {
 	ticker := timex.NewFakeTicker()
-	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
+	tw, _ := NewTimingWheelWithTicker(testStep, 10, func(k, v any) {
 	}, ticker)
-	defer tw.Stop()
 	tw.SetTimer("first", 3, testStep*4)
 	tw.SetTimer("second", 5, testStep*7)
 	tw.SetTimer("third", 7, testStep*7)
@@ -37,7 +36,7 @@ func TestTimingWheel_Drain(t *testing.T) {
 	var lock sync.Mutex
 	var wg sync.WaitGroup
 	wg.Add(3)
-	tw.Drain(func(key, value interface{}) {
+	tw.Drain(func(key, value any) {
 		lock.Lock()
 		defer lock.Unlock()
 		keys = append(keys, key.(string))
@@ -51,17 +50,48 @@ func TestTimingWheel_Drain(t *testing.T) {
 	assert.EqualValues(t, []string{"first", "second", "third"}, keys)
 	assert.EqualValues(t, []int{3, 5, 7}, vals)
 	var count int
-	tw.Drain(func(key, value interface{}) {
+	tw.Drain(func(key, value any) {
 		count++
 	})
 	time.Sleep(time.Millisecond * 100)
 	assert.Equal(t, 0, count)
+	tw.Stop()
+	assert.Equal(t, ErrClosed, tw.Drain(func(key, value any) {}))
+}
+
+func TestTimingWheel_SetTimerAfterDrain(t *testing.T) {
+	ticker := timex.NewFakeTicker()
+	var count int32
+	tw, _ := NewTimingWheelWithTicker(testStep, 10, func(key, value any) {
+		assert.Equal(t, "foo", key)
+		assert.Equal(t, 2, value)
+		atomic.AddInt32(&count, 1)
+		ticker.Done()
+	}, ticker)
+	defer tw.Stop()
+
+	tw.SetTimer("foo", 1, testStep*4)
+	drained := make(chan lang.PlaceholderType)
+	tw.Drain(func(key, value any) {
+		assert.Equal(t, "foo", key)
+		assert.Equal(t, 1, value)
+		close(drained)
+	})
+	<-drained
+
+	tw.SetTimer("foo", 2, testStep*4)
+	for i := 0; i < 4; i++ {
+		ticker.Tick()
+	}
+
+	assert.Nil(t, ticker.Wait(waitTime))
+	assert.Equal(t, int32(1), atomic.LoadInt32(&count))
 }
 
 func TestTimingWheel_SetTimerSoon(t *testing.T) {
 	run := syncx.NewAtomicBool()
 	ticker := timex.NewFakeTicker()
-	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
+	tw, _ := NewTimingWheelWithTicker(testStep, 10, func(k, v any) {
 		assert.True(t, run.CompareAndSwap(false, true))
 		assert.Equal(t, "any", k)
 		assert.Equal(t, 3, v.(int))
@@ -77,7 +107,7 @@ func TestTimingWheel_SetTimerSoon(t *testing.T) {
 func TestTimingWheel_SetTimerTwice(t *testing.T) {
 	run := syncx.NewAtomicBool()
 	ticker := timex.NewFakeTicker()
-	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
+	tw, _ := NewTimingWheelWithTicker(testStep, 10, func(k, v any) {
 		assert.True(t, run.CompareAndSwap(false, true))
 		assert.Equal(t, "any", k)
 		assert.Equal(t, 5, v.(int))
@@ -95,23 +125,29 @@ func TestTimingWheel_SetTimerTwice(t *testing.T) {
 
 func TestTimingWheel_SetTimerWrongDelay(t *testing.T) {
 	ticker := timex.NewFakeTicker()
-	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {}, ticker)
+	tw, _ := NewTimingWheelWithTicker(testStep, 10, func(k, v any) {}, ticker)
 	defer tw.Stop()
 	assert.NotPanics(t, func() {
 		tw.SetTimer("any", 3, -testStep)
 	})
 }
 
+func TestTimingWheel_SetTimerAfterClose(t *testing.T) {
+	ticker := timex.NewFakeTicker()
+	tw, _ := NewTimingWheelWithTicker(testStep, 10, func(k, v any) {}, ticker)
+	tw.Stop()
+	assert.Equal(t, ErrClosed, tw.SetTimer("any", 3, testStep))
+}
+
 func TestTimingWheel_MoveTimer(t *testing.T) {
 	run := syncx.NewAtomicBool()
 	ticker := timex.NewFakeTicker()
-	tw, _ := newTimingWheelWithClock(testStep, 3, func(k, v interface{}) {
+	tw, _ := NewTimingWheelWithTicker(testStep, 3, func(k, v any) {
 		assert.True(t, run.CompareAndSwap(false, true))
 		assert.Equal(t, "any", k)
 		assert.Equal(t, 3, v.(int))
 		ticker.Done()
 	}, ticker)
-	defer tw.Stop()
 	tw.SetTimer("any", 3, testStep*4)
 	tw.MoveTimer("any", testStep*7)
 	tw.MoveTimer("any", -testStep)
@@ -125,12 +161,14 @@ func TestTimingWheel_MoveTimer(t *testing.T) {
 	}
 	assert.Nil(t, ticker.Wait(waitTime))
 	assert.True(t, run.True())
+	tw.Stop()
+	assert.Equal(t, ErrClosed, tw.MoveTimer("any", time.Millisecond))
 }
 
 func TestTimingWheel_MoveTimerSoon(t *testing.T) {
 	run := syncx.NewAtomicBool()
 	ticker := timex.NewFakeTicker()
-	tw, _ := newTimingWheelWithClock(testStep, 3, func(k, v interface{}) {
+	tw, _ := NewTimingWheelWithTicker(testStep, 3, func(k, v any) {
 		assert.True(t, run.CompareAndSwap(false, true))
 		assert.Equal(t, "any", k)
 		assert.Equal(t, 3, v.(int))
@@ -146,7 +184,7 @@ func TestTimingWheel_MoveTimerSoon(t *testing.T) {
 func TestTimingWheel_MoveTimerEarlier(t *testing.T) {
 	run := syncx.NewAtomicBool()
 	ticker := timex.NewFakeTicker()
-	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
+	tw, _ := NewTimingWheelWithTicker(testStep, 10, func(k, v any) {
 		assert.True(t, run.CompareAndSwap(false, true))
 		assert.Equal(t, "any", k)
 		assert.Equal(t, 3, v.(int))
@@ -164,7 +202,7 @@ func TestTimingWheel_MoveTimerEarlier(t *testing.T) {
 
 func TestTimingWheel_RemoveTimer(t *testing.T) {
 	ticker := timex.NewFakeTicker()
-	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {}, ticker)
+	tw, _ := NewTimingWheelWithTicker(testStep, 10, func(k, v any) {}, ticker)
 	tw.SetTimer("any", 3, testStep)
 	assert.NotPanics(t, func() {
 		tw.RemoveTimer("any")
@@ -175,6 +213,7 @@ func TestTimingWheel_RemoveTimer(t *testing.T) {
 		ticker.Tick()
 	}
 	tw.Stop()
+	assert.Equal(t, ErrClosed, tw.RemoveTimer("any"))
 }
 
 func TestTimingWheel_SetTimer(t *testing.T) {
@@ -226,7 +265,7 @@ func TestTimingWheel_SetTimer(t *testing.T) {
 			}
 			var actual int32
 			done := make(chan lang.PlaceholderType)
-			tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
+			tw, err := NewTimingWheelWithTicker(testStep, test.slots, func(key, value any) {
 				assert.Equal(t, 1, key.(int))
 				assert.Equal(t, 2, value.(int))
 				actual = atomic.LoadInt32(&count)
@@ -307,7 +346,7 @@ func TestTimingWheel_SetAndMoveThenStart(t *testing.T) {
 			}
 			var actual int32
 			done := make(chan lang.PlaceholderType)
-			tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
+			tw, err := NewTimingWheelWithTicker(testStep, test.slots, func(key, value any) {
 				actual = atomic.LoadInt32(&count)
 				close(done)
 			}, ticker)
@@ -395,7 +434,7 @@ func TestTimingWheel_SetAndMoveTwice(t *testing.T) {
 			}
 			var actual int32
 			done := make(chan lang.PlaceholderType)
-			tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
+			tw, err := NewTimingWheelWithTicker(testStep, test.slots, func(key, value any) {
 				actual = atomic.LoadInt32(&count)
 				close(done)
 			}, ticker)
@@ -476,7 +515,7 @@ func TestTimingWheel_ElapsedAndSet(t *testing.T) {
 			}
 			var actual int32
 			done := make(chan lang.PlaceholderType)
-			tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
+			tw, err := NewTimingWheelWithTicker(testStep, test.slots, func(key, value any) {
 				actual = atomic.LoadInt32(&count)
 				close(done)
 			}, ticker)
@@ -567,7 +606,7 @@ func TestTimingWheel_ElapsedAndSetThenMove(t *testing.T) {
 			}
 			var actual int32
 			done := make(chan lang.PlaceholderType)
-			tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
+			tw, err := NewTimingWheelWithTicker(testStep, test.slots, func(key, value any) {
 				actual = atomic.LoadInt32(&count)
 				close(done)
 			}, ticker)
@@ -594,10 +633,186 @@ func TestTimingWheel_ElapsedAndSetThenMove(t *testing.T) {
 	}
 }
 
+func TestMoveAndRemoveTask(t *testing.T) {
+	ticker := timex.NewFakeTicker()
+	tick := func(v int) {
+		for i := 0; i < v; i++ {
+			ticker.Tick()
+		}
+	}
+	var keys []int
+	tw, _ := NewTimingWheelWithTicker(testStep, 10, func(k, v any) {
+		assert.Equal(t, "any", k)
+		assert.Equal(t, 3, v.(int))
+		keys = append(keys, v.(int))
+		ticker.Done()
+	}, ticker)
+	defer tw.Stop()
+	tw.SetTimer("any", 3, testStep*8)
+	tick(6)
+	tw.MoveTimer("any", testStep*7)
+	tick(3)
+	tw.RemoveTimer("any")
+	tick(30)
+	time.Sleep(time.Millisecond)
+	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()
 
-	tw, _ := NewTimingWheel(time.Second, 100, func(k, v interface{}) {})
+	tw, _ := NewTimingWheel(time.Second, 100, func(k, v any) {})
 	for i := 0; i < b.N; i++ {
 		tw.SetTimer(i, i, time.Second)
 		tw.SetTimer(b.N+i, b.N+i, time.Second)

core/color/color.go

@@ -0,0 +1,73 @@
+package color
+
+import "github.com/fatih/color"
+
+const (
+	// NoColor is no color for both foreground and background.
+	NoColor Color = iota
+	// FgBlack is the foreground color black.
+	FgBlack
+	// FgRed is the foreground color red.
+	FgRed
+	// FgGreen is the foreground color green.
+	FgGreen
+	// FgYellow is the foreground color yellow.
+	FgYellow
+	// FgBlue is the foreground color blue.
+	FgBlue
+	// FgMagenta is the foreground color magenta.
+	FgMagenta
+	// FgCyan is the foreground color cyan.
+	FgCyan
+	// FgWhite is the foreground color white.
+	FgWhite
+
+	// BgBlack is the background color black.
+	BgBlack
+	// BgRed is the background color red.
+	BgRed
+	// BgGreen is the background color green.
+	BgGreen
+	// BgYellow is the background color yellow.
+	BgYellow
+	// BgBlue is the background color blue.
+	BgBlue
+	// BgMagenta is the background color magenta.
+	BgMagenta
+	// BgCyan is the background color cyan.
+	BgCyan
+	// BgWhite is the background color white.
+	BgWhite
+)
+
+var colors = map[Color][]color.Attribute{
+	FgBlack:   {color.FgBlack, color.Bold},
+	FgRed:     {color.FgRed, color.Bold},
+	FgGreen:   {color.FgGreen, color.Bold},
+	FgYellow:  {color.FgYellow, color.Bold},
+	FgBlue:    {color.FgBlue, color.Bold},
+	FgMagenta: {color.FgMagenta, color.Bold},
+	FgCyan:    {color.FgCyan, color.Bold},
+	FgWhite:   {color.FgWhite, color.Bold},
+	BgBlack:   {color.BgBlack, color.FgHiWhite, color.Bold},
+	BgRed:     {color.BgRed, color.FgHiWhite, color.Bold},
+	BgGreen:   {color.BgGreen, color.FgHiWhite, color.Bold},
+	BgYellow:  {color.BgHiYellow, color.FgHiBlack, color.Bold},
+	BgBlue:    {color.BgBlue, color.FgHiWhite, color.Bold},
+	BgMagenta: {color.BgMagenta, color.FgHiWhite, color.Bold},
+	BgCyan:    {color.BgCyan, color.FgHiWhite, color.Bold},
+	BgWhite:   {color.BgHiWhite, color.FgHiBlack, color.Bold},
+}
+
+type Color uint32
+
+// WithColor returns a string with the given color applied.
+func WithColor(text string, colour Color) string {
+	c := color.New(colors[colour]...)
+	return c.Sprint(text)
+}
+
+// WithColorPadding returns a string with the given color applied with leading and trailing spaces.
+func WithColorPadding(text string, colour Color) string {
+	return WithColor(" "+text+" ", colour)
+}

core/color/color_test.go

@@ -0,0 +1,17 @@
+package color
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWithColor(t *testing.T) {
+	output := WithColor("Hello", BgRed)
+	assert.Equal(t, "Hello", output)
+}
+
+func TestWithColorPadding(t *testing.T) {
+	output := WithColorPadding("Hello", BgRed)
+	assert.Equal(t, " Hello ", output)
+}

core/conf/config.go

@@ -2,39 +2,378 @@ package conf
 
 import (
 	"fmt"
-	"io/ioutil"
 	"log"
+	"os"
 	"path"
+	"reflect"
+	"strings"
 
-	"github.com/tal-tech/go-zero/core/mapping"
+	"github.com/zeromicro/go-zero/core/jsonx"
+	"github.com/zeromicro/go-zero/core/mapping"
+	"github.com/zeromicro/go-zero/internal/encoding"
 )
 
-var loaders = map[string]func([]byte, interface{}) error{
-	".json": LoadConfigFromJsonBytes,
-	".yaml": LoadConfigFromYamlBytes,
-	".yml":  LoadConfigFromYamlBytes,
-}
+const (
+	jsonTagKey = "json"
+	jsonTagSep = ','
+)
 
-func LoadConfig(file string, v interface{}) error {
-	if content, err := ioutil.ReadFile(file); err != nil {
-		return err
-	} else if loader, ok := loaders[path.Ext(file)]; ok {
-		return loader(content, v)
-	} else {
-		return fmt.Errorf("unrecoginized file type: %s", file)
+var (
+	fillDefaultUnmarshaler = mapping.NewUnmarshaler(jsonTagKey, mapping.WithDefault())
+	loaders                = map[string]func([]byte, any) error{
+		".json":  LoadFromJsonBytes,
+		".json5": LoadFromJson5Bytes,
+		".toml":  LoadFromTomlBytes,
+		".yaml":  LoadFromYamlBytes,
+		".yml":   LoadFromYamlBytes,
 	}
+)
+
+// children and mapField should not be both filled.
+// named fields and map cannot be bound to the same field name.
+type fieldInfo struct {
+	children map[string]*fieldInfo
+	mapField *fieldInfo
 }
 
-func LoadConfigFromJsonBytes(content []byte, v interface{}) error {
-	return mapping.UnmarshalJsonBytes(content, v)
+// FillDefault fills the default values for the given v,
+// and the premise is that the value of v must be guaranteed to be empty.
+func FillDefault(v any) error {
+	return fillDefaultUnmarshaler.Unmarshal(map[string]any{}, v)
 }
 
-func LoadConfigFromYamlBytes(content []byte, v interface{}) error {
-	return mapping.UnmarshalYamlBytes(content, v)
+// Load loads config into v from file, .json, .json5, .toml, .yaml and .yml are acceptable.
+func Load(file string, v any, opts ...Option) error {
+	content, err := os.ReadFile(file)
+	if err != nil {
+		return err
+	}
+
+	loader, ok := loaders[strings.ToLower(path.Ext(file))]
+	if !ok {
+		return fmt.Errorf("unrecognized file type: %s", file)
+	}
+
+	var opt options
+	for _, o := range opts {
+		o(&opt)
+	}
+
+	if opt.env {
+		return loader([]byte(os.ExpandEnv(string(content))), v)
+	}
+
+	return loader(content, v)
 }
 
-func MustLoad(path string, v interface{}) {
-	if err := LoadConfig(path, v); err != nil {
+// LoadConfig loads config into v from file, .json, .json5, .toml, .yaml and .yml are acceptable.
+// Deprecated: use Load instead.
+func LoadConfig(file string, v any, opts ...Option) error {
+	return Load(file, v, opts...)
+}
+
+// LoadFromJsonBytes loads config into v from content json bytes.
+func LoadFromJsonBytes(content []byte, v any) error {
+	info, err := buildFieldsInfo(reflect.TypeOf(v), "")
+	if err != nil {
+		return err
+	}
+
+	var m map[string]any
+	if err = jsonx.Unmarshal(content, &m); err != nil {
+		return err
+	}
+
+	lowerCaseKeyMap := toLowerCaseKeyMap(m, info)
+
+	if err = mapping.UnmarshalJsonMap(lowerCaseKeyMap, v,
+		mapping.WithCanonicalKeyFunc(toLowerCase)); err != nil {
+		return err
+	}
+
+	return validate(v)
+}
+
+// LoadConfigFromJsonBytes loads config into v from content json bytes.
+// Deprecated: use LoadFromJsonBytes instead.
+func LoadConfigFromJsonBytes(content []byte, v any) error {
+	return LoadFromJsonBytes(content, v)
+}
+
+// LoadFromTomlBytes loads config into v from content toml bytes.
+func LoadFromTomlBytes(content []byte, v any) error {
+	b, err := encoding.TomlToJson(content)
+	if err != nil {
+		return err
+	}
+
+	return LoadFromJsonBytes(b, v)
+}
+
+// LoadFromYamlBytes loads config into v from content yaml bytes.
+func LoadFromYamlBytes(content []byte, v any) error {
+	b, err := encoding.YamlToJson(content)
+	if err != nil {
+		return err
+	}
+
+	return LoadFromJsonBytes(b, v)
+}
+
+// LoadFromJson5Bytes loads config into v from content json5 bytes.
+func LoadFromJson5Bytes(content []byte, v any) error {
+	b, err := encoding.Json5ToJson(content)
+	if err != nil {
+		return err
+	}
+
+	return LoadFromJsonBytes(b, v)
+}
+
+// LoadConfigFromYamlBytes loads config into v from content yaml bytes.
+// Deprecated: use LoadFromYamlBytes instead.
+func LoadConfigFromYamlBytes(content []byte, v any) error {
+	return LoadFromYamlBytes(content, v)
+}
+
+// MustLoad loads config into v from path, exits on error.
+func MustLoad(path string, v any, opts ...Option) {
+	if err := Load(path, v, opts...); err != nil {
 		log.Fatalf("error: config file %s, %s", path, err.Error())
 	}
 }
+
+func addOrMergeFields(info *fieldInfo, key string, child *fieldInfo, fullName string) error {
+	if prev, ok := info.children[key]; ok {
+		if child.mapField != nil {
+			return newConflictKeyError(fullName)
+		}
+
+		if err := mergeFields(prev, child.children, fullName); err != nil {
+			return err
+		}
+	} else {
+		info.children[key] = child
+	}
+
+	return nil
+}
+
+func buildAnonymousFieldInfo(info *fieldInfo, lowerCaseName string, ft reflect.Type, fullName string) error {
+	switch ft.Kind() {
+	case reflect.Struct:
+		fields, err := buildFieldsInfo(ft, fullName)
+		if err != nil {
+			return err
+		}
+
+		for k, v := range fields.children {
+			if err = addOrMergeFields(info, k, v, fullName); err != nil {
+				return err
+			}
+		}
+	case reflect.Map:
+		elemField, err := buildFieldsInfo(mapping.Deref(ft.Elem()), fullName)
+		if err != nil {
+			return err
+		}
+
+		if _, ok := info.children[lowerCaseName]; ok {
+			return newConflictKeyError(fullName)
+		}
+
+		info.children[lowerCaseName] = &fieldInfo{
+			children: make(map[string]*fieldInfo),
+			mapField: elemField,
+		}
+	default:
+		if _, ok := info.children[lowerCaseName]; ok {
+			return newConflictKeyError(fullName)
+		}
+
+		info.children[lowerCaseName] = &fieldInfo{
+			children: make(map[string]*fieldInfo),
+		}
+	}
+
+	return nil
+}
+
+func buildFieldsInfo(tp reflect.Type, fullName string) (*fieldInfo, error) {
+	tp = mapping.Deref(tp)
+
+	switch tp.Kind() {
+	case reflect.Struct:
+		return buildStructFieldsInfo(tp, fullName)
+	case reflect.Array, reflect.Slice, reflect.Map:
+		return buildFieldsInfo(mapping.Deref(tp.Elem()), fullName)
+	case reflect.Chan, reflect.Func:
+		return nil, fmt.Errorf("unsupported type: %s, fullName: %s", tp.Kind(), fullName)
+	default:
+		return &fieldInfo{
+			children: make(map[string]*fieldInfo),
+		}, nil
+	}
+}
+
+func buildNamedFieldInfo(info *fieldInfo, lowerCaseName string, ft reflect.Type, fullName string) error {
+	var finfo *fieldInfo
+	var err error
+
+	switch ft.Kind() {
+	case reflect.Struct:
+		finfo, err = buildFieldsInfo(ft, fullName)
+		if err != nil {
+			return err
+		}
+	case reflect.Array, reflect.Slice:
+		finfo, err = buildFieldsInfo(ft.Elem(), fullName)
+		if err != nil {
+			return err
+		}
+	case reflect.Map:
+		elemInfo, err := buildFieldsInfo(mapping.Deref(ft.Elem()), fullName)
+		if err != nil {
+			return err
+		}
+
+		finfo = &fieldInfo{
+			children: make(map[string]*fieldInfo),
+			mapField: elemInfo,
+		}
+	default:
+		finfo, err = buildFieldsInfo(ft, fullName)
+		if err != nil {
+			return err
+		}
+	}
+
+	return addOrMergeFields(info, lowerCaseName, finfo, fullName)
+}
+
+func buildStructFieldsInfo(tp reflect.Type, fullName string) (*fieldInfo, error) {
+	info := &fieldInfo{
+		children: make(map[string]*fieldInfo),
+	}
+
+	for i := 0; i < tp.NumField(); i++ {
+		field := tp.Field(i)
+		if !field.IsExported() {
+			continue
+		}
+
+		name := getTagName(field)
+		lowerCaseName := toLowerCase(name)
+		ft := mapping.Deref(field.Type)
+		// flatten anonymous fields
+		if field.Anonymous {
+			if err := buildAnonymousFieldInfo(info, lowerCaseName, ft,
+				getFullName(fullName, lowerCaseName)); err != nil {
+				return nil, err
+			}
+		} else if err := buildNamedFieldInfo(info, lowerCaseName, ft,
+			getFullName(fullName, lowerCaseName)); err != nil {
+			return nil, err
+		}
+	}
+
+	return info, nil
+}
+
+// getTagName get the tag name of the given field, if no tag name, use file.Name.
+// field.Name is returned on tags like `json:""` and `json:",optional"`.
+func getTagName(field reflect.StructField) string {
+	if tag, ok := field.Tag.Lookup(jsonTagKey); ok {
+		if pos := strings.IndexByte(tag, jsonTagSep); pos >= 0 {
+			tag = tag[:pos]
+		}
+
+		tag = strings.TrimSpace(tag)
+		if len(tag) > 0 {
+			return tag
+		}
+	}
+
+	return field.Name
+}
+
+func mergeFields(prev *fieldInfo, children map[string]*fieldInfo, fullName string) error {
+	if len(prev.children) == 0 || len(children) == 0 {
+		return newConflictKeyError(fullName)
+	}
+
+	// merge fields
+	for k, v := range children {
+		if _, ok := prev.children[k]; ok {
+			return newConflictKeyError(fullName)
+		}
+
+		prev.children[k] = v
+	}
+
+	return nil
+}
+
+func toLowerCase(s string) string {
+	return strings.ToLower(s)
+}
+
+func toLowerCaseInterface(v any, info *fieldInfo) any {
+	switch vv := v.(type) {
+	case map[string]any:
+		return toLowerCaseKeyMap(vv, info)
+	case []any:
+		arr := make([]any, 0, len(vv))
+		for _, vvv := range vv {
+			arr = append(arr, toLowerCaseInterface(vvv, info))
+		}
+		return arr
+	default:
+		return v
+	}
+}
+
+func toLowerCaseKeyMap(m map[string]any, info *fieldInfo) map[string]any {
+	res := make(map[string]any)
+
+	for k, v := range m {
+		ti, ok := info.children[k]
+		if ok {
+			res[k] = toLowerCaseInterface(v, ti)
+			continue
+		}
+
+		lk := toLowerCase(k)
+		if ti, ok = info.children[lk]; ok {
+			res[lk] = toLowerCaseInterface(v, ti)
+		} else if info.mapField != nil {
+			res[k] = toLowerCaseInterface(v, info.mapField)
+		} else if vv, ok := v.(map[string]any); ok {
+			res[k] = toLowerCaseKeyMap(vv, info)
+		} else {
+			res[k] = v
+		}
+	}
+
+	return res
+}
+
+type conflictKeyError struct {
+	key string
+}
+
+func newConflictKeyError(key string) conflictKeyError {
+	return conflictKeyError{key: key}
+}
+
+func (e conflictKeyError) Error() string {
+	return fmt.Sprintf("conflict key %s, pay attention to anonymous fields", e.key)
+}
+
+func getFullName(parent, child string) string {
+	if len(parent) == 0 {
+		return child
+	}
+
+	return parent + "." + child
+}

core/conf/options.go

@@ -0,0 +1,17 @@
+package conf
+
+type (
+	// Option defines the method to customize the config options.
+	Option func(opt *options)
+
+	options struct {
+		env bool
+	}
+)
+
+// UseEnv customizes the config to use environment variables.
+func UseEnv() Option {
+	return func(opt *options) {
+		opt.env = true
+	}
+}

core/conf/properties.go

@@ -2,16 +2,16 @@ package conf
 
 import (
 	"fmt"
+	"os"
 	"strconv"
 	"strings"
 	"sync"
 
-	"github.com/tal-tech/go-zero/core/iox"
+	"github.com/zeromicro/go-zero/core/iox"
 )
 
 // PropertyError represents a configuration error message.
 type PropertyError struct {
-	error
 	message string
 }
 
@@ -30,17 +30,22 @@ type mapBasedProperties struct {
 	lock       sync.RWMutex
 }
 
-// Loads the properties into a properties configuration instance.
+// LoadProperties loads the properties into a properties configuration instance.
 // Returns an error that indicates if there was a problem loading the configuration.
-func LoadProperties(filename string) (Properties, error) {
+func LoadProperties(filename string, opts ...Option) (Properties, error) {
 	lines, err := iox.ReadTextLines(filename, iox.WithoutBlank(), iox.OmitWithPrefix("#"))
 	if err != nil {
 		return nil, err
 	}
 
+	var opt options
+	for _, o := range opts {
+		o(&opt)
+	}
+
 	raw := make(map[string]string)
 	for i := range lines {
-		pair := strings.Split(lines[i], "=")
+		pair := strings.SplitN(lines[i], "=", 2)
 		if len(pair) != 2 {
 			// invalid property format
 			return nil, &PropertyError{
@@ -50,7 +55,11 @@ func LoadProperties(filename string) (Properties, error) {
 
 		key := strings.TrimSpace(pair[0])
 		value := strings.TrimSpace(pair[1])
-		raw[key] = value
+		if opt.env {
+			raw[key] = os.ExpandEnv(value)
+		} else {
+			raw[key] = value
+		}
 	}
 
 	return &mapBasedProperties{
@@ -87,7 +96,7 @@ func (config *mapBasedProperties) SetInt(key string, value int) {
 	config.lock.Unlock()
 }
 
-// Dumps the configuration internal map into a string.
+// ToString dumps the configuration internal map into a string.
 func (config *mapBasedProperties) ToString() string {
 	config.lock.RLock()
 	ret := fmt.Sprintf("%s", config.properties)
@@ -96,12 +105,12 @@ func (config *mapBasedProperties) ToString() string {
 	return ret
 }
 
-// Returns the error message.
+// Error returns the error message.
 func (configError *PropertyError) Error() string {
 	return configError.message
 }
 
-// Builds a new properties configuration structure
+// NewProperties builds a new properties configuration structure.
 func NewProperties() Properties {
 	return &mapBasedProperties{
 		properties: make(map[string]string),

core/conf/properties_test.go

@@ -5,7 +5,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/fs"
+	"github.com/zeromicro/go-zero/core/fs"
 )
 
 func TestProperties(t *testing.T) {
@@ -24,6 +24,52 @@ func TestProperties(t *testing.T) {
 	assert.Equal(t, "test", props.GetString("app.name"))
 	assert.Equal(t, "app", props.GetString("app.program"))
 	assert.Equal(t, 5, props.GetInt("app.threads"))
+
+	val := props.ToString()
+	assert.Contains(t, val, "app.name")
+	assert.Contains(t, val, "app.program")
+	assert.Contains(t, val, "app.threads")
+}
+
+func TestPropertiesEnv(t *testing.T) {
+	text := `app.name = test
+
+    app.program=app
+
+	app.env1 = ${FOO}
+	app.env2 = $none
+
+    # this is comment
+    app.threads = 5`
+	tmpfile, err := fs.TempFilenameWithText(text)
+	assert.Nil(t, err)
+	defer os.Remove(tmpfile)
+
+	t.Setenv("FOO", "2")
+
+	props, err := LoadProperties(tmpfile, UseEnv())
+	assert.Nil(t, err)
+	assert.Equal(t, "test", props.GetString("app.name"))
+	assert.Equal(t, "app", props.GetString("app.program"))
+	assert.Equal(t, 5, props.GetInt("app.threads"))
+	assert.Equal(t, "2", props.GetString("app.env1"))
+	assert.Equal(t, "", props.GetString("app.env2"))
+
+	val := props.ToString()
+	assert.Contains(t, val, "app.name")
+	assert.Contains(t, val, "app.program")
+	assert.Contains(t, val, "app.threads")
+	assert.Contains(t, val, "app.env1")
+	assert.Contains(t, val, "app.env2")
+}
+
+func TestLoadProperties_badContent(t *testing.T) {
+	filename, err := fs.TempFilenameWithText("hello")
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+	_, err = LoadProperties(filename)
+	assert.NotNil(t, err)
+	assert.True(t, len(err.Error()) > 0)
 }
 
 func TestSetString(t *testing.T) {
@@ -46,3 +92,70 @@ func TestLoadBadFile(t *testing.T) {
 	_, err := LoadProperties("nosuchfile")
 	assert.NotNil(t, err)
 }
+
+func TestProperties_valueWithEqualSymbols(t *testing.T) {
+	text := `# test with equal symbols in value
+	db.url=postgres://localhost:5432/db?param=value
+	math.equation=a=b=c
+	base64.data=SGVsbG8=World=Test=
+	url.with.params=http://example.com?foo=bar&baz=qux
+	empty.value=
+	key.with.space = value = with = equals`
+	tmpfile, err := fs.TempFilenameWithText(text)
+	assert.Nil(t, err)
+	defer os.Remove(tmpfile)
+
+	props, err := LoadProperties(tmpfile)
+	assert.Nil(t, err)
+	assert.Equal(t, "postgres://localhost:5432/db?param=value", props.GetString("db.url"))
+	assert.Equal(t, "a=b=c", props.GetString("math.equation"))
+	assert.Equal(t, "SGVsbG8=World=Test=", props.GetString("base64.data"))
+	assert.Equal(t, "http://example.com?foo=bar&baz=qux", props.GetString("url.with.params"))
+	assert.Equal(t, "", props.GetString("empty.value"))
+	assert.Equal(t, "value = with = equals", props.GetString("key.with.space"))
+}
+
+func TestProperties_edgeCases(t *testing.T) {
+	tests := []struct {
+		name    string
+		content string
+		wantErr bool
+		errMsg  string
+	}{
+		{
+			name:    "no equal sign",
+			content: "invalid line without equal",
+			wantErr: true,
+		},
+		{
+			name:    "only equal sign",
+			content: "=",
+			wantErr: false, // "=" 会被解析为空 key 和空 value，len(pair) == 2，是合法的
+		},
+		{
+			name:    "empty key",
+			content: "=value",
+			wantErr: false, // 空 key 也会被 trim，但 len(pair) == 2 所以不会报错
+		},
+		{
+			name:    "equal at end",
+			content: "key.name=",
+			wantErr: false, // 空 value 是合法的
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tmpfile, err := fs.TempFilenameWithText(tt.content)
+			assert.Nil(t, err)
+			defer os.Remove(tmpfile)
+
+			_, err = LoadProperties(tmpfile)
+			if tt.wantErr {
+				assert.NotNil(t, err, "expected error for case: %s", tt.name)
+			} else {
+				assert.Nil(t, err, "unexpected error for case: %s", tt.name)
+			}
+		})
+	}
+}

core/conf/readme.md

@@ -0,0 +1,58 @@
+## How to use
+
+1. Define a config structure, like below:
+
+```go
+type RestfulConf struct {
+  ServiceName  string        `json:",env=SERVICE_NAME"`  // read from env automatically
+	Host         string        `json:",default=0.0.0.0"`
+	Port         int
+	LogMode      string        `json:",options=[file,console]"`
+	Verbose      bool          `json:",optional"`
+	MaxConns     int           `json:",default=10000"`
+	MaxBytes     int64         `json:",default=1048576"`
+	Timeout      time.Duration `json:",default=3s"`
+	CpuThreshold int64         `json:",default=900,range=[0:1000)"`
+}
+```
+
+2. Write the yaml, toml or json config file:
+
+- yaml example
+
+```yaml
+# most fields are optional or have default values
+port: 8080
+logMode: console
+# you can use env settings
+maxBytes: ${MAX_BYTES}
+```
+
+- toml example
+
+```toml
+# most fields are optional or have default values
+port = 8_080
+logMode = "console"
+# you can use env settings
+maxBytes = "${MAX_BYTES}"
+```
+
+3. Load the config from a file:
+
+```go
+// exit on error
+var config RestfulConf
+conf.MustLoad(configFile, &config)
+
+// or handle the error on your own
+var config RestfulConf
+if err := conf.Load(configFile, &config); err != nil {
+  log.Fatal(err)
+}
+
+// enable reading from environments
+var config RestfulConf
+conf.MustLoad(configFile, &config, conf.UseEnv())
+```
+

core/conf/validate.go

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

core/conf/validate_test.go

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

core/configcenter/configurator.go

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

core/configcenter/configurator_test.go

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

core/configcenter/subscriber/etcd.go

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

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

core/configcenter/unmarshaler.go

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

core/configcenter/unmarshaler_test.go

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

core/contextx/deadline.go

@@ -1,17 +0,0 @@
-package contextx
-
-import (
-	"context"
-	"time"
-)
-
-func ShrinkDeadline(ctx context.Context, timeout time.Duration) (context.Context, func()) {
-	if deadline, ok := ctx.Deadline(); ok {
-		leftTime := time.Until(deadline)
-		if leftTime < timeout {
-			timeout = leftTime
-		}
-	}
-
-	return context.WithDeadline(ctx, time.Now().Add(timeout))
-}

core/contextx/deadline_test.go

@@ -1,31 +0,0 @@
-package contextx
-
-import (
-	"context"
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/assert"
-)
-
-func TestShrinkDeadlineLess(t *testing.T) {
-	deadline := time.Now().Add(time.Second)
-	ctx, cancel := context.WithDeadline(context.Background(), deadline)
-	defer cancel()
-	ctx, cancel = ShrinkDeadline(ctx, time.Minute)
-	defer cancel()
-	dl, ok := ctx.Deadline()
-	assert.True(t, ok)
-	assert.Equal(t, deadline, dl)
-}
-
-func TestShrinkDeadlineMore(t *testing.T) {
-	deadline := time.Now().Add(time.Minute)
-	ctx, cancel := context.WithDeadline(context.Background(), deadline)
-	defer cancel()
-	ctx, cancel = ShrinkDeadline(ctx, time.Second)
-	defer cancel()
-	dl, ok := ctx.Deadline()
-	assert.True(t, ok)
-	assert.True(t, dl.Before(deadline))
-}

core/contextx/unmarshaler.go

@@ -3,7 +3,7 @@ package contextx
 import (
 	"context"
 
-	"github.com/tal-tech/go-zero/core/mapping"
+	"github.com/zeromicro/go-zero/core/mapping"
 )
 
 const contextTagKey = "ctx"
@@ -14,12 +14,13 @@ type contextValuer struct {
 	context.Context
 }
 
-func (cv contextValuer) Value(key string) (interface{}, bool) {
+func (cv contextValuer) Value(key string) (any, bool) {
 	v := cv.Context.Value(key)
 	return v, v != nil
 }
 
-func For(ctx context.Context, v interface{}) error {
+// For unmarshals ctx into v.
+func For(ctx context.Context, v any) error {
 	return unmarshaler.UnmarshalValuer(contextValuer{
 		Context: ctx,
 	}, v)

core/contextx/unmarshaler_test.go

@@ -47,9 +47,11 @@ func TestUnmarshalContextWithMissing(t *testing.T) {
 		Name string `ctx:"name"`
 		Age  int    `ctx:"age"`
 	}
+	type name string
+	const PersonNameKey name = "name"
 
 	ctx := context.Background()
-	ctx = context.WithValue(ctx, "name", "kevin")
+	ctx = context.WithValue(ctx, PersonNameKey, "kevin")
 
 	var person Person
 	err := For(ctx, &person)

core/contextx/valueonlycontext.go

@@ -21,6 +21,7 @@ func (valueOnlyContext) Err() error {
 	return nil
 }
 
+// ValueOnlyFrom takes all values from the given ctx, without deadline and error control.
 func ValueOnlyFrom(ctx context.Context) context.Context {
 	return valueOnlyContext{
 		Context: ctx,

core/contextx/valueonlycontext_test.go

@@ -9,7 +9,9 @@ import (
 )
 
 func TestContextCancel(t *testing.T) {
-	c := context.WithValue(context.Background(), "key", "value")
+	type key string
+	var nameKey key = "name"
+	c := context.WithValue(context.Background(), nameKey, "value")
 	c1, cancel := context.WithCancel(c)
 	o := ValueOnlyFrom(c1)
 	c2, cancel2 := context.WithCancel(o)

core/discov/accountregistry.go

@@ -0,0 +1,14 @@
+package discov
+
+import "github.com/zeromicro/go-zero/core/discov/internal"
+
+// RegisterAccount registers the username/password to the given etcd cluster.
+func RegisterAccount(endpoints []string, user, pass string) {
+	internal.AddAccount(endpoints, user, pass)
+}
+
+// RegisterTLS registers the CertFile/CertKeyFile/CACertFile to the given etcd.
+func RegisterTLS(endpoints []string, certFile, certKeyFile, caFile string,
+	insecureSkipVerify bool) error {
+	return internal.AddTLS(endpoints, certFile, certKeyFile, caFile, insecureSkipVerify)
+}

core/discov/accountregistry_test.go

@@ -0,0 +1,21 @@
+package discov
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/discov/internal"
+	"github.com/zeromicro/go-zero/core/stringx"
+)
+
+func TestRegisterAccount(t *testing.T) {
+	endpoints := []string{
+		"localhost:2379",
+	}
+	user := "foo" + stringx.Rand()
+	RegisterAccount(endpoints, user, "bar")
+	account, ok := internal.GetAccount(endpoints)
+	assert.True(t, ok)
+	assert.Equal(t, user, account.User)
+	assert.Equal(t, "bar", account.Pass)
+}

core/discov/clients.go

@@ -4,7 +4,7 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/tal-tech/go-zero/core/discov/internal"
+	"github.com/zeromicro/go-zero/core/discov/internal"
 )
 
 const (
@@ -14,6 +14,7 @@ const (
 
 const timeToLive int64 = 10
 
+// TimeToLive is seconds to live in etcd.
 var TimeToLive = timeToLive
 
 func extract(etcdKey string, index int) (string, bool) {

core/discov/clients_test.go

@@ -5,7 +5,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/discov/internal"
+	"github.com/zeromicro/go-zero/core/discov/internal"
 )
 
 var mockLock sync.Mutex
@@ -28,6 +28,9 @@ func TestExtract(t *testing.T) {
 
 	_, ok = extract("any", -1)
 	assert.False(t, ok)
+
+	_, ok = extract("any", 10)
+	assert.False(t, ok)
 }
 
 func TestMakeKey(t *testing.T) {

core/discov/config.go

@@ -2,16 +2,47 @@ package discov
 
 import "errors"
 
+var (
+	// errEmptyEtcdHosts indicates that etcd hosts are empty.
+	errEmptyEtcdHosts = errors.New("empty etcd hosts")
+	// errEmptyEtcdKey indicates that etcd key is empty.
+	errEmptyEtcdKey = errors.New("empty etcd key")
+)
+
+// EtcdConf is the config item with the given key on etcd.
 type EtcdConf struct {
-	Hosts []string
-	Key   string
+	Hosts              []string
+	Key                string
+	ID                 int64  `json:",optional"`
+	User               string `json:",optional"`
+	Pass               string `json:",optional"`
+	CertFile           string `json:",optional"`
+	CertKeyFile        string `json:",optional=CertFile"`
+	CACertFile         string `json:",optional=CertFile"`
+	InsecureSkipVerify bool   `json:",optional"`
 }
 
+// HasAccount returns if account provided.
+func (c EtcdConf) HasAccount() bool {
+	return len(c.User) > 0 && len(c.Pass) > 0
+}
+
+// HasID returns if ID provided.
+func (c EtcdConf) HasID() bool {
+	return c.ID > 0
+}
+
+// HasTLS returns if TLS CertFile/CertKeyFile/CACertFile are provided.
+func (c EtcdConf) HasTLS() bool {
+	return len(c.CertFile) > 0 && len(c.CertKeyFile) > 0 && len(c.CACertFile) > 0
+}
+
+// Validate validates c.
 func (c EtcdConf) Validate() error {
 	if len(c.Hosts) == 0 {
-		return errors.New("empty etcd hosts")
+		return errEmptyEtcdHosts
 	} else if len(c.Key) == 0 {
-		return errors.New("empty etcd key")
+		return errEmptyEtcdKey
 	} else {
 		return nil
 	}

core/discov/config_test.go

@@ -44,3 +44,126 @@ func TestConfig(t *testing.T) {
 		}
 	}
 }
+
+func TestEtcdConf_HasAccount(t *testing.T) {
+	tests := []struct {
+		EtcdConf
+		hasAccount bool
+	}{
+		{
+			EtcdConf: EtcdConf{
+				Hosts: []string{"any"},
+				Key:   "key",
+			},
+			hasAccount: false,
+		},
+		{
+			EtcdConf: EtcdConf{
+				Hosts: []string{"any"},
+				Key:   "key",
+				User:  "foo",
+			},
+			hasAccount: false,
+		},
+		{
+			EtcdConf: EtcdConf{
+				Hosts: []string{"any"},
+				Key:   "key",
+				User:  "foo",
+				Pass:  "bar",
+			},
+			hasAccount: true,
+		},
+	}
+
+	for _, test := range tests {
+		assert.Equal(t, test.hasAccount, test.EtcdConf.HasAccount())
+	}
+}
+
+func TestEtcdConf_HasID(t *testing.T) {
+	tests := []struct {
+		EtcdConf
+		hasServerID bool
+	}{
+		{
+			EtcdConf: EtcdConf{
+				Hosts: []string{"any"},
+				ID:    -1,
+			},
+			hasServerID: false,
+		},
+		{
+			EtcdConf: EtcdConf{
+				Hosts: []string{"any"},
+				ID:    0,
+			},
+			hasServerID: false,
+		},
+		{
+			EtcdConf: EtcdConf{
+				Hosts: []string{"any"},
+				ID:    10000,
+			},
+			hasServerID: true,
+		},
+	}
+
+	for _, test := range tests {
+		assert.Equal(t, test.hasServerID, test.EtcdConf.HasID())
+	}
+}
+
+func TestEtcdConf_HasTLS(t *testing.T) {
+	tests := []struct {
+		name string
+		conf EtcdConf
+		want bool
+	}{
+		{
+			name: "empty config",
+			conf: EtcdConf{},
+			want: false,
+		},
+		{
+			name: "missing CertFile",
+			conf: EtcdConf{
+				CertKeyFile: "key",
+				CACertFile:  "ca",
+			},
+			want: false,
+		},
+		{
+			name: "missing CertKeyFile",
+			conf: EtcdConf{
+				CertFile:   "cert",
+				CACertFile: "ca",
+			},
+			want: false,
+		},
+		{
+			name: "missing CACertFile",
+			conf: EtcdConf{
+				CertFile:    "cert",
+				CertKeyFile: "key",
+			},
+			want: false,
+		},
+		{
+			name: "valid config",
+			conf: EtcdConf{
+				CertFile:    "cert",
+				CertKeyFile: "key",
+				CACertFile:  "ca",
+			},
+			want: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := tt.conf.HasTLS()
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}

core/discov/facade.go

@@ -1,47 +0,0 @@
-package discov
-
-import (
-	"github.com/tal-tech/go-zero/core/discov/internal"
-	"github.com/tal-tech/go-zero/core/logx"
-)
-
-type (
-	Facade struct {
-		endpoints []string
-		registry  *internal.Registry
-	}
-
-	FacadeListener interface {
-		OnAdd(key, val string)
-		OnDelete(key string)
-	}
-)
-
-func NewFacade(endpoints []string) Facade {
-	return Facade{
-		endpoints: endpoints,
-		registry:  internal.GetRegistry(),
-	}
-}
-
-func (f Facade) Client() internal.EtcdClient {
-	conn, err := f.registry.GetConn(f.endpoints)
-	logx.Must(err)
-	return conn
-}
-
-func (f Facade) Monitor(key string, l FacadeListener) {
-	f.registry.Monitor(f.endpoints, key, listenerAdapter{l})
-}
-
-type listenerAdapter struct {
-	l FacadeListener
-}
-
-func (la listenerAdapter) OnAdd(kv internal.KV) {
-	la.l.OnAdd(kv.Key, kv.Val)
-}
-
-func (la listenerAdapter) OnDelete(kv internal.KV) {
-	la.l.OnDelete(kv.Key)
-}

core/discov/internal/accountmanager.go

@@ -0,0 +1,75 @@
+package internal
+
+import (
+	"crypto/tls"
+	"crypto/x509"
+	"os"
+	"sync"
+)
+
+var (
+	accounts   = make(map[string]Account)
+	tlsConfigs = make(map[string]*tls.Config)
+	lock       sync.RWMutex
+)
+
+// Account holds the username/password for an etcd cluster.
+type Account struct {
+	User string
+	Pass string
+}
+
+// AddAccount adds the username/password for the given etcd cluster.
+func AddAccount(endpoints []string, user, pass string) {
+	lock.Lock()
+	defer lock.Unlock()
+
+	accounts[getClusterKey(endpoints)] = Account{
+		User: user,
+		Pass: pass,
+	}
+}
+
+// AddTLS adds the tls cert files for the given etcd cluster.
+func AddTLS(endpoints []string, certFile, certKeyFile, caFile string, insecureSkipVerify bool) error {
+	cert, err := tls.LoadX509KeyPair(certFile, certKeyFile)
+	if err != nil {
+		return err
+	}
+
+	caData, err := os.ReadFile(caFile)
+	if err != nil {
+		return err
+	}
+
+	pool := x509.NewCertPool()
+	pool.AppendCertsFromPEM(caData)
+
+	lock.Lock()
+	defer lock.Unlock()
+	tlsConfigs[getClusterKey(endpoints)] = &tls.Config{
+		Certificates:       []tls.Certificate{cert},
+		RootCAs:            pool,
+		InsecureSkipVerify: insecureSkipVerify,
+	}
+
+	return nil
+}
+
+// GetAccount gets the username/password for the given etcd cluster.
+func GetAccount(endpoints []string) (Account, bool) {
+	lock.RLock()
+	defer lock.RUnlock()
+
+	account, ok := accounts[getClusterKey(endpoints)]
+	return account, ok
+}
+
+// GetTLS gets the tls config for the given etcd cluster.
+func GetTLS(endpoints []string) (*tls.Config, bool) {
+	lock.RLock()
+	defer lock.RUnlock()
+
+	cfg, ok := tlsConfigs[getClusterKey(endpoints)]
+	return cfg, ok
+}

core/discov/internal/accountmanager_test.go

@@ -0,0 +1,138 @@
+package internal
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/stringx"
+)
+
+const (
+	certContent = `-----BEGIN CERTIFICATE-----
+MIIDazCCAlOgAwIBAgIUEg9GVO2oaPn+YSmiqmFIuAo10WIwDQYJKoZIhvcNAQEM
+BQAwRTELMAkGA1UEBhMCQVUxEzARBgNVBAgMClNvbWUtU3RhdGUxITAfBgNVBAoM
+GEludGVybmV0IFdpZGdpdHMgUHR5IEx0ZDAgFw0yMzAzMTExMzIxMjNaGA8yMTIz
+MDIxNTEzMjEyM1owRTELMAkGA1UEBhMCQVUxEzARBgNVBAgMClNvbWUtU3RhdGUx
+ITAfBgNVBAoMGEludGVybmV0IFdpZGdpdHMgUHR5IEx0ZDCCASIwDQYJKoZIhvcN
+AQEBBQADggEPADCCAQoCggEBALplXlWsIf0O/IgnIplmiZHKGnxyfyufyE2FBRNk
+OofRqbKuPH8GNqbkvZm7N29fwTDAQ+mViAggCkDht4hOzoWJMA7KYJt8JnTSWL48
+M1lcrpc9DL2gszC/JF/FGvyANbBtLklkZPFBGdHUX14pjrT937wqPtm+SqUHSvRT
+B7bmwmm2drRcmhpVm98LSlV7uQ2EgnJgsLjBPITKUejLmVLHfgX0RwQ2xIpX9pS4
+FCe1BTacwl2gGp7Mje7y4Mfv3o0ArJW6Tuwbjx59ZXwb1KIP71b7bT04AVS8ZeYO
+UMLKKuB5UR9x9Rn6cLXOTWBpcMVyzDgrAFLZjnE9LPUolZMCAwEAAaNRME8wHwYD
+VR0jBBgwFoAUeW8w8pmhncbRgTsl48k4/7wnfx8wCQYDVR0TBAIwADALBgNVHQ8E
+BAMCBPAwFAYDVR0RBA0wC4IJbG9jYWxob3N0MA0GCSqGSIb3DQEBDAUAA4IBAQAI
+y9xaoS88CLPBsX6mxfcTAFVfGNTRW9VN9Ng1cCnUR+YGoXGM/l+qP4f7p8ocdGwK
+iYZErVTzXYIn+D27//wpY3klJk3gAnEUBT3QRkStBw7XnpbeZ2oPBK+cmDnCnZPS
+BIF1wxPX7vIgaxs5Zsdqwk3qvZ4Djr2wP7LabNWTLSBKgQoUY45Liw6pffLwcGF9
+UKlu54bvGze2SufISCR3ib+I+FLvqpvJhXToZWYb/pfI/HccuCL1oot1x8vx6DQy
+U+TYxlZsKS5mdNxAX3dqEkEMsgEi+g/tzDPXJImfeCGGBhIOXLm8SRypiuGdEbc9
+xkWYxRPegajuEZGvCqVs
+-----END CERTIFICATE-----`
+	keyContent = `-----BEGIN RSA PRIVATE KEY-----
+MIIEowIBAAKCAQEAumVeVawh/Q78iCcimWaJkcoafHJ/K5/ITYUFE2Q6h9Gpsq48
+fwY2puS9mbs3b1/BMMBD6ZWICCAKQOG3iE7OhYkwDspgm3wmdNJYvjwzWVyulz0M
+vaCzML8kX8Ua/IA1sG0uSWRk8UEZ0dRfXimOtP3fvCo+2b5KpQdK9FMHtubCabZ2
+tFyaGlWb3wtKVXu5DYSCcmCwuME8hMpR6MuZUsd+BfRHBDbEilf2lLgUJ7UFNpzC
+XaAansyN7vLgx+/ejQCslbpO7BuPHn1lfBvUog/vVvttPTgBVLxl5g5Qwsoq4HlR
+H3H1Gfpwtc5NYGlwxXLMOCsAUtmOcT0s9SiVkwIDAQABAoIBAD5meTJNMgO55Kjg
+ESExxpRcCIno+tHr5+6rvYtEXqPheOIsmmwb9Gfi4+Z3WpOaht5/Pz0Ppj6yGzyl
+U//6AgGKb+BDuBvVcDpjwPnOxZIBCSHwejdxeQu0scSuA97MPS0XIAvJ5FEv7ijk
+5Bht6SyGYURpECltHygoTNuGgGqmO+McCJRLE9L09lTBI6UQ/JQwWJqSr7wx6iPU
+M1Ze/srIV+7cyEPu6i0DGjS1gSQKkX68Lqn1w6oE290O+OZvleO0gZ02fLDWCZke
+aeD9+EU/Pw+rqm3H6o0szOFIpzhRp41FUdW9sybB3Yp3u7c/574E+04Z/e30LMKs
+TCtE1QECgYEA3K7KIpw0NH2HXL5C3RHcLmr204xeBfS70riBQQuVUgYdmxak2ima
+80RInskY8hRhSGTg0l+VYIH8cmjcUyqMSOELS5XfRH99r4QPiK8AguXg80T4VumY
+W3Pf+zEC2ssgP/gYthV0g0Xj5m2QxktOF9tRw5nkg739ZR4dI9lm/iECgYEA2Dnf
+uwEDGqHiQRF6/fh5BG/nGVMvrefkqx6WvTJQ3k/M/9WhxB+lr/8yH46TuS8N2b29
+FoTf3Mr9T7pr/PWkOPzoY3P56nYbKU8xSwCim9xMzhBMzj8/N9ukJvXy27/VOz56
+eQaKqnvdXNGtPJrIMDGHps2KKWlKLyAlapzjVTMCgYAA/W++tACv85g13EykfT4F
+n0k4LbsGP9DP4zABQLIMyiY72eAncmRVjwrcW36XJ2xATOONTgx3gF3HjZzfaqNy
+eD/6uNNllUTVEryXGmHgNHPL45VRnn6memCY2eFvZdXhM5W4y2PYaunY0MkDercA
++GTngbs6tBF88KOk04bYwQKBgFl68cRgsdkmnwwQYNaTKfmVGYzYaQXNzkqmWPko
+xmCJo6tHzC7ubdG8iRCYHzfmahPuuj6EdGPZuSRyYFgJi5Ftz/nAN+84OxtIQ3zn
+YWOgskQgaLh9YfsKsQ7Sf1NDOsnOnD5TX7UXl07fEpLe9vNCvAFiU8e5Y9LGudU5
+4bYTAoGBAMdX3a3bXp4cZvXNBJ/QLVyxC6fP1Q4haCR1Od3m+T00Jth2IX2dk/fl
+p6xiJT1av5JtYabv1dFKaXOS5s1kLGGuCCSKpkvFZm826aQ2AFm0XGqEQDLeei5b
+A52Kpy/YJ+RkG4BTFtAooFq6DmA0cnoP6oPvG2h6XtDJwDTPInJb
+-----END RSA PRIVATE KEY-----`
+	caContent = `-----BEGIN CERTIFICATE-----
+MIIDbTCCAlWgAwIBAgIUBJvFoCowKich7MMfseJ+DYzzirowDQYJKoZIhvcNAQEM
+BQAwRTELMAkGA1UEBhMCQVUxEzARBgNVBAgMClNvbWUtU3RhdGUxITAfBgNVBAoM
+GEludGVybmV0IFdpZGdpdHMgUHR5IEx0ZDAgFw0yMzAzMTExMzIxMDNaGA8yMTIz
+MDIxNTEzMjEwM1owRTELMAkGA1UEBhMCQVUxEzARBgNVBAgMClNvbWUtU3RhdGUx
+ITAfBgNVBAoMGEludGVybmV0IFdpZGdpdHMgUHR5IEx0ZDCCASIwDQYJKoZIhvcN
+AQEBBQADggEPADCCAQoCggEBAO4to2YMYj0bxgr2FCiweSTSFuPx33zSw2x/s9Wf
+OR41bm2DFsyYT5f3sOIKlXZEdLmOKty2e3ho3yC0EyNpVHdykkkHT3aDI17quZax
+kYi/URqqtl1Z08A22txolc04hAZisg2BypGi3vql81UW1t3zyloGnJoIAeXR9uca
+ljP6Bk3bwsxoVBLi1JtHrO0hHLQaeHmKhAyrys06X0LRdn7Px48yRZlt6FaLSa8X
+YiRM0G44bVy/h6BkoQjMYGwVmCVk6zjJ9U7ZPFqdnDMNxAfR+hjDnYodqdLDMTTR
+1NPVrnEnNwFx0AMLvgt/ba/45vZCEAmSZnFXFAJJcM7ai9ECAwEAAaNTMFEwHQYD
+VR0OBBYEFHlvMPKZoZ3G0YE7JePJOP+8J38fMB8GA1UdIwQYMBaAFHlvMPKZoZ3G
+0YE7JePJOP+8J38fMA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZIhvcNAQEMBQADggEB
+AMX8dNulADOo9uQgBMyFb9TVra7iY0zZjzv4GY5XY7scd52n6CnfAPvYBBDnTr/O
+BgNp5jaujb4+9u/2qhV3f9n+/3WOb2CmPehBgVSzlXqHeQ9lshmgwZPeem2T+8Tm
+Nnc/xQnsUfCFszUDxpkr55+aLVM22j02RWqcZ4q7TAaVYL+kdFVMc8FoqG/0ro6A
+BjE/Qn0Nn7ciX1VUjDt8l+k7ummPJTmzdi6i6E4AwO9dzrGNgGJ4aWL8cC6xYcIX
+goVIRTFeONXSDno/oPjWHpIPt7L15heMpKBHNuzPkKx2YVqPHE5QZxWfS+Lzgx+Q
+E2oTTM0rYKOZ8p6000mhvKI=
+-----END CERTIFICATE-----`
+)
+
+func TestAccount(t *testing.T) {
+	endpoints := []string{
+		"192.168.0.2:2379",
+		"192.168.0.3:2379",
+		"192.168.0.4:2379",
+	}
+	username := "foo" + stringx.Rand()
+	password := "bar"
+	anotherPassword := "any"
+
+	_, ok := GetAccount(endpoints)
+	assert.False(t, ok)
+
+	AddAccount(endpoints, username, password)
+	account, ok := GetAccount(endpoints)
+	assert.True(t, ok)
+	assert.Equal(t, username, account.User)
+	assert.Equal(t, password, account.Pass)
+
+	AddAccount(endpoints, username, anotherPassword)
+	account, ok = GetAccount(endpoints)
+	assert.True(t, ok)
+	assert.Equal(t, username, account.User)
+	assert.Equal(t, anotherPassword, account.Pass)
+}
+
+func TestTLSMethods(t *testing.T) {
+	certFile := createTempFile(t, []byte(certContent))
+	defer os.Remove(certFile)
+	keyFile := createTempFile(t, []byte(keyContent))
+	defer os.Remove(keyFile)
+	caFile := createTempFile(t, []byte(caContent))
+	defer os.Remove(caFile)
+
+	assert.NoError(t, AddTLS([]string{"foo"}, certFile, keyFile, caFile, false))
+	cfg, ok := GetTLS([]string{"foo"})
+	assert.True(t, ok)
+	assert.NotNil(t, cfg)
+
+	assert.Error(t, AddTLS([]string{"bar"}, "bad-file", keyFile, caFile, false))
+	assert.Error(t, AddTLS([]string{"bar"}, certFile, keyFile, "bad-file", false))
+}
+
+func createTempFile(t *testing.T, body []byte) string {
+	tmpFile, err := os.CreateTemp(os.TempDir(), "go-unit-*.tmp")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	tmpFile.Close()
+	if err = os.WriteFile(tmpFile.Name(), body, os.ModePerm); err != nil {
+		t.Fatal(err)
+	}
+
+	return tmpFile.Name()
+}

core/discov/internal/etcdclient.go

@@ -1,13 +1,15 @@
 //go:generate mockgen -package internal -destination etcdclient_mock.go -source etcdclient.go EtcdClient
+
 package internal
 
 import (
 	"context"
 
-	"go.etcd.io/etcd/clientv3"
+	clientv3 "go.etcd.io/etcd/client/v3"
 	"google.golang.org/grpc"
 )
 
+// EtcdClient interface represents an etcd client.
 type EtcdClient interface {
 	ActiveConnection() *grpc.ClientConn
 	Close() error

core/discov/internal/etcdclient_mock.go

@@ -1,41 +1,48 @@
 // 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
 
 import (
 	context "context"
-	gomock "github.com/golang/mock/gomock"
-	clientv3 "go.etcd.io/etcd/clientv3"
-	grpc "google.golang.org/grpc"
 	reflect "reflect"
+
+	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")
@@ -43,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")
@@ -57,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")
@@ -71,16 +78,16 @@ 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 := []interface{}{ctx, key}
+	varargs := []any{ctx, key}
 	for _, a := range opts {
 		varargs = append(varargs, a)
 	}
@@ -90,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
-func (mr *MockEtcdClientMockRecorder) Get(ctx, key interface{}, opts ...interface{}) *gomock.Call {
+// 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([]interface{}{ctx, key}, opts...)
+	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)
@@ -106,13 +113,13 @@ func (m *MockEtcdClient) Grant(ctx context.Context, ttl int64) (*clientv3.LeaseG
 	return ret0, ret1
 }
 
-// Grant indicates an expected call of Grant
-func (mr *MockEtcdClientMockRecorder) Grant(ctx, ttl interface{}) *gomock.Call {
+// 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)
@@ -121,16 +128,16 @@ func (m *MockEtcdClient) KeepAlive(ctx context.Context, id clientv3.LeaseID) (<-
 	return ret0, ret1
 }
 
-// KeepAlive indicates an expected call of KeepAlive
-func (mr *MockEtcdClientMockRecorder) KeepAlive(ctx, id interface{}) *gomock.Call {
+// 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 := []interface{}{ctx, key, val}
+	varargs := []any{ctx, key, val}
 	for _, a := range opts {
 		varargs = append(varargs, a)
 	}
@@ -140,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
-func (mr *MockEtcdClientMockRecorder) Put(ctx, key, val interface{}, opts ...interface{}) *gomock.Call {
+// 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([]interface{}{ctx, key, val}, opts...)
+	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)
@@ -156,16 +163,16 @@ func (m *MockEtcdClient) Revoke(ctx context.Context, id clientv3.LeaseID) (*clie
 	return ret0, ret1
 }
 
-// Revoke indicates an expected call of Revoke
-func (mr *MockEtcdClientMockRecorder) Revoke(ctx, id interface{}) *gomock.Call {
+// 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 := []interface{}{ctx, key}
+	varargs := []any{ctx, key}
 	for _, a := range opts {
 		varargs = append(varargs, a)
 	}
@@ -174,9 +181,9 @@ func (m *MockEtcdClient) Watch(ctx context.Context, key string, opts ...clientv3
 	return ret0
 }
 
-// Watch indicates an expected call of Watch
-func (mr *MockEtcdClientMockRecorder) Watch(ctx, key interface{}, opts ...interface{}) *gomock.Call {
+// 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([]interface{}{ctx, key}, opts...)
+	varargs := append([]any{ctx, key}, opts...)
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Watch", reflect.TypeOf((*MockEtcdClient)(nil).Watch), varargs...)
 }

core/discov/internal/listener.go

@@ -1,5 +1,6 @@
 package internal
 
+// Listener interface wraps the OnUpdate method.
 type Listener interface {
-	OnUpdate(keys []string, values []string, newKey string)
+	OnUpdate(keys, values []string, newKey string)
 }

core/discov/internal/registry.go

@@ -2,6 +2,7 @@ package internal
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"sort"
@@ -9,74 +10,181 @@ import (
 	"sync"
 	"time"
 
-	"github.com/tal-tech/go-zero/core/contextx"
-	"github.com/tal-tech/go-zero/core/lang"
-	"github.com/tal-tech/go-zero/core/logx"
-	"github.com/tal-tech/go-zero/core/syncx"
-	"github.com/tal-tech/go-zero/core/threading"
-	"go.etcd.io/etcd/clientv3"
+	"github.com/zeromicro/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/logc"
+	"github.com/zeromicro/go-zero/core/mathx"
+	"github.com/zeromicro/go-zero/core/syncx"
+	"github.com/zeromicro/go-zero/core/threading"
+	"go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
+	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
+const coolDownDeviation = 0.05
+
 var (
-	registryInstance = Registry{
+	registry = Registry{
 		clusters: make(map[string]*cluster),
 	}
-	connManager = syncx.NewResourceManager()
+	connManager      = syncx.NewResourceManager()
+	coolDownUnstable = mathx.NewUnstable(coolDownDeviation)
+	errClosed        = errors.New("etcd monitor chan has been closed")
 )
 
+// A Registry is a registry that manages the etcd client connections.
 type Registry struct {
 	clusters map[string]*cluster
-	lock     sync.Mutex
+	lock     sync.RWMutex
 }
 
+// GetRegistry returns a global Registry.
 func GetRegistry() *Registry {
-	return &registryInstance
+	return &registry
 }
 
-func (r *Registry) getCluster(endpoints []string) *cluster {
-	clusterKey := getClusterKey(endpoints)
-	r.lock.Lock()
-	defer r.lock.Unlock()
-	c, ok := r.clusters[clusterKey]
-	if !ok {
-		c = newCluster(endpoints)
-		r.clusters[clusterKey] = c
+// GetConn returns an etcd client connection associated with given endpoints.
+func (r *Registry) GetConn(endpoints []string) (EtcdClient, error) {
+	c, _ := r.getOrCreateCluster(endpoints)
+	return c.getClient()
+}
+
+// Monitor monitors the key on given etcd endpoints, notify with the given UpdateListener.
+func (r *Registry) Monitor(endpoints []string, key string, exactMatch bool, l UpdateListener) error {
+	wkey := watchKey{
+		key:        key,
+		exactMatch: exactMatch,
 	}
 
-	return c
+	c, exists := r.getOrCreateCluster(endpoints)
+	// if exists, the existing values should be updated to the listener.
+	if exists {
+		c.lock.Lock()
+		watcher, ok := c.watchers[wkey]
+		if ok {
+			watcher.listeners = append(watcher.listeners, l)
+		}
+		c.lock.Unlock()
+
+		if ok {
+			kvs := c.getCurrent(wkey)
+			for _, kv := range kvs {
+				l.OnAdd(kv)
+			}
+
+			return nil
+		}
+	}
+
+	return c.monitor(wkey, l)
 }
 
-func (r *Registry) GetConn(endpoints []string) (EtcdClient, error) {
-	return r.getCluster(endpoints).getClient()
+func (r *Registry) Unmonitor(endpoints []string, key string, exactMatch bool, l UpdateListener) {
+	c, exists := r.getCluster(endpoints)
+	if !exists {
+		return
+	}
+
+	wkey := watchKey{
+		key:        key,
+		exactMatch: exactMatch,
+	}
+
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	watcher, ok := c.watchers[wkey]
+	if !ok {
+		return
+	}
+
+	for i, listener := range watcher.listeners {
+		if listener == l {
+			watcher.listeners = append(watcher.listeners[:i], watcher.listeners[i+1:]...)
+			break
+		}
+	}
+
+	if len(watcher.listeners) == 0 {
+		if watcher.cancel != nil {
+			watcher.cancel()
+		}
+		delete(c.watchers, wkey)
+	}
 }
 
-func (r *Registry) Monitor(endpoints []string, key string, l UpdateListener) error {
-	return r.getCluster(endpoints).monitor(key, l)
+func (r *Registry) getCluster(endpoints []string) (*cluster, bool) {
+	clusterKey := getClusterKey(endpoints)
+
+	r.lock.RLock()
+	c, ok := r.clusters[clusterKey]
+	r.lock.RUnlock()
+
+	return c, ok
 }
 
-type cluster struct {
-	endpoints  []string
-	key        string
-	values     map[string]map[string]string
-	listeners  map[string][]UpdateListener
-	watchGroup *threading.RoutineGroup
-	done       chan lang.PlaceholderType
-	lock       sync.Mutex
+func (r *Registry) getOrCreateCluster(endpoints []string) (c *cluster, exists bool) {
+	c, exists = r.getCluster(endpoints)
+	if !exists {
+		clusterKey := getClusterKey(endpoints)
+
+		r.lock.Lock()
+		defer r.lock.Unlock()
+
+		// double-check locking
+		c, exists = r.clusters[clusterKey]
+		if !exists {
+			c = newCluster(endpoints)
+			r.clusters[clusterKey] = c
+		}
+	}
+
+	return
 }
 
+type (
+	watchKey struct {
+		key        string
+		exactMatch bool
+	}
+
+	watchValue struct {
+		listeners []UpdateListener
+		values    map[string]string
+		cancel    context.CancelFunc
+	}
+
+	cluster struct {
+		endpoints  []string
+		key        string
+		watchers   map[watchKey]*watchValue
+		watchGroup *threading.RoutineGroup
+		done       chan lang.PlaceholderType
+		lock       sync.RWMutex
+	}
+)
+
 func newCluster(endpoints []string) *cluster {
 	return &cluster{
 		endpoints:  endpoints,
 		key:        getClusterKey(endpoints),
-		values:     make(map[string]map[string]string),
-		listeners:  make(map[string][]UpdateListener),
+		watchers:   make(map[watchKey]*watchValue),
 		watchGroup: threading.NewRoutineGroup(),
 		done:       make(chan lang.PlaceholderType),
 	}
 }
 
-func (c *cluster) context(cli EtcdClient) context.Context {
-	return contextx.ValueOnlyFrom(cli.Ctx())
+func (c *cluster) addListener(key watchKey, l UpdateListener) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	watcher, ok := c.watchers[key]
+	if ok {
+		watcher.listeners = append(watcher.listeners, l)
+		return
+	}
+
+	val := newWatchValue()
+	val.listeners = []UpdateListener{l}
+	c.watchers[key] = val
 }
 
 func (c *cluster) getClient() (EtcdClient, error) {
@@ -90,42 +198,43 @@ func (c *cluster) getClient() (EtcdClient, error) {
 	return val.(EtcdClient), nil
 }
 
-func (c *cluster) handleChanges(key string, kvs []KV) {
-	var add []KV
-	var remove []KV
-	c.lock.Lock()
-	listeners := append([]UpdateListener(nil), c.listeners[key]...)
-	vals, ok := c.values[key]
+func (c *cluster) getCurrent(key watchKey) []KV {
+	c.lock.RLock()
+	defer c.lock.RUnlock()
+
+	watcher, ok := c.watchers[key]
 	if !ok {
-		add = kvs
-		vals = make(map[string]string)
-		for _, kv := range kvs {
-			vals[kv.Key] = kv.Val
-		}
-		c.values[key] = vals
-	} else {
-		m := make(map[string]string)
-		for _, kv := range kvs {
-			m[kv.Key] = kv.Val
-		}
-		for k, v := range vals {
-			if val, ok := m[k]; !ok || v != val {
-				remove = append(remove, KV{
-					Key: k,
-					Val: v,
-				})
-			}
-		}
-		for k, v := range m {
-			if val, ok := vals[k]; !ok || v != val {
-				add = append(add, KV{
-					Key: k,
-					Val: v,
-				})
-			}
-		}
-		c.values[key] = m
+		return nil
 	}
+
+	kvs := make([]KV, 0, len(watcher.values))
+	for k, v := range watcher.values {
+		kvs = append(kvs, KV{
+			Key: k,
+			Val: v,
+		})
+	}
+
+	return kvs
+}
+
+func (c *cluster) handleChanges(key watchKey, kvs []KV) {
+	c.lock.Lock()
+	watcher, ok := c.watchers[key]
+	if !ok {
+		c.lock.Unlock()
+		return
+	}
+
+	listeners := append([]UpdateListener(nil), watcher.listeners...)
+	// watcher.values cannot be nil
+	vals := watcher.values
+	newVals := make(map[string]string, len(kvs)+len(vals))
+	for _, kv := range kvs {
+		newVals[kv.Key] = kv.Val
+	}
+	add, remove := calculateChanges(vals, newVals)
+	watcher.values = newVals
 	c.lock.Unlock()
 
 	for _, kv := range add {
@@ -140,31 +249,43 @@ func (c *cluster) handleChanges(key string, kvs []KV) {
 	}
 }
 
-func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
-	c.lock.Lock()
-	listeners := append([]UpdateListener(nil), c.listeners[key]...)
-	c.lock.Unlock()
+func (c *cluster) handleWatchEvents(ctx context.Context, key watchKey, events []*clientv3.Event) {
+	c.lock.RLock()
+	watcher, ok := c.watchers[key]
+	if !ok {
+		c.lock.RUnlock()
+		return
+	}
+
+	listeners := append([]UpdateListener(nil), watcher.listeners...)
+	c.lock.RUnlock()
 
 	for _, ev := range events {
 		switch ev.Type {
 		case clientv3.EventTypePut:
+			evKey := string(ev.Kv.Key)
+			evVal := string(ev.Kv.Value)
 			c.lock.Lock()
-			if vals, ok := c.values[key]; ok {
-				vals[string(ev.Kv.Key)] = string(ev.Kv.Value)
-			} else {
-				c.values[key] = map[string]string{string(ev.Kv.Key): string(ev.Kv.Value)}
-			}
+			oldVal, exists := watcher.values[evKey]
+			watcher.values[evKey] = evVal
 			c.lock.Unlock()
+			if exists && oldVal == evVal {
+				// duplicate PUT with same value, skip to prevent unbounded growth
+				continue
+			}
+			if exists {
+				// key moved to a new value, notify delete of old entry first
+				for _, l := range listeners {
+					l.OnDelete(KV{Key: evKey, Val: oldVal})
+				}
+			}
 			for _, l := range listeners {
-				l.OnAdd(KV{
-					Key: string(ev.Kv.Key),
-					Val: string(ev.Kv.Value),
-				})
+				l.OnAdd(KV{Key: evKey, Val: evVal})
 			}
 		case clientv3.EventTypeDelete:
-			if vals, ok := c.values[key]; ok {
-				delete(vals, string(ev.Kv.Key))
-			}
+			c.lock.Lock()
+			delete(watcher.values, string(ev.Kv.Key))
+			c.lock.Unlock()
 			for _, l := range listeners {
 				l.OnDelete(KV{
 					Key: string(ev.Kv.Key),
@@ -172,52 +293,54 @@ func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
 				})
 			}
 		default:
-			logx.Errorf("Unknown event type: %v", ev.Type)
+			logc.Errorf(ctx, "Unknown event type: %v", ev.Type)
 		}
 	}
 }
 
-func (c *cluster) load(cli EtcdClient, key string) {
+func (c *cluster) load(cli EtcdClient, key watchKey) int64 {
 	var resp *clientv3.GetResponse
 	for {
 		var err error
-		ctx, cancel := context.WithTimeout(c.context(cli), RequestTimeout)
-		resp, err = cli.Get(ctx, makeKeyPrefix(key), clientv3.WithPrefix())
+		ctx, cancel := context.WithTimeout(cli.Ctx(), RequestTimeout)
+		if key.exactMatch {
+			resp, err = cli.Get(ctx, key.key)
+		} else {
+			resp, err = cli.Get(ctx, makeKeyPrefix(key.key), clientv3.WithPrefix())
+		}
+
 		cancel()
 		if err == nil {
 			break
 		}
 
-		logx.Error(err)
-		time.Sleep(coolDownInterval)
+		logc.Errorf(cli.Ctx(), "%s, key: %s, exactMatch: %t", err.Error(), key.key, key.exactMatch)
+		time.Sleep(coolDownUnstable.AroundDuration(coolDownInterval))
 	}
 
-	var kvs []KV
-	c.lock.Lock()
+	kvs := make([]KV, 0, len(resp.Kvs))
 	for _, ev := range resp.Kvs {
 		kvs = append(kvs, KV{
 			Key: string(ev.Key),
 			Val: string(ev.Value),
 		})
 	}
-	c.lock.Unlock()
 
 	c.handleChanges(key, kvs)
+
+	return resp.Header.Revision
 }
 
-func (c *cluster) monitor(key string, l UpdateListener) error {
-	c.lock.Lock()
-	c.listeners[key] = append(c.listeners[key], l)
-	c.lock.Unlock()
-
+func (c *cluster) monitor(key watchKey, l UpdateListener) error {
 	cli, err := c.getClient()
 	if err != nil {
 		return err
 	}
 
-	c.load(cli, key)
+	c.addListener(key, l)
+	rev := c.load(cli, key)
 	c.watchGroup.Run(func() {
-		c.watch(cli, key)
+		c.watch(cli, key, rev)
 	})
 
 	return nil
@@ -236,50 +359,106 @@ func (c *cluster) newClient() (EtcdClient, error) {
 
 func (c *cluster) reload(cli EtcdClient) {
 	c.lock.Lock()
+	// cancel the previous watches
 	close(c.done)
 	c.watchGroup.Wait()
+	keys := make([]watchKey, 0, len(c.watchers))
+	for wk, wval := range c.watchers {
+		keys = append(keys, wk)
+		if wval.cancel != nil {
+			wval.cancel()
+		}
+	}
+
 	c.done = make(chan lang.PlaceholderType)
 	c.watchGroup = threading.NewRoutineGroup()
-	var keys []string
-	for k := range c.listeners {
-		keys = append(keys, k)
-	}
 	c.lock.Unlock()
 
+	// start new watches
 	for _, key := range keys {
 		k := key
 		c.watchGroup.Run(func() {
-			c.load(cli, k)
-			c.watch(cli, k)
+			rev := c.load(cli, k)
+			c.watch(cli, k, rev)
 		})
 	}
 }
 
-func (c *cluster) watch(cli EtcdClient, key string) {
-	rch := cli.Watch(clientv3.WithRequireLeader(c.context(cli)), makeKeyPrefix(key), clientv3.WithPrefix())
+func (c *cluster) watch(cli EtcdClient, key watchKey, rev int64) {
+	for {
+		err := c.watchStream(cli, key, rev)
+		if err == nil {
+			return
+		}
+
+		if rev != 0 && errors.Is(err, rpctypes.ErrCompacted) {
+			logc.Errorf(cli.Ctx(), "etcd watch stream has been compacted, try to reload, rev %d", rev)
+			rev = c.load(cli, key)
+		}
+
+		// log the error and retry with cooldown to prevent CPU/disk exhaustion
+		logc.Error(cli.Ctx(), err)
+		time.Sleep(coolDownUnstable.AroundDuration(coolDownInterval))
+	}
+}
+
+func (c *cluster) watchStream(cli EtcdClient, key watchKey, rev int64) error {
+	ctx, rch := c.setupWatch(cli, key, rev)
+
 	for {
 		select {
 		case wresp, ok := <-rch:
 			if !ok {
-				logx.Error("etcd monitor chan has been closed")
-				return
+				return errClosed
 			}
 			if wresp.Canceled {
-				logx.Error("etcd monitor chan has been canceled")
-				return
+				return fmt.Errorf("etcd monitor chan has been canceled, error: %w", wresp.Err())
 			}
 			if wresp.Err() != nil {
-				logx.Error(fmt.Sprintf("etcd monitor chan error: %v", wresp.Err()))
-				return
+				return fmt.Errorf("etcd monitor chan error: %w", wresp.Err())
 			}
 
-			c.handleWatchEvents(key, wresp.Events)
+			c.handleWatchEvents(ctx, key, wresp.Events)
+		case <-ctx.Done():
+			return nil
 		case <-c.done:
-			return
+			return nil
 		}
 	}
 }
 
+func (c *cluster) setupWatch(cli EtcdClient, key watchKey, rev int64) (context.Context, clientv3.WatchChan) {
+	var (
+		rch  clientv3.WatchChan
+		ops  []clientv3.OpOption
+		wkey = key.key
+	)
+
+	if !key.exactMatch {
+		wkey = makeKeyPrefix(key.key)
+		ops = append(ops, clientv3.WithPrefix())
+	}
+	if rev != 0 {
+		ops = append(ops, clientv3.WithRev(rev+1))
+	}
+
+	ctx, cancel := context.WithCancel(cli.Ctx())
+
+	c.lock.Lock()
+	if watcher, ok := c.watchers[key]; ok {
+		watcher.cancel = cancel
+	} else {
+		val := newWatchValue()
+		val.cancel = cancel
+		c.watchers[key] = val
+	}
+	c.lock.Unlock()
+
+	rch = cli.Watch(clientv3.WithRequireLeader(ctx), wkey, ops...)
+
+	return ctx, rch
+}
+
 func (c *cluster) watchConnState(cli EtcdClient) {
 	watcher := newStateWatcher()
 	watcher.addListener(func() {
@@ -288,15 +467,46 @@ func (c *cluster) watchConnState(cli EtcdClient) {
 	watcher.watch(cli.ActiveConnection())
 }
 
+// DialClient dials an etcd cluster with given endpoints.
 func DialClient(endpoints []string) (EtcdClient, error) {
-	return clientv3.New(clientv3.Config{
-		Endpoints:            endpoints,
-		AutoSyncInterval:     autoSyncInterval,
-		DialTimeout:          DialTimeout,
-		DialKeepAliveTime:    dialKeepAliveTime,
-		DialKeepAliveTimeout: DialTimeout,
-		RejectOldCluster:     true,
-	})
+	cfg := clientv3.Config{
+		Endpoints:           endpoints,
+		AutoSyncInterval:    autoSyncInterval,
+		DialTimeout:         DialTimeout,
+		RejectOldCluster:    true,
+		PermitWithoutStream: true,
+	}
+	if account, ok := GetAccount(endpoints); ok {
+		cfg.Username = account.User
+		cfg.Password = account.Pass
+	}
+	if tlsCfg, ok := GetTLS(endpoints); ok {
+		cfg.TLS = tlsCfg
+	}
+
+	return clientv3.New(cfg)
+}
+
+func calculateChanges(oldVals, newVals map[string]string) (add, remove []KV) {
+	for k, v := range newVals {
+		if val, ok := oldVals[k]; !ok || v != val {
+			add = append(add, KV{
+				Key: k,
+				Val: v,
+			})
+		}
+	}
+
+	for k, v := range oldVals {
+		if val, ok := newVals[k]; !ok || v != val {
+			remove = append(remove, KV{
+				Key: k,
+				Val: v,
+			})
+		}
+	}
+
+	return add, remove
 }
 
 func getClusterKey(endpoints []string) string {
@@ -307,3 +517,10 @@ func getClusterKey(endpoints []string) string {
 func makeKeyPrefix(key string) string {
 	return fmt.Sprintf("%s%c", key, Delimiter)
 }
+
+// newWatchValue returns a watchValue that make sure values are not nil.
+func newWatchValue() *watchValue {
+	return &watchValue{
+		values: make(map[string]string),
+	}
+}

core/discov/internal/registry_test.go

@@ -2,16 +2,22 @@ package internal
 
 import (
 	"context"
+	"os"
 	"sync"
 	"testing"
+	"time"
 
-	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/contextx"
-	"github.com/tal-tech/go-zero/core/logx"
-	"github.com/tal-tech/go-zero/core/stringx"
-	"go.etcd.io/etcd/clientv3"
-	"go.etcd.io/etcd/mvcc/mvccpb"
+	"github.com/zeromicro/go-zero/core/contextx"
+	"github.com/zeromicro/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/stringx"
+	"github.com/zeromicro/go-zero/core/threading"
+	"go.etcd.io/etcd/api/v3/etcdserverpb"
+	"go.etcd.io/etcd/api/v3/mvccpb"
+	clientv3 "go.etcd.io/etcd/client/v3"
+	"go.etcd.io/etcd/client/v3/mock/mockserver"
+	"go.uber.org/mock/gomock"
 )
 
 var mockLock sync.Mutex
@@ -32,9 +38,10 @@ func setMockClient(cli EtcdClient) func() {
 }
 
 func TestGetCluster(t *testing.T) {
-	c1 := GetRegistry().getCluster([]string{"first"})
-	c2 := GetRegistry().getCluster([]string{"second"})
-	c3 := GetRegistry().getCluster([]string{"first"})
+	AddAccount([]string{"first"}, "foo", "bar")
+	c1, _ := GetRegistry().getOrCreateCluster([]string{"first"})
+	c2, _ := GetRegistry().getOrCreateCluster([]string{"second"})
+	c3, _ := GetRegistry().getOrCreateCluster([]string{"first"})
 	assert.Equal(t, c1, c3)
 	assert.NotEqual(t, c1, c2)
 }
@@ -44,6 +51,36 @@ func TestGetClusterKey(t *testing.T) {
 		getClusterKey([]string{"remotehost:5678", "localhost:1234"}))
 }
 
+func TestUnmonitor(t *testing.T) {
+	t.Run("no listener", func(t *testing.T) {
+		reg := &Registry{
+			clusters: map[string]*cluster{},
+		}
+		assert.NotPanics(t, func() {
+			reg.Unmonitor([]string{"any"}, "any", false, nil)
+		})
+	})
+
+	t.Run("no value", func(t *testing.T) {
+		reg := &Registry{
+			clusters: map[string]*cluster{
+				"any": {
+					watchers: map[watchKey]*watchValue{
+						{
+							key: "any",
+						}: {
+							values: map[string]string{},
+						},
+					},
+				},
+			},
+		}
+		assert.NotPanics(t, func() {
+			reg.Unmonitor([]string{"any"}, "another", false, nil)
+		})
+	})
+}
+
 func TestCluster_HandleChanges(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	l := NewMockUpdateListener(ctrl)
@@ -72,8 +109,14 @@ func TestCluster_HandleChanges(t *testing.T) {
 		Val: "4",
 	})
 	c := newCluster([]string{"any"})
-	c.listeners["any"] = []UpdateListener{l}
-	c.handleChanges("any", []KV{
+	key := watchKey{
+		key:        "any",
+		exactMatch: false,
+	}
+	c.watchers[key] = &watchValue{
+		listeners: []UpdateListener{l},
+	}
+	c.handleChanges(key, []KV{
 		{
 			Key: "first",
 			Val: "1",
@@ -86,8 +129,8 @@ func TestCluster_HandleChanges(t *testing.T) {
 	assert.EqualValues(t, map[string]string{
 		"first":  "1",
 		"second": "2",
-	}, c.values["any"])
-	c.handleChanges("any", []KV{
+	}, c.watchers[key].values)
+	c.handleChanges(key, []KV{
 		{
 			Key: "third",
 			Val: "3",
@@ -100,7 +143,7 @@ func TestCluster_HandleChanges(t *testing.T) {
 	assert.EqualValues(t, map[string]string{
 		"third":  "3",
 		"fourth": "4",
-	}, c.values["any"])
+	}, c.watchers[key].values)
 }
 
 func TestCluster_Load(t *testing.T) {
@@ -110,6 +153,7 @@ func TestCluster_Load(t *testing.T) {
 	restore := setMockClient(cli)
 	defer restore()
 	cli.EXPECT().Get(gomock.Any(), "any/", gomock.Any()).Return(&clientv3.GetResponse{
+		Header: &etcdserverpb.ResponseHeader{},
 		Kvs: []*mvccpb.KeyValue{
 			{
 				Key:   []byte("hello"),
@@ -119,9 +163,11 @@ func TestCluster_Load(t *testing.T) {
 	}, nil)
 	cli.EXPECT().Ctx().Return(context.Background())
 	c := &cluster{
-		values: make(map[string]map[string]string),
+		watchers: make(map[watchKey]*watchValue),
 	}
-	c.load(cli, "any")
+	c.load(cli, watchKey{
+		key: "any",
+	})
 }
 
 func TestCluster_Watch(t *testing.T) {
@@ -153,20 +199,25 @@ func TestCluster_Watch(t *testing.T) {
 			var wg sync.WaitGroup
 			wg.Add(1)
 			c := &cluster{
-				listeners: make(map[string][]UpdateListener),
-				values:    make(map[string]map[string]string),
+				watchers: make(map[watchKey]*watchValue),
+			}
+			key := watchKey{
+				key: "any",
 			}
 			listener := NewMockUpdateListener(ctrl)
-			c.listeners["any"] = []UpdateListener{listener}
+			c.watchers[key] = &watchValue{
+				listeners: []UpdateListener{listener},
+				values:    make(map[string]string),
+			}
 			listener.EXPECT().OnAdd(gomock.Any()).Do(func(kv KV) {
 				assert.Equal(t, "hello", kv.Key)
 				assert.Equal(t, "world", kv.Val)
 				wg.Done()
 			}).MaxTimes(1)
-			listener.EXPECT().OnDelete(gomock.Any()).Do(func(_ interface{}) {
+			listener.EXPECT().OnDelete(gomock.Any()).Do(func(_ any) {
 				wg.Done()
 			}).MaxTimes(1)
-			go c.watch(cli, "any")
+			go c.watch(cli, key, 0)
 			ch <- clientv3.WatchResponse{
 				Events: []*clientv3.Event{
 					{
@@ -202,17 +253,113 @@ func TestClusterWatch_RespFailures(t *testing.T) {
 			restore := setMockClient(cli)
 			defer restore()
 			ch := make(chan clientv3.WatchResponse)
-			cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch)
+			cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch).AnyTimes()
 			cli.EXPECT().Ctx().Return(context.Background()).AnyTimes()
-			c := new(cluster)
+			c := &cluster{
+				watchers: make(map[watchKey]*watchValue),
+			}
+			c.done = make(chan lang.PlaceholderType)
 			go func() {
 				ch <- resp
+				close(c.done)
 			}()
-			c.watch(cli, "any")
+			key := watchKey{
+				key: "any",
+			}
+			c.watch(cli, key, 0)
 		})
 	}
 }
 
+func TestCluster_getCurrent(t *testing.T) {
+	t.Run("no value", func(t *testing.T) {
+		c := &cluster{
+			watchers: map[watchKey]*watchValue{
+				{
+					key: "any",
+				}: {
+					values: map[string]string{},
+				},
+			},
+		}
+		assert.Nil(t, c.getCurrent(watchKey{
+			key: "another",
+		}))
+	})
+}
+
+func TestCluster_handleWatchEvents(t *testing.T) {
+	t.Run("no value", func(t *testing.T) {
+		c := &cluster{
+			watchers: map[watchKey]*watchValue{
+				{
+					key: "any",
+				}: {
+					values: map[string]string{},
+				},
+			},
+		}
+		assert.NotPanics(t, func() {
+			c.handleWatchEvents(context.Background(), watchKey{
+				key: "another",
+			}, nil)
+		})
+	})
+}
+
+func TestCluster_addListener(t *testing.T) {
+	t.Run("has listener", func(t *testing.T) {
+		c := &cluster{
+			watchers: map[watchKey]*watchValue{
+				{
+					key: "any",
+				}: {
+					listeners: make([]UpdateListener, 0),
+				},
+			},
+		}
+		assert.NotPanics(t, func() {
+			c.addListener(watchKey{
+				key: "any",
+			}, nil)
+		})
+	})
+
+	t.Run("no listener", func(t *testing.T) {
+		c := &cluster{
+			watchers: map[watchKey]*watchValue{
+				{
+					key: "any",
+				}: {
+					listeners: make([]UpdateListener, 0),
+				},
+			},
+		}
+		assert.NotPanics(t, func() {
+			c.addListener(watchKey{
+				key: "another",
+			}, nil)
+		})
+	})
+}
+
+func TestCluster_reload(t *testing.T) {
+	c := &cluster{
+		watchers:   map[watchKey]*watchValue{},
+		watchGroup: threading.NewRoutineGroup(),
+		done:       make(chan lang.PlaceholderType),
+	}
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	cli := NewMockEtcdClient(ctrl)
+	restore := setMockClient(cli)
+	defer restore()
+	assert.NotPanics(t, func() {
+		c.reload(cli)
+	})
+}
+
 func TestClusterWatch_CloseChan(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
@@ -220,13 +367,19 @@ func TestClusterWatch_CloseChan(t *testing.T) {
 	restore := setMockClient(cli)
 	defer restore()
 	ch := make(chan clientv3.WatchResponse)
-	cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch)
+	cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch).AnyTimes()
 	cli.EXPECT().Ctx().Return(context.Background()).AnyTimes()
-	c := new(cluster)
+	c := &cluster{
+		watchers: make(map[watchKey]*watchValue),
+	}
+	c.done = make(chan lang.PlaceholderType)
 	go func() {
 		close(ch)
+		close(c.done)
 	}()
-	c.watch(cli, "any")
+	c.watch(cli, watchKey{
+		key: "any",
+	}, 0)
 }
 
 func TestValueOnlyContext(t *testing.T) {
@@ -234,3 +387,211 @@ func TestValueOnlyContext(t *testing.T) {
 	ctx.Done()
 	assert.Nil(t, ctx.Err())
 }
+
+func TestDialClient(t *testing.T) {
+	svr, err := mockserver.StartMockServers(1)
+	assert.NoError(t, err)
+	svr.StartAt(0)
+
+	certFile := createTempFile(t, []byte(certContent))
+	defer os.Remove(certFile)
+	keyFile := createTempFile(t, []byte(keyContent))
+	defer os.Remove(keyFile)
+	caFile := createTempFile(t, []byte(caContent))
+	defer os.Remove(caFile)
+
+	endpoints := []string{svr.Servers[0].Address}
+	AddAccount(endpoints, "foo", "bar")
+	assert.NoError(t, AddTLS(endpoints, certFile, keyFile, caFile, false))
+
+	old := DialTimeout
+	DialTimeout = time.Millisecond
+	defer func() {
+		DialTimeout = old
+	}()
+	_, err = DialClient(endpoints)
+	assert.Error(t, err)
+}
+
+func TestRegistry_Monitor(t *testing.T) {
+	svr, err := mockserver.StartMockServers(1)
+	assert.NoError(t, err)
+	svr.StartAt(0)
+
+	endpoints := []string{svr.Servers[0].Address}
+	GetRegistry().lock.Lock()
+	GetRegistry().clusters = map[string]*cluster{
+		getClusterKey(endpoints): {
+			watchers: map[watchKey]*watchValue{
+				{
+					key:        "foo",
+					exactMatch: true,
+				}: {
+					values: map[string]string{
+						"bar": "baz",
+					},
+				},
+			},
+		},
+	}
+	GetRegistry().lock.Unlock()
+	assert.Error(t, GetRegistry().Monitor(endpoints, "foo", false, new(mockListener)))
+}
+
+func TestRegistry_Unmonitor(t *testing.T) {
+	svr, err := mockserver.StartMockServers(1)
+	assert.NoError(t, err)
+	svr.StartAt(0)
+
+	_, cancel := context.WithCancel(context.Background())
+	endpoints := []string{svr.Servers[0].Address}
+	GetRegistry().lock.Lock()
+	GetRegistry().clusters = map[string]*cluster{
+		getClusterKey(endpoints): {
+			watchers: map[watchKey]*watchValue{
+				{
+					key:        "foo",
+					exactMatch: true,
+				}: {
+					values: map[string]string{
+						"bar": "baz",
+					},
+					cancel: cancel,
+				},
+			},
+		},
+	}
+	GetRegistry().lock.Unlock()
+	l := new(mockListener)
+	assert.NoError(t, GetRegistry().Monitor(endpoints, "foo", true, l))
+	watchVals := GetRegistry().clusters[getClusterKey(endpoints)].watchers[watchKey{
+		key:        "foo",
+		exactMatch: true,
+	}]
+	assert.Equal(t, 1, len(watchVals.listeners))
+	GetRegistry().Unmonitor(endpoints, "foo", true, l)
+	watchVals = GetRegistry().clusters[getClusterKey(endpoints)].watchers[watchKey{
+		key:        "foo",
+		exactMatch: true,
+	}]
+	assert.Nil(t, watchVals)
+}
+
+// TestCluster_ConcurrentMonitor tests the race condition fix in setupWatch
+// This test specifically covers the scenario from issue #5394 where:
+// - addListener() writes to the watchers map (with lock)
+// - setupWatch() reads from the watchers map (now with lock after fix)
+// Running with -race flag will detect any race conditions
+func TestCluster_ConcurrentMonitor(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	cli := NewMockEtcdClient(ctrl)
+	cli.EXPECT().Ctx().Return(context.Background()).AnyTimes()
+	cli.EXPECT().Watch(gomock.Any(), gomock.Any(), gomock.Any()).Return(make(chan clientv3.WatchResponse)).AnyTimes()
+
+	c := &cluster{
+		endpoints:  []string{"localhost:2379"},
+		key:        "test-cluster",
+		watchers:   make(map[watchKey]*watchValue),
+		watchGroup: threading.NewRoutineGroup(),
+		done:       make(chan lang.PlaceholderType),
+		lock:       sync.RWMutex{},
+	}
+
+	// Spawn multiple concurrent operations that simulate the race condition:
+	// - Some goroutines call addListener (write to map)
+	// - Some goroutines call setupWatch (read from map)
+	var wg sync.WaitGroup
+	numGoroutines := 20
+	wg.Add(numGoroutines)
+
+	keys := []watchKey{
+		{key: "key-0", exactMatch: false},
+		{key: "key-1", exactMatch: false},
+		{key: "key-2", exactMatch: false},
+	}
+
+	for i := 0; i < numGoroutines; i++ {
+		idx := i
+		go func() {
+			defer wg.Done()
+			key := keys[idx%len(keys)]
+
+			if idx%2 == 0 {
+				// Half the goroutines add listeners (write operation)
+				c.addListener(key, &mockListener{})
+			} else {
+				// Half the goroutines setup watches (read operation)
+				_, _ = c.setupWatch(cli, key, 0)
+			}
+		}()
+	}
+
+	// Wait for all goroutines to complete
+	wg.Wait()
+
+	// Verify that watchers were correctly added
+	c.lock.RLock()
+	assert.True(t, len(c.watchers) > 0, "watchers should be added")
+	for _, watcher := range c.watchers {
+		assert.NotNil(t, watcher, "watcher should not be nil")
+	}
+	c.lock.RUnlock()
+
+	// Clean up
+	close(c.done)
+}
+
+func TestCluster_handleWatchEvents_DuplicatePut(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	listener := NewMockUpdateListener(ctrl)
+	// OnAdd must be called exactly once despite two PUT events with the same key+value.
+	listener.EXPECT().OnAdd(KV{Key: "hello", Val: "world"}).Times(1)
+
+	c := newCluster([]string{"any"})
+	key := watchKey{key: "any"}
+	c.watchers[key] = &watchValue{
+		listeners: []UpdateListener{listener},
+		values:    make(map[string]string),
+	}
+	events := []*clientv3.Event{
+		{Type: clientv3.EventTypePut, Kv: &mvccpb.KeyValue{Key: []byte("hello"), Value: []byte("world")}},
+		{Type: clientv3.EventTypePut, Kv: &mvccpb.KeyValue{Key: []byte("hello"), Value: []byte("world")}},
+	}
+	c.handleWatchEvents(context.Background(), key, events)
+}
+
+func TestCluster_handleWatchEvents_ValueChange(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	listener := NewMockUpdateListener(ctrl)
+	gomock.InOrder(
+		listener.EXPECT().OnAdd(KV{Key: "hello", Val: "world1"}),
+		listener.EXPECT().OnDelete(KV{Key: "hello", Val: "world1"}),
+		listener.EXPECT().OnAdd(KV{Key: "hello", Val: "world2"}),
+	)
+
+	c := newCluster([]string{"any"})
+	key := watchKey{key: "any"}
+	c.watchers[key] = &watchValue{
+		listeners: []UpdateListener{listener},
+		values:    make(map[string]string),
+	}
+	c.handleWatchEvents(context.Background(), key, []*clientv3.Event{
+		{Type: clientv3.EventTypePut, Kv: &mvccpb.KeyValue{Key: []byte("hello"), Value: []byte("world1")}},
+	})
+	c.handleWatchEvents(context.Background(), key, []*clientv3.Event{
+		{Type: clientv3.EventTypePut, Kv: &mvccpb.KeyValue{Key: []byte("hello"), Value: []byte("world2")}},
+	})
+}
+
+type mockListener struct {
+}
+
+func (m *mockListener) OnAdd(_ KV) {
+}
+
+func (m *mockListener) OnDelete(_ KV) {
+}

core/discov/internal/statewatcher.go

@@ -1,4 +1,5 @@
 //go:generate mockgen -package internal -destination statewatcher_mock.go -source statewatcher.go etcdConn
+
 package internal
 
 import (
@@ -18,7 +19,8 @@ type (
 		disconnected bool
 		currentState connectivity.State
 		listeners    []func()
-		lock         sync.Mutex
+		// lock only guards listeners, because only listens can be accessed by other goroutines.
+		lock sync.Mutex
 	}
 )
 
@@ -32,27 +34,33 @@ func (sw *stateWatcher) addListener(l func()) {
 	sw.lock.Unlock()
 }
 
+func (sw *stateWatcher) notifyListeners() {
+	sw.lock.Lock()
+	defer sw.lock.Unlock()
+
+	for _, l := range sw.listeners {
+		l()
+	}
+}
+
+func (sw *stateWatcher) updateState(conn etcdConn) {
+	sw.currentState = conn.GetState()
+	switch sw.currentState {
+	case connectivity.TransientFailure, connectivity.Shutdown:
+		sw.disconnected = true
+	case connectivity.Ready:
+		if sw.disconnected {
+			sw.disconnected = false
+			sw.notifyListeners()
+		}
+	}
+}
+
 func (sw *stateWatcher) watch(conn etcdConn) {
 	sw.currentState = conn.GetState()
 	for {
 		if conn.WaitForStateChange(context.Background(), sw.currentState) {
-			newState := conn.GetState()
-			sw.lock.Lock()
-			sw.currentState = newState
-
-			switch newState {
-			case connectivity.TransientFailure, connectivity.Shutdown:
-				sw.disconnected = true
-			case connectivity.Ready:
-				if sw.disconnected {
-					sw.disconnected = false
-					for _, l := range sw.listeners {
-						l()
-					}
-				}
-			}
-
-			sw.lock.Unlock()
+			sw.updateState(conn)
 		}
 	}
 }

core/discov/internal/statewatcher_mock.go

@@ -1,40 +1,47 @@
 // 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
 
 import (
 	context "context"
-	gomock "github.com/golang/mock/gomock"
-	connectivity "google.golang.org/grpc/connectivity"
 	reflect "reflect"
+
+	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")
@@ -42,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)
@@ -56,8 +63,8 @@ func (m *MocketcdConn) WaitForStateChange(ctx context.Context, sourceState conne
 	return ret0
 }
 
-// WaitForStateChange indicates an expected call of WaitForStateChange
-func (mr *MocketcdConnMockRecorder) WaitForStateChange(ctx, sourceState interface{}) *gomock.Call {
+// 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.go

@@ -1,12 +1,16 @@
 //go:generate mockgen -package internal -destination updatelistener_mock.go -source updatelistener.go UpdateListener
+
 package internal
 
 type (
+	// A KV is used to store an etcd entry with key and value.
 	KV struct {
 		Key string
 		Val string
 	}
 
+	// UpdateListener wraps the OnAdd and OnDelete methods.
+	// The implementation should be thread-safe and idempotent.
 	UpdateListener interface {
 		OnAdd(kv KV)
 		OnDelete(kv KV)

core/discov/internal/updatelistener_mock.go

@@ -1,57 +1,64 @@
 // 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
 
 import (
-	gomock "github.com/golang/mock/gomock"
 	reflect "reflect"
+
+	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
-func (mr *MockUpdateListenerMockRecorder) OnAdd(kv interface{}) *gomock.Call {
+// 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
-func (mr *MockUpdateListenerMockRecorder) OnDelete(kv interface{}) *gomock.Call {
+// 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)
 }