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

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

.codecov.yml

@@ -1,4 +1,13 @@
+coverage:
+  status:
+    patch: true
+    project: false # disabled because project coverage is not stable
+comment:
+  layout: "flags, files"
+  behavior: once
+  require_changes: true
 ignore:
-  - "doc"
-  - "example"
   - "tools"
+  - "**/mock"
+  - "**/*_mock.go"
+  - "**/*test"

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

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

.github/workflows/go.yml

@@ -7,28 +7,59 @@ 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@v4
 
-    - 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@v5
+        with:
+          go-version: '1.20'
+          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@v4
+
+  test-win:
+    name: Windows
+    runs-on: windows-latest
+    steps:
+      - name: Checkout codebase
+        uses: actions/checkout@v4
+
+      - name: Set up Go 1.x
+        uses: actions/setup-go@v5
+        with:
+          # use 1.20 to guarantee Go 1.20 compatibility
+          go-version: '1.20'
+          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/issue-translator.yml

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

.github/workflows/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@v9
+        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@v4
+      - 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.19.13.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,19 @@
+name: reviewdog
+on: [pull_request]
+jobs:
+  staticcheck:
+    name: runner / staticcheck
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - 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_on_error: true
+          # Set staticcheck flags
+          staticcheck_flags: -checks=inherit,-SA1019,-SA1029,-SA5008

.gitignore

@@ -10,12 +10,21 @@
 !*/
 !api
 
-.idea
+# ignore
+**/.idea
+**/.vscode
 **/.DS_Store
 **/logs
+**/adhoc
+**/coverage.txt
+
+# 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 create 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,16 +1,17 @@
 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/mathx"
+	"github.com/zeromicro/go-zero/core/proc"
+	"github.com/zeromicro/go-zero/core/stat"
+	"github.com/zeromicro/go-zero/core/stringx"
 )
 
 const (
@@ -18,53 +19,74 @@ const (
 	timeFormat        = "15:04:05"
 )
 
-// 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 +102,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 +130,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 +232,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 +243,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 +262,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(timeFormat), reason)
 	ew.index = (ew.index + 1) % numHistoryReasons
 	ew.count = mathx.MinInt(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,10 +9,12 @@ import (
 )
 
 func TestNopBreaker(t *testing.T) {
-	b := newNoOpBreaker()
-	assert.Equal(t, noOpBreakerName, b.Name())
+	b := NopBreaker()
+	assert.Equal(t, nopBreakerName, b.Name())
 	p, err := b.Allow()
 	assert.Nil(t, err)
+	p, err = b.AllowCtx(context.Background())
+	assert.Nil(t, err)
 	p.Accept()
 	for i := 0; i < 1000; i++ {
 		p, err := b.Allow()
@@ -21,18 +24,34 @@ func TestNopBreaker(t *testing.T) {
 	assert.Nil(t, b.Do(func() error {
 		return nil
 	}))
+	assert.Nil(t, b.DoCtx(context.Background(), func() error {
+		return nil
+	}))
 	assert.Nil(t, b.DoWithAcceptable(func() error {
 		return nil
 	}, defaultAcceptable))
+	assert.Nil(t, b.DoWithAcceptableCtx(context.Background(), func() error {
+		return nil
+	}, defaultAcceptable))
 	errDummy := errors.New("any")
 	assert.Equal(t, errDummy, b.DoWithFallback(func() error {
 		return errDummy
 	}, func(err error) error {
 		return nil
 	}))
+	assert.Equal(t, errDummy, b.DoWithFallbackCtx(context.Background(), func() error {
+		return errDummy
+	}, func(err error) error {
+		return nil
+	}))
 	assert.Equal(t, errDummy, b.DoWithFallbackAcceptable(func() error {
 		return errDummy
 	}, func(err error) error {
 		return nil
 	}, defaultAcceptable))
+	assert.Equal(t, errDummy, b.DoWithFallbackAcceptableCtx(context.Background(), func() error {
+		return errDummy
+	}, func(err error) error {
+		return nil
+	}, defaultAcceptable))
 }

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

@@ -7,9 +7,10 @@ import (
 	"encoding/base64"
 	"errors"
 
-	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logx"
 )
 
+// ErrPaddingSize indicates bad padding size.
 var ErrPaddingSize = errors.New("padding size error")
 
 type ecb struct {
@@ -26,13 +27,16 @@ func newECB(b cipher.Block) *ecb {
 
 type ecbEncrypter ecb
 
+// NewECBEncrypter returns an ECB encrypter.
 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")
@@ -52,15 +56,18 @@ func (x *ecbEncrypter) CryptBlocks(dst, src []byte) {
 
 type ecbDecrypter ecb
 
+// NewECBDecrypter returns an ECB decrypter.
 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")
@@ -78,6 +85,7 @@ func (x *ecbDecrypter) CryptBlocks(dst, src []byte) {
 	}
 }
 
+// EcbDecrypt decrypts src with the given key.
 func EcbDecrypt(key, src []byte) ([]byte, error) {
 	block, err := aes.NewCipher(key)
 	if err != nil {
@@ -92,6 +100,8 @@ func EcbDecrypt(key, src []byte) ([]byte, error) {
 	return pkcs5Unpadding(decrypted, decrypter.BlockSize())
 }
 
+// EcbDecryptBase64 decrypts base64 encoded src with the given base64 encoded key.
+// The returned string is also base64 encoded.
 func EcbDecryptBase64(key, src string) (string, error) {
 	keyBytes, err := getKeyBytes(key)
 	if err != nil {
@@ -111,6 +121,7 @@ func EcbDecryptBase64(key, src string) (string, error) {
 	return base64.StdEncoding.EncodeToString(decryptedBytes), nil
 }
 
+// EcbEncrypt encrypts src with the given key.
 func EcbEncrypt(key, src []byte) ([]byte, error) {
 	block, err := aes.NewCipher(key)
 	if err != nil {
@@ -126,6 +137,8 @@ func EcbEncrypt(key, src []byte) ([]byte, error) {
 	return crypted, nil
 }
 
+// EcbEncryptBase64 encrypts base64 encoded src with the given base64 encoded key.
+// The returned string is also base64 encoded.
 func EcbEncryptBase64(key, src string) (string, error) {
 	keyBytes, err := getKeyBytes(key)
 	if err != nil {
@@ -146,15 +159,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 {

core/codec/aesecb_test.go

@@ -0,0 +1,100 @@
+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 enough block, just nil
+	assert.Nil(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)
+	encrypter.CryptBlocks(dst, val)
+	for _, b := range dst {
+		assert.Equal(t, byte(0), b)
+	}
+
+	dst = make([]byte, 8)
+	encrypter.CryptBlocks(dst, valLong)
+	for _, b := range dst {
+		assert.Equal(t, byte(0), b)
+	}
+
+	dst = make([]byte, 8)
+	decrypter.CryptBlocks(dst, val)
+	for _, b := range dst {
+		assert.Equal(t, byte(0), b)
+	}
+
+	dst = make([]byte, 8)
+	decrypter.CryptBlocks(dst, valLong)
+	for _, b := range dst {
+		assert.Equal(t, byte(0), b)
+	}
+
+	_, 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))
+}

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,11 +23,13 @@ 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
@@ -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,17 @@ 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("")))
+}

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

@@ -7,21 +7,26 @@ import (
 	"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 +46,9 @@ type (
 	}
 )
 
+// NewRsaDecrypter returns a RsaDecrypter with the given file.
 func NewRsaDecrypter(file string) (RsaDecrypter, error) {
-	content, err := ioutil.ReadFile(file)
+	content, err := os.ReadFile(file)
 	if err != nil {
 		return nil, err
 	}
@@ -84,6 +90,7 @@ func (r *rsaDecrypter) DecryptBase64(input string) ([]byte, error) {
 	return r.Decrypt(base64Decoded)
 }
 
+// NewRsaEncrypter returns a RsaEncrypter with the given key.
 func NewRsaEncrypter(key []byte) (RsaEncrypter, error) {
 	block, _ := pem.Decode(key)
 	if block == nil {

core/codec/rsa_test.go

@@ -2,10 +2,11 @@ package codec
 
 import (
 	"encoding/base64"
+	"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 +42,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)

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,6 +75,7 @@ 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)
@@ -80,7 +84,8 @@ func (c *Cache) Del(key string) {
 	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 +96,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 +117,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 +150,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 +181,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 +190,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,25 @@ package collection
 
 import "sync"
 
+// A Queue is a FIFO queue.
 type Queue struct {
 	lock     sync.Mutex
-	elements []interface{}
+	elements []any
 	size     int
 	head     int
 	tail     int
 	count    int
 }
 
+// NewQueue returns a Queue object.
 func NewQueue(size int) *Queue {
 	return &Queue{
-		elements: make([]interface{}, size),
+		elements: make([]any, size),
 		size:     size,
 	}
 }
 
+// Empty checks if q is empty.
 func (q *Queue) Empty() bool {
 	q.lock.Lock()
 	empty := q.count == 0
@@ -26,12 +29,13 @@ 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)
+		nodes := make([]any, len(q.elements)+q.size)
 		copy(nodes, q.elements[q.head:])
 		copy(nodes[len(q.elements)-q.head:], q.elements[:q.head])
 		q.head = 0
@@ -44,7 +48,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()
 

core/collection/fifo_test.go

@@ -61,3 +61,41 @@ func TestPutMore(t *testing.T) {
 		assert.Equal(t, string(element), string(body.([]byte)))
 	}
 }
+
+func TestPutMoreWithHeaderNotZero(t *testing.T) {
+	elements := [][]byte{
+		[]byte("hello"),
+		[]byte("world"),
+		[]byte("again"),
+	}
+	queue := NewQueue(4)
+	for i := range elements {
+		queue.Put(elements[i])
+	}
+
+	// take 1
+	body, ok := queue.Take()
+	assert.True(t, ok)
+	element, ok := body.([]byte)
+	assert.True(t, ok)
+	assert.Equal(t, element, []byte("hello"))
+
+	// put more
+	queue.Put([]byte("b4"))
+	queue.Put([]byte("b5")) // will store in elements[0]
+	queue.Put([]byte("b6")) // cause expansion
+
+	results := [][]byte{
+		[]byte("world"),
+		[]byte("again"),
+		[]byte("b4"),
+		[]byte("b5"),
+		[]byte("b6"),
+	}
+
+	for _, element := range results {
+		body, ok := queue.Take()
+		assert.True(t, ok)
+		assert.Equal(t, string(element), string(body.([]byte)))
+	}
+}

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

core/collection/rollingwindow_test.go

@@ -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)
@@ -108,10 +120,12 @@ func TestRollingWindowReduce(t *testing.T) {
 func TestRollingWindowBucketTimeBoundary(t *testing.T) {
 	const size = 3
 	interval := time.Millisecond * 30
-	r := NewRollingWindow(size, interval)
+	r := NewRollingWindow[float64, *Bucket[float64]](func() *Bucket[float64] {
+		return new(Bucket[float64])
+	}, size, interval)
 	listBuckets := func() []float64 {
 		var buckets []float64
-		r.Reduce(func(b *Bucket) {
+		r.Reduce(func(b *Bucket[float64]) {
 			buckets = append(buckets, b.Sum)
 		})
 		return buckets
@@ -138,8 +152,10 @@ func TestRollingWindowBucketTimeBoundary(t *testing.T) {
 
 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 {
@@ -157,7 +173,7 @@ func TestRollingWindowDataRace(t *testing.T) {
 			case <-stop:
 				return
 			default:
-				r.Reduce(func(b *Bucket) {})
+				r.Reduce(func(b *Bucket[float64]) {})
 			}
 		}
 	}()

core/collection/safemap.go

@@ -14,19 +14,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++
@@ -40,33 +44,54 @@ func (m *SafeMap) Del(key interface{}) {
 		}
 		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{})
+		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 +105,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,8 +1,8 @@
 package collection
 
 import (
-	"github.com/tal-tech/go-zero/core/lang"
-	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/logx"
 )
 
 const (
@@ -17,61 +17,70 @@ const (
 
 // Set is not thread-safe, for concurrent use, make sure to use it with synchronization.
 type Set struct {
-	data map[interface{}]lang.PlaceholderType
+	data map[any]lang.PlaceholderType
 	tp   int
 }
 
+// NewSet returns a managed Set, can only put the values with the same type.
 func NewSet() *Set {
 	return &Set{
-		data: make(map[interface{}]lang.PlaceholderType),
+		data: make(map[any]lang.PlaceholderType),
 		tp:   untyped,
 	}
 }
 
+// NewUnmanagedSet returns an unmanaged Set, which can put values with different types.
 func NewUnmanagedSet() *Set {
 	return &Set{
-		data: make(map[interface{}]lang.PlaceholderType),
+		data: make(map[any]lang.PlaceholderType),
 		tp:   unmanaged,
 	}
 }
 
-func (s *Set) Add(i ...interface{}) {
+// Add adds i into s.
+func (s *Set) Add(i ...any) {
 	for _, each := range i {
 		s.add(each)
 	}
 }
 
+// AddInt adds int values ii into s.
 func (s *Set) AddInt(ii ...int) {
 	for _, each := range ii {
 		s.add(each)
 	}
 }
 
+// AddInt64 adds int64 values ii into s.
 func (s *Set) AddInt64(ii ...int64) {
 	for _, each := range ii {
 		s.add(each)
 	}
 }
 
+// AddUint adds uint values ii into s.
 func (s *Set) AddUint(ii ...uint) {
 	for _, each := range ii {
 		s.add(each)
 	}
 }
 
+// AddUint64 adds uint64 values ii into s.
 func (s *Set) AddUint64(ii ...uint64) {
 	for _, each := range ii {
 		s.add(each)
 	}
 }
 
+// AddStr adds string values ss into s.
 func (s *Set) AddStr(ss ...string) {
 	for _, each := range ss {
 		s.add(each)
 	}
 }
 
-func (s *Set) Contains(i interface{}) bool {
+// Contains checks if i is in s.
+func (s *Set) Contains(i any) bool {
 	if len(s.data) == 0 {
 		return false
 	}
@@ -81,8 +90,9 @@ func (s *Set) Contains(i interface{}) bool {
 	return ok
 }
 
-func (s *Set) Keys() []interface{} {
-	var keys []interface{}
+// Keys returns the keys in s.
+func (s *Set) Keys() []any {
+	var keys []any
 
 	for key := range s.data {
 		keys = append(keys, key)
@@ -91,13 +101,12 @@ func (s *Set) Keys() []interface{} {
 	return keys
 }
 
+// KeysInt returns the int keys in s.
 func (s *Set) KeysInt() []int {
 	var keys []int
 
 	for key := range s.data {
-		if intKey, ok := key.(int); !ok {
-			continue
-		} else {
+		if intKey, ok := key.(int); ok {
 			keys = append(keys, intKey)
 		}
 	}
@@ -105,13 +114,12 @@ func (s *Set) KeysInt() []int {
 	return keys
 }
 
+// KeysInt64 returns int64 keys in s.
 func (s *Set) KeysInt64() []int64 {
 	var keys []int64
 
 	for key := range s.data {
-		if intKey, ok := key.(int64); !ok {
-			continue
-		} else {
+		if intKey, ok := key.(int64); ok {
 			keys = append(keys, intKey)
 		}
 	}
@@ -119,13 +127,12 @@ func (s *Set) KeysInt64() []int64 {
 	return keys
 }
 
+// KeysUint returns uint keys in s.
 func (s *Set) KeysUint() []uint {
 	var keys []uint
 
 	for key := range s.data {
-		if intKey, ok := key.(uint); !ok {
-			continue
-		} else {
+		if intKey, ok := key.(uint); ok {
 			keys = append(keys, intKey)
 		}
 	}
@@ -133,13 +140,12 @@ func (s *Set) KeysUint() []uint {
 	return keys
 }
 
+// KeysUint64 returns uint64 keys in s.
 func (s *Set) KeysUint64() []uint64 {
 	var keys []uint64
 
 	for key := range s.data {
-		if intKey, ok := key.(uint64); !ok {
-			continue
-		} else {
+		if intKey, ok := key.(uint64); ok {
 			keys = append(keys, intKey)
 		}
 	}
@@ -147,13 +153,12 @@ func (s *Set) KeysUint64() []uint64 {
 	return keys
 }
 
+// KeysStr returns string keys in s.
 func (s *Set) KeysStr() []string {
 	var keys []string
 
 	for key := range s.data {
-		if strKey, ok := key.(string); !ok {
-			continue
-		} else {
+		if strKey, ok := key.(string); ok {
 			keys = append(keys, strKey)
 		}
 	}
@@ -161,16 +166,18 @@ func (s *Set) KeysStr() []string {
 	return keys
 }
 
-func (s *Set) Remove(i interface{}) {
+// Remove removes i from s.
+func (s *Set) Remove(i any) {
 	s.validate(i)
 	delete(s.data, i)
 }
 
+// Count returns the number of items in s.
 func (s *Set) Count() int {
 	return len(s.data)
 }
 
-func (s *Set) add(i interface{}) {
+func (s *Set) add(i any) {
 	switch s.tp {
 	case unmanaged:
 		// do nothing
@@ -182,7 +189,7 @@ func (s *Set) add(i interface{}) {
 	s.data[i] = lang.Placeholder
 }
 
-func (s *Set) setType(i interface{}) {
+func (s *Set) setType(i any) {
 	// s.tp can only be untyped here
 	switch i.(type) {
 	case int:
@@ -198,7 +205,7 @@ func (s *Set) setType(i interface{}) {
 	}
 }
 
-func (s *Set) validate(i interface{}) {
+func (s *Set) validate(i any) {
 	if s.tp == unmanaged {
 		return
 	}
@@ -206,23 +213,23 @@ func (s *Set) validate(i interface{}) {
 	switch i.(type) {
 	case int:
 		if s.tp != intType {
-			logx.Errorf("Error: element is int, but set contains elements with type %d", s.tp)
+			logx.Errorf("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)
+			logx.Errorf("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)
+			logx.Errorf("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)
+			logx.Errorf("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)
+			logx.Errorf("element is string, but set contains elements with type %d", s.tp)
 		}
 	}
 }

core/collection/set_test.go

@@ -5,7 +5,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logx"
 )
 
 func init() {
@@ -13,7 +13,7 @@ func init() {
 }
 
 func BenchmarkRawSet(b *testing.B) {
-	m := make(map[interface{}]struct{})
+	m := make(map[any]struct{})
 	for i := 0; i < b.N; i++ {
 		m[i] = struct{}{}
 		_ = m[i]
@@ -39,7 +39,7 @@ func BenchmarkSet(b *testing.B) {
 func TestAdd(t *testing.T) {
 	// given
 	set := NewUnmanagedSet()
-	values := []interface{}{1, 2, 3}
+	values := []any{1, 2, 3}
 
 	// when
 	set.Add(values...)
@@ -135,7 +135,7 @@ func TestContainsUnmanagedWithoutElements(t *testing.T) {
 func TestRemove(t *testing.T) {
 	// given
 	set := NewSet()
-	set.Add([]interface{}{1, 2, 3}...)
+	set.Add([]any{1, 2, 3}...)
 
 	// when
 	set.Remove(2)
@@ -147,7 +147,7 @@ func TestRemove(t *testing.T) {
 func TestCount(t *testing.T) {
 	// given
 	set := NewSet()
-	set.Add([]interface{}{1, 2, 3}...)
+	set.Add([]any{1, 2, 3}...)
 
 	// then
 	assert.Equal(t, set.Count(), 3)
@@ -198,5 +198,5 @@ func TestSetType(t *testing.T) {
 	set.add(1)
 	set.add("2")
 	vals := set.Keys()
-	assert.ElementsMatch(t, []interface{}{1, "2"}, vals)
+	assert.ElementsMatch(t, []any{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,48 +96,73 @@ 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; {
@@ -143,7 +179,7 @@ func (tw *TimingWheel) drainAll(fn func(key, value interface{})) {
 	}
 }
 
-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 +232,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

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,19 @@ 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_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 +78,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 +96,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 +132,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 +155,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 +173,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 +184,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 +236,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 +317,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 +405,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 +486,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 +577,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)
@@ -602,7 +612,7 @@ func TestMoveAndRemoveTask(t *testing.T) {
 		}
 	}
 	var keys []int
-	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
+	tw, _ := NewTimingWheelWithTicker(testStep, 10, func(k, v any) {
 		assert.Equal(t, "any", k)
 		assert.Equal(t, 3, v.(int))
 		keys = append(keys, v.(int))
@@ -622,7 +632,7 @@ func TestMoveAndRemoveTask(t *testing.T) {
 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,40 +2,360 @@ 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([]byte(os.ExpandEnv(string(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,
+		".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, .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, .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)
+
+	return mapping.UnmarshalJsonMap(lowerCaseKeyMap, v, mapping.WithCanonicalKeyFunc(toLowerCase))
+}
+
+// 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)
+}
+
+// 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:
+		return buildFieldsInfo(mapping.Deref(tp.Elem()), fullName)
+	case reflect.Chan, reflect.Func:
+		return nil, fmt.Errorf("unsupported type: %s", tp.Kind())
+	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:
+		var arr []any
+		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 {
+			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 strings.Join([]string{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,14 +30,19 @@ 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], "=")
@@ -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) {
@@ -31,6 +31,38 @@ func TestProperties(t *testing.T) {
 	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)

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

@@ -6,10 +6,11 @@ package internal
 
 import (
 	context "context"
-	gomock "github.com/golang/mock/gomock"
-	clientv3 "go.etcd.io/etcd/clientv3"
-	grpc "google.golang.org/grpc"
 	reflect "reflect"
+
+	gomock "github.com/golang/mock/gomock"
+	clientv3 "go.etcd.io/etcd/client/v3"
+	grpc "google.golang.org/grpc"
 )
 
 // MockEtcdClient is a mock of EtcdClient interface
@@ -80,7 +81,7 @@ func (mr *MockEtcdClientMockRecorder) Ctx() *gomock.Call {
 // 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)
 	}
@@ -91,9 +92,9 @@ func (m *MockEtcdClient) Get(ctx context.Context, key string, opts ...clientv3.O
 }
 
 // Get indicates an expected call of Get
-func (mr *MockEtcdClientMockRecorder) Get(ctx, key interface{}, opts ...interface{}) *gomock.Call {
+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...)
 }
 
@@ -107,7 +108,7 @@ func (m *MockEtcdClient) Grant(ctx context.Context, ttl int64) (*clientv3.LeaseG
 }
 
 // Grant indicates an expected call of Grant
-func (mr *MockEtcdClientMockRecorder) Grant(ctx, ttl interface{}) *gomock.Call {
+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)
 }
@@ -122,7 +123,7 @@ func (m *MockEtcdClient) KeepAlive(ctx context.Context, id clientv3.LeaseID) (<-
 }
 
 // KeepAlive indicates an expected call of KeepAlive
-func (mr *MockEtcdClientMockRecorder) KeepAlive(ctx, id interface{}) *gomock.Call {
+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)
 }
@@ -130,7 +131,7 @@ func (mr *MockEtcdClientMockRecorder) KeepAlive(ctx, id interface{}) *gomock.Cal
 // 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)
 	}
@@ -141,9 +142,9 @@ func (m *MockEtcdClient) Put(ctx context.Context, key, val string, opts ...clien
 }
 
 // Put indicates an expected call of Put
-func (mr *MockEtcdClientMockRecorder) Put(ctx, key, val interface{}, opts ...interface{}) *gomock.Call {
+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...)
 }
 
@@ -157,7 +158,7 @@ func (m *MockEtcdClient) Revoke(ctx context.Context, id clientv3.LeaseID) (*clie
 }
 
 // Revoke indicates an expected call of Revoke
-func (mr *MockEtcdClientMockRecorder) Revoke(ctx, id interface{}) *gomock.Call {
+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)
 }
@@ -165,7 +166,7 @@ func (mr *MockEtcdClientMockRecorder) Revoke(ctx, id interface{}) *gomock.Call {
 // 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)
 	}
@@ -175,8 +176,8 @@ func (m *MockEtcdClient) Watch(ctx context.Context, key string, opts ...clientv3
 }
 
 // Watch indicates an expected call of Watch
-func (mr *MockEtcdClientMockRecorder) Watch(ctx, key interface{}, opts ...interface{}) *gomock.Call {
+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,49 +10,72 @@ 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/contextx"
+	"github.com/zeromicro/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/logx"
+	"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"
 )
 
 var (
-	registryInstance = Registry{
+	registry = Registry{
 		clusters: make(map[string]*cluster),
 	}
 	connManager = syncx.NewResourceManager()
+	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.getCluster(endpoints)
+	return c.getClient()
+}
+
+// Monitor monitors the key on given etcd endpoints, notify with the given UpdateListener.
+func (r *Registry) Monitor(endpoints []string, key string, l UpdateListener) error {
+	c, exists := r.getCluster(endpoints)
+	// if exists, the existing values should be updated to the listener.
+	if exists {
+		kvs := c.getCurrent(key)
+		for _, kv := range kvs {
+			l.OnAdd(kv)
+		}
 	}
 
-	return c
+	return c.monitor(key, l)
 }
 
-func (r *Registry) GetConn(endpoints []string) (EtcdClient, error) {
-	return r.getCluster(endpoints).getClient()
-}
+func (r *Registry) getCluster(endpoints []string) (c *cluster, exists bool) {
+	clusterKey := getClusterKey(endpoints)
+	r.lock.RLock()
+	c, exists = r.clusters[clusterKey]
+	r.lock.RUnlock()
 
-func (r *Registry) Monitor(endpoints []string, key string, l UpdateListener) error {
-	return r.getCluster(endpoints).monitor(key, l)
+	if !exists {
+		r.lock.Lock()
+		defer r.lock.Unlock()
+		// double-check locking
+		c, exists = r.clusters[clusterKey]
+		if !exists {
+			c = newCluster(endpoints)
+			r.clusters[clusterKey] = c
+		}
+	}
+
+	return
 }
 
 type cluster struct {
@@ -61,7 +85,7 @@ type cluster struct {
 	listeners  map[string][]UpdateListener
 	watchGroup *threading.RoutineGroup
 	done       chan lang.PlaceholderType
-	lock       sync.Mutex
+	lock       sync.RWMutex
 }
 
 func newCluster(endpoints []string) *cluster {
@@ -90,9 +114,25 @@ func (c *cluster) getClient() (EtcdClient, error) {
 	return val.(EtcdClient), nil
 }
 
+func (c *cluster) getCurrent(key string) []KV {
+	c.lock.RLock()
+	defer c.lock.RUnlock()
+
+	var kvs []KV
+	for k, v := range c.values[key] {
+		kvs = append(kvs, KV{
+			Key: k,
+			Val: v,
+		})
+	}
+
+	return kvs
+}
+
 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]
@@ -141,9 +181,9 @@ func (c *cluster) handleChanges(key string, kvs []KV) {
 }
 
 func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
-	c.lock.Lock()
+	c.lock.RLock()
 	listeners := append([]UpdateListener(nil), c.listeners[key]...)
-	c.lock.Unlock()
+	c.lock.RUnlock()
 
 	for _, ev := range events {
 		switch ev.Type {
@@ -162,9 +202,11 @@ func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
 				})
 			}
 		case clientv3.EventTypeDelete:
+			c.lock.Lock()
 			if vals, ok := c.values[key]; ok {
 				delete(vals, string(ev.Kv.Key))
 			}
+			c.lock.Unlock()
 			for _, l := range listeners {
 				l.OnDelete(KV{
 					Key: string(ev.Kv.Key),
@@ -177,7 +219,7 @@ func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
 	}
 }
 
-func (c *cluster) load(cli EtcdClient, key string) {
+func (c *cluster) load(cli EtcdClient, key string) int64 {
 	var resp *clientv3.GetResponse
 	for {
 		var err error
@@ -188,21 +230,21 @@ func (c *cluster) load(cli EtcdClient, key string) {
 			break
 		}
 
-		logx.Error(err)
+		logx.Errorf("%s, key is %s", err.Error(), key)
 		time.Sleep(coolDownInterval)
 	}
 
 	var kvs []KV
-	c.lock.Lock()
 	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 {
@@ -215,9 +257,9 @@ func (c *cluster) monitor(key string, l UpdateListener) error {
 		return err
 	}
 
-	c.load(cli, key)
+	rev := c.load(cli, key)
 	c.watchGroup.Run(func() {
-		c.watch(cli, key)
+		c.watch(cli, key, rev)
 	})
 
 	return nil
@@ -249,33 +291,55 @@ func (c *cluster) reload(cli EtcdClient) {
 	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 string, rev int64) {
+	for {
+		err := c.watchStream(cli, key, rev)
+		if err == nil {
+			return
+		}
+
+		if rev != 0 && errors.Is(err, rpctypes.ErrCompacted) {
+			logx.Errorf("etcd watch stream has been compacted, try to reload, rev %d", rev)
+			rev = c.load(cli, key)
+		}
+
+		// log the error and retry
+		logx.Error(err)
+	}
+}
+
+func (c *cluster) watchStream(cli EtcdClient, key string, rev int64) error {
+	var rch clientv3.WatchChan
+	if rev != 0 {
+		rch = cli.Watch(clientv3.WithRequireLeader(c.context(cli)), makeKeyPrefix(key),
+			clientv3.WithPrefix(), clientv3.WithRev(rev+1))
+	} else {
+		rch = cli.Watch(clientv3.WithRequireLeader(c.context(cli)), makeKeyPrefix(key),
+			clientv3.WithPrefix())
+	}
+
 	for {
 		select {
 		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)
 		case <-c.done:
-			return
+			return nil
 		}
 	}
 }
@@ -288,15 +352,24 @@ 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 getClusterKey(endpoints []string) string {

core/discov/internal/registry_test.go

@@ -2,16 +2,21 @@ 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"
+	"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"
 )
 
 var mockLock sync.Mutex
@@ -32,9 +37,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().getCluster([]string{"first"})
+	c2, _ := GetRegistry().getCluster([]string{"second"})
+	c3, _ := GetRegistry().getCluster([]string{"first"})
 	assert.Equal(t, c1, c3)
 	assert.NotEqual(t, c1, c2)
 }
@@ -110,6 +116,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"),
@@ -163,10 +170,10 @@ func TestCluster_Watch(t *testing.T) {
 				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, "any", 0)
 			ch <- clientv3.WatchResponse{
 				Events: []*clientv3.Event{
 					{
@@ -202,13 +209,15 @@ 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.done = make(chan lang.PlaceholderType)
 			go func() {
 				ch <- resp
+				close(c.done)
 			}()
-			c.watch(cli, "any")
+			c.watch(cli, "any", 0)
 		})
 	}
 }
@@ -220,13 +229,15 @@ 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.done = make(chan lang.PlaceholderType)
 	go func() {
 		close(ch)
+		close(c.done)
 	}()
-	c.watch(cli, "any")
+	c.watch(cli, "any", 0)
 }
 
 func TestValueOnlyContext(t *testing.T) {
@@ -234,3 +245,58 @@ 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): {
+			listeners: map[string][]UpdateListener{},
+			values: map[string]map[string]string{
+				"foo": {
+					"bar": "baz",
+				},
+			},
+		},
+	}
+	GetRegistry().lock.Unlock()
+	assert.Error(t, GetRegistry().Monitor(endpoints, "foo", new(mockListener)))
+}
+
+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 (

core/discov/internal/statewatcher_mock.go

@@ -6,9 +6,10 @@ package internal
 
 import (
 	context "context"
+	reflect "reflect"
+
 	gomock "github.com/golang/mock/gomock"
 	connectivity "google.golang.org/grpc/connectivity"
-	reflect "reflect"
 )
 
 // MocketcdConn is a mock of etcdConn interface
@@ -57,7 +58,7 @@ func (m *MocketcdConn) WaitForStateChange(ctx context.Context, sourceState conne
 }
 
 // WaitForStateChange indicates an expected call of WaitForStateChange
-func (mr *MocketcdConnMockRecorder) WaitForStateChange(ctx, sourceState interface{}) *gomock.Call {
+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/updatelistener.go

@@ -1,12 +1,15 @@
 //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.
 	UpdateListener interface {
 		OnAdd(kv KV)
 		OnDelete(kv KV)

core/discov/internal/updatelistener_mock.go

@@ -5,8 +5,9 @@
 package internal
 
 import (
-	gomock "github.com/golang/mock/gomock"
 	reflect "reflect"
+
+	gomock "github.com/golang/mock/gomock"
 )
 
 // MockUpdateListener is a mock of UpdateListener interface
@@ -39,7 +40,7 @@ func (m *MockUpdateListener) OnAdd(kv KV) {
 }
 
 // OnAdd indicates an expected call of OnAdd
-func (mr *MockUpdateListenerMockRecorder) OnAdd(kv interface{}) *gomock.Call {
+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)
 }
@@ -51,7 +52,7 @@ func (m *MockUpdateListener) OnDelete(kv KV) {
 }
 
 // OnDelete indicates an expected call of OnDelete
-func (mr *MockUpdateListenerMockRecorder) OnDelete(kv interface{}) *gomock.Call {
+func (mr *MockUpdateListenerMockRecorder) OnDelete(kv any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnDelete", reflect.TypeOf((*MockUpdateListener)(nil).OnDelete), kv)
 }

core/discov/internal/vars.go

@@ -3,17 +3,21 @@ package internal
 import "time"
 
 const (
+	// Delimiter is a separator that separates the etcd path.
+	Delimiter = '/'
+
 	autoSyncInterval   = time.Minute
 	coolDownInterval   = time.Second
 	dialTimeout        = 5 * time.Second
-	dialKeepAliveTime  = 5 * time.Second
 	requestTimeout     = 3 * time.Second
-	Delimiter          = '/'
 	endpointsSeparator = ","
 )
 
 var (
-	DialTimeout    = dialTimeout
+	// DialTimeout is the dial timeout.
+	DialTimeout = dialTimeout
+	// RequestTimeout is the request timeout.
 	RequestTimeout = requestTimeout
-	NewClient      = DialClient
+	// NewClient is used to create etcd clients.
+	NewClient = DialClient
 )

core/discov/kubernetes/etcd-statefulset.yaml

@@ -0,0 +1,64 @@
+apiVersion: apps/v1
+kind: StatefulSet
+metadata:
+  name: "etcd"
+  namespace: discov
+  labels:
+    app: "etcd"
+spec:
+  serviceName: "etcd"
+  replicas: 5
+  template:
+    metadata:
+      name: "etcd"
+      labels:
+        app: "etcd"
+    spec:
+      volumes:
+      - name: etcd-pvc
+        persistentVolumeClaim:
+          claimName: etcd-pvc
+      containers:
+      - name: "etcd"
+        image: quay.io/coreos/etcd:latest
+        ports:
+        - containerPort: 2379
+          name: client
+        - containerPort: 2380
+          name: peer
+        env:
+        - name: CLUSTER_SIZE
+          value: "5"
+        - name: SET_NAME
+          value: "etcd"
+        - name: VOLNAME
+          valueFrom:
+            fieldRef:
+              apiVersion: v1
+              fieldPath: metadata.name
+        volumeMounts:
+        - name: etcd-pvc
+          mountPath: /var/lib/etcd
+          subPathExpr: $(VOLNAME)  # data mounted respectively in each pod
+        command:
+          - "/bin/sh"
+          - "-ecx"
+          - |
+
+            chmod 700 /var/lib/etcd
+
+            IP=$(hostname -i)
+            PEERS=""
+            for i in $(seq 0 $((${CLUSTER_SIZE} - 1))); do
+                PEERS="${PEERS}${PEERS:+,}${SET_NAME}-${i}=http://${SET_NAME}-${i}.${SET_NAME}:2380"
+            done
+            exec etcd --name ${HOSTNAME} \
+              --listen-peer-urls http://0.0.0.0:2380 \
+              --listen-client-urls http://0.0.0.0:2379 \
+              --advertise-client-urls http://${HOSTNAME}.${SET_NAME}.discov:2379 \
+              --initial-advertise-peer-urls http://${HOSTNAME}.${SET_NAME}:2380 \
+              --initial-cluster ${PEERS} \
+              --initial-cluster-state new \
+              --logger zap \
+              --data-dir /var/lib/etcd \
+              --auto-compaction-retention 1

core/discov/publisher.go

@@ -1,18 +1,22 @@
 package discov
 
 import (
-	"github.com/tal-tech/go-zero/core/discov/internal"
-	"github.com/tal-tech/go-zero/core/lang"
-	"github.com/tal-tech/go-zero/core/logx"
-	"github.com/tal-tech/go-zero/core/proc"
-	"github.com/tal-tech/go-zero/core/syncx"
-	"github.com/tal-tech/go-zero/core/threading"
-	"go.etcd.io/etcd/clientv3"
+	"time"
+
+	"github.com/zeromicro/go-zero/core/discov/internal"
+	"github.com/zeromicro/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/proc"
+	"github.com/zeromicro/go-zero/core/syncx"
+	"github.com/zeromicro/go-zero/core/threading"
+	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
 type (
-	PublisherOption func(client *Publisher)
+	// PubOption defines the method to customize a Publisher.
+	PubOption func(client *Publisher)
 
+	// A Publisher can be used to publish the value to an etcd cluster on the given key.
 	Publisher struct {
 		endpoints  []string
 		key        string
@@ -26,7 +30,11 @@ type (
 	}
 )
 
-func NewPublisher(endpoints []string, key, value string, opts ...PublisherOption) *Publisher {
+// NewPublisher returns a Publisher.
+// endpoints is the hosts of the etcd cluster.
+// key:value are a pair to be published.
+// opts are used to customize the Publisher.
+func NewPublisher(endpoints []string, key, value string, opts ...PubOption) *Publisher {
 	publisher := &Publisher{
 		endpoints:  endpoints,
 		key:        key,
@@ -43,13 +51,9 @@ func NewPublisher(endpoints []string, key, value string, opts ...PublisherOption
 	return publisher
 }
 
+// KeepAlive keeps key:value alive.
 func (p *Publisher) KeepAlive() error {
-	cli, err := internal.GetRegistry().GetConn(p.endpoints)
-	if err != nil {
-		return err
-	}
-
-	p.lease, err = p.register(cli)
+	cli, err := p.doRegister()
 	if err != nil {
 		return err
 	}
@@ -61,18 +65,58 @@ func (p *Publisher) KeepAlive() error {
 	return p.keepAliveAsync(cli)
 }
 
+// Pause pauses the renewing of key:value.
 func (p *Publisher) Pause() {
 	p.pauseChan <- lang.Placeholder
 }
 
+// Resume resumes the renewing of key:value.
 func (p *Publisher) Resume() {
 	p.resumeChan <- lang.Placeholder
 }
 
+// Stop stops the renewing and revokes the registration.
 func (p *Publisher) Stop() {
 	p.quit.Close()
 }
 
+func (p *Publisher) doKeepAlive() error {
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+
+	for range ticker.C {
+		select {
+		case <-p.quit.Done():
+			return nil
+		default:
+			cli, err := p.doRegister()
+			if err != nil {
+				logx.Errorf("etcd publisher doRegister: %s", err.Error())
+				break
+			}
+
+			if err := p.keepAliveAsync(cli); err != nil {
+				logx.Errorf("etcd publisher keepAliveAsync: %s", err.Error())
+				break
+			}
+
+			return nil
+		}
+	}
+
+	return nil
+}
+
+func (p *Publisher) doRegister() (internal.EtcdClient, error) {
+	cli, err := internal.GetRegistry().GetConn(p.endpoints)
+	if err != nil {
+		return nil, err
+	}
+
+	p.lease, err = p.register(cli)
+	return cli, err
+}
+
 func (p *Publisher) keepAliveAsync(cli internal.EtcdClient) error {
 	ch, err := cli.KeepAlive(cli.Ctx(), p.lease)
 	if err != nil {
@@ -85,8 +129,8 @@ func (p *Publisher) keepAliveAsync(cli internal.EtcdClient) error {
 			case _, ok := <-ch:
 				if !ok {
 					p.revoke(cli)
-					if err := p.KeepAlive(); err != nil {
-						logx.Errorf("KeepAlive: %s", err.Error())
+					if err := p.doKeepAlive(); err != nil {
+						logx.Errorf("etcd publisher KeepAlive: %s", err.Error())
 					}
 					return
 				}
@@ -95,8 +139,8 @@ func (p *Publisher) keepAliveAsync(cli internal.EtcdClient) error {
 				p.revoke(cli)
 				select {
 				case <-p.resumeChan:
-					if err := p.KeepAlive(); err != nil {
-						logx.Errorf("KeepAlive: %s", err.Error())
+					if err := p.doKeepAlive(); err != nil {
+						logx.Errorf("etcd publisher KeepAlive: %s", err.Error())
 					}
 					return
 				case <-p.quit.Done():
@@ -131,12 +175,27 @@ func (p *Publisher) register(client internal.EtcdClient) (clientv3.LeaseID, erro
 
 func (p *Publisher) revoke(cli internal.EtcdClient) {
 	if _, err := cli.Revoke(cli.Ctx(), p.lease); err != nil {
-		logx.Error(err)
+		logx.Errorf("etcd publisher revoke: %s", err.Error())
 	}
 }
 
-func WithId(id int64) PublisherOption {
+// WithId customizes a Publisher with the id.
+func WithId(id int64) PubOption {
 	return func(publisher *Publisher) {
 		publisher.id = id
 	}
 }
+
+// WithPubEtcdAccount provides the etcd username/password.
+func WithPubEtcdAccount(user, pass string) PubOption {
+	return func(pub *Publisher) {
+		RegisterAccount(pub.endpoints, user, pass)
+	}
+}
+
+// WithPubEtcdTLS provides the etcd CertFile/CertKeyFile/CACertFile.
+func WithPubEtcdTLS(certFile, certKeyFile, caFile string, insecureSkipVerify bool) PubOption {
+	return func(pub *Publisher) {
+		logx.Must(RegisterTLS(pub.endpoints, certFile, certKeyFile, caFile, insecureSkipVerify))
+	}
+}

core/discov/publisher_test.go

@@ -1,15 +1,98 @@
 package discov
 
 import (
+	"context"
 	"errors"
+	"net"
+	"os"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/discov/internal"
-	"github.com/tal-tech/go-zero/core/logx"
-	"go.etcd.io/etcd/clientv3"
+	"github.com/zeromicro/go-zero/core/discov/internal"
+	"github.com/zeromicro/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/stringx"
+	clientv3 "go.etcd.io/etcd/client/v3"
+	"golang.org/x/net/http2"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/manual"
+)
+
+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 init() {
@@ -28,12 +111,13 @@ func TestPublisher_register(t *testing.T) {
 		ID: id,
 	}, nil)
 	cli.EXPECT().Put(gomock.Any(), makeEtcdKey("thekey", id), "thevalue", gomock.Any())
-	pub := NewPublisher(nil, "thekey", "thevalue")
+	pub := NewPublisher(nil, "thekey", "thevalue",
+		WithPubEtcdAccount(stringx.Rand(), "bar"))
 	_, err := pub.register(cli)
 	assert.Nil(t, err)
 }
 
-func TestPublisher_registerWithId(t *testing.T) {
+func TestPublisher_registerWithOptions(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 	const id = 2
@@ -45,7 +129,15 @@ func TestPublisher_registerWithId(t *testing.T) {
 		ID: 1,
 	}, nil)
 	cli.EXPECT().Put(gomock.Any(), makeEtcdKey("thekey", id), "thevalue", gomock.Any())
-	pub := NewPublisher(nil, "thekey", "thevalue", WithId(id))
+
+	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)
+	pub := NewPublisher(nil, "thekey", "thevalue", WithId(id),
+		WithPubEtcdTLS(certFile, keyFile, caFile, true))
 	_, err := pub.register(cli)
 	assert.Nil(t, err)
 }
@@ -121,7 +213,7 @@ func TestPublisher_keepAliveAsyncQuit(t *testing.T) {
 	cli.EXPECT().KeepAlive(gomock.Any(), id)
 	var wg sync.WaitGroup
 	wg.Add(1)
-	cli.EXPECT().Revoke(gomock.Any(), id).Do(func(_, _ interface{}) {
+	cli.EXPECT().Revoke(gomock.Any(), id).Do(func(_, _ any) {
 		wg.Done()
 	})
 	pub := NewPublisher(nil, "thekey", "thevalue")
@@ -143,7 +235,7 @@ func TestPublisher_keepAliveAsyncPause(t *testing.T) {
 	pub := NewPublisher(nil, "thekey", "thevalue")
 	var wg sync.WaitGroup
 	wg.Add(1)
-	cli.EXPECT().Revoke(gomock.Any(), id).Do(func(_, _ interface{}) {
+	cli.EXPECT().Revoke(gomock.Any(), id).Do(func(_, _ any) {
 		pub.Stop()
 		wg.Done()
 	})
@@ -152,3 +244,105 @@ func TestPublisher_keepAliveAsyncPause(t *testing.T) {
 	pub.Pause()
 	wg.Wait()
 }
+
+func TestPublisher_Resume(t *testing.T) {
+	publisher := new(Publisher)
+	publisher.resumeChan = make(chan lang.PlaceholderType)
+	go func() {
+		publisher.Resume()
+	}()
+	go func() {
+		time.Sleep(time.Minute)
+		t.Fail()
+	}()
+	<-publisher.resumeChan
+}
+
+func TestPublisher_keepAliveAsync(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	const id clientv3.LeaseID = 1
+	conn := createMockConn(t)
+	defer conn.Close()
+	cli := internal.NewMockEtcdClient(ctrl)
+	cli.EXPECT().ActiveConnection().Return(conn).AnyTimes()
+	cli.EXPECT().Close()
+	defer cli.Close()
+	cli.ActiveConnection()
+	restore := setMockClient(cli)
+	defer restore()
+	cli.EXPECT().Ctx().AnyTimes()
+	cli.EXPECT().KeepAlive(gomock.Any(), id)
+	cli.EXPECT().Grant(gomock.Any(), timeToLive).Return(&clientv3.LeaseGrantResponse{
+		ID: 1,
+	}, nil)
+	cli.EXPECT().Put(gomock.Any(), makeEtcdKey("thekey", int64(id)), "thevalue", gomock.Any())
+	var wg sync.WaitGroup
+	wg.Add(1)
+	cli.EXPECT().Revoke(gomock.Any(), id).Do(func(_, _ any) {
+		wg.Done()
+	})
+	pub := NewPublisher([]string{"the-endpoint"}, "thekey", "thevalue")
+	pub.lease = id
+	assert.Nil(t, pub.KeepAlive())
+	pub.Stop()
+	wg.Wait()
+}
+
+func createMockConn(t *testing.T) *grpc.ClientConn {
+	lis, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatalf("Error while listening. Err: %v", err)
+	}
+	defer lis.Close()
+	lisAddr := resolver.Address{Addr: lis.Addr().String()}
+	lisDone := make(chan struct{})
+	dialDone := make(chan struct{})
+	// 1st listener accepts the connection and then does nothing
+	go func() {
+		defer close(lisDone)
+		conn, err := lis.Accept()
+		if err != nil {
+			t.Errorf("Error while accepting. Err: %v", err)
+			return
+		}
+		framer := http2.NewFramer(conn, conn)
+		if err := framer.WriteSettings(http2.Setting{}); err != nil {
+			t.Errorf("Error while writing settings. Err: %v", err)
+			return
+		}
+		<-dialDone // Close conn only after dial returns.
+	}()
+
+	r := manual.NewBuilderWithScheme("whatever")
+	r.InitialState(resolver.State{Addresses: []resolver.Address{lisAddr}})
+	client, err := grpc.DialContext(context.Background(), r.Scheme()+":///test.server",
+		grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithResolvers(r))
+	close(dialDone)
+	if err != nil {
+		t.Fatalf("Dial failed. Err: %v", err)
+	}
+
+	timeout := time.After(1 * time.Second)
+	select {
+	case <-timeout:
+		t.Fatal("timed out waiting for server to finish")
+	case <-lisDone:
+	}
+
+	return client
+}
+
+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/subscriber.go

@@ -4,31 +4,36 @@ import (
 	"sync"
 	"sync/atomic"
 
-	"github.com/tal-tech/go-zero/core/discov/internal"
-	"github.com/tal-tech/go-zero/core/syncx"
+	"github.com/zeromicro/go-zero/core/discov/internal"
+	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/syncx"
 )
 
 type (
-	subOptions struct {
-		exclusive bool
-	}
-
-	SubOption func(opts *subOptions)
+	// SubOption defines the method to customize a Subscriber.
+	SubOption func(sub *Subscriber)
 
+	// A Subscriber is used to subscribe the given key on an etcd cluster.
 	Subscriber struct {
-		items *container
+		endpoints []string
+		exclusive bool
+		items     *container
 	}
 )
 
+// NewSubscriber returns a Subscriber.
+// endpoints is the hosts of the etcd cluster.
+// key is the key to subscribe.
+// opts are used to customize the Subscriber.
 func NewSubscriber(endpoints []string, key string, opts ...SubOption) (*Subscriber, error) {
-	var subOpts subOptions
-	for _, opt := range opts {
-		opt(&subOpts)
-	}
-
 	sub := &Subscriber{
-		items: newContainer(subOpts.exclusive),
+		endpoints: endpoints,
 	}
+	for _, opt := range opts {
+		opt(sub)
+	}
+	sub.items = newContainer(sub.exclusive)
+
 	if err := internal.GetRegistry().Monitor(endpoints, key, sub.items); err != nil {
 		return nil, err
 	}
@@ -36,19 +41,35 @@ func NewSubscriber(endpoints []string, key string, opts ...SubOption) (*Subscrib
 	return sub, nil
 }
 
+// AddListener adds listener to s.
 func (s *Subscriber) AddListener(listener func()) {
 	s.items.addListener(listener)
 }
 
+// Values returns all the subscription values.
 func (s *Subscriber) Values() []string {
 	return s.items.getValues()
 }
 
-// exclusive means that key value can only be 1:1,
+// Exclusive means that key value can only be 1:1,
 // which means later added value will remove the keys associated with the same value previously.
 func Exclusive() SubOption {
-	return func(opts *subOptions) {
-		opts.exclusive = true
+	return func(sub *Subscriber) {
+		sub.exclusive = true
+	}
+}
+
+// WithSubEtcdAccount provides the etcd username/password.
+func WithSubEtcdAccount(user, pass string) SubOption {
+	return func(sub *Subscriber) {
+		RegisterAccount(sub.endpoints, user, pass)
+	}
+}
+
+// WithSubEtcdTLS provides the etcd CertFile/CertKeyFile/CACertFile.
+func WithSubEtcdTLS(certFile, certKeyFile, caFile string, insecureSkipVerify bool) SubOption {
+	return func(sub *Subscriber) {
+		logx.Must(RegisterTLS(sub.endpoints, certFile, certKeyFile, caFile, insecureSkipVerify))
 	}
 }
 
@@ -100,9 +121,9 @@ func (c *container) addKv(key, value string) ([]string, bool) {
 
 	if early {
 		return previous, true
-	} else {
-		return nil, false
 	}
+
+	return nil, false
 }
 
 func (c *container) addListener(listener func()) {

core/discov/subscriber_test.go

@@ -1,10 +1,12 @@
 package discov
 
 import (
+	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/discov/internal"
+	"github.com/zeromicro/go-zero/core/discov/internal"
+	"github.com/zeromicro/go-zero/core/stringx"
 )
 
 const (
@@ -198,3 +200,28 @@ func TestContainer(t *testing.T) {
 		}
 	}
 }
+
+func TestSubscriber(t *testing.T) {
+	sub := new(Subscriber)
+	Exclusive()(sub)
+	sub.items = newContainer(sub.exclusive)
+	var count int32
+	sub.AddListener(func() {
+		atomic.AddInt32(&count, 1)
+	})
+	sub.items.notifyChange()
+	assert.Empty(t, sub.Values())
+	assert.Equal(t, int32(1), atomic.LoadInt32(&count))
+}
+
+func TestWithSubEtcdAccount(t *testing.T) {
+	endpoints := []string{"localhost:2379"}
+	user := stringx.Rand()
+	WithSubEtcdAccount(user, "bar")(&Subscriber{
+		endpoints: endpoints,
+	})
+	account, ok := internal.GetAccount(endpoints)
+	assert.True(t, ok)
+	assert.Equal(t, user, account.User)
+	assert.Equal(t, "bar", account.Pass)
+}

core/errorx/atomicerror.go

@@ -2,14 +2,19 @@ package errorx
 
 import "sync/atomic"
 
+// AtomicError defines an atomic error.
 type AtomicError struct {
 	err atomic.Value // error
 }
 
+// Set sets the error.
 func (ae *AtomicError) Set(err error) {
-	ae.err.Store(err)
+	if err != nil {
+		ae.err.Store(err)
+	}
 }
 
+// Load returns the error.
 func (ae *AtomicError) Load() error {
 	if v := ae.err.Load(); v != nil {
 		return v.(error)

core/errorx/atomicerror_test.go

@@ -17,6 +17,15 @@ func TestAtomicError(t *testing.T) {
 	assert.Equal(t, errDummy, err.Load())
 }
 
+func TestAtomicErrorSetNil(t *testing.T) {
+	var (
+		errNil error
+		err    AtomicError
+	)
+	err.Set(errNil)
+	assert.Equal(t, errNil, err.Load())
+}
+
 func TestAtomicErrorNil(t *testing.T) {
 	var err AtomicError
 	assert.Nil(t, err.Load())

core/errorx/batcherror.go

@@ -1,22 +1,37 @@
 package errorx
 
-import "bytes"
+import (
+	"bytes"
+	"sync"
+)
 
 type (
+	// A BatchError is an error that can hold multiple errors.
 	BatchError struct {
 		errs errorArray
+		lock sync.Mutex
 	}
 
 	errorArray []error
 )
 
-func (be *BatchError) Add(err error) {
-	if err != nil {
-		be.errs = append(be.errs, err)
+// Add adds errs to be, nil errors are ignored.
+func (be *BatchError) Add(errs ...error) {
+	be.lock.Lock()
+	defer be.lock.Unlock()
+
+	for _, err := range errs {
+		if err != nil {
+			be.errs = append(be.errs, err)
+		}
 	}
 }
 
+// Err returns an error that represents all errors.
 func (be *BatchError) Err() error {
+	be.lock.Lock()
+	defer be.lock.Unlock()
+
 	switch len(be.errs) {
 	case 0:
 		return nil
@@ -27,10 +42,15 @@ func (be *BatchError) Err() error {
 	}
 }
 
+// NotNil checks if any error inside.
 func (be *BatchError) NotNil() bool {
+	be.lock.Lock()
+	defer be.lock.Unlock()
+
 	return len(be.errs) > 0
 }
 
+// Error returns a string that represents inside errors.
 func (ea errorArray) Error() string {
 	var buf bytes.Buffer
 

core/errorx/batcherror_test.go

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

core/errorx/callchain.go

@@ -1,5 +1,6 @@
 package errorx
 
+// Chain runs funs one by one until an error occurred.
 func Chain(fns ...func() error) error {
 	for _, fn := range fns {
 		if err := fn(); err != nil {