Fix pg model generation without tag (#1407)

1. fix pg model struct haven't tag
2. add pg command test from datasource

.codecov.yml

@@ -1,4 +1,3 @@
+comment: false
 ignore:
-  - "doc"
-  - "example"
-  - "tools"
+  - "tools"

.github/FUNDING.yml

@@ -0,0 +1,12 @@
+# These are supported funding model platforms
+
+github: # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]
+patreon: # Replace with a single Patreon username
+open_collective: # Replace with a single Open Collective username
+ko_fi: # Replace with a single Ko-fi username
+tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel
+community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry
+liberapay: # Replace with a single Liberapay username
+issuehunt: # Replace with a single IssueHunt username
+otechie: # Replace with a single Otechie username
+custom: https://gitee.com/kevwan/static/raw/master/images/sponsor.jpg

.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/workflows/go.yml

@@ -15,7 +15,7 @@ jobs:
     - name: Set up Go 1.x
       uses: actions/setup-go@v2
       with:
-        go-version: ^1.13
+        go-version: ^1.14
       id: go
 
     - name: Check out code into the Go module directory
@@ -25,10 +25,14 @@ jobs:
       run: |
         go get -v -t -d ./...
 
+    - name: Lint
+      run: |
+        go vet -stdmethods=false $(go list ./...)
+        go install mvdan.cc/gofumpt@latest
+        test -z "$(gofumpt -s -l -extra .)" || echo "Please run 'gofumpt -l -w -extra .'"
+
     - name: Test
       run: go test -race -coverprofile=coverage.txt -covermode=atomic ./...
 
     - name: Codecov
-      uses: codecov/codecov-action@v1.0.6
-      with:
-        token: ${{secrets.CODECOV_TOKEN}}
+      uses: codecov/codecov-action@v2

.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@v3
+        with:
+          days-before-issue-stale: 30
+          days-before-issue-close: 14
+          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/reviewdog.yml

@@ -0,0 +1,19 @@
+name: reviewdog
+on: [pull_request]
+jobs:
+  staticcheck:
+    name: runner / staticcheck
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - 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,10 +10,14 @@
 !*/
 !api
 
+# ignore
 .idea
 **/.DS_Store
 **/logs
 
+# for test purpose
+adhoc
+
 # gitlab ci
 .cache
 

CONTRIBUTING.md

@@ -0,0 +1,102 @@
+# Contributing
+
+Welcome to go-zero!
+
+-   [Before you get started](#before-you-get-started)
+    -   [Code of Conduct](#code-of-conduct)
+    -   [Community Expectations](#community-expectations)
+-   [Getting started](#getting-started)
+-   [Your First Contribution](#your-first-contribution)
+    -   [Find something to work on](#find-something-to-work-on)
+        -   [Find a good first topic](#find-a-good-first-topic)
+        -   [Work on an Issue](#work-on-an-issue)
+        -   [File an Issue](#file-an-issue)
+-   [Contributor Workflow](#contributor-workflow)
+    -   [Creating Pull Requests](#creating-pull-requests)
+    -   [Code Review](#code-review)
+    -   [Testing](#testing)
+
+# Before you get started
+
+## Code of Conduct
+
+Please make sure to read and observe our [Code of Conduct](/code-of-conduct.md).
+
+## Community Expectations
+
+go-zero is a community project driven by its community which strives to promote a healthy, friendly and productive environment.
+go-zero is a web and rpc framework written in Go. It's born to ensure the stability of the busy sites with resilient design. Builtin goctl greatly improves the development productivity.
+
+# Getting started
+
+- Fork the repository on GitHub.
+- Make your changes on your fork repository.
+- Submit a PR.
+
+
+# Your First Contribution
+
+We will help you to contribute in different areas like filing issues, developing features, fixing critical bugs and
+getting your work reviewed and merged.
+
+If you have questions about the development process,
+feel free to [file an issue](https://github.com/tal-tech/go-zero/issues/new/choose).
+
+## Find something to work on
+
+We are always in need of help, be it fixing documentation, reporting bugs or writing some code.
+Look at places where you feel best coding practices aren't followed, code refactoring is needed or tests are missing.
+Here is how you get started.
+
+### Find a good first topic
+
+[go-zero](https://github.com/tal-tech/go-zero) has beginner-friendly issues that provide a good first issue.
+For example, [go-zero](https://github.com/tal-tech/go-zero) has
+[help wanted](https://github.com/tal-tech/go-zero/issues?q=is%3Aopen+is%3Aissue+label%3A%22help+wanted%22) and
+[good first issue](https://github.com/tal-tech/go-zero/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22)
+labels for issues that should not need deep knowledge of the system.
+We can help new contributors who wish to work on such issues.
+
+Another good way to contribute is to find a documentation improvement, such as a missing/broken link.
+Please see [Contributing](#contributing) below for the workflow.
+
+#### Work on an issue
+
+When you are willing to take on an issue, just reply on the issue. The maintainer will assign it to you.
+
+### File an Issue
+
+While we encourage everyone to contribute code, it is also appreciated when someone reports an issue.
+
+Please follow the prompted submission guidelines while opening an issue.
+
+# Contributor Workflow
+
+Please do not ever hesitate to ask a question or send a pull request.
+
+This is a rough outline of what a contributor's workflow looks like:
+
+- Create a topic branch from where to base the contribution. This is usually master.
+- Make commits of logical units.
+- Push changes in a topic branch to a personal fork of the repository.
+- Submit a pull request to [go-zero](https://github.com/tal-tech/go-zero).
+
+## Creating Pull Requests
+
+Pull requests are often called simply "PR".
+go-zero generally follows the standard [github pull request](https://help.github.com/articles/about-pull-requests/) process.
+To submit a proposed change, please develop the code/fix and add new test cases.
+After that, run these local verifications before submitting pull request to predict the pass or
+fail of continuous integration.
+
+* Format the code with `gofmt`
+* Run the test with data race enabled `go test -race ./...`
+
+## Code Review
+
+To make it easier for your PR to receive reviews, consider the reviewers will need you to:
+
+* follow [good coding guidelines](https://github.com/golang/go/wiki/CodeReviewComments).
+* write [good commit messages](https://chris.beams.io/posts/git-commit/).
+* break large changes into a logical series of smaller patches which individually make easily understandable changes, and in aggregate solve a broader issue.
+

ROADMAP.md

@@ -0,0 +1,28 @@
+# go-zero Roadmap
+
+This document defines a high level roadmap for go-zero development and upcoming releases.
+Community and contributor involvement is vital for successfully implementing all desired items for each release.
+We hope that the items listed below will inspire further engagement from the community to keep go-zero progressing and shipping exciting and valuable features.
+
+## 2021 Q2
+- [x] Support service discovery through K8S client api
+- [x] Log full sql statements for easier sql problem solving
+
+## 2021 Q3
+- [x] Support `goctl model pg` to support PostgreSQL code generation
+- [x] Adapt builtin tracing mechanism to opentracing solutions
+
+## 2021 Q4
+- [x] Support `username/password` authentication in ETCD
+- [x] Support `SSL/TLS` in ETCD
+- [x] Support `SSL/TLS` in `zRPC`
+- [x] Support `TLS` in redis connections
+- [x] Support `goctl bug` to report bugs conveniently
+
+## 2022
+- [ ] Support `goctl mock` command to start a mocking server with given `.api` file
+- [ ] Add `httpx.Client` with governance, like circuit breaker etc.
+- [ ] Support `goctl doctor` command to report potential issues for given service
+- [ ] Support `context` in redis related methods for timeout and tracing
+- [ ] Support `context` in sql related methods for timeout and tracing
+- [ ] Support `context` in mongodb related methods for timeout and tracing

code-of-conduct.md

@@ -0,0 +1,76 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to make participation in our project and
+our community a harassment-free experience for everyone, regardless of age, body
+size, disability, ethnicity, sex characteristics, gender identity and expression,
+level of experience, education, socio-economic status, nationality, personal
+appearance, race, religion, or sexual identity and orientation.
+
+## Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+  advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+  address, without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers 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, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+## Scope
+
+This Code of Conduct applies within all project spaces, and it also applies when
+an individual is representing the project or its community in public spaces.
+Examples of representing a project or community include using an official
+project e-mail address, posting via an official social media account, or acting
+as an appointed representative at an online or offline event. Representation of
+a project may be further defined and clarified by project maintainers.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the project team at [INSERT EMAIL ADDRESS]. All
+complaints will be reviewed and investigated and will result in a response that
+is deemed necessary and appropriate to the circumstances. The project team is
+obligated to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
+available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see
+https://www.contributor-covenant.org/faq

core/breaker/breaker_test.go

@@ -122,8 +122,7 @@ func BenchmarkGoogleBreaker(b *testing.B) {
 	}
 }
 
-type mockedPromise struct {
-}
+type mockedPromise struct{}
 
 func (m *mockedPromise) Accept() {
 }

core/breaker/googlebreaker.go

@@ -94,7 +94,7 @@ func (b *googleBreaker) markFailure() {
 	b.stat.Add(0)
 }
 
-func (b *googleBreaker) history() (accepts int64, total int64) {
+func (b *googleBreaker) history() (accepts, total int64) {
 	b.stat.Reduce(func(b *collection.Bucket) {
 		accepts += int64(b.Sum)
 		total += b.Count

core/codec/aesecb_test.go

@@ -40,7 +40,7 @@ func TestAesEcbBase64(t *testing.T) {
 		// more than 32 chars
 		badKey2 = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
 	)
-	var key = []byte("q4t7w!z%C*F-JaNdRgUjXn2r5u8x/A?D")
+	key := []byte("q4t7w!z%C*F-JaNdRgUjXn2r5u8x/A?D")
 	b64Key := base64.StdEncoding.EncodeToString(key)
 	b64Val := base64.StdEncoding.EncodeToString([]byte(val))
 	_, err := EcbEncryptBase64(badKey1, val)

core/collection/cache.go

@@ -34,7 +34,7 @@ type (
 		expire         time.Duration
 		timingWheel    *TimingWheel
 		lruCache       lru
-		barrier        syncx.SharedCalls
+		barrier        syncx.SingleFlight
 		unstableExpiry mathx.Unstable
 		stats          *cacheStat
 	}
@@ -46,7 +46,7 @@ func NewCache(expire time.Duration, opts ...CacheOption) (*Cache, error) {
 		data:           make(map[string]interface{}),
 		expire:         expire,
 		lruCache:       emptyLruCache,
-		barrier:        syncx.NewSharedCalls(),
+		barrier:        syncx.NewSingleFlight(),
 		unstableExpiry: mathx.NewUnstable(expiryDeviation),
 	}
 

core/collection/rollingwindow_test.go

@@ -139,7 +139,7 @@ func TestRollingWindowBucketTimeBoundary(t *testing.T) {
 func TestRollingWindowDataRace(t *testing.T) {
 	const size = 3
 	r := NewRollingWindow(size, duration)
-	var stop = make(chan bool)
+	stop := make(chan bool)
 	go func() {
 		for {
 			select {

core/collection/set.go

@@ -106,9 +106,7 @@ 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)
 		}
 	}
@@ -121,9 +119,7 @@ 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)
 		}
 	}
@@ -136,9 +132,7 @@ 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)
 		}
 	}
@@ -151,9 +145,7 @@ 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)
 		}
 	}
@@ -166,9 +158,7 @@ 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)
 		}
 	}

core/collection/timingwheel.go

@@ -151,7 +151,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

core/conf/config.go

@@ -25,7 +25,7 @@ func LoadConfig(file string, v interface{}, opts ...Option) error {
 
 	loader, ok := loaders[path.Ext(file)]
 	if !ok {
-		return fmt.Errorf("unrecoginized file type: %s", file)
+		return fmt.Errorf("unrecognized file type: %s", file)
 	}
 
 	var opt options

core/contextx/deadline.go

@@ -1,19 +0,0 @@
-package contextx
-
-import (
-	"context"
-	"time"
-)
-
-// ShrinkDeadline returns a new Context with proper deadline base on the given ctx and timeout.
-// And returns a cancel function as well.
-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_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_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/tal-tech/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/tal-tech/go-zero/core/discov/internal"
+	"github.com/tal-tech/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/config.go

@@ -4,8 +4,24 @@ import "errors"
 
 // EtcdConf is the config item with the given key on etcd.
 type EtcdConf struct {
-	Hosts []string
-	Key   string
+	Hosts              []string
+	Key                string
+	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
+}
+
+// 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.

core/discov/config_test.go

@@ -44,3 +44,39 @@ 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())
+	}
+}

core/discov/internal/accountmanager.go

@@ -0,0 +1,75 @@
+package internal
+
+import (
+	"crypto/tls"
+	"crypto/x509"
+	"io/ioutil"
+	"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 := ioutil.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,34 @@
+package internal
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/stringx"
+)
+
+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)
+}

core/discov/internal/etcdclient.go

@@ -5,7 +5,7 @@ package internal
 import (
 	"context"
 
-	"go.etcd.io/etcd/clientv3"
+	clientv3 "go.etcd.io/etcd/client/v3"
 	"google.golang.org/grpc"
 )
 

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

core/discov/internal/listener.go

@@ -2,5 +2,5 @@ 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

@@ -14,7 +14,7 @@ import (
 	"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"
+	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
 var (
@@ -37,25 +37,35 @@ func GetRegistry() *Registry {
 
 // GetConn returns an etcd client connection associated with given endpoints.
 func (r *Registry) GetConn(endpoints []string) (EtcdClient, error) {
-	return r.getCluster(endpoints).getClient()
+	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 {
-	return r.getCluster(endpoints).monitor(key, l)
+	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.monitor(key, l)
 }
 
-func (r *Registry) getCluster(endpoints []string) *cluster {
+func (r *Registry) getCluster(endpoints []string) (c *cluster, exists bool) {
 	clusterKey := getClusterKey(endpoints)
 	r.lock.Lock()
 	defer r.lock.Unlock()
-	c, ok := r.clusters[clusterKey]
-	if !ok {
+	c, exists = r.clusters[clusterKey]
+	if !exists {
 		c = newCluster(endpoints)
 		r.clusters[clusterKey] = c
 	}
 
-	return c
+	return
 }
 
 type cluster struct {
@@ -94,6 +104,21 @@ func (c *cluster) getClient() (EtcdClient, error) {
 	return val.(EtcdClient), nil
 }
 
+func (c *cluster) getCurrent(key string) []KV {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	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
@@ -197,14 +222,12 @@ func (c *cluster) load(cli EtcdClient, key string) {
 	}
 
 	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)
 }
@@ -260,26 +283,34 @@ func (c *cluster) reload(cli EtcdClient) {
 }
 
 func (c *cluster) watch(cli EtcdClient, key string) {
+	for {
+		if c.watchStream(cli, key) {
+			return
+		}
+	}
+}
+
+func (c *cluster) watchStream(cli EtcdClient, key string) bool {
 	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 false
 			}
 			if wresp.Canceled {
-				logx.Error("etcd monitor chan has been canceled")
-				return
+				logx.Errorf("etcd monitor chan has been canceled, error: %v", wresp.Err())
+				return false
 			}
 			if wresp.Err() != nil {
 				logx.Error(fmt.Sprintf("etcd monitor chan error: %v", wresp.Err()))
-				return
+				return false
 			}
 
 			c.handleWatchEvents(key, wresp.Events)
 		case <-c.done:
-			return
+			return true
 		}
 	}
 }
@@ -294,14 +325,23 @@ func (c *cluster) watchConnState(cli EtcdClient) {
 
 // DialClient dials an etcd cluster with given endpoints.
 func DialClient(endpoints []string) (EtcdClient, error) {
-	return clientv3.New(clientv3.Config{
+	cfg := clientv3.Config{
 		Endpoints:            endpoints,
 		AutoSyncInterval:     autoSyncInterval,
 		DialTimeout:          DialTimeout,
 		DialKeepAliveTime:    dialKeepAliveTime,
 		DialKeepAliveTimeout: DialTimeout,
 		RejectOldCluster:     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

@@ -8,10 +8,11 @@ import (
 	"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/lang"
 	"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"
+	"go.etcd.io/etcd/api/v3/mvccpb"
+	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
 var mockLock sync.Mutex
@@ -32,9 +33,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)
 }
@@ -202,11 +204,13 @@ 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")
 		})
@@ -220,11 +224,13 @@ 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")
 }

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

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

core/discov/publisher.go

@@ -7,12 +7,12 @@ import (
 	"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"
+	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
 type (
-	// PublisherOption defines the method to customize a Publisher.
-	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 {
@@ -32,7 +32,7 @@ type (
 // 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 ...PublisherOption) *Publisher {
+func NewPublisher(endpoints []string, key, value string, opts ...PubOption) *Publisher {
 	publisher := &Publisher{
 		endpoints:  endpoints,
 		key:        key,
@@ -146,8 +146,22 @@ func (p *Publisher) revoke(cli internal.EtcdClient) {
 }
 
 // WithId customizes a Publisher with the id.
-func WithId(id int64) PublisherOption {
+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

@@ -11,7 +11,8 @@ 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"
-	"go.etcd.io/etcd/clientv3"
+	"github.com/tal-tech/go-zero/core/stringx"
+	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
 func init() {
@@ -30,7 +31,8 @@ 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)
 }

core/discov/subscriber.go

@@ -5,20 +5,19 @@ import (
 	"sync/atomic"
 
 	"github.com/tal-tech/go-zero/core/discov/internal"
+	"github.com/tal-tech/go-zero/core/logx"
 	"github.com/tal-tech/go-zero/core/syncx"
 )
 
 type (
-	subOptions struct {
-		exclusive bool
-	}
-
 	// SubOption defines the method to customize a Subscriber.
-	SubOption func(opts *subOptions)
+	SubOption func(sub *Subscriber)
 
 	// A Subscriber is used to subscribe the given key on a etcd cluster.
 	Subscriber struct {
-		items *container
+		endpoints []string
+		exclusive bool
+		items     *container
 	}
 )
 
@@ -27,14 +26,14 @@ type (
 // 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
 	}
@@ -55,8 +54,22 @@ func (s *Subscriber) Values() []string {
 // 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))
 	}
 }
 

core/discov/subscriber_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/tal-tech/go-zero/core/discov/internal"
+	"github.com/tal-tech/go-zero/core/stringx"
 )
 
 const (
@@ -201,11 +202,9 @@ func TestContainer(t *testing.T) {
 }
 
 func TestSubscriber(t *testing.T) {
-	var opt subOptions
-	Exclusive()(&opt)
-
 	sub := new(Subscriber)
-	sub.items = newContainer(opt.exclusive)
+	Exclusive()(sub)
+	sub.items = newContainer(sub.exclusive)
 	var count int32
 	sub.AddListener(func() {
 		atomic.AddInt32(&count, 1)
@@ -214,3 +213,15 @@ func TestSubscriber(t *testing.T) {
 	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

@@ -9,7 +9,9 @@ type AtomicError struct {
 
 // 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.

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

@@ -8,7 +8,7 @@ import (
 )
 
 func TestChain(t *testing.T) {
-	var errDummy = errors.New("dummy")
+	errDummy := errors.New("dummy")
 	assert.Nil(t, Chain(func() error {
 		return nil
 	}, func() error {

core/fs/files+polyfill.go

@@ -1,3 +1,4 @@
+//go:build windows
 // +build windows
 
 package fs

core/fs/files.go

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

core/fx/fn_test.go

@@ -1,354 +0,0 @@
-package fx
-
-import (
-	"io/ioutil"
-	"log"
-	"runtime"
-	"sync"
-	"sync/atomic"
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/stringx"
-)
-
-func TestBuffer(t *testing.T) {
-	const N = 5
-	var count int32
-	var wait sync.WaitGroup
-	wait.Add(1)
-	From(func(source chan<- interface{}) {
-		ticker := time.NewTicker(10 * time.Millisecond)
-		defer ticker.Stop()
-
-		for i := 0; i < 2*N; i++ {
-			select {
-			case source <- i:
-				atomic.AddInt32(&count, 1)
-			case <-ticker.C:
-				wait.Done()
-				return
-			}
-		}
-	}).Buffer(N).ForAll(func(pipe <-chan interface{}) {
-		wait.Wait()
-		// why N+1, because take one more to wait for sending into the channel
-		assert.Equal(t, int32(N+1), atomic.LoadInt32(&count))
-	})
-}
-
-func TestBufferNegative(t *testing.T) {
-	var result int
-	Just(1, 2, 3, 4).Buffer(-1).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
-		for item := range pipe {
-			result += item.(int)
-		}
-		return result, nil
-	})
-	assert.Equal(t, 10, result)
-}
-
-func TestCount(t *testing.T) {
-	tests := []struct {
-		name     string
-		elements []interface{}
-	}{
-		{
-			name: "no elements with nil",
-		},
-		{
-			name:     "no elements",
-			elements: []interface{}{},
-		},
-		{
-			name:     "1 element",
-			elements: []interface{}{1},
-		},
-		{
-			name:     "multiple elements",
-			elements: []interface{}{1, 2, 3},
-		},
-	}
-
-	for _, test := range tests {
-		t.Run(test.name, func(t *testing.T) {
-			val := Just(test.elements...).Count()
-			assert.Equal(t, len(test.elements), val)
-		})
-	}
-}
-
-func TestDone(t *testing.T) {
-	var count int32
-	Just(1, 2, 3).Walk(func(item interface{}, pipe chan<- interface{}) {
-		time.Sleep(time.Millisecond * 100)
-		atomic.AddInt32(&count, int32(item.(int)))
-	}).Done()
-	assert.Equal(t, int32(6), count)
-}
-
-func TestJust(t *testing.T) {
-	var result int
-	Just(1, 2, 3, 4).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
-		for item := range pipe {
-			result += item.(int)
-		}
-		return result, nil
-	})
-	assert.Equal(t, 10, result)
-}
-
-func TestDistinct(t *testing.T) {
-	var result int
-	Just(4, 1, 3, 2, 3, 4).Distinct(func(item interface{}) interface{} {
-		return item
-	}).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
-		for item := range pipe {
-			result += item.(int)
-		}
-		return result, nil
-	})
-	assert.Equal(t, 10, result)
-}
-
-func TestFilter(t *testing.T) {
-	var result int
-	Just(1, 2, 3, 4).Filter(func(item interface{}) bool {
-		return item.(int)%2 == 0
-	}).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
-		for item := range pipe {
-			result += item.(int)
-		}
-		return result, nil
-	})
-	assert.Equal(t, 6, result)
-}
-
-func TestForAll(t *testing.T) {
-	var result int
-	Just(1, 2, 3, 4).Filter(func(item interface{}) bool {
-		return item.(int)%2 == 0
-	}).ForAll(func(pipe <-chan interface{}) {
-		for item := range pipe {
-			result += item.(int)
-		}
-	})
-	assert.Equal(t, 6, result)
-}
-
-func TestGroup(t *testing.T) {
-	var groups [][]int
-	Just(10, 11, 20, 21).Group(func(item interface{}) interface{} {
-		v := item.(int)
-		return v / 10
-	}).ForEach(func(item interface{}) {
-		v := item.([]interface{})
-		var group []int
-		for _, each := range v {
-			group = append(group, each.(int))
-		}
-		groups = append(groups, group)
-	})
-
-	assert.Equal(t, 2, len(groups))
-	for _, group := range groups {
-		assert.Equal(t, 2, len(group))
-		assert.True(t, group[0]/10 == group[1]/10)
-	}
-}
-
-func TestHead(t *testing.T) {
-	var result int
-	Just(1, 2, 3, 4).Head(2).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
-		for item := range pipe {
-			result += item.(int)
-		}
-		return result, nil
-	})
-	assert.Equal(t, 3, result)
-}
-
-func TestHeadZero(t *testing.T) {
-	assert.Panics(t, func() {
-		Just(1, 2, 3, 4).Head(0).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
-			return nil, nil
-		})
-	})
-}
-
-func TestHeadMore(t *testing.T) {
-	var result int
-	Just(1, 2, 3, 4).Head(6).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
-		for item := range pipe {
-			result += item.(int)
-		}
-		return result, nil
-	})
-	assert.Equal(t, 10, result)
-}
-
-func TestMap(t *testing.T) {
-	log.SetOutput(ioutil.Discard)
-
-	tests := []struct {
-		mapper MapFunc
-		expect int
-	}{
-		{
-			mapper: func(item interface{}) interface{} {
-				v := item.(int)
-				return v * v
-			},
-			expect: 30,
-		},
-		{
-			mapper: func(item interface{}) interface{} {
-				v := item.(int)
-				if v%2 == 0 {
-					return 0
-				}
-				return v * v
-			},
-			expect: 10,
-		},
-		{
-			mapper: func(item interface{}) interface{} {
-				v := item.(int)
-				if v%2 == 0 {
-					panic(v)
-				}
-				return v * v
-			},
-			expect: 10,
-		},
-	}
-
-	// Map(...) works even WithWorkers(0)
-	for i, test := range tests {
-		t.Run(stringx.Rand(), func(t *testing.T) {
-			var result int
-			var workers int
-			if i%2 == 0 {
-				workers = 0
-			} else {
-				workers = runtime.NumCPU()
-			}
-			From(func(source chan<- interface{}) {
-				for i := 1; i < 5; i++ {
-					source <- i
-				}
-			}).Map(test.mapper, WithWorkers(workers)).Reduce(
-				func(pipe <-chan interface{}) (interface{}, error) {
-					for item := range pipe {
-						result += item.(int)
-					}
-					return result, nil
-				})
-
-			assert.Equal(t, test.expect, result)
-		})
-	}
-}
-
-func TestMerge(t *testing.T) {
-	Just(1, 2, 3, 4).Merge().ForEach(func(item interface{}) {
-		assert.ElementsMatch(t, []interface{}{1, 2, 3, 4}, item.([]interface{}))
-	})
-}
-
-func TestParallelJust(t *testing.T) {
-	var count int32
-	Just(1, 2, 3).Parallel(func(item interface{}) {
-		time.Sleep(time.Millisecond * 100)
-		atomic.AddInt32(&count, int32(item.(int)))
-	}, UnlimitedWorkers())
-	assert.Equal(t, int32(6), count)
-}
-
-func TestReverse(t *testing.T) {
-	Just(1, 2, 3, 4).Reverse().Merge().ForEach(func(item interface{}) {
-		assert.ElementsMatch(t, []interface{}{4, 3, 2, 1}, item.([]interface{}))
-	})
-}
-
-func TestSort(t *testing.T) {
-	var prev int
-	Just(5, 3, 7, 1, 9, 6, 4, 8, 2).Sort(func(a, b interface{}) bool {
-		return a.(int) < b.(int)
-	}).ForEach(func(item interface{}) {
-		next := item.(int)
-		assert.True(t, prev < next)
-		prev = next
-	})
-}
-
-func TestSplit(t *testing.T) {
-	assert.Panics(t, func() {
-		Just(1, 2, 3, 4, 5, 6, 7, 8, 9, 10).Split(0).Done()
-	})
-	var chunks [][]interface{}
-	Just(1, 2, 3, 4, 5, 6, 7, 8, 9, 10).Split(4).ForEach(func(item interface{}) {
-		chunk := item.([]interface{})
-		chunks = append(chunks, chunk)
-	})
-	assert.EqualValues(t, [][]interface{}{
-		{1, 2, 3, 4},
-		{5, 6, 7, 8},
-		{9, 10},
-	}, chunks)
-}
-
-func TestTail(t *testing.T) {
-	var result int
-	Just(1, 2, 3, 4).Tail(2).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
-		for item := range pipe {
-			result += item.(int)
-		}
-		return result, nil
-	})
-	assert.Equal(t, 7, result)
-}
-
-func TestTailZero(t *testing.T) {
-	assert.Panics(t, func() {
-		Just(1, 2, 3, 4).Tail(0).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
-			return nil, nil
-		})
-	})
-}
-
-func TestWalk(t *testing.T) {
-	var result int
-	Just(1, 2, 3, 4, 5).Walk(func(item interface{}, pipe chan<- interface{}) {
-		if item.(int)%2 != 0 {
-			pipe <- item
-		}
-	}, UnlimitedWorkers()).ForEach(func(item interface{}) {
-		result += item.(int)
-	})
-	assert.Equal(t, 9, result)
-}
-
-func BenchmarkMapReduce(b *testing.B) {
-	b.ReportAllocs()
-
-	mapper := func(v interface{}) interface{} {
-		return v.(int64) * v.(int64)
-	}
-	reducer := func(input <-chan interface{}) (interface{}, error) {
-		var result int64
-		for v := range input {
-			result += v.(int64)
-		}
-		return result, nil
-	}
-
-	for i := 0; i < b.N; i++ {
-		From(func(input chan<- interface{}) {
-			for j := 0; j < 2; j++ {
-				input <- int64(j)
-			}
-		}).Map(mapper).Reduce(reducer)
-	}
-}

core/fx/retry.go

@@ -15,7 +15,7 @@ type (
 
 // DoWithRetry runs fn, and retries if failed. Default to retry 3 times.
 func DoWithRetry(fn func() error, opts ...RetryOption) error {
-	var options = newRetryOptions()
+	options := newRetryOptions()
 	for _, opt := range opts {
 		opt(options)
 	}

core/fx/retry_test.go

@@ -30,7 +30,7 @@ func TestRetry(t *testing.T) {
 		return errors.New("any")
 	}))
 
-	var total = 2 * defaultRetryTimes
+	total := 2 * defaultRetryTimes
 	times = 0
 	assert.Nil(t, DoWithRetry(func() error {
 		times++

core/fx/stream.go

@@ -49,6 +49,11 @@ type (
 	}
 )
 
+// Concat returns a concatenated Stream.
+func Concat(s Stream, others ...Stream) Stream {
+	return s.Concat(others...)
+}
+
 // From constructs a Stream from the given GenerateFunc.
 func From(generate GenerateFunc) Stream {
 	source := make(chan interface{})
@@ -79,16 +84,46 @@ func Range(source <-chan interface{}) Stream {
 	}
 }
 
+// AllMach returns whether all elements of this stream match the provided predicate.
+// May not evaluate the predicate on all elements if not necessary for determining the result.
+// If the stream is empty then true is returned and the predicate is not evaluated.
+func (s Stream) AllMach(predicate func(item interface{}) bool) bool {
+	for item := range s.source {
+		if !predicate(item) {
+			// make sure the former goroutine not block, and current func returns fast.
+			go drain(s.source)
+			return false
+		}
+	}
+
+	return true
+}
+
+// AnyMach returns whether any elements of this stream match the provided predicate.
+// May not evaluate the predicate on all elements if not necessary for determining the result.
+// If the stream is empty then false is returned and the predicate is not evaluated.
+func (s Stream) AnyMach(predicate func(item interface{}) bool) bool {
+	for item := range s.source {
+		if predicate(item) {
+			// make sure the former goroutine not block, and current func returns fast.
+			go drain(s.source)
+			return true
+		}
+	}
+
+	return false
+}
+
 // Buffer buffers the items into a queue with size n.
 // It can balance the producer and the consumer if their processing throughput don't match.
-func (p Stream) Buffer(n int) Stream {
+func (s Stream) Buffer(n int) Stream {
 	if n < 0 {
 		n = 0
 	}
 
 	source := make(chan interface{}, n)
 	go func() {
-		for item := range p.source {
+		for item := range s.source {
 			source <- item
 		}
 		close(source)
@@ -97,23 +132,51 @@ func (p Stream) Buffer(n int) Stream {
 	return Range(source)
 }
 
+// Concat returns a Stream that concatenated other streams
+func (s Stream) Concat(others ...Stream) Stream {
+	source := make(chan interface{})
+
+	go func() {
+		group := threading.NewRoutineGroup()
+		group.Run(func() {
+			for item := range s.source {
+				source <- item
+			}
+		})
+
+		for _, each := range others {
+			each := each
+			group.Run(func() {
+				for item := range each.source {
+					source <- item
+				}
+			})
+		}
+
+		group.Wait()
+		close(source)
+	}()
+
+	return Range(source)
+}
+
 // Count counts the number of elements in the result.
-func (p Stream) Count() (count int) {
-	for range p.source {
+func (s Stream) Count() (count int) {
+	for range s.source {
 		count++
 	}
 	return
 }
 
 // Distinct removes the duplicated items base on the given KeyFunc.
-func (p Stream) Distinct(fn KeyFunc) Stream {
+func (s Stream) Distinct(fn KeyFunc) Stream {
 	source := make(chan interface{})
 
 	threading.GoSafe(func() {
 		defer close(source)
 
 		keys := make(map[interface{}]lang.PlaceholderType)
-		for item := range p.source {
+		for item := range s.source {
 			key := fn(item)
 			if _, ok := keys[key]; !ok {
 				source <- item
@@ -126,36 +189,48 @@ func (p Stream) Distinct(fn KeyFunc) Stream {
 }
 
 // Done waits all upstreaming operations to be done.
-func (p Stream) Done() {
-	for range p.source {
-	}
+func (s Stream) Done() {
+	drain(s.source)
 }
 
 // Filter filters the items by the given FilterFunc.
-func (p Stream) Filter(fn FilterFunc, opts ...Option) Stream {
-	return p.Walk(func(item interface{}, pipe chan<- interface{}) {
+func (s Stream) Filter(fn FilterFunc, opts ...Option) Stream {
+	return s.Walk(func(item interface{}, pipe chan<- interface{}) {
 		if fn(item) {
 			pipe <- item
 		}
 	}, opts...)
 }
 
+// First returns the first item, nil if no items.
+func (s Stream) First() interface{} {
+	for item := range s.source {
+		// make sure the former goroutine not block, and current func returns fast.
+		go drain(s.source)
+		return item
+	}
+
+	return nil
+}
+
 // ForAll handles the streaming elements from the source and no later streams.
-func (p Stream) ForAll(fn ForAllFunc) {
-	fn(p.source)
+func (s Stream) ForAll(fn ForAllFunc) {
+	fn(s.source)
+	// avoid goroutine leak on fn not consuming all items.
+	go drain(s.source)
 }
 
 // ForEach seals the Stream with the ForEachFunc on each item, no successive operations.
-func (p Stream) ForEach(fn ForEachFunc) {
-	for item := range p.source {
+func (s Stream) ForEach(fn ForEachFunc) {
+	for item := range s.source {
 		fn(item)
 	}
 }
 
 // Group groups the elements into different groups based on their keys.
-func (p Stream) Group(fn KeyFunc) Stream {
+func (s Stream) Group(fn KeyFunc) Stream {
 	groups := make(map[interface{}][]interface{})
-	for item := range p.source {
+	for item := range s.source {
 		key := fn(item)
 		groups[key] = append(groups[key], item)
 	}
@@ -172,7 +247,7 @@ func (p Stream) Group(fn KeyFunc) Stream {
 }
 
 // Head returns the first n elements in p.
-func (p Stream) Head(n int64) Stream {
+func (s Stream) Head(n int64) Stream {
 	if n < 1 {
 		panic("n must be greater than 0")
 	}
@@ -180,18 +255,21 @@ func (p Stream) Head(n int64) Stream {
 	source := make(chan interface{})
 
 	go func() {
-		for item := range p.source {
+		for item := range s.source {
 			n--
 			if n >= 0 {
 				source <- item
 			}
 			if n == 0 {
 				// let successive method go ASAP even we have more items to skip
-				// why we don't just break the loop, because if break,
-				// this former goroutine will block forever, which will cause goroutine leak.
 				close(source)
+				// why we don't just break the loop, and drain to consume all items.
+				// because if breaks, this former goroutine will block forever,
+				// which will cause goroutine leak.
+				drain(s.source)
 			}
 		}
+		// not enough items in s.source, but we need to let successive method to go ASAP.
 		if n > 0 {
 			close(source)
 		}
@@ -200,17 +278,24 @@ func (p Stream) Head(n int64) Stream {
 	return Range(source)
 }
 
+// Last returns the last item, or nil if no items.
+func (s Stream) Last() (item interface{}) {
+	for item = range s.source {
+	}
+	return
+}
+
 // Map converts each item to another corresponding item, which means it's a 1:1 model.
-func (p Stream) Map(fn MapFunc, opts ...Option) Stream {
-	return p.Walk(func(item interface{}, pipe chan<- interface{}) {
+func (s Stream) Map(fn MapFunc, opts ...Option) Stream {
+	return s.Walk(func(item interface{}, pipe chan<- interface{}) {
 		pipe <- fn(item)
 	}, opts...)
 }
 
 // Merge merges all the items into a slice and generates a new stream.
-func (p Stream) Merge() Stream {
+func (s Stream) Merge() Stream {
 	var items []interface{}
-	for item := range p.source {
+	for item := range s.source {
 		items = append(items, item)
 	}
 
@@ -221,22 +306,37 @@ func (p Stream) Merge() Stream {
 	return Range(source)
 }
 
+// NoneMatch returns whether all elements of this stream don't match the provided predicate.
+// May not evaluate the predicate on all elements if not necessary for determining the result.
+// If the stream is empty then true is returned and the predicate is not evaluated.
+func (s Stream) NoneMatch(predicate func(item interface{}) bool) bool {
+	for item := range s.source {
+		if predicate(item) {
+			// make sure the former goroutine not block, and current func returns fast.
+			go drain(s.source)
+			return false
+		}
+	}
+
+	return true
+}
+
 // Parallel applies the given ParallelFunc to each item concurrently with given number of workers.
-func (p Stream) Parallel(fn ParallelFunc, opts ...Option) {
-	p.Walk(func(item interface{}, pipe chan<- interface{}) {
+func (s Stream) Parallel(fn ParallelFunc, opts ...Option) {
+	s.Walk(func(item interface{}, pipe chan<- interface{}) {
 		fn(item)
 	}, opts...).Done()
 }
 
 // Reduce is a utility method to let the caller deal with the underlying channel.
-func (p Stream) Reduce(fn ReduceFunc) (interface{}, error) {
-	return fn(p.source)
+func (s Stream) Reduce(fn ReduceFunc) (interface{}, error) {
+	return fn(s.source)
 }
 
 // Reverse reverses the elements in the stream.
-func (p Stream) Reverse() Stream {
+func (s Stream) Reverse() Stream {
 	var items []interface{}
-	for item := range p.source {
+	for item := range s.source {
 		items = append(items, item)
 	}
 	// reverse, official method
@@ -248,10 +348,36 @@ func (p Stream) Reverse() Stream {
 	return Just(items...)
 }
 
+// Skip returns a Stream that skips size elements.
+func (s Stream) Skip(n int64) Stream {
+	if n < 0 {
+		panic("n must not be negative")
+	}
+	if n == 0 {
+		return s
+	}
+
+	source := make(chan interface{})
+
+	go func() {
+		for item := range s.source {
+			n--
+			if n >= 0 {
+				continue
+			} else {
+				source <- item
+			}
+		}
+		close(source)
+	}()
+
+	return Range(source)
+}
+
 // Sort sorts the items from the underlying source.
-func (p Stream) Sort(less LessFunc) Stream {
+func (s Stream) Sort(less LessFunc) Stream {
 	var items []interface{}
-	for item := range p.source {
+	for item := range s.source {
 		items = append(items, item)
 	}
 	sort.Slice(items, func(i, j int) bool {
@@ -263,7 +389,7 @@ func (p Stream) Sort(less LessFunc) Stream {
 
 // Split splits the elements into chunk with size up to n,
 // might be less than n on tailing elements.
-func (p Stream) Split(n int) Stream {
+func (s Stream) Split(n int) Stream {
 	if n < 1 {
 		panic("n should be greater than 0")
 	}
@@ -271,7 +397,7 @@ func (p Stream) Split(n int) Stream {
 	source := make(chan interface{})
 	go func() {
 		var chunk []interface{}
-		for item := range p.source {
+		for item := range s.source {
 			chunk = append(chunk, item)
 			if len(chunk) == n {
 				source <- chunk
@@ -288,7 +414,7 @@ func (p Stream) Split(n int) Stream {
 }
 
 // Tail returns the last n elements in p.
-func (p Stream) Tail(n int64) Stream {
+func (s Stream) Tail(n int64) Stream {
 	if n < 1 {
 		panic("n should be greater than 0")
 	}
@@ -297,7 +423,7 @@ func (p Stream) Tail(n int64) Stream {
 
 	go func() {
 		ring := collection.NewRing(int(n))
-		for item := range p.source {
+		for item := range s.source {
 			ring.Add(item)
 		}
 		for _, item := range ring.Take() {
@@ -310,31 +436,28 @@ func (p Stream) Tail(n int64) Stream {
 }
 
 // Walk lets the callers handle each item, the caller may write zero, one or more items base on the given item.
-func (p Stream) Walk(fn WalkFunc, opts ...Option) Stream {
+func (s Stream) Walk(fn WalkFunc, opts ...Option) Stream {
 	option := buildOptions(opts...)
 	if option.unlimitedWorkers {
-		return p.walkUnlimited(fn, option)
+		return s.walkUnlimited(fn, option)
 	}
 
-	return p.walkLimited(fn, option)
+	return s.walkLimited(fn, option)
 }
 
-func (p Stream) walkLimited(fn WalkFunc, option *rxOptions) Stream {
+func (s Stream) walkLimited(fn WalkFunc, option *rxOptions) Stream {
 	pipe := make(chan interface{}, option.workers)
 
 	go func() {
 		var wg sync.WaitGroup
 		pool := make(chan lang.PlaceholderType, option.workers)
 
-		for {
+		for item := range s.source {
+			// important, used in another goroutine
+			val := item
 			pool <- lang.Placeholder
-			item, ok := <-p.source
-			if !ok {
-				<-pool
-				break
-			}
-
 			wg.Add(1)
+
 			// better to safely run caller defined method
 			threading.GoSafe(func() {
 				defer func() {
@@ -342,7 +465,7 @@ func (p Stream) walkLimited(fn WalkFunc, option *rxOptions) Stream {
 					<-pool
 				}()
 
-				fn(item, pipe)
+				fn(val, pipe)
 			})
 		}
 
@@ -353,23 +476,20 @@ func (p Stream) walkLimited(fn WalkFunc, option *rxOptions) Stream {
 	return Range(pipe)
 }
 
-func (p Stream) walkUnlimited(fn WalkFunc, option *rxOptions) Stream {
-	pipe := make(chan interface{}, defaultWorkers)
+func (s Stream) walkUnlimited(fn WalkFunc, option *rxOptions) Stream {
+	pipe := make(chan interface{}, option.workers)
 
 	go func() {
 		var wg sync.WaitGroup
 
-		for {
-			item, ok := <-p.source
-			if !ok {
-				break
-			}
-
+		for item := range s.source {
+			// important, used in another goroutine
+			val := item
 			wg.Add(1)
 			// better to safely run caller defined method
 			threading.GoSafe(func() {
 				defer wg.Done()
-				fn(item, pipe)
+				fn(val, pipe)
 			})
 		}
 
@@ -380,14 +500,14 @@ func (p Stream) walkUnlimited(fn WalkFunc, option *rxOptions) Stream {
 	return Range(pipe)
 }
 
-// UnlimitedWorkers lets the caller to use as many workers as the tasks.
+// UnlimitedWorkers lets the caller use as many workers as the tasks.
 func UnlimitedWorkers() Option {
 	return func(opts *rxOptions) {
 		opts.unlimitedWorkers = true
 	}
 }
 
-// WithWorkers lets the caller to customize the concurrent workers.
+// WithWorkers lets the caller customize the concurrent workers.
 func WithWorkers(workers int) Option {
 	return func(opts *rxOptions) {
 		if workers < minWorkers {
@@ -398,6 +518,7 @@ func WithWorkers(workers int) Option {
 	}
 }
 
+// buildOptions returns a rxOptions with given customizations.
 func buildOptions(opts ...Option) *rxOptions {
 	options := newOptions()
 	for _, opt := range opts {
@@ -407,6 +528,13 @@ func buildOptions(opts ...Option) *rxOptions {
 	return options
 }
 
+// drain drains the given channel.
+func drain(channel <-chan interface{}) {
+	for range channel {
+	}
+}
+
+// newOptions returns a default rxOptions.
 func newOptions() *rxOptions {
 	return &rxOptions{
 		workers: defaultWorkers,

core/fx/stream_test.go

@@ -0,0 +1,571 @@
+package fx
+
+import (
+	"io/ioutil"
+	"log"
+	"math/rand"
+	"reflect"
+	"runtime"
+	"sort"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/stringx"
+)
+
+func TestBuffer(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		const N = 5
+		var count int32
+		var wait sync.WaitGroup
+		wait.Add(1)
+		From(func(source chan<- interface{}) {
+			ticker := time.NewTicker(10 * time.Millisecond)
+			defer ticker.Stop()
+
+			for i := 0; i < 2*N; i++ {
+				select {
+				case source <- i:
+					atomic.AddInt32(&count, 1)
+				case <-ticker.C:
+					wait.Done()
+					return
+				}
+			}
+		}).Buffer(N).ForAll(func(pipe <-chan interface{}) {
+			wait.Wait()
+			// why N+1, because take one more to wait for sending into the channel
+			assert.Equal(t, int32(N+1), atomic.LoadInt32(&count))
+		})
+	})
+}
+
+func TestBufferNegative(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var result int
+		Just(1, 2, 3, 4).Buffer(-1).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
+			for item := range pipe {
+				result += item.(int)
+			}
+			return result, nil
+		})
+		assert.Equal(t, 10, result)
+	})
+}
+
+func TestCount(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		tests := []struct {
+			name     string
+			elements []interface{}
+		}{
+			{
+				name: "no elements with nil",
+			},
+			{
+				name:     "no elements",
+				elements: []interface{}{},
+			},
+			{
+				name:     "1 element",
+				elements: []interface{}{1},
+			},
+			{
+				name:     "multiple elements",
+				elements: []interface{}{1, 2, 3},
+			},
+		}
+
+		for _, test := range tests {
+			t.Run(test.name, func(t *testing.T) {
+				val := Just(test.elements...).Count()
+				assert.Equal(t, len(test.elements), val)
+			})
+		}
+	})
+}
+
+func TestDone(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var count int32
+		Just(1, 2, 3).Walk(func(item interface{}, pipe chan<- interface{}) {
+			time.Sleep(time.Millisecond * 100)
+			atomic.AddInt32(&count, int32(item.(int)))
+		}).Done()
+		assert.Equal(t, int32(6), count)
+	})
+}
+
+func TestJust(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var result int
+		Just(1, 2, 3, 4).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
+			for item := range pipe {
+				result += item.(int)
+			}
+			return result, nil
+		})
+		assert.Equal(t, 10, result)
+	})
+}
+
+func TestDistinct(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var result int
+		Just(4, 1, 3, 2, 3, 4).Distinct(func(item interface{}) interface{} {
+			return item
+		}).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
+			for item := range pipe {
+				result += item.(int)
+			}
+			return result, nil
+		})
+		assert.Equal(t, 10, result)
+	})
+}
+
+func TestFilter(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var result int
+		Just(1, 2, 3, 4).Filter(func(item interface{}) bool {
+			return item.(int)%2 == 0
+		}).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
+			for item := range pipe {
+				result += item.(int)
+			}
+			return result, nil
+		})
+		assert.Equal(t, 6, result)
+	})
+}
+
+func TestFirst(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		assert.Nil(t, Just().First())
+		assert.Equal(t, "foo", Just("foo").First())
+		assert.Equal(t, "foo", Just("foo", "bar").First())
+	})
+}
+
+func TestForAll(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var result int
+		Just(1, 2, 3, 4).Filter(func(item interface{}) bool {
+			return item.(int)%2 == 0
+		}).ForAll(func(pipe <-chan interface{}) {
+			for item := range pipe {
+				result += item.(int)
+			}
+		})
+		assert.Equal(t, 6, result)
+	})
+}
+
+func TestGroup(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var groups [][]int
+		Just(10, 11, 20, 21).Group(func(item interface{}) interface{} {
+			v := item.(int)
+			return v / 10
+		}).ForEach(func(item interface{}) {
+			v := item.([]interface{})
+			var group []int
+			for _, each := range v {
+				group = append(group, each.(int))
+			}
+			groups = append(groups, group)
+		})
+
+		assert.Equal(t, 2, len(groups))
+		for _, group := range groups {
+			assert.Equal(t, 2, len(group))
+			assert.True(t, group[0]/10 == group[1]/10)
+		}
+	})
+}
+
+func TestHead(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var result int
+		Just(1, 2, 3, 4).Head(2).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
+			for item := range pipe {
+				result += item.(int)
+			}
+			return result, nil
+		})
+		assert.Equal(t, 3, result)
+	})
+}
+
+func TestHeadZero(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		assert.Panics(t, func() {
+			Just(1, 2, 3, 4).Head(0).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
+				return nil, nil
+			})
+		})
+	})
+}
+
+func TestHeadMore(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var result int
+		Just(1, 2, 3, 4).Head(6).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
+			for item := range pipe {
+				result += item.(int)
+			}
+			return result, nil
+		})
+		assert.Equal(t, 10, result)
+	})
+}
+
+func TestLast(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		goroutines := runtime.NumGoroutine()
+		assert.Nil(t, Just().Last())
+		assert.Equal(t, "foo", Just("foo").Last())
+		assert.Equal(t, "bar", Just("foo", "bar").Last())
+		// let scheduler schedule first
+		runtime.Gosched()
+		assert.Equal(t, goroutines, runtime.NumGoroutine())
+	})
+}
+
+func TestMap(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		log.SetOutput(ioutil.Discard)
+
+		tests := []struct {
+			mapper MapFunc
+			expect int
+		}{
+			{
+				mapper: func(item interface{}) interface{} {
+					v := item.(int)
+					return v * v
+				},
+				expect: 30,
+			},
+			{
+				mapper: func(item interface{}) interface{} {
+					v := item.(int)
+					if v%2 == 0 {
+						return 0
+					}
+					return v * v
+				},
+				expect: 10,
+			},
+			{
+				mapper: func(item interface{}) interface{} {
+					v := item.(int)
+					if v%2 == 0 {
+						panic(v)
+					}
+					return v * v
+				},
+				expect: 10,
+			},
+		}
+
+		// Map(...) works even WithWorkers(0)
+		for i, test := range tests {
+			t.Run(stringx.Rand(), func(t *testing.T) {
+				var result int
+				var workers int
+				if i%2 == 0 {
+					workers = 0
+				} else {
+					workers = runtime.NumCPU()
+				}
+				From(func(source chan<- interface{}) {
+					for i := 1; i < 5; i++ {
+						source <- i
+					}
+				}).Map(test.mapper, WithWorkers(workers)).Reduce(
+					func(pipe <-chan interface{}) (interface{}, error) {
+						for item := range pipe {
+							result += item.(int)
+						}
+						return result, nil
+					})
+
+				assert.Equal(t, test.expect, result)
+			})
+		}
+	})
+}
+
+func TestMerge(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		Just(1, 2, 3, 4).Merge().ForEach(func(item interface{}) {
+			assert.ElementsMatch(t, []interface{}{1, 2, 3, 4}, item.([]interface{}))
+		})
+	})
+}
+
+func TestParallelJust(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var count int32
+		Just(1, 2, 3).Parallel(func(item interface{}) {
+			time.Sleep(time.Millisecond * 100)
+			atomic.AddInt32(&count, int32(item.(int)))
+		}, UnlimitedWorkers())
+		assert.Equal(t, int32(6), count)
+	})
+}
+
+func TestReverse(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		Just(1, 2, 3, 4).Reverse().Merge().ForEach(func(item interface{}) {
+			assert.ElementsMatch(t, []interface{}{4, 3, 2, 1}, item.([]interface{}))
+		})
+	})
+}
+
+func TestSort(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var prev int
+		Just(5, 3, 7, 1, 9, 6, 4, 8, 2).Sort(func(a, b interface{}) bool {
+			return a.(int) < b.(int)
+		}).ForEach(func(item interface{}) {
+			next := item.(int)
+			assert.True(t, prev < next)
+			prev = next
+		})
+	})
+}
+
+func TestSplit(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		assert.Panics(t, func() {
+			Just(1, 2, 3, 4, 5, 6, 7, 8, 9, 10).Split(0).Done()
+		})
+		var chunks [][]interface{}
+		Just(1, 2, 3, 4, 5, 6, 7, 8, 9, 10).Split(4).ForEach(func(item interface{}) {
+			chunk := item.([]interface{})
+			chunks = append(chunks, chunk)
+		})
+		assert.EqualValues(t, [][]interface{}{
+			{1, 2, 3, 4},
+			{5, 6, 7, 8},
+			{9, 10},
+		}, chunks)
+	})
+}
+
+func TestTail(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var result int
+		Just(1, 2, 3, 4).Tail(2).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
+			for item := range pipe {
+				result += item.(int)
+			}
+			return result, nil
+		})
+		assert.Equal(t, 7, result)
+	})
+}
+
+func TestTailZero(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		assert.Panics(t, func() {
+			Just(1, 2, 3, 4).Tail(0).Reduce(func(pipe <-chan interface{}) (interface{}, error) {
+				return nil, nil
+			})
+		})
+	})
+}
+
+func TestWalk(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		var result int
+		Just(1, 2, 3, 4, 5).Walk(func(item interface{}, pipe chan<- interface{}) {
+			if item.(int)%2 != 0 {
+				pipe <- item
+			}
+		}, UnlimitedWorkers()).ForEach(func(item interface{}) {
+			result += item.(int)
+		})
+		assert.Equal(t, 9, result)
+	})
+}
+
+func TestStream_AnyMach(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		assetEqual(t, false, Just(1, 2, 3).AnyMach(func(item interface{}) bool {
+			return item.(int) == 4
+		}))
+		assetEqual(t, false, Just(1, 2, 3).AnyMach(func(item interface{}) bool {
+			return item.(int) == 0
+		}))
+		assetEqual(t, true, Just(1, 2, 3).AnyMach(func(item interface{}) bool {
+			return item.(int) == 2
+		}))
+		assetEqual(t, true, Just(1, 2, 3).AnyMach(func(item interface{}) bool {
+			return item.(int) == 2
+		}))
+	})
+}
+
+func TestStream_AllMach(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		assetEqual(
+			t, true, Just(1, 2, 3).AllMach(func(item interface{}) bool {
+				return true
+			}),
+		)
+		assetEqual(
+			t, false, Just(1, 2, 3).AllMach(func(item interface{}) bool {
+				return false
+			}),
+		)
+		assetEqual(
+			t, false, Just(1, 2, 3).AllMach(func(item interface{}) bool {
+				return item.(int) == 1
+			}),
+		)
+	})
+}
+
+func TestStream_NoneMatch(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		assetEqual(
+			t, true, Just(1, 2, 3).NoneMatch(func(item interface{}) bool {
+				return false
+			}),
+		)
+		assetEqual(
+			t, false, Just(1, 2, 3).NoneMatch(func(item interface{}) bool {
+				return true
+			}),
+		)
+		assetEqual(
+			t, true, Just(1, 2, 3).NoneMatch(func(item interface{}) bool {
+				return item.(int) == 4
+			}),
+		)
+	})
+}
+
+func TestConcat(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		a1 := []interface{}{1, 2, 3}
+		a2 := []interface{}{4, 5, 6}
+		s1 := Just(a1...)
+		s2 := Just(a2...)
+		stream := Concat(s1, s2)
+		var items []interface{}
+		for item := range stream.source {
+			items = append(items, item)
+		}
+		sort.Slice(items, func(i, j int) bool {
+			return items[i].(int) < items[j].(int)
+		})
+		ints := make([]interface{}, 0)
+		ints = append(ints, a1...)
+		ints = append(ints, a2...)
+		assetEqual(t, ints, items)
+	})
+}
+
+func TestStream_Skip(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		assetEqual(t, 3, Just(1, 2, 3, 4).Skip(1).Count())
+		assetEqual(t, 1, Just(1, 2, 3, 4).Skip(3).Count())
+		assetEqual(t, 4, Just(1, 2, 3, 4).Skip(0).Count())
+		equal(t, Just(1, 2, 3, 4).Skip(3), []interface{}{4})
+		assert.Panics(t, func() {
+			Just(1, 2, 3, 4).Skip(-1)
+		})
+	})
+}
+
+func TestStream_Concat(t *testing.T) {
+	runCheckedTest(t, func(t *testing.T) {
+		stream := Just(1).Concat(Just(2), Just(3))
+		var items []interface{}
+		for item := range stream.source {
+			items = append(items, item)
+		}
+		sort.Slice(items, func(i, j int) bool {
+			return items[i].(int) < items[j].(int)
+		})
+		assetEqual(t, []interface{}{1, 2, 3}, items)
+
+		just := Just(1)
+		equal(t, just.Concat(just), []interface{}{1})
+	})
+}
+
+func BenchmarkParallelMapReduce(b *testing.B) {
+	b.ReportAllocs()
+
+	mapper := func(v interface{}) interface{} {
+		return v.(int64) * v.(int64)
+	}
+	reducer := func(input <-chan interface{}) (interface{}, error) {
+		var result int64
+		for v := range input {
+			result += v.(int64)
+		}
+		return result, nil
+	}
+	b.ResetTimer()
+	From(func(input chan<- interface{}) {
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				input <- int64(rand.Int())
+			}
+		})
+	}).Map(mapper).Reduce(reducer)
+}
+
+func BenchmarkMapReduce(b *testing.B) {
+	b.ReportAllocs()
+
+	mapper := func(v interface{}) interface{} {
+		return v.(int64) * v.(int64)
+	}
+	reducer := func(input <-chan interface{}) (interface{}, error) {
+		var result int64
+		for v := range input {
+			result += v.(int64)
+		}
+		return result, nil
+	}
+	b.ResetTimer()
+	From(func(input chan<- interface{}) {
+		for i := 0; i < b.N; i++ {
+			input <- int64(rand.Int())
+		}
+	}).Map(mapper).Reduce(reducer)
+}
+
+func assetEqual(t *testing.T, except, data interface{}) {
+	if !reflect.DeepEqual(except, data) {
+		t.Errorf(" %v, want %v", data, except)
+	}
+}
+
+func equal(t *testing.T, stream Stream, data []interface{}) {
+	items := make([]interface{}, 0)
+	for item := range stream.source {
+		items = append(items, item)
+	}
+	if !reflect.DeepEqual(items, data) {
+		t.Errorf(" %v, want %v", items, data)
+	}
+}
+
+func runCheckedTest(t *testing.T, fn func(t *testing.T)) {
+	goroutines := runtime.NumGoroutine()
+	fn(t)
+	// let scheduler schedule first
+	time.Sleep(time.Millisecond)
+	assert.True(t, runtime.NumGoroutine() <= goroutines)
+}

core/fx/timeout.go

@@ -2,9 +2,10 @@ package fx
 
 import (
 	"context"
+	"fmt"
+	"runtime/debug"
+	"strings"
 	"time"
-
-	"github.com/tal-tech/go-zero/core/contextx"
 )
 
 var (
@@ -23,19 +24,20 @@ func DoWithTimeout(fn func() error, timeout time.Duration, opts ...DoOption) err
 	for _, opt := range opts {
 		parentCtx = opt()
 	}
-	ctx, cancel := contextx.ShrinkDeadline(parentCtx, timeout)
+	ctx, cancel := context.WithTimeout(parentCtx, timeout)
 	defer cancel()
 
-	done := make(chan error)
+	// create channel with buffer size 1 to avoid goroutine leak
+	done := make(chan error, 1)
 	panicChan := make(chan interface{}, 1)
 	go func() {
 		defer func() {
 			if p := recover(); p != nil {
-				panicChan <- p
+				// attach call stack to avoid missing in different goroutine
+				panicChan <- fmt.Sprintf("%+v\n\n%s", p, strings.TrimSpace(string(debug.Stack())))
 			}
 		}()
 		done <- fn()
-		close(done)
 	}()
 
 	select {

core/hash/consistenthash.go

@@ -83,7 +83,7 @@ func (h *ConsistentHash) AddWithReplicas(node interface{}, replicas int) {
 		h.ring[hash] = append(h.ring[hash], node)
 	}
 
-	sort.Slice(h.keys, func(i int, j int) bool {
+	sort.Slice(h.keys, func(i, j int) bool {
 		return h.keys[i] < h.keys[j]
 	})
 }
@@ -140,7 +140,7 @@ func (h *ConsistentHash) Remove(node interface{}) {
 		index := sort.Search(len(h.keys), func(i int) bool {
 			return h.keys[i] >= hash
 		})
-		if index < len(h.keys) {
+		if index < len(h.keys) && h.keys[index] == hash {
 			h.keys = append(h.keys[:index], h.keys[index+1:]...)
 		}
 		h.removeRingNode(hash, nodeRepr)

core/hash/consistenthash_test.go

@@ -74,12 +74,12 @@ func TestConsistentHashIncrementalTransfer(t *testing.T) {
 		laterCh := create()
 		laterCh.AddWithWeight(node, 10*(i+1))
 
-		for i := 0; i < requestSize; i++ {
-			key, ok := laterCh.Get(requestSize + i)
+		for j := 0; j < requestSize; j++ {
+			key, ok := laterCh.Get(requestSize + j)
 			assert.True(t, ok)
 			assert.NotNil(t, key)
 			value := key.(string)
-			assert.True(t, value == keys[i] || value == node)
+			assert.True(t, value == keys[j] || value == node)
 		}
 	}
 }

core/hash/hash_test.go

@@ -20,6 +20,11 @@ func TestMd5(t *testing.T) {
 	assert.Equal(t, md5Digest, actual)
 }
 
+func TestMd5Hex(t *testing.T) {
+	actual := Md5Hex([]byte(text))
+	assert.Equal(t, md5Digest, actual)
+}
+
 func BenchmarkHashFnv(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		h := fnv.New32()

core/jsontype/time_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/globalsign/mgo/bson"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -106,3 +107,20 @@ func TestMilliTime_UnmarshalJSON(t *testing.T) {
 		})
 	}
 }
+
+func TestUnmarshalWithError(t *testing.T) {
+	var mt MilliTime
+	assert.NotNil(t, mt.UnmarshalJSON([]byte("hello")))
+}
+
+func TestSetBSON(t *testing.T) {
+	data, err := bson.Marshal(time.Now())
+	assert.Nil(t, err)
+
+	var raw bson.Raw
+	assert.Nil(t, bson.Unmarshal(data, &raw))
+
+	var mt MilliTime
+	assert.Nil(t, mt.SetBSON(raw))
+	assert.NotNil(t, mt.SetBSON(bson.Raw{}))
+}

core/lang/lang.go

@@ -4,8 +4,8 @@ package lang
 var Placeholder PlaceholderType
 
 type (
-	// GenericType can be used to hold any type.
-	GenericType = interface{}
+	// AnyType can be used to hold any type.
+	AnyType = interface{}
 	// PlaceholderType represents a placeholder type.
 	PlaceholderType = struct{}
 )

core/limit/tokenlimit_test.go

@@ -24,7 +24,7 @@ func TestTokenLimit_Rescue(t *testing.T) {
 		rate  = 5
 		burst = 10
 	)
-	l := NewTokenLimiter(rate, burst, redis.NewRedis(s.Addr(), redis.NodeType), "tokenlimit")
+	l := NewTokenLimiter(rate, burst, redis.New(s.Addr()), "tokenlimit")
 	s.Close()
 
 	var allowed int

core/load/adaptiveshedder.go

@@ -31,6 +31,8 @@ var (
 
 	// default to be enabled
 	enabled = syncx.ForAtomicBool(true)
+	// default to be enabled
+	logEnabled = syncx.ForAtomicBool(true)
 	// make it a variable for unit test
 	systemOverloadChecker = func(cpuThreshold int64) bool {
 		return stat.CpuUsage() >= cpuThreshold
@@ -80,6 +82,11 @@ func Disable() {
 	enabled.Set(false)
 }
 
+// DisableLog disables the stat logs for load shedding.
+func DisableLog() {
+	logEnabled.Set(false)
+}
+
 // NewAdaptiveShedder returns an adaptive shedder.
 // opts can be used to customize the Shedder.
 func NewAdaptiveShedder(opts ...ShedderOption) Shedder {
@@ -168,7 +175,7 @@ func (as *adaptiveShedder) maxPass() int64 {
 }
 
 func (as *adaptiveShedder) minRt() float64 {
-	var result = defaultMinRt
+	result := defaultMinRt
 
 	as.rtCounter.Reduce(func(b *collection.Bucket) {
 		if b.Count <= 0 {

core/load/adaptiveshedder_test.go

@@ -25,6 +25,7 @@ func init() {
 }
 
 func TestAdaptiveShedder(t *testing.T) {
+	DisableLog()
 	shedder := NewAdaptiveShedder(WithWindow(bucketDuration), WithBuckets(buckets), WithCpuThreshold(100))
 	var wg sync.WaitGroup
 	var drop int64
@@ -201,7 +202,7 @@ func BenchmarkAdaptiveShedder_Allow(b *testing.B) {
 	logx.Disable()
 
 	bench := func(b *testing.B) {
-		var shedder = NewAdaptiveShedder()
+		shedder := NewAdaptiveShedder()
 		proba := mathx.NewProba()
 		for i := 0; i < 6000; i++ {
 			p, err := shedder.Allow()

core/load/nopshedder.go

@@ -1,7 +1,6 @@
 package load
 
-type nopShedder struct {
-}
+type nopShedder struct{}
 
 func newNopShedder() Shedder {
 	return nopShedder{}
@@ -11,8 +10,7 @@ func (s nopShedder) Allow() (Promise, error) {
 	return nopPromise{}, nil
 }
 
-type nopPromise struct {
-}
+type nopPromise struct{}
 
 func (p nopPromise) Pass() {
 }

core/load/sheddingstat.go

@@ -48,6 +48,25 @@ func (s *SheddingStat) IncrementDrop() {
 	atomic.AddInt64(&s.drop, 1)
 }
 
+func (s *SheddingStat) loop(c <-chan time.Time) {
+	for range c {
+		st := s.reset()
+
+		if !logEnabled.True() {
+			continue
+		}
+
+		c := stat.CpuUsage()
+		if st.Drop == 0 {
+			logx.Statf("(%s) shedding_stat [1m], cpu: %d, total: %d, pass: %d, drop: %d",
+				s.name, c, st.Total, st.Pass, st.Drop)
+		} else {
+			logx.Statf("(%s) shedding_stat_drop [1m], cpu: %d, total: %d, pass: %d, drop: %d",
+				s.name, c, st.Total, st.Pass, st.Drop)
+		}
+	}
+}
+
 func (s *SheddingStat) reset() snapshot {
 	return snapshot{
 		Total: atomic.SwapInt64(&s.total, 0),
@@ -59,15 +78,6 @@ func (s *SheddingStat) reset() snapshot {
 func (s *SheddingStat) run() {
 	ticker := time.NewTicker(time.Minute)
 	defer ticker.Stop()
-	for range ticker.C {
-		c := stat.CpuUsage()
-		st := s.reset()
-		if st.Drop == 0 {
-			logx.Statf("(%s) shedding_stat [1m], cpu: %d, total: %d, pass: %d, drop: %d",
-				s.name, c, st.Total, st.Pass, st.Drop)
-		} else {
-			logx.Statf("(%s) shedding_stat_drop [1m], cpu: %d, total: %d, pass: %d, drop: %d",
-				s.name, c, st.Total, st.Pass, st.Drop)
-		}
-	}
+
+	s.loop(ticker.C)
 }

core/load/sheddingstat_test.go

@@ -2,6 +2,7 @@ package load
 
 import (
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -22,3 +23,32 @@ func TestSheddingStat(t *testing.T) {
 	assert.Equal(t, int64(5), result.Pass)
 	assert.Equal(t, int64(7), result.Drop)
 }
+
+func TestLoopTrue(t *testing.T) {
+	ch := make(chan time.Time, 1)
+	ch <- time.Now()
+	close(ch)
+	st := new(SheddingStat)
+	logEnabled.Set(true)
+	st.loop(ch)
+}
+
+func TestLoopTrueAndDrop(t *testing.T) {
+	ch := make(chan time.Time, 1)
+	ch <- time.Now()
+	close(ch)
+	st := new(SheddingStat)
+	st.IncrementDrop()
+	logEnabled.Set(true)
+	st.loop(ch)
+}
+
+func TestLoopFalseAndDrop(t *testing.T) {
+	ch := make(chan time.Time, 1)
+	ch <- time.Now()
+	close(ch)
+	st := new(SheddingStat)
+	st.IncrementDrop()
+	logEnabled.Set(false)
+	st.loop(ch)
+}

core/logx/config.go

@@ -4,6 +4,7 @@ package logx
 type LogConf struct {
 	ServiceName         string `json:",optional"`
 	Mode                string `json:",default=console,options=console|file|volume"`
+	TimeFormat          string `json:",optional"`
 	Path                string `json:",default=logs"`
 	Level               string `json:",default=info,options=info|error|severe"`
 	Compress            bool   `json:",optional"`

core/logx/durationlogger.go

@@ -20,49 +20,69 @@ func WithDuration(d time.Duration) Logger {
 }
 
 func (l *durationLogger) Error(v ...interface{}) {
-	if shouldLog(ErrorLevel) {
+	if shallLog(ErrorLevel) {
 		l.write(errorLog, levelError, formatWithCaller(fmt.Sprint(v...), durationCallerDepth))
 	}
 }
 
 func (l *durationLogger) Errorf(format string, v ...interface{}) {
-	if shouldLog(ErrorLevel) {
+	if shallLog(ErrorLevel) {
 		l.write(errorLog, levelError, formatWithCaller(fmt.Sprintf(format, v...), durationCallerDepth))
 	}
 }
 
+func (l *durationLogger) Errorv(v interface{}) {
+	if shallLog(ErrorLevel) {
+		l.write(errorLog, levelError, v)
+	}
+}
+
 func (l *durationLogger) Info(v ...interface{}) {
-	if shouldLog(InfoLevel) {
+	if shallLog(InfoLevel) {
 		l.write(infoLog, levelInfo, fmt.Sprint(v...))
 	}
 }
 
 func (l *durationLogger) Infof(format string, v ...interface{}) {
-	if shouldLog(InfoLevel) {
+	if shallLog(InfoLevel) {
 		l.write(infoLog, levelInfo, fmt.Sprintf(format, v...))
 	}
 }
 
+func (l *durationLogger) Infov(v interface{}) {
+	if shallLog(InfoLevel) {
+		l.write(infoLog, levelInfo, v)
+	}
+}
+
 func (l *durationLogger) Slow(v ...interface{}) {
-	if shouldLog(ErrorLevel) {
+	if shallLog(ErrorLevel) {
 		l.write(slowLog, levelSlow, fmt.Sprint(v...))
 	}
 }
 
 func (l *durationLogger) Slowf(format string, v ...interface{}) {
-	if shouldLog(ErrorLevel) {
+	if shallLog(ErrorLevel) {
 		l.write(slowLog, levelSlow, fmt.Sprintf(format, v...))
 	}
 }
 
+func (l *durationLogger) Slowv(v interface{}) {
+	if shallLog(ErrorLevel) {
+		l.write(slowLog, levelSlow, v)
+	}
+}
+
 func (l *durationLogger) WithDuration(duration time.Duration) Logger {
 	l.Duration = timex.ReprOfDuration(duration)
 	return l
 }
 
-func (l *durationLogger) write(writer io.Writer, level, content string) {
-	l.Timestamp = getTimestamp()
-	l.Level = level
-	l.Content = content
-	outputJson(writer, logEntry(*l))
+func (l *durationLogger) write(writer io.Writer, level string, val interface{}) {
+	outputJson(writer, &durationLogger{
+		Timestamp: getTimestamp(),
+		Level:     level,
+		Content:   val,
+		Duration:  l.Duration,
+	})
 }

core/logx/durationlogger_test.go

@@ -23,6 +23,13 @@ func TestWithDurationErrorf(t *testing.T) {
 	assert.True(t, strings.Contains(builder.String(), "duration"), builder.String())
 }
 
+func TestWithDurationErrorv(t *testing.T) {
+	var builder strings.Builder
+	log.SetOutput(&builder)
+	WithDuration(time.Second).Errorv("foo")
+	assert.True(t, strings.Contains(builder.String(), "duration"), builder.String())
+}
+
 func TestWithDurationInfo(t *testing.T) {
 	var builder strings.Builder
 	log.SetOutput(&builder)
@@ -37,6 +44,13 @@ func TestWithDurationInfof(t *testing.T) {
 	assert.True(t, strings.Contains(builder.String(), "duration"), builder.String())
 }
 
+func TestWithDurationInfov(t *testing.T) {
+	var builder strings.Builder
+	log.SetOutput(&builder)
+	WithDuration(time.Second).Infov("foo")
+	assert.True(t, strings.Contains(builder.String(), "duration"), builder.String())
+}
+
 func TestWithDurationSlow(t *testing.T) {
 	var builder strings.Builder
 	log.SetOutput(&builder)
@@ -50,3 +64,10 @@ func TestWithDurationSlowf(t *testing.T) {
 	WithDuration(time.Second).WithDuration(time.Hour).Slowf("foo")
 	assert.True(t, strings.Contains(builder.String(), "duration"), builder.String())
 }
+
+func TestWithDurationSlowv(t *testing.T) {
+	var builder strings.Builder
+	log.SetOutput(&builder)
+	WithDuration(time.Second).WithDuration(time.Hour).Slowv("foo")
+	assert.True(t, strings.Contains(builder.String(), "duration"), builder.String())
+}

core/logx/limitedexecutor_test.go

@@ -0,0 +1,62 @@
+package logx
+
+import (
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/timex"
+)
+
+func TestLimitedExecutor_logOrDiscard(t *testing.T) {
+	tests := []struct {
+		name      string
+		threshold time.Duration
+		lastTime  time.Duration
+		discarded uint32
+		executed  bool
+	}{
+		{
+			name:     "nil executor",
+			executed: true,
+		},
+		{
+			name:      "regular",
+			threshold: time.Hour,
+			lastTime:  timex.Now(),
+			discarded: 10,
+			executed:  false,
+		},
+		{
+			name:      "slow",
+			threshold: time.Duration(1),
+			lastTime:  -1000,
+			discarded: 10,
+			executed:  true,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			executor := newLimitedExecutor(0)
+			executor.threshold = test.threshold
+			executor.discarded = test.discarded
+			executor.lastTime.Set(test.lastTime)
+
+			var run int32
+			executor.logOrDiscard(func() {
+				atomic.AddInt32(&run, 1)
+			})
+			if test.executed {
+				assert.Equal(t, int32(1), atomic.LoadInt32(&run))
+			} else {
+				assert.Equal(t, int32(0), atomic.LoadInt32(&run))
+				assert.Equal(t, test.discarded+1, atomic.LoadUint32(&executor.discarded))
+			}
+		})
+	}
+}

core/logx/logs.go

@@ -32,8 +32,6 @@ const (
 )
 
 const (
-	timeFormat = "2006-01-02T15:04:05.000Z07"
-
 	accessFilename = "access.log"
 	errorFilename  = "error.log"
 	severeFilename = "severe.log"
@@ -64,14 +62,17 @@ var (
 	// ErrLogServiceNameNotSet is an error that indicates that the service name is not set.
 	ErrLogServiceNameNotSet = errors.New("log service name must be set")
 
+	timeFormat   = "2006-01-02T15:04:05.000Z07"
 	writeConsole bool
 	logLevel     uint32
-	infoLog      io.WriteCloser
-	errorLog     io.WriteCloser
-	severeLog    io.WriteCloser
-	slowLog      io.WriteCloser
-	statLog      io.WriteCloser
-	stackLog     io.Writer
+	// use uint32 for atomic operations
+	disableStat uint32
+	infoLog     io.WriteCloser
+	errorLog    io.WriteCloser
+	severeLog   io.WriteCloser
+	slowLog     io.WriteCloser
+	statLog     io.WriteCloser
+	stackLog    io.Writer
 
 	once        sync.Once
 	initialized uint32
@@ -80,10 +81,10 @@ var (
 
 type (
 	logEntry struct {
-		Timestamp string `json:"@timestamp"`
-		Level     string `json:"level"`
-		Duration  string `json:"duration,omitempty"`
-		Content   string `json:"content"`
+		Timestamp string      `json:"@timestamp"`
+		Level     string      `json:"level"`
+		Duration  string      `json:"duration,omitempty"`
+		Content   interface{} `json:"content"`
 	}
 
 	logOptions struct {
@@ -99,10 +100,13 @@ type (
 	Logger interface {
 		Error(...interface{})
 		Errorf(string, ...interface{})
+		Errorv(interface{})
 		Info(...interface{})
 		Infof(string, ...interface{})
+		Infov(interface{})
 		Slow(...interface{})
 		Slowf(string, ...interface{})
+		Slowv(interface{})
 		WithDuration(time.Duration) Logger
 	}
 )
@@ -117,6 +121,10 @@ func MustSetup(c LogConf) {
 // we need to allow different service frameworks to initialize logx respectively.
 // the same logic for SetUp
 func SetUp(c LogConf) error {
+	if len(c.TimeFormat) > 0 {
+		timeFormat = c.TimeFormat
+	}
+
 	switch c.Mode {
 	case consoleMode:
 		setupWithConsole(c)
@@ -130,7 +138,7 @@ func SetUp(c LogConf) error {
 
 // Alert alerts v in alert level, and the message is written to error log.
 func Alert(v string) {
-	output(errorLog, levelAlert, v)
+	outputText(errorLog, levelAlert, v)
 }
 
 // Close closes the logging.
@@ -192,24 +200,29 @@ func Disable() {
 	})
 }
 
+// DisableStat disables the stat logs.
+func DisableStat() {
+	atomic.StoreUint32(&disableStat, 1)
+}
+
 // Error writes v into error log.
 func Error(v ...interface{}) {
 	ErrorCaller(1, v...)
 }
 
-// Errorf writes v with format into error log.
-func Errorf(format string, v ...interface{}) {
-	ErrorCallerf(1, format, v...)
-}
-
 // ErrorCaller writes v with context into error log.
 func ErrorCaller(callDepth int, v ...interface{}) {
-	errorSync(fmt.Sprint(v...), callDepth+callerInnerDepth)
+	errorTextSync(fmt.Sprint(v...), callDepth+callerInnerDepth)
 }
 
 // ErrorCallerf writes v with context in format into error log.
 func ErrorCallerf(callDepth int, format string, v ...interface{}) {
-	errorSync(fmt.Sprintf(format, v...), callDepth+callerInnerDepth)
+	errorTextSync(fmt.Errorf(format, v...).Error(), callDepth+callerInnerDepth)
+}
+
+// Errorf writes v with format into error log.
+func Errorf(format string, v ...interface{}) {
+	ErrorCallerf(1, format, v...)
 }
 
 // ErrorStack writes v along with call stack into error log.
@@ -224,14 +237,25 @@ func ErrorStackf(format string, v ...interface{}) {
 	stackSync(fmt.Sprintf(format, v...))
 }
 
+// Errorv writes v into error log with json content.
+// No call stack attached, because not elegant to pack the messages.
+func Errorv(v interface{}) {
+	errorAnySync(v)
+}
+
 // Info writes v into access log.
 func Info(v ...interface{}) {
-	infoSync(fmt.Sprint(v...))
+	infoTextSync(fmt.Sprint(v...))
 }
 
 // Infof writes v with format into access log.
 func Infof(format string, v ...interface{}) {
-	infoSync(fmt.Sprintf(format, v...))
+	infoTextSync(fmt.Sprintf(format, v...))
+}
+
+// Infov writes v into access log with json content.
+func Infov(v interface{}) {
+	infoAnySync(v)
 }
 
 // Must checks if err is nil, otherwise logs the err and exits.
@@ -239,7 +263,7 @@ func Must(err error) {
 	if err != nil {
 		msg := formatWithCaller(err.Error(), 3)
 		log.Print(msg)
-		output(severeLog, levelFatal, msg)
+		outputText(severeLog, levelFatal, msg)
 		os.Exit(1)
 	}
 }
@@ -261,12 +285,17 @@ func Severef(format string, v ...interface{}) {
 
 // Slow writes v into slow log.
 func Slow(v ...interface{}) {
-	slowSync(fmt.Sprint(v...))
+	slowTextSync(fmt.Sprint(v...))
 }
 
 // Slowf writes v with format into slow log.
 func Slowf(format string, v ...interface{}) {
-	slowSync(fmt.Sprintf(format, v...))
+	slowTextSync(fmt.Sprintf(format, v...))
+}
+
+// Slowv writes v into slow log with json content.
+func Slowv(v interface{}) {
+	slowAnySync(v)
 }
 
 // Stat writes v into stat log.
@@ -309,8 +338,14 @@ func createOutput(path string) (io.WriteCloser, error) {
 		options.gzipEnabled), options.gzipEnabled)
 }
 
-func errorSync(msg string, callDepth int) {
-	if shouldLog(ErrorLevel) {
+func errorAnySync(v interface{}) {
+	if shallLog(ErrorLevel) {
+		outputAny(errorLog, levelError, v)
+	}
+}
+
+func errorTextSync(msg string, callDepth int) {
+	if shallLog(ErrorLevel) {
 		outputError(errorLog, msg, callDepth)
 	}
 }
@@ -359,13 +394,28 @@ func handleOptions(opts []LogOption) {
 	}
 }
 
-func infoSync(msg string) {
-	if shouldLog(InfoLevel) {
-		output(infoLog, levelInfo, msg)
+func infoAnySync(val interface{}) {
+	if shallLog(InfoLevel) {
+		outputAny(infoLog, levelInfo, val)
 	}
 }
 
-func output(writer io.Writer, level, msg string) {
+func infoTextSync(msg string) {
+	if shallLog(InfoLevel) {
+		outputText(infoLog, levelInfo, msg)
+	}
+}
+
+func outputAny(writer io.Writer, level string, val interface{}) {
+	info := logEntry{
+		Timestamp: getTimestamp(),
+		Level:     level,
+		Content:   val,
+	}
+	outputJson(writer, info)
+}
+
+func outputText(writer io.Writer, level, msg string) {
 	info := logEntry{
 		Timestamp: getTimestamp(),
 		Level:     level,
@@ -376,7 +426,7 @@ func output(writer io.Writer, level, msg string) {
 
 func outputError(writer io.Writer, msg string, callDepth int) {
 	content := formatWithCaller(msg, callDepth)
-	output(writer, levelError, content)
+	outputText(writer, levelError, content)
 }
 
 func outputJson(writer io.Writer, info interface{}) {
@@ -478,30 +528,40 @@ func setupWithVolume(c LogConf) error {
 }
 
 func severeSync(msg string) {
-	if shouldLog(SevereLevel) {
-		output(severeLog, levelSevere, fmt.Sprintf("%s\n%s", msg, string(debug.Stack())))
+	if shallLog(SevereLevel) {
+		outputText(severeLog, levelSevere, fmt.Sprintf("%s\n%s", msg, string(debug.Stack())))
 	}
 }
 
-func shouldLog(level uint32) bool {
+func shallLog(level uint32) bool {
 	return atomic.LoadUint32(&logLevel) <= level
 }
 
-func slowSync(msg string) {
-	if shouldLog(ErrorLevel) {
-		output(slowLog, levelSlow, msg)
+func shallLogStat() bool {
+	return atomic.LoadUint32(&disableStat) == 0
+}
+
+func slowAnySync(v interface{}) {
+	if shallLog(ErrorLevel) {
+		outputAny(slowLog, levelSlow, v)
+	}
+}
+
+func slowTextSync(msg string) {
+	if shallLog(ErrorLevel) {
+		outputText(slowLog, levelSlow, msg)
 	}
 }
 
 func stackSync(msg string) {
-	if shouldLog(ErrorLevel) {
-		output(stackLog, levelError, fmt.Sprintf("%s\n%s", msg, string(debug.Stack())))
+	if shallLog(ErrorLevel) {
+		outputText(stackLog, levelError, fmt.Sprintf("%s\n%s", msg, string(debug.Stack())))
 	}
 }
 
 func statSync(msg string) {
-	if shouldLog(InfoLevel) {
-		output(statLog, levelStat, msg)
+	if shallLogStat() && shallLog(InfoLevel) {
+		outputText(statLog, levelStat, msg)
 	}
 }
 

core/logx/logs_test.go

@@ -2,6 +2,7 @@ package logx
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -92,6 +93,30 @@ func TestStructedLogAlert(t *testing.T) {
 	})
 }
 
+func TestStructedLogError(t *testing.T) {
+	doTestStructedLog(t, levelError, func(writer io.WriteCloser) {
+		errorLog = writer
+	}, func(v ...interface{}) {
+		Error(v...)
+	})
+}
+
+func TestStructedLogErrorf(t *testing.T) {
+	doTestStructedLog(t, levelError, func(writer io.WriteCloser) {
+		errorLog = writer
+	}, func(v ...interface{}) {
+		Errorf("%s", fmt.Sprint(v...))
+	})
+}
+
+func TestStructedLogErrorv(t *testing.T) {
+	doTestStructedLog(t, levelError, func(writer io.WriteCloser) {
+		errorLog = writer
+	}, func(v ...interface{}) {
+		Errorv(fmt.Sprint(v...))
+	})
+}
+
 func TestStructedLogInfo(t *testing.T) {
 	doTestStructedLog(t, levelInfo, func(writer io.WriteCloser) {
 		infoLog = writer
@@ -100,6 +125,22 @@ func TestStructedLogInfo(t *testing.T) {
 	})
 }
 
+func TestStructedLogInfof(t *testing.T) {
+	doTestStructedLog(t, levelInfo, func(writer io.WriteCloser) {
+		infoLog = writer
+	}, func(v ...interface{}) {
+		Infof("%s", fmt.Sprint(v...))
+	})
+}
+
+func TestStructedLogInfov(t *testing.T) {
+	doTestStructedLog(t, levelInfo, func(writer io.WriteCloser) {
+		infoLog = writer
+	}, func(v ...interface{}) {
+		Infov(fmt.Sprint(v...))
+	})
+}
+
 func TestStructedLogSlow(t *testing.T) {
 	doTestStructedLog(t, levelSlow, func(writer io.WriteCloser) {
 		slowLog = writer
@@ -116,6 +157,14 @@ func TestStructedLogSlowf(t *testing.T) {
 	})
 }
 
+func TestStructedLogSlowv(t *testing.T) {
+	doTestStructedLog(t, levelSlow, func(writer io.WriteCloser) {
+		slowLog = writer
+	}, func(v ...interface{}) {
+		Slowv(fmt.Sprint(v...))
+	})
+}
+
 func TestStructedLogStat(t *testing.T) {
 	doTestStructedLog(t, levelStat, func(writer io.WriteCloser) {
 		statLog = writer
@@ -194,6 +243,16 @@ func TestSetLevelWithDuration(t *testing.T) {
 	assert.Equal(t, 0, writer.builder.Len())
 }
 
+func TestErrorfWithWrappedError(t *testing.T) {
+	SetLevel(ErrorLevel)
+	const message = "there"
+	writer := new(mockWriter)
+	errorLog = writer
+	atomic.StoreUint32(&initialized, 1)
+	Errorf("hello %w", errors.New(message))
+	assert.True(t, strings.Contains(writer.builder.String(), "hello there"))
+}
+
 func TestMustNil(t *testing.T) {
 	Must(nil)
 }
@@ -246,6 +305,17 @@ func TestDisable(t *testing.T) {
 	assert.Nil(t, Close())
 }
 
+func TestDisableStat(t *testing.T) {
+	DisableStat()
+
+	const message = "hello there"
+	writer := new(mockWriter)
+	statLog = writer
+	atomic.StoreUint32(&initialized, 1)
+	Stat(message)
+	assert.Equal(t, 0, writer.builder.Len())
+}
+
 func TestWithGzip(t *testing.T) {
 	fn := WithGzip()
 	var opt logOptions
@@ -357,7 +427,9 @@ func doTestStructedLog(t *testing.T, level string, setup func(writer io.WriteClo
 		t.Error(err)
 	}
 	assert.Equal(t, level, entry.Level)
-	assert.True(t, strings.Contains(entry.Content, message))
+	val, ok := entry.Content.(string)
+	assert.True(t, ok)
+	assert.True(t, strings.Contains(val, message))
 }
 
 func testSetLevelTwiceWithMode(t *testing.T, mode string) {

core/logx/rotatelogger.go

@@ -22,8 +22,8 @@ const (
 	dateFormat      = "2006-01-02"
 	hoursPerDay     = 24
 	bufferSize      = 100
-	defaultDirMode  = 0755
-	defaultFileMode = 0600
+	defaultDirMode  = 0o755
+	defaultFileMode = 0o600
 )
 
 // ErrLogFileClosed is an error that indicates the log file is already closed.
@@ -47,7 +47,6 @@ type (
 		done     chan lang.PlaceholderType
 		rule     RotateRule
 		compress bool
-		keepDays int
 		// can't use threading.RoutineGroup because of cycle import
 		waitGroup sync.WaitGroup
 		closeOnce sync.Once

core/logx/rotatelogger_test.go

@@ -3,6 +3,7 @@ package logx
 import (
 	"os"
 	"path/filepath"
+	"syscall"
 	"testing"
 	"time"
 
@@ -97,7 +98,13 @@ func TestRotateLoggerRotate(t *testing.T) {
 		}()
 	}
 	err = logger.rotate()
-	assert.Nil(t, err)
+	switch v := err.(type) {
+	case *os.LinkError:
+		// avoid rename error on docker container
+		assert.Equal(t, syscall.EXDEV, v.Err)
+	default:
+		assert.Nil(t, err)
+	}
 }
 
 func TestRotateLoggerWrite(t *testing.T) {

core/logx/syslog_test.go

@@ -44,5 +44,5 @@ func captureOutput(f func()) string {
 func getContent(jsonStr string) string {
 	var entry logEntry
 	json.Unmarshal([]byte(jsonStr), &entry)
-	return entry.Content
+	return entry.Content.(string)
 }

core/logx/tracelogger.go

@@ -7,7 +7,7 @@ import (
 	"time"
 
 	"github.com/tal-tech/go-zero/core/timex"
-	"github.com/tal-tech/go-zero/core/trace/tracespec"
+	"go.opentelemetry.io/otel/trace"
 )
 
 type traceLogger struct {
@@ -18,53 +18,75 @@ type traceLogger struct {
 }
 
 func (l *traceLogger) Error(v ...interface{}) {
-	if shouldLog(ErrorLevel) {
+	if shallLog(ErrorLevel) {
 		l.write(errorLog, levelError, formatWithCaller(fmt.Sprint(v...), durationCallerDepth))
 	}
 }
 
 func (l *traceLogger) Errorf(format string, v ...interface{}) {
-	if shouldLog(ErrorLevel) {
+	if shallLog(ErrorLevel) {
 		l.write(errorLog, levelError, formatWithCaller(fmt.Sprintf(format, v...), durationCallerDepth))
 	}
 }
 
+func (l *traceLogger) Errorv(v interface{}) {
+	if shallLog(ErrorLevel) {
+		l.write(errorLog, levelError, v)
+	}
+}
+
 func (l *traceLogger) Info(v ...interface{}) {
-	if shouldLog(InfoLevel) {
+	if shallLog(InfoLevel) {
 		l.write(infoLog, levelInfo, fmt.Sprint(v...))
 	}
 }
 
 func (l *traceLogger) Infof(format string, v ...interface{}) {
-	if shouldLog(InfoLevel) {
+	if shallLog(InfoLevel) {
 		l.write(infoLog, levelInfo, fmt.Sprintf(format, v...))
 	}
 }
 
+func (l *traceLogger) Infov(v interface{}) {
+	if shallLog(InfoLevel) {
+		l.write(infoLog, levelInfo, v)
+	}
+}
+
 func (l *traceLogger) Slow(v ...interface{}) {
-	if shouldLog(ErrorLevel) {
+	if shallLog(ErrorLevel) {
 		l.write(slowLog, levelSlow, fmt.Sprint(v...))
 	}
 }
 
 func (l *traceLogger) Slowf(format string, v ...interface{}) {
-	if shouldLog(ErrorLevel) {
+	if shallLog(ErrorLevel) {
 		l.write(slowLog, levelSlow, fmt.Sprintf(format, v...))
 	}
 }
 
+func (l *traceLogger) Slowv(v interface{}) {
+	if shallLog(ErrorLevel) {
+		l.write(slowLog, levelSlow, v)
+	}
+}
+
 func (l *traceLogger) WithDuration(duration time.Duration) Logger {
 	l.Duration = timex.ReprOfDuration(duration)
 	return l
 }
 
-func (l *traceLogger) write(writer io.Writer, level, content string) {
-	l.Timestamp = getTimestamp()
-	l.Level = level
-	l.Content = content
-	l.Trace = traceIdFromContext(l.ctx)
-	l.Span = spanIdFromContext(l.ctx)
-	outputJson(writer, l)
+func (l *traceLogger) write(writer io.Writer, level string, val interface{}) {
+	outputJson(writer, &traceLogger{
+		logEntry: logEntry{
+			Timestamp: getTimestamp(),
+			Level:     level,
+			Duration:  l.Duration,
+			Content:   val,
+		},
+		Trace: traceIdFromContext(l.ctx),
+		Span:  spanIdFromContext(l.ctx),
+	})
 }
 
 // WithContext sets ctx to log, for keeping tracing information.
@@ -75,19 +97,19 @@ func WithContext(ctx context.Context) Logger {
 }
 
 func spanIdFromContext(ctx context.Context) string {
-	t, ok := ctx.Value(tracespec.TracingKey).(tracespec.Trace)
-	if !ok {
-		return ""
+	spanCtx := trace.SpanContextFromContext(ctx)
+	if spanCtx.HasSpanID() {
+		return spanCtx.SpanID().String()
 	}
 
-	return t.SpanId()
+	return ""
 }
 
 func traceIdFromContext(ctx context.Context) string {
-	t, ok := ctx.Value(tracespec.TracingKey).(tracespec.Trace)
-	if !ok {
-		return ""
+	spanCtx := trace.SpanContextFromContext(ctx)
+	if spanCtx.HasTraceID() {
+		return spanCtx.TraceID().String()
 	}
 
-	return t.TraceId()
+	return ""
 }

core/logx/tracelogger_test.go

@@ -9,71 +9,102 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/trace/tracespec"
+	"go.opentelemetry.io/otel"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 )
 
 const (
-	mockTraceID = "mock-trace-id"
-	mockSpanID  = "mock-span-id"
+	traceKey = "trace"
+	spanKey  = "span"
 )
 
-var mock tracespec.Trace = new(mockTrace)
-
 func TestTraceLog(t *testing.T) {
 	var buf mockWriter
 	atomic.StoreUint32(&initialized, 1)
-	ctx := context.WithValue(context.Background(), tracespec.TracingKey, mock)
+	otp := otel.GetTracerProvider()
+	tp := sdktrace.NewTracerProvider(sdktrace.WithSampler(sdktrace.AlwaysSample()))
+	otel.SetTracerProvider(tp)
+	defer otel.SetTracerProvider(otp)
+
+	ctx, _ := tp.Tracer("foo").Start(context.Background(), "bar")
 	WithContext(ctx).(*traceLogger).write(&buf, levelInfo, testlog)
-	assert.True(t, strings.Contains(buf.String(), mockTraceID))
-	assert.True(t, strings.Contains(buf.String(), mockSpanID))
+	assert.True(t, strings.Contains(buf.String(), traceKey))
+	assert.True(t, strings.Contains(buf.String(), spanKey))
 }
 
 func TestTraceError(t *testing.T) {
 	var buf mockWriter
 	atomic.StoreUint32(&initialized, 1)
 	errorLog = newLogWriter(log.New(&buf, "", flags))
-	ctx := context.WithValue(context.Background(), tracespec.TracingKey, mock)
+	otp := otel.GetTracerProvider()
+	tp := sdktrace.NewTracerProvider(sdktrace.WithSampler(sdktrace.AlwaysSample()))
+	otel.SetTracerProvider(tp)
+	defer otel.SetTracerProvider(otp)
+
+	ctx, _ := tp.Tracer("foo").Start(context.Background(), "bar")
 	l := WithContext(ctx).(*traceLogger)
 	SetLevel(InfoLevel)
 	l.WithDuration(time.Second).Error(testlog)
-	assert.True(t, strings.Contains(buf.String(), mockTraceID))
-	assert.True(t, strings.Contains(buf.String(), mockSpanID))
+	assert.True(t, strings.Contains(buf.String(), traceKey))
+	assert.True(t, strings.Contains(buf.String(), spanKey))
 	buf.Reset()
 	l.WithDuration(time.Second).Errorf(testlog)
-	assert.True(t, strings.Contains(buf.String(), mockTraceID))
-	assert.True(t, strings.Contains(buf.String(), mockSpanID))
+	assert.True(t, strings.Contains(buf.String(), traceKey))
+	assert.True(t, strings.Contains(buf.String(), spanKey))
+	buf.Reset()
+	l.WithDuration(time.Second).Errorv(testlog)
+	assert.True(t, strings.Contains(buf.String(), traceKey))
+	assert.True(t, strings.Contains(buf.String(), spanKey))
 }
 
 func TestTraceInfo(t *testing.T) {
 	var buf mockWriter
 	atomic.StoreUint32(&initialized, 1)
 	infoLog = newLogWriter(log.New(&buf, "", flags))
-	ctx := context.WithValue(context.Background(), tracespec.TracingKey, mock)
+	otp := otel.GetTracerProvider()
+	tp := sdktrace.NewTracerProvider(sdktrace.WithSampler(sdktrace.AlwaysSample()))
+	otel.SetTracerProvider(tp)
+	defer otel.SetTracerProvider(otp)
+
+	ctx, _ := tp.Tracer("foo").Start(context.Background(), "bar")
 	l := WithContext(ctx).(*traceLogger)
 	SetLevel(InfoLevel)
 	l.WithDuration(time.Second).Info(testlog)
-	assert.True(t, strings.Contains(buf.String(), mockTraceID))
-	assert.True(t, strings.Contains(buf.String(), mockSpanID))
+	assert.True(t, strings.Contains(buf.String(), traceKey))
+	assert.True(t, strings.Contains(buf.String(), spanKey))
 	buf.Reset()
 	l.WithDuration(time.Second).Infof(testlog)
-	assert.True(t, strings.Contains(buf.String(), mockTraceID))
-	assert.True(t, strings.Contains(buf.String(), mockSpanID))
+	assert.True(t, strings.Contains(buf.String(), traceKey))
+	assert.True(t, strings.Contains(buf.String(), spanKey))
+	buf.Reset()
+	l.WithDuration(time.Second).Infov(testlog)
+	assert.True(t, strings.Contains(buf.String(), traceKey))
+	assert.True(t, strings.Contains(buf.String(), spanKey))
 }
 
 func TestTraceSlow(t *testing.T) {
 	var buf mockWriter
 	atomic.StoreUint32(&initialized, 1)
 	slowLog = newLogWriter(log.New(&buf, "", flags))
-	ctx := context.WithValue(context.Background(), tracespec.TracingKey, mock)
+	otp := otel.GetTracerProvider()
+	tp := sdktrace.NewTracerProvider(sdktrace.WithSampler(sdktrace.AlwaysSample()))
+	otel.SetTracerProvider(tp)
+	defer otel.SetTracerProvider(otp)
+
+	ctx, _ := tp.Tracer("foo").Start(context.Background(), "bar")
 	l := WithContext(ctx).(*traceLogger)
 	SetLevel(InfoLevel)
 	l.WithDuration(time.Second).Slow(testlog)
-	assert.True(t, strings.Contains(buf.String(), mockTraceID))
-	assert.True(t, strings.Contains(buf.String(), mockSpanID))
+	assert.True(t, strings.Contains(buf.String(), traceKey))
+	assert.True(t, strings.Contains(buf.String(), spanKey))
 	buf.Reset()
 	l.WithDuration(time.Second).Slowf(testlog)
-	assert.True(t, strings.Contains(buf.String(), mockTraceID))
-	assert.True(t, strings.Contains(buf.String(), mockSpanID))
+	assert.True(t, strings.Contains(buf.String(), traceKey))
+	assert.True(t, strings.Contains(buf.String(), spanKey))
+	buf.Reset()
+	l.WithDuration(time.Second).Slowv(testlog)
+	assert.True(t, strings.Contains(buf.String(), traceKey))
+	assert.True(t, strings.Contains(buf.String(), spanKey))
 }
 
 func TestTraceWithoutContext(t *testing.T) {
@@ -83,34 +114,10 @@ func TestTraceWithoutContext(t *testing.T) {
 	l := WithContext(context.Background()).(*traceLogger)
 	SetLevel(InfoLevel)
 	l.WithDuration(time.Second).Info(testlog)
-	assert.False(t, strings.Contains(buf.String(), mockTraceID))
-	assert.False(t, strings.Contains(buf.String(), mockSpanID))
+	assert.False(t, strings.Contains(buf.String(), traceKey))
+	assert.False(t, strings.Contains(buf.String(), spanKey))
 	buf.Reset()
 	l.WithDuration(time.Second).Infof(testlog)
-	assert.False(t, strings.Contains(buf.String(), mockTraceID))
-	assert.False(t, strings.Contains(buf.String(), mockSpanID))
-}
-
-type mockTrace struct{}
-
-func (t mockTrace) TraceId() string {
-	return mockTraceID
-}
-
-func (t mockTrace) SpanId() string {
-	return mockSpanID
-}
-
-func (t mockTrace) Finish() {
-}
-
-func (t mockTrace) Fork(ctx context.Context, serviceName, operationName string) (context.Context, tracespec.Trace) {
-	return nil, nil
-}
-
-func (t mockTrace) Follow(ctx context.Context, serviceName, operationName string) (context.Context, tracespec.Trace) {
-	return nil, nil
-}
-
-func (t mockTrace) Visit(fn func(key string, val string) bool) {
+	assert.False(t, strings.Contains(buf.String(), traceKey))
+	assert.False(t, strings.Contains(buf.String(), spanKey))
 }

core/mapping/jsonunmarshaler.go

@@ -15,6 +15,11 @@ func UnmarshalJsonBytes(content []byte, v interface{}) error {
 	return unmarshalJsonBytes(content, v, jsonUnmarshaler)
 }
 
+// UnmarshalJsonMap unmarshals content from m into v.
+func UnmarshalJsonMap(m map[string]interface{}, v interface{}) error {
+	return jsonUnmarshaler.Unmarshal(m, v)
+}
+
 // UnmarshalJsonReader unmarshals content from reader into v.
 func UnmarshalJsonReader(reader io.Reader, v interface{}) error {
 	return unmarshalJsonReader(reader, v, jsonUnmarshaler)

core/mapping/jsonunmarshaler_test.go

@@ -871,3 +871,50 @@ func TestUnmarshalReaderError(t *testing.T) {
 	assert.NotNil(t, err)
 	assert.True(t, strings.Contains(err.Error(), payload))
 }
+
+func TestUnmarshalMap(t *testing.T) {
+	t.Run("nil map and valid", func(t *testing.T) {
+		var m map[string]interface{}
+		var v struct {
+			Any string `json:",optional"`
+		}
+
+		err := UnmarshalJsonMap(m, &v)
+		assert.Nil(t, err)
+		assert.True(t, len(v.Any) == 0)
+	})
+
+	t.Run("empty map but not valid", func(t *testing.T) {
+		m := map[string]interface{}{}
+		var v struct {
+			Any string
+		}
+
+		err := UnmarshalJsonMap(m, &v)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("empty map and valid", func(t *testing.T) {
+		m := map[string]interface{}{}
+		var v struct {
+			Any string `json:",optional"`
+		}
+
+		err := UnmarshalJsonMap(m, &v)
+		assert.Nil(t, err)
+		assert.True(t, len(v.Any) == 0)
+	})
+
+	t.Run("valid map", func(t *testing.T) {
+		m := map[string]interface{}{
+			"Any": "foo",
+		}
+		var v struct {
+			Any string
+		}
+
+		err := UnmarshalJsonMap(m, &v)
+		assert.Nil(t, err)
+		assert.Equal(t, "foo", v.Any)
+	})
+}

core/mapping/unmarshaler.go

@@ -7,7 +7,6 @@ import (
 	"reflect"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/tal-tech/go-zero/core/jsonx"
@@ -25,15 +24,17 @@ var (
 	errValueNotSettable = errors.New("value is not settable")
 	errValueNotStruct   = errors.New("value type is not struct")
 	keyUnmarshaler      = NewUnmarshaler(defaultKeyName)
-	cacheKeys           atomic.Value
-	cacheKeysLock       sync.Mutex
 	durationType        = reflect.TypeOf(time.Duration(0))
+	cacheKeys           map[string][]string
+	cacheKeysLock       sync.Mutex
+	defaultCache        map[string]interface{}
+	defaultCacheLock    sync.Mutex
 	emptyMap            = map[string]interface{}{}
 	emptyValue          = reflect.ValueOf(lang.Placeholder)
 )
 
 type (
-	// A Unmarshaler is used to unmarshal with given tag key.
+	// Unmarshaler is used to unmarshal with given tag key.
 	Unmarshaler struct {
 		key  string
 		opts unmarshalOptions
@@ -43,14 +44,14 @@ type (
 	UnmarshalOption func(*unmarshalOptions)
 
 	unmarshalOptions struct {
-		fromString bool
+		fromString   bool
+		canonicalKey func(key string) string
 	}
-
-	keyCache map[string][]string
 )
 
 func init() {
-	cacheKeys.Store(make(keyCache))
+	cacheKeys = make(map[string][]string)
+	defaultCache = make(map[string]interface{})
 }
 
 // NewUnmarshaler returns a Unmarshaler.
@@ -206,6 +207,8 @@ func (u *Unmarshaler) processFieldNotFromString(field reflect.StructField, value
 	switch {
 	case valueKind == reflect.Map && typeKind == reflect.Struct:
 		return u.processFieldStruct(field, value, mapValue, fullName)
+	case valueKind == reflect.Map && typeKind == reflect.Map:
+		return u.fillMap(field, value, mapValue)
 	case valueKind == reflect.String && typeKind == reflect.Slice:
 		return u.fillSliceFromString(fieldType, value, mapValue)
 	case valueKind == reflect.String && derefedFieldType == durationType:
@@ -229,7 +232,7 @@ func (u *Unmarshaler) processFieldPrimitive(field reflect.StructField, value ref
 	default:
 		switch v := mapValue.(type) {
 		case json.Number:
-			return u.processFieldPrimitiveWithJsonNumber(field, value, v, opts, fullName)
+			return u.processFieldPrimitiveWithJSONNumber(field, value, v, opts, fullName)
 		default:
 			if typeKind == valueKind {
 				if err := validateValueInOptions(opts.options(), mapValue); err != nil {
@@ -244,7 +247,7 @@ func (u *Unmarshaler) processFieldPrimitive(field reflect.StructField, value ref
 	return newTypeMismatchError(fullName)
 }
 
-func (u *Unmarshaler) processFieldPrimitiveWithJsonNumber(field reflect.StructField, value reflect.Value,
+func (u *Unmarshaler) processFieldPrimitiveWithJSONNumber(field reflect.StructField, value reflect.Value,
 	v json.Number, opts *fieldOptionsWithContext, fullName string) error {
 	fieldType := field.Type
 	fieldKind := fieldType.Kind()
@@ -323,7 +326,11 @@ func (u *Unmarshaler) processNamedField(field reflect.StructField, value reflect
 	}
 
 	fullName = join(fullName, key)
-	mapValue, hasValue := getValue(m, key)
+	canonicalKey := key
+	if u.opts.canonicalKey != nil {
+		canonicalKey = u.opts.canonicalKey(key)
+	}
+	mapValue, hasValue := getValue(m, canonicalKey)
 	if hasValue {
 		return u.processNamedFieldWithValue(field, value, mapValue, key, opts, fullName)
 	}
@@ -381,7 +388,13 @@ func (u *Unmarshaler) processNamedFieldWithoutValue(field reflect.StructField, v
 		if derefedType == durationType {
 			return fillDurationValue(fieldKind, value, defaultValue)
 		}
-		return setValue(fieldKind, value, defaultValue)
+
+		switch fieldKind {
+		case reflect.Array, reflect.Slice:
+			return u.fillSliceWithDefault(derefedType, value, defaultValue)
+		default:
+			return setValue(fieldKind, value, defaultValue)
+		}
 	}
 
 	switch fieldKind {
@@ -457,6 +470,10 @@ func (u *Unmarshaler) fillSlice(fieldType reflect.Type, value reflect.Value, map
 			} else {
 				conv.Index(i).Set(target.Elem())
 			}
+		case reflect.Slice:
+			if err := u.fillSlice(dereffedBaseType, conv.Index(i), ithValue); err != nil {
+				return err
+			}
 		default:
 			if err := u.fillSliceValue(conv, i, dereffedBaseKind, ithValue); err != nil {
 				return err
@@ -491,22 +508,58 @@ func (u *Unmarshaler) fillSliceFromString(fieldType reflect.Type, value reflect.
 	return nil
 }
 
-func (u *Unmarshaler) fillSliceValue(slice reflect.Value, index int, baseKind reflect.Kind, value interface{}) error {
+func (u *Unmarshaler) fillSliceValue(slice reflect.Value, index int,
+	baseKind reflect.Kind, value interface{}) error {
+	ithVal := slice.Index(index)
 	switch v := value.(type) {
 	case json.Number:
-		return setValue(baseKind, slice.Index(index), v.String())
+		return setValue(baseKind, ithVal, v.String())
 	default:
 		// don't need to consider the difference between int, int8, int16, int32, int64,
 		// uint, uint8, uint16, uint32, uint64, because they're handled as json.Number.
-		if slice.Index(index).Kind() != reflect.TypeOf(value).Kind() {
+		if ithVal.Kind() == reflect.Ptr {
+			baseType := Deref(ithVal.Type())
+			if baseType.Kind() != reflect.TypeOf(value).Kind() {
+				return errTypeMismatch
+			}
+
+			target := reflect.New(baseType).Elem()
+			target.Set(reflect.ValueOf(value))
+			ithVal.Set(target.Addr())
+			return nil
+		}
+
+		if ithVal.Kind() != reflect.TypeOf(value).Kind() {
 			return errTypeMismatch
 		}
 
-		slice.Index(index).Set(reflect.ValueOf(value))
+		ithVal.Set(reflect.ValueOf(value))
 		return nil
 	}
 }
 
+func (u *Unmarshaler) fillSliceWithDefault(derefedType reflect.Type, value reflect.Value,
+	defaultValue string) error {
+	baseFieldType := Deref(derefedType.Elem())
+	baseFieldKind := baseFieldType.Kind()
+	defaultCacheLock.Lock()
+	slice, ok := defaultCache[defaultValue]
+	defaultCacheLock.Unlock()
+	if !ok {
+		if baseFieldKind == reflect.String {
+			slice = parseGroupedSegments(defaultValue)
+		} else if err := jsonx.UnmarshalFromString(defaultValue, &slice); err != nil {
+			return err
+		}
+
+		defaultCacheLock.Lock()
+		defaultCache[defaultValue] = slice
+		defaultCacheLock.Unlock()
+	}
+
+	return u.fillSlice(derefedType, value, slice)
+}
+
 func (u *Unmarshaler) generateMap(keyType, elemType reflect.Type, mapValue interface{}) (reflect.Value, error) {
 	mapType := reflect.MapOf(keyType, elemType)
 	valueType := reflect.TypeOf(mapValue)
@@ -562,6 +615,8 @@ func (u *Unmarshaler) generateMap(keyType, elemType reflect.Type, mapValue inter
 			targetValue.SetMapIndex(key, innerValue)
 		default:
 			switch v := keythData.(type) {
+			case bool:
+				targetValue.SetMapIndex(key, reflect.ValueOf(v))
 			case string:
 				targetValue.SetMapIndex(key, reflect.ValueOf(v))
 			case json.Number:
@@ -604,6 +659,13 @@ func WithStringValues() UnmarshalOption {
 	}
 }
 
+// WithCanonicalKeyFunc customizes a Unmarshaler with Canonical Key func
+func WithCanonicalKeyFunc(f func(string) string) UnmarshalOption {
+	return func(opt *unmarshalOptions) {
+		opt.canonicalKey = f
+	}
+}
+
 func fillDurationValue(fieldKind reflect.Kind, value reflect.Value, dur string) error {
 	d, err := time.ParseDuration(dur)
 	if err != nil {
@@ -691,20 +753,6 @@ func getValueWithChainedKeys(m Valuer, keys []string) (interface{}, bool) {
 	return nil, false
 }
 
-func insertKeys(key string, cache []string) {
-	cacheKeysLock.Lock()
-	defer cacheKeysLock.Unlock()
-
-	keys := cacheKeys.Load().(keyCache)
-	// copy the contents into the new map, to guarantee the old map is immutable
-	newKeys := make(keyCache)
-	for k, v := range keys {
-		newKeys[k] = v
-	}
-	newKeys[key] = cache
-	cacheKeys.Store(newKeys)
-}
-
 func join(elem ...string) string {
 	var builder strings.Builder
 
@@ -735,15 +783,19 @@ func newTypeMismatchError(name string) error {
 }
 
 func readKeys(key string) []string {
-	cache := cacheKeys.Load().(keyCache)
-	if keys, ok := cache[key]; ok {
+	cacheKeysLock.Lock()
+	keys, ok := cacheKeys[key]
+	cacheKeysLock.Unlock()
+	if ok {
 		return keys
 	}
 
-	keys := strings.FieldsFunc(key, func(c rune) bool {
+	keys = strings.FieldsFunc(key, func(c rune) bool {
 		return c == delimiter
 	})
-	insertKeys(key, keys)
+	cacheKeysLock.Lock()
+	cacheKeys[key] = keys
+	cacheKeysLock.Unlock()
 
 	return keys
 }

core/mapping/unmarshaler_test.go

@@ -3,6 +3,7 @@ package mapping
 import (
 	"encoding/json"
 	"strconv"
+	"strings"
 	"testing"
 	"time"
 
@@ -197,6 +198,66 @@ func TestUnmarshalIntWithDefault(t *testing.T) {
 	assert.Equal(t, 1, in.Int)
 }
 
+func TestUnmarshalBoolSliceWithDefault(t *testing.T) {
+	type inner struct {
+		Bools []bool `key:"bools,default=[true,false]"`
+	}
+
+	var in inner
+	assert.Nil(t, UnmarshalKey(nil, &in))
+	assert.ElementsMatch(t, []bool{true, false}, in.Bools)
+}
+
+func TestUnmarshalIntSliceWithDefault(t *testing.T) {
+	type inner struct {
+		Ints []int `key:"ints,default=[1,2,3]"`
+	}
+
+	var in inner
+	assert.Nil(t, UnmarshalKey(nil, &in))
+	assert.ElementsMatch(t, []int{1, 2, 3}, in.Ints)
+}
+
+func TestUnmarshalIntSliceWithDefaultHasSpaces(t *testing.T) {
+	type inner struct {
+		Ints []int `key:"ints,default=[1, 2, 3]"`
+	}
+
+	var in inner
+	assert.Nil(t, UnmarshalKey(nil, &in))
+	assert.ElementsMatch(t, []int{1, 2, 3}, in.Ints)
+}
+
+func TestUnmarshalFloatSliceWithDefault(t *testing.T) {
+	type inner struct {
+		Floats []float32 `key:"floats,default=[1.1,2.2,3.3]"`
+	}
+
+	var in inner
+	assert.Nil(t, UnmarshalKey(nil, &in))
+	assert.ElementsMatch(t, []float32{1.1, 2.2, 3.3}, in.Floats)
+}
+
+func TestUnmarshalStringSliceWithDefault(t *testing.T) {
+	type inner struct {
+		Strs []string `key:"strs,default=[foo,bar,woo]"`
+	}
+
+	var in inner
+	assert.Nil(t, UnmarshalKey(nil, &in))
+	assert.ElementsMatch(t, []string{"foo", "bar", "woo"}, in.Strs)
+}
+
+func TestUnmarshalStringSliceWithDefaultHasSpaces(t *testing.T) {
+	type inner struct {
+		Strs []string `key:"strs,default=[foo, bar, woo]"`
+	}
+
+	var in inner
+	assert.Nil(t, UnmarshalKey(nil, &in))
+	assert.ElementsMatch(t, []string{"foo", "bar", "woo"}, in.Strs)
+}
+
 func TestUnmarshalUint(t *testing.T) {
 	type inner struct {
 		Uint          uint   `key:"uint"`
@@ -752,7 +813,7 @@ func TestUnmarshalJsonNumberInt64(t *testing.T) {
 	for i := 0; i <= maxUintBitsToTest; i++ {
 		var intValue int64 = 1 << uint(i)
 		strValue := strconv.FormatInt(intValue, 10)
-		var number = json.Number(strValue)
+		number := json.Number(strValue)
 		m := map[string]interface{}{
 			"ID": number,
 		}
@@ -768,7 +829,7 @@ func TestUnmarshalJsonNumberUint64(t *testing.T) {
 	for i := 0; i <= maxUintBitsToTest; i++ {
 		var intValue uint64 = 1 << uint(i)
 		strValue := strconv.FormatUint(intValue, 10)
-		var number = json.Number(strValue)
+		number := json.Number(strValue)
 		m := map[string]interface{}{
 			"ID": number,
 		}
@@ -784,7 +845,7 @@ func TestUnmarshalJsonNumberUint64Ptr(t *testing.T) {
 	for i := 0; i <= maxUintBitsToTest; i++ {
 		var intValue uint64 = 1 << uint(i)
 		strValue := strconv.FormatUint(intValue, 10)
-		var number = json.Number(strValue)
+		number := json.Number(strValue)
 		m := map[string]interface{}{
 			"ID": number,
 		}
@@ -860,10 +921,12 @@ func TestUnmarshalSliceOfStruct(t *testing.T) {
 func TestUnmarshalWithStringOptionsCorrect(t *testing.T) {
 	type inner struct {
 		Value   string `key:"value,options=first|second"`
+		Foo     string `key:"foo,options=[bar,baz]"`
 		Correct string `key:"correct,options=1|2"`
 	}
 	m := map[string]interface{}{
 		"value":   "first",
+		"foo":     "bar",
 		"correct": "2",
 	}
 
@@ -871,6 +934,7 @@ func TestUnmarshalWithStringOptionsCorrect(t *testing.T) {
 	ast := assert.New(t)
 	ast.Nil(UnmarshalKey(m, &in))
 	ast.Equal("first", in.Value)
+	ast.Equal("bar", in.Foo)
 	ast.Equal("2", in.Correct)
 }
 
@@ -942,6 +1006,22 @@ func TestUnmarshalStringOptionsWithStringOptionsIncorrect(t *testing.T) {
 	ast.NotNil(unmarshaler.Unmarshal(m, &in))
 }
 
+func TestUnmarshalStringOptionsWithStringOptionsIncorrectGrouped(t *testing.T) {
+	type inner struct {
+		Value   string `key:"value,options=[first,second]"`
+		Correct string `key:"correct,options=1|2"`
+	}
+	m := map[string]interface{}{
+		"value":   "third",
+		"correct": "2",
+	}
+
+	var in inner
+	unmarshaler := NewUnmarshaler(defaultKeyName, WithStringValues())
+	ast := assert.New(t)
+	ast.NotNil(unmarshaler.Unmarshal(m, &in))
+}
+
 func TestUnmarshalWithStringOptionsIncorrect(t *testing.T) {
 	type inner struct {
 		Value     string `key:"value,options=first|second"`
@@ -2480,3 +2560,66 @@ func BenchmarkUnmarshal(b *testing.B) {
 		UnmarshalKey(data, &an)
 	}
 }
+
+func TestUnmarshalJsonReaderMultiArray(t *testing.T) {
+	payload := `{"a": "133", "b": [["add", "cccd"], ["eeee"]]}`
+	var res struct {
+		A string     `json:"a"`
+		B [][]string `json:"b"`
+	}
+	reader := strings.NewReader(payload)
+	err := UnmarshalJsonReader(reader, &res)
+	assert.Nil(t, err)
+	assert.Equal(t, 2, len(res.B))
+}
+
+func TestUnmarshalJsonReaderPtrMultiArray(t *testing.T) {
+	payload := `{"a": "133", "b": [["add", "cccd"], ["eeee"]]}`
+	var res struct {
+		A string      `json:"a"`
+		B [][]*string `json:"b"`
+	}
+	reader := strings.NewReader(payload)
+	err := UnmarshalJsonReader(reader, &res)
+	assert.Nil(t, err)
+	assert.Equal(t, 2, len(res.B))
+	assert.Equal(t, 2, len(res.B[0]))
+}
+
+func TestUnmarshalJsonReaderPtrArray(t *testing.T) {
+	payload := `{"a": "133", "b": ["add", "cccd", "eeee"]}`
+	var res struct {
+		A string    `json:"a"`
+		B []*string `json:"b"`
+	}
+	reader := strings.NewReader(payload)
+	err := UnmarshalJsonReader(reader, &res)
+	assert.Nil(t, err)
+	assert.Equal(t, 3, len(res.B))
+}
+
+func TestUnmarshalJsonWithoutKey(t *testing.T) {
+	payload := `{"A": "1", "B": "2"}`
+	var res struct {
+		A string `json:""`
+		B string `json:","`
+	}
+	reader := strings.NewReader(payload)
+	err := UnmarshalJsonReader(reader, &res)
+	assert.Nil(t, err)
+	assert.Equal(t, "1", res.A)
+	assert.Equal(t, "2", res.B)
+}
+
+func BenchmarkDefaultValue(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		var a struct {
+			Ints []int    `json:"ints,default=[1,2,3]"`
+			Strs []string `json:"strs,default=[foo,bar,baz]"`
+		}
+		_ = UnmarshalJsonMap(nil, &a)
+		if len(a.Strs) != 3 || len(a.Ints) != 3 {
+			b.Fatal("failed")
+		}
+	}
+}

core/mapping/utils.go

@@ -14,13 +14,19 @@ import (
 )
 
 const (
-	defaultOption   = "default"
-	stringOption    = "string"
-	optionalOption  = "optional"
-	optionsOption   = "options"
-	rangeOption     = "range"
-	optionSeparator = "|"
-	equalToken      = "="
+	defaultOption      = "default"
+	stringOption       = "string"
+	optionalOption     = "optional"
+	optionsOption      = "options"
+	rangeOption        = "range"
+	optionSeparator    = "|"
+	equalToken         = "="
+	escapeChar         = '\\'
+	leftBracket        = '('
+	rightBracket       = ')'
+	leftSquareBracket  = '['
+	rightSquareBracket = ']'
+	segmentSeparator   = ','
 )
 
 var (
@@ -118,7 +124,7 @@ func convertType(kind reflect.Kind, str string) (interface{}, error) {
 }
 
 func doParseKeyAndOptions(field reflect.StructField, value string) (string, *fieldOptions, error) {
-	segments := strings.Split(value, ",")
+	segments := parseSegments(value)
 	key := strings.TrimSpace(segments[0])
 	options := segments[1:]
 
@@ -170,12 +176,44 @@ func implicitValueRequiredStruct(tag string, tp reflect.Type) (bool, error) {
 	return false, nil
 }
 
+func isLeftInclude(b byte) (bool, error) {
+	switch b {
+	case '[':
+		return true, nil
+	case '(':
+		return false, nil
+	default:
+		return false, errNumberRange
+	}
+}
+
+func isRightInclude(b byte) (bool, error) {
+	switch b {
+	case ']':
+		return true, nil
+	case ')':
+		return false, nil
+	default:
+		return false, errNumberRange
+	}
+}
+
 func maybeNewValue(field reflect.StructField, value reflect.Value) {
 	if field.Type.Kind() == reflect.Ptr && value.IsNil() {
 		value.Set(reflect.New(value.Type().Elem()))
 	}
 }
 
+func parseGroupedSegments(val string) []string {
+	val = strings.TrimLeftFunc(val, func(r rune) bool {
+		return r == leftBracket || r == leftSquareBracket
+	})
+	val = strings.TrimRightFunc(val, func(r rune) bool {
+		return r == rightBracket || r == rightSquareBracket
+	})
+	return parseSegments(val)
+}
+
 // don't modify returned fieldOptions, it's cached and shared among different calls.
 func parseKeyAndOptions(tagName string, field reflect.StructField) (string, *fieldOptions, error) {
 	value := field.Tag.Get(tagName)
@@ -211,14 +249,9 @@ func parseNumberRange(str string) (*numberRange, error) {
 		return nil, errNumberRange
 	}
 
-	var leftInclude bool
-	switch str[0] {
-	case '[':
-		leftInclude = true
-	case '(':
-		leftInclude = false
-	default:
-		return nil, errNumberRange
+	leftInclude, err := isLeftInclude(str[0])
+	if err != nil {
+		return nil, err
 	}
 
 	str = str[1:]
@@ -226,14 +259,9 @@ func parseNumberRange(str string) (*numberRange, error) {
 		return nil, errNumberRange
 	}
 
-	var rightInclude bool
-	switch str[len(str)-1] {
-	case ']':
-		rightInclude = true
-	case ')':
-		rightInclude = false
-	default:
-		return nil, errNumberRange
+	rightInclude, err := isRightInclude(str[len(str)-1])
+	if err != nil {
+		return nil, err
 	}
 
 	str = str[:len(str)-1]
@@ -274,7 +302,7 @@ func parseNumberRange(str string) (*numberRange, error) {
 	}, nil
 }
 
-func parseOption(fieldOpts *fieldOptions, fieldName string, option string) error {
+func parseOption(fieldOpts *fieldOptions, fieldName, option string) error {
 	switch {
 	case option == stringOption:
 		fieldOpts.FromString = true
@@ -297,7 +325,7 @@ func parseOption(fieldOpts *fieldOptions, fieldName string, option string) error
 			return fmt.Errorf("field %s has wrong options", fieldName)
 		}
 
-		fieldOpts.Options = strings.Split(segs[1], optionSeparator)
+		fieldOpts.Options = parseOptions(segs[1])
 	case strings.HasPrefix(option, defaultOption):
 		segs := strings.Split(option, equalToken)
 		if len(segs) != 2 {
@@ -322,6 +350,69 @@ func parseOption(fieldOpts *fieldOptions, fieldName string, option string) error
 	return nil
 }
 
+// parseOptions parses the given options in tag.
+// for example: `json:"name,options=foo|bar"` or `json:"name,options=[foo,bar]"`
+func parseOptions(val string) []string {
+	if len(val) == 0 {
+		return nil
+	}
+
+	if val[0] == leftSquareBracket {
+		return parseGroupedSegments(val)
+	}
+
+	return strings.Split(val, optionSeparator)
+}
+
+func parseSegments(val string) []string {
+	var segments []string
+	var escaped, grouped bool
+	var buf strings.Builder
+
+	for _, ch := range val {
+		if escaped {
+			buf.WriteRune(ch)
+			escaped = false
+			continue
+		}
+
+		switch ch {
+		case segmentSeparator:
+			if grouped {
+				buf.WriteRune(ch)
+			} else {
+				// need to trim spaces, but we cannot ignore empty string,
+				// because the first segment stands for the key might be empty.
+				// if ignored, the later tag will be used as the key.
+				segments = append(segments, strings.TrimSpace(buf.String()))
+				buf.Reset()
+			}
+		case escapeChar:
+			if grouped {
+				buf.WriteRune(ch)
+			} else {
+				escaped = true
+			}
+		case leftBracket, leftSquareBracket:
+			buf.WriteRune(ch)
+			grouped = true
+		case rightBracket, rightSquareBracket:
+			buf.WriteRune(ch)
+			grouped = false
+		default:
+			buf.WriteRune(ch)
+		}
+	}
+
+	last := strings.TrimSpace(buf.String())
+	// ignore last empty string
+	if len(last) > 0 {
+		segments = append(segments, last)
+	}
+
+	return segments
+}
+
 func reprOfValue(val reflect.Value) string {
 	switch vt := val.Interface().(type) {
 	case bool:

core/mapping/utils_test.go

@@ -16,8 +16,8 @@ type Foo struct {
 }
 
 func TestDeferInt(t *testing.T) {
-	var i = 1
-	var s = "hello"
+	i := 1
+	s := "hello"
 	number := struct {
 		f float64
 	}{
@@ -90,6 +90,82 @@ func TestParseKeyAndOptionWithTagAndOption(t *testing.T) {
 	assert.True(t, options.FromString)
 }
 
+func TestParseSegments(t *testing.T) {
+	tests := []struct {
+		input  string
+		expect []string
+	}{
+		{
+			input:  "",
+			expect: []string{},
+		},
+		{
+			input:  ",",
+			expect: []string{""},
+		},
+		{
+			input:  "foo,",
+			expect: []string{"foo"},
+		},
+		{
+			input: ",foo",
+			// the first empty string cannot be ignored, it's the key.
+			expect: []string{"", "foo"},
+		},
+		{
+			input:  "foo",
+			expect: []string{"foo"},
+		},
+		{
+			input:  "foo,bar",
+			expect: []string{"foo", "bar"},
+		},
+		{
+			input:  "foo,bar,baz",
+			expect: []string{"foo", "bar", "baz"},
+		},
+		{
+			input:  "foo,options=a|b",
+			expect: []string{"foo", "options=a|b"},
+		},
+		{
+			input:  "foo,bar,default=[baz,qux]",
+			expect: []string{"foo", "bar", "default=[baz,qux]"},
+		},
+		{
+			input:  "foo,bar,options=[baz,qux]",
+			expect: []string{"foo", "bar", "options=[baz,qux]"},
+		},
+		{
+			input:  `foo\,bar,options=[baz,qux]`,
+			expect: []string{`foo,bar`, "options=[baz,qux]"},
+		},
+		{
+			input:  `foo,bar,options=\[baz,qux]`,
+			expect: []string{"foo", "bar", "options=[baz", "qux]"},
+		},
+		{
+			input:  `foo,bar,options=[baz\,qux]`,
+			expect: []string{"foo", "bar", `options=[baz\,qux]`},
+		},
+		{
+			input:  `foo\,bar,options=[baz,qux],default=baz`,
+			expect: []string{`foo,bar`, "options=[baz,qux]", "default=baz"},
+		},
+		{
+			input:  `foo\,bar,options=[baz,qux, quux],default=[qux, baz]`,
+			expect: []string{`foo,bar`, "options=[baz,qux, quux]", "default=[qux, baz]"},
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.input, func(t *testing.T) {
+			assert.ElementsMatch(t, test.expect, parseSegments(test.input))
+		})
+	}
+}
+
 func TestValidatePtrWithNonPtr(t *testing.T) {
 	var foo string
 	rve := reflect.ValueOf(foo)
@@ -209,6 +285,12 @@ func TestRepr(t *testing.T) {
 			newMockPtr(),
 			"mockptr",
 		},
+		{
+			&mockOpacity{
+				val: 1,
+			},
+			"{1}",
+		},
 		{
 			true,
 			"true",

core/mapping/yamlunmarshaler_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"k8s.io/utils/io"
 )
 
 func TestUnmarshalYamlBytes(t *testing.T) {
@@ -18,6 +19,22 @@ func TestUnmarshalYamlBytes(t *testing.T) {
 	assert.Equal(t, "liao", c.Name)
 }
 
+func TestUnmarshalYamlBytesErrorInput(t *testing.T) {
+	var c struct {
+		Name string
+	}
+	content := []byte(`liao`)
+	assert.NotNil(t, UnmarshalYamlBytes(content, &c))
+}
+
+func TestUnmarshalYamlBytesEmptyInput(t *testing.T) {
+	var c struct {
+		Name string
+	}
+	content := []byte(``)
+	assert.NotNil(t, UnmarshalYamlBytes(content, &c))
+}
+
 func TestUnmarshalYamlBytesOptional(t *testing.T) {
 	var c struct {
 		Name string
@@ -918,3 +935,82 @@ func TestUnmarshalYamlReaderError(t *testing.T) {
 	err := UnmarshalYamlReader(reader, &v)
 	assert.NotNil(t, err)
 }
+
+func TestUnmarshalYamlBadReader(t *testing.T) {
+	var v struct {
+		Any string
+	}
+
+	err := UnmarshalYamlReader(new(badReader), &v)
+	assert.NotNil(t, err)
+}
+
+func TestUnmarshalYamlMapBool(t *testing.T) {
+	text := `machine:
+  node1: true
+  node2: true
+  node3: true
+`
+	var v struct {
+		Machine map[string]bool `json:"machine,optional"`
+	}
+	reader := strings.NewReader(text)
+	assert.Nil(t, UnmarshalYamlReader(reader, &v))
+	assert.True(t, v.Machine["node1"])
+	assert.True(t, v.Machine["node2"])
+	assert.True(t, v.Machine["node3"])
+}
+
+func TestUnmarshalYamlMapInt(t *testing.T) {
+	text := `machine:
+  node1: 1
+  node2: 2
+  node3: 3
+`
+	var v struct {
+		Machine map[string]int `json:"machine,optional"`
+	}
+	reader := strings.NewReader(text)
+	assert.Nil(t, UnmarshalYamlReader(reader, &v))
+	assert.Equal(t, 1, v.Machine["node1"])
+	assert.Equal(t, 2, v.Machine["node2"])
+	assert.Equal(t, 3, v.Machine["node3"])
+}
+
+func TestUnmarshalYamlMapByte(t *testing.T) {
+	text := `machine:
+  node1: 1
+  node2: 2
+  node3: 3
+`
+	var v struct {
+		Machine map[string]byte `json:"machine,optional"`
+	}
+	reader := strings.NewReader(text)
+	assert.Nil(t, UnmarshalYamlReader(reader, &v))
+	assert.Equal(t, byte(1), v.Machine["node1"])
+	assert.Equal(t, byte(2), v.Machine["node2"])
+	assert.Equal(t, byte(3), v.Machine["node3"])
+}
+
+func TestUnmarshalYamlMapRune(t *testing.T) {
+	text := `machine:
+  node1: 1
+  node2: 2
+  node3: 3
+`
+	var v struct {
+		Machine map[string]rune `json:"machine,optional"`
+	}
+	reader := strings.NewReader(text)
+	assert.Nil(t, UnmarshalYamlReader(reader, &v))
+	assert.Equal(t, rune(1), v.Machine["node1"])
+	assert.Equal(t, rune(2), v.Machine["node2"])
+	assert.Equal(t, rune(3), v.Machine["node3"])
+}
+
+type badReader struct{}
+
+func (b *badReader) Read(p []byte) (n int, err error) {
+	return 0, io.ErrLimitReached
+}

core/mathx/unstable.go

@@ -36,7 +36,7 @@ func (u Unstable) AroundDuration(base time.Duration) time.Duration {
 	return val
 }
 
-// AroundInt returns a randome int64 with given base and deviation.
+// AroundInt returns a random int64 with given base and deviation.
 func (u Unstable) AroundInt(base int64) int64 {
 	u.lock.Lock()
 	val := int64((1 + u.deviation - 2*u.deviation*u.r.Float64()) * float64(base))

core/metric/counter.go

@@ -12,7 +12,7 @@ type (
 	// CounterVec interface represents a counter vector.
 	CounterVec interface {
 		// Inc increments labels.
-		Inc(lables ...string)
+		Inc(labels ...string)
 		// Add adds labels with v.
 		Add(v float64, labels ...string)
 		close() bool
@@ -50,8 +50,8 @@ func (cv *promCounterVec) Inc(labels ...string) {
 	cv.counter.WithLabelValues(labels...).Inc()
 }
 
-func (cv *promCounterVec) Add(v float64, lables ...string) {
-	cv.counter.WithLabelValues(lables...).Add(v)
+func (cv *promCounterVec) Add(v float64, labels ...string) {
+	cv.counter.WithLabelValues(labels...).Add(v)
 }
 
 func (cv *promCounterVec) close() bool {

core/metric/gauge.go

@@ -20,7 +20,7 @@ type (
 		close() bool
 	}
 
-	promGuageVec struct {
+	promGaugeVec struct {
 		gauge *prom.GaugeVec
 	}
 )
@@ -39,7 +39,7 @@ func NewGaugeVec(cfg *GaugeVecOpts) GaugeVec {
 			Help:      cfg.Help,
 		}, cfg.Labels)
 	prom.MustRegister(vec)
-	gv := &promGuageVec{
+	gv := &promGaugeVec{
 		gauge: vec,
 	}
 	proc.AddShutdownListener(func() {
@@ -49,18 +49,18 @@ func NewGaugeVec(cfg *GaugeVecOpts) GaugeVec {
 	return gv
 }
 
-func (gv *promGuageVec) Inc(labels ...string) {
+func (gv *promGaugeVec) Inc(labels ...string) {
 	gv.gauge.WithLabelValues(labels...).Inc()
 }
 
-func (gv *promGuageVec) Add(v float64, lables ...string) {
-	gv.gauge.WithLabelValues(lables...).Add(v)
+func (gv *promGaugeVec) Add(v float64, labels ...string) {
+	gv.gauge.WithLabelValues(labels...).Add(v)
 }
 
-func (gv *promGuageVec) Set(v float64, lables ...string) {
-	gv.gauge.WithLabelValues(lables...).Set(v)
+func (gv *promGaugeVec) Set(v float64, labels ...string) {
+	gv.gauge.WithLabelValues(labels...).Set(v)
 }
 
-func (gv *promGuageVec) close() bool {
+func (gv *promGaugeVec) close() bool {
 	return prom.Unregister(gv.gauge)
 }

core/metric/gauge_test.go

@@ -29,7 +29,7 @@ func TestGaugeInc(t *testing.T) {
 		Labels:    []string{"path"},
 	})
 	defer gaugeVec.close()
-	gv, _ := gaugeVec.(*promGuageVec)
+	gv, _ := gaugeVec.(*promGaugeVec)
 	gv.Inc("/users")
 	gv.Inc("/users")
 	r := testutil.ToFloat64(gv.gauge)
@@ -45,7 +45,7 @@ func TestGaugeAdd(t *testing.T) {
 		Labels:    []string{"path"},
 	})
 	defer gaugeVec.close()
-	gv, _ := gaugeVec.(*promGuageVec)
+	gv, _ := gaugeVec.(*promGaugeVec)
 	gv.Add(-10, "/classroom")
 	gv.Add(30, "/classroom")
 	r := testutil.ToFloat64(gv.gauge)
@@ -61,7 +61,7 @@ func TestGaugeSet(t *testing.T) {
 		Labels:    []string{"path"},
 	})
 	gaugeVec.close()
-	gv, _ := gaugeVec.(*promGuageVec)
+	gv, _ := gaugeVec.(*promGaugeVec)
 	gv.Set(666, "/users")
 	r := testutil.ToFloat64(gv.gauge)
 	assert.Equal(t, float64(666), r)

core/metric/histogram.go

@@ -19,7 +19,7 @@ type (
 	// A HistogramVec interface represents a histogram vector.
 	HistogramVec interface {
 		// Observe adds observation v to labels.
-		Observe(v int64, lables ...string)
+		Observe(v int64, labels ...string)
 		close() bool
 	}
 

core/mr/mapreduce.go

@@ -1,13 +1,13 @@
 package mr
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"sync"
 
 	"github.com/tal-tech/go-zero/core/errorx"
 	"github.com/tal-tech/go-zero/core/lang"
-	"github.com/tal-tech/go-zero/core/syncx"
 	"github.com/tal-tech/go-zero/core/threading"
 )
 
@@ -43,6 +43,7 @@ type (
 	Option func(opts *mapReduceOptions)
 
 	mapReduceOptions struct {
+		ctx     context.Context
 		workers int
 	}
 
@@ -93,16 +94,17 @@ func Map(generate GenerateFunc, mapper MapFunc, opts ...Option) chan interface{}
 	options := buildOptions(opts...)
 	source := buildSource(generate)
 	collector := make(chan interface{}, options.workers)
-	done := syncx.NewDoneChan()
+	done := make(chan lang.PlaceholderType)
 
-	go executeMappers(mapper, source, collector, done.Done(), options.workers)
+	go executeMappers(options.ctx, mapper, source, collector, done, options.workers)
 
 	return collector
 }
 
 // MapReduce maps all elements generated from given generate func,
-// and reduces the output elemenets with given reducer.
-func MapReduce(generate GenerateFunc, mapper MapperFunc, reducer ReducerFunc, opts ...Option) (interface{}, error) {
+// and reduces the output elements with given reducer.
+func MapReduce(generate GenerateFunc, mapper MapperFunc, reducer ReducerFunc,
+	opts ...Option) (interface{}, error) {
 	source := buildSource(generate)
 	return MapReduceWithSource(source, mapper, reducer, opts...)
 }
@@ -112,14 +114,20 @@ func MapReduceWithSource(source <-chan interface{}, mapper MapperFunc, reducer R
 	opts ...Option) (interface{}, error) {
 	options := buildOptions(opts...)
 	output := make(chan interface{})
+	defer func() {
+		for range output {
+			panic("more than one element written in reducer")
+		}
+	}()
+
 	collector := make(chan interface{}, options.workers)
-	done := syncx.NewDoneChan()
-	writer := newGuardedWriter(output, done.Done())
+	done := make(chan lang.PlaceholderType)
+	writer := newGuardedWriter(options.ctx, output, done)
 	var closeOnce sync.Once
 	var retErr errorx.AtomicError
 	finish := func() {
 		closeOnce.Do(func() {
-			done.Close()
+			close(done)
 			close(output)
 		})
 	}
@@ -136,19 +144,21 @@ func MapReduceWithSource(source <-chan interface{}, mapper MapperFunc, reducer R
 
 	go func() {
 		defer func() {
+			drain(collector)
+
 			if r := recover(); r != nil {
 				cancel(fmt.Errorf("%v", r))
 			} else {
 				finish()
 			}
 		}()
+
 		reducer(collector, writer, cancel)
-		drain(collector)
 	}()
 
-	go executeMappers(func(item interface{}, w Writer) {
+	go executeMappers(options.ctx, func(item interface{}, w Writer) {
 		mapper(item, w, cancel)
-	}, source, collector, done.Done(), options.workers)
+	}, source, collector, done, options.workers)
 
 	value, ok := <-output
 	if err := retErr.Load(); err != nil {
@@ -165,7 +175,6 @@ func MapReduceWithSource(source <-chan interface{}, mapper MapperFunc, reducer R
 func MapReduceVoid(generate GenerateFunc, mapper MapperFunc, reducer VoidReducerFunc, opts ...Option) error {
 	_, err := MapReduce(generate, mapper, func(input <-chan interface{}, writer Writer, cancel func(error)) {
 		reducer(input, cancel)
-		drain(input)
 		// We need to write a placeholder to let MapReduce to continue on reducer done,
 		// otherwise, all goroutines are waiting. The placeholder will be discarded by MapReduce.
 		writer.Write(lang.Placeholder)
@@ -180,6 +189,13 @@ func MapVoid(generate GenerateFunc, mapper VoidMapFunc, opts ...Option) {
 	}, opts...))
 }
 
+// WithContext customizes a mapreduce processing accepts a given ctx.
+func WithContext(ctx context.Context) Option {
+	return func(opts *mapReduceOptions) {
+		opts.ctx = ctx
+	}
+}
+
 // WithWorkers customizes a mapreduce processing with given workers.
 func WithWorkers(workers int) Option {
 	return func(opts *mapReduceOptions) {
@@ -217,8 +233,8 @@ func drain(channel <-chan interface{}) {
 	}
 }
 
-func executeMappers(mapper MapFunc, input <-chan interface{}, collector chan<- interface{},
-	done <-chan lang.PlaceholderType, workers int) {
+func executeMappers(ctx context.Context, mapper MapFunc, input <-chan interface{},
+	collector chan<- interface{}, done <-chan lang.PlaceholderType, workers int) {
 	var wg sync.WaitGroup
 	defer func() {
 		wg.Wait()
@@ -226,9 +242,11 @@ func executeMappers(mapper MapFunc, input <-chan interface{}, collector chan<- i
 	}()
 
 	pool := make(chan lang.PlaceholderType, workers)
-	writer := newGuardedWriter(collector, done)
+	writer := newGuardedWriter(ctx, collector, done)
 	for {
 		select {
+		case <-ctx.Done():
+			return
 		case <-done:
 			return
 		case pool <- lang.Placeholder:
@@ -254,6 +272,7 @@ func executeMappers(mapper MapFunc, input <-chan interface{}, collector chan<- i
 
 func newOptions() *mapReduceOptions {
 	return &mapReduceOptions{
+		ctx:     context.Background(),
 		workers: defaultWorkers,
 	}
 }
@@ -268,12 +287,15 @@ func once(fn func(error)) func(error) {
 }
 
 type guardedWriter struct {
+	ctx     context.Context
 	channel chan<- interface{}
 	done    <-chan lang.PlaceholderType
 }
 
-func newGuardedWriter(channel chan<- interface{}, done <-chan lang.PlaceholderType) guardedWriter {
+func newGuardedWriter(ctx context.Context, channel chan<- interface{},
+	done <-chan lang.PlaceholderType) guardedWriter {
 	return guardedWriter{
+		ctx:     ctx,
 		channel: channel,
 		done:    done,
 	}
@@ -281,6 +303,8 @@ func newGuardedWriter(channel chan<- interface{}, done <-chan lang.PlaceholderTy
 
 func (gw guardedWriter) Write(v interface{}) {
 	select {
+	case <-gw.ctx.Done():
+		return
 	case <-gw.done:
 		return
 	default:

core/mr/mapreduce_test.go

@@ -1,6 +1,7 @@
 package mr
 
 import (
+	"context"
 	"errors"
 	"io/ioutil"
 	"log"
@@ -202,6 +203,22 @@ func TestMapReduce(t *testing.T) {
 	}
 }
 
+func TestMapReduceWithReduerWriteMoreThanOnce(t *testing.T) {
+	assert.Panics(t, func() {
+		MapReduce(func(source chan<- interface{}) {
+			for i := 0; i < 10; i++ {
+				source <- i
+			}
+		}, func(item interface{}, writer Writer, cancel func(error)) {
+			writer.Write(item)
+		}, func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
+			drain(pipe)
+			writer.Write("one")
+			writer.Write("two")
+		})
+	})
+}
+
 func TestMapReduceVoid(t *testing.T) {
 	var value uint32
 	tests := []struct {
@@ -394,6 +411,50 @@ func TestMapReduceWithoutReducerWrite(t *testing.T) {
 	assert.Nil(t, res)
 }
 
+func TestMapReduceVoidPanicInReducer(t *testing.T) {
+	const message = "foo"
+	var done syncx.AtomicBool
+	err := MapReduceVoid(func(source chan<- interface{}) {
+		for i := 0; i < defaultWorkers*2; i++ {
+			source <- i
+		}
+		done.Set(true)
+	}, func(item interface{}, writer Writer, cancel func(error)) {
+		i := item.(int)
+		writer.Write(i)
+	}, func(pipe <-chan interface{}, cancel func(error)) {
+		panic(message)
+	}, WithWorkers(1))
+	assert.NotNil(t, err)
+	assert.Equal(t, message, err.Error())
+	assert.True(t, done.True())
+}
+
+func TestMapReduceWithContext(t *testing.T) {
+	var done syncx.AtomicBool
+	var result []int
+	ctx, cancel := context.WithCancel(context.Background())
+	err := MapReduceVoid(func(source chan<- interface{}) {
+		for i := 0; i < defaultWorkers*2; i++ {
+			source <- i
+		}
+		done.Set(true)
+	}, func(item interface{}, writer Writer, c func(error)) {
+		i := item.(int)
+		if i == defaultWorkers/2 {
+			cancel()
+		}
+		writer.Write(i)
+	}, func(pipe <-chan interface{}, cancel func(error)) {
+		for item := range pipe {
+			i := item.(int)
+			result = append(result, i)
+		}
+	}, WithContext(ctx))
+	assert.NotNil(t, err)
+	assert.Equal(t, ErrReduceNoOutput, err)
+}
+
 func BenchmarkMapReduce(b *testing.B) {
 	b.ReportAllocs()
 

core/mr/readme-cn.md

@@ -0,0 +1,89 @@
+# mapreduce
+
+[English](readme.md) | 简体中文
+
+## 为什么需要 MapReduce
+
+在实际的业务场景中我们常常需要从不同的 rpc 服务中获取相应属性来组装成复杂对象。
+
+比如要查询商品详情：
+
+1. 商品服务-查询商品属性
+2. 库存服务-查询库存属性
+3. 价格服务-查询价格属性
+4. 营销服务-查询营销属性
+
+如果是串行调用的话响应时间会随着 rpc 调用次数呈线性增长，所以我们要优化性能一般会将串行改并行。
+
+简单的场景下使用 `WaitGroup` 也能够满足需求，但是如果我们需要对 rpc 调用返回的数据进行校验、数据加工转换、数据汇总呢？继续使用 `WaitGroup` 就有点力不从心了，go 的官方库中并没有这种工具（java 中提供了 CompleteFuture），我们依据 MapReduce 架构思想实现了进程内的数据批处理 MapReduce 并发工具类。
+
+## 设计思路
+
+我们尝试把自己代入到作者的角色梳理一下并发工具可能的业务场景：
+
+1. 查询商品详情：支持并发调用多个服务来组合产品属性，支持调用错误可以立即结束。
+2. 商品详情页自动推荐用户卡券：支持并发校验卡券，校验失败自动剔除，返回全部卡券。
+
+以上实际都是在进行对输入数据进行处理最后输出清洗后的数据，针对数据处理有个非常经典的异步模式：生产者消费者模式。于是我们可以抽象一下数据批处理的生命周期，大致可以分为三个阶段：
+
+<img src="https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/mapreduce-serial-cn.png" width="500">
+
+1. 数据生产 generate
+2. 数据加工 mapper
+3. 数据聚合 reducer
+
+其中数据生产是不可或缺的阶段，数据加工、数据聚合是可选阶段，数据生产与加工支持并发调用，数据聚合基本属于纯内存操作单协程即可。
+
+再来思考一下不同阶段之间数据应该如何流转，既然不同阶段的数据处理都是由不同 goroutine 执行的，那么很自然的可以考虑采用 channel 来实现 goroutine 之间的通信。
+
+<img src="https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/mapreduce-cn.png" width="500">
+
+
+如何实现随时终止流程呢？
+
+`goroutine` 中监听一个全局的结束 `channel` 和调用方提供的 `ctx` 就行。
+
+## 简单示例
+
+并行求平方和（不要嫌弃示例简单，只是模拟并发）
+
+```go
+package main
+
+import (
+    "fmt"
+    "log"
+
+    "github.com/tal-tech/go-zero/core/mr"
+)
+
+func main() {
+    val, err := mr.MapReduce(func(source chan<- interface{}) {
+        // generator
+        for i := 0; i < 10; i++ {
+            source <- i
+        }
+    }, func(item interface{}, writer mr.Writer, cancel func(error)) {
+        // mapper
+        i := item.(int)
+        writer.Write(i * i)
+    }, func(pipe <-chan interface{}, writer mr.Writer, cancel func(error)) {
+        // reducer
+        var sum int
+        for i := range pipe {
+            sum += i.(int)
+        }
+        writer.Write(sum)
+    })
+    if err != nil {
+        log.Fatal(err)
+    }
+    fmt.Println("result:", val)
+}
+```
+
+更多示例：[https://github.com/zeromicro/zero-examples/tree/main/mapreduce](https://github.com/zeromicro/zero-examples/tree/main/mapreduce)
+
+## 欢迎 star！⭐
+
+如果你正在使用或者觉得这个项目对你有帮助，请 **star** 支持，感谢！

core/mr/readme.md

@@ -0,0 +1,90 @@
+<img align="right" width="150px" src="https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/go-zero.png">
+
+# mapreduce
+
+English | [简体中文](readme-cn.md)
+
+## Why MapReduce is needed
+
+In practical business scenarios we often need to get the corresponding properties from different rpc services to assemble complex objects.
+
+For example, to query product details.
+
+1. product service - query product attributes
+2. inventory service - query inventory properties
+3. price service - query price attributes
+4. marketing service - query marketing properties
+
+If it is a serial call, the response time will increase linearly with the number of rpc calls, so we will generally change serial to parallel to optimize response time.
+
+Simple scenarios using `WaitGroup` can also meet the needs, but what if we need to check the data returned by the rpc call, data processing, data aggregation? The official go library does not have such a tool (CompleteFuture is provided in java), so we implemented an in-process data batching MapReduce concurrent tool based on the MapReduce architecture.
+
+## Design ideas
+
+Let's try to put ourselves in the author's shoes and sort out the possible business scenarios for the concurrency tool:
+
+1. querying product details: supporting concurrent calls to multiple services to combine product attributes, and supporting call errors that can be ended immediately.
+2. automatic recommendation of user card coupons on product details page: support concurrently verifying card coupons, automatically rejecting them if they fail, and returning all of them.
+
+The above is actually processing the input data and finally outputting the cleaned data. There is a very classic asynchronous pattern for data processing: the producer-consumer pattern. So we can abstract the life cycle of data batch processing, which can be roughly divided into three phases.
+
+<img src="https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/mapreduce-serial-en.png" width="500">
+
+1. data production generate
+2. data processing mapper
+3. data aggregation reducer
+
+Data producing is an indispensable stage, data processing and data aggregation are optional stages, data producing and processing support concurrent calls, data aggregation is basically a pure memory operation, so a single concurrent process can do it.
+
+Since different stages of data processing are performed by different goroutines, it is natural to consider the use of channel to achieve communication between goroutines.
+
+<img src="https://raw.githubusercontent.com/zeromicro/zero-doc/main/doc/images/mapreduce-en.png" width="500">
+
+How can I terminate the process at any time?
+
+It's simple, just receive from a  channel or the given context in the goroutine.
+
+## A simple example
+
+Calculate the sum of squares, simulating the concurrency.
+
+```go
+package main
+
+import (
+    "fmt"
+    "log"
+
+    "github.com/tal-tech/go-zero/core/mr"
+)
+
+func main() {
+    val, err := mr.MapReduce(func(source chan<- interface{}) {
+        // generator
+        for i := 0; i < 10; i++ {
+            source <- i
+        }
+    }, func(item interface{}, writer mr.Writer, cancel func(error)) {
+        // mapper
+        i := item.(int)
+        writer.Write(i * i)
+    }, func(pipe <-chan interface{}, writer mr.Writer, cancel func(error)) {
+        // reducer
+        var sum int
+        for i := range pipe {
+            sum += i.(int)
+        }
+        writer.Write(sum)
+    })
+    if err != nil {
+        log.Fatal(err)
+    }
+    fmt.Println("result:", val)
+}
+```
+
+More examples: [https://github.com/zeromicro/zero-examples/tree/main/mapreduce](https://github.com/zeromicro/zero-examples/tree/main/mapreduce)
+
+## Give a Star! ⭐
+
+If you like or are using this project to learn or start your solution, please give it a star. Thanks!

core/proc/goroutines+polyfill.go

@@ -1,3 +1,4 @@
+//go:build windows
 // +build windows
 
 package proc

core/proc/goroutines.go

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

core/proc/profile+polyfill.go

@@ -1,3 +1,4 @@
+//go:build windows
 // +build windows
 
 package proc

core/proc/profile.go

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

core/proc/shutdown+polyfill.go

@@ -1,3 +1,4 @@
+//go:build windows
 // +build windows
 
 package proc
@@ -14,5 +15,5 @@ func AddWrapUpListener(fn func()) func() {
 	return fn
 }
 
-func SetTimeoutToForceQuit(duration time.Duration) {
+func SetTimeToForceQuit(duration time.Duration) {
 }

core/proc/shutdown.go

@@ -1,3 +1,4 @@
+//go:build linux || darwin
 // +build linux darwin
 
 package proc
@@ -10,6 +11,7 @@ import (
 	"time"
 
 	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/threading"
 )
 
 const (
@@ -45,10 +47,10 @@ func gracefulStop(signals chan os.Signal) {
 	signal.Stop(signals)
 
 	logx.Info("Got signal SIGTERM, shutting down...")
-	wrapUpListeners.notifyListeners()
+	go wrapUpListeners.notifyListeners()
 
 	time.Sleep(wrapUpTime)
-	shutdownListeners.notifyListeners()
+	go shutdownListeners.notifyListeners()
 
 	time.Sleep(delayTimeBeforeForceQuit - wrapUpTime)
 	logx.Infof("Still alive after %v, going to force kill the process...", delayTimeBeforeForceQuit)
@@ -80,7 +82,9 @@ func (lm *listenerManager) notifyListeners() {
 	lm.lock.Lock()
 	defer lm.lock.Unlock()
 
+	group := threading.NewRoutineGroup()
 	for _, listener := range lm.listeners {
-		listener()
+		group.RunSafe(listener)
 	}
+	group.Wait()
 }

core/proc/signals+polyfill.go

@@ -0,0 +1,10 @@
+//go:build windows
+// +build windows
+
+package proc
+
+import "context"
+
+func Done() <-chan struct{} {
+	return context.Background().Done()
+}

core/proc/signals.go

@@ -1,3 +1,4 @@
+//go:build linux || darwin
 // +build linux darwin
 
 package proc
@@ -12,6 +13,8 @@ import (
 
 const timeFormat = "0102150405"
 
+var done = make(chan struct{})
+
 func init() {
 	go func() {
 		var profiler Stopper
@@ -33,6 +36,13 @@ func init() {
 					profiler = nil
 				}
 			case syscall.SIGTERM:
+				select {
+				case <-done:
+					// already closed
+				default:
+					close(done)
+				}
+
 				gracefulStop(signals)
 			default:
 				logx.Error("Got unregistered signal:", v)
@@ -40,3 +50,8 @@ func init() {
 		}
 	}()
 }
+
+// Done returns the channel that notifies the process quitting.
+func Done() <-chan struct{} {
+	return done
+}

core/proc/signals_test.go

@@ -0,0 +1,16 @@
+package proc
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDone(t *testing.T) {
+	select {
+	case <-Done():
+		assert.Fail(t, "should run")
+	default:
+	}
+	assert.NotNil(t, Done())
+}

core/prometheus/agent.go

@@ -7,18 +7,28 @@ import (
 
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/syncx"
 	"github.com/tal-tech/go-zero/core/threading"
 )
 
-var once sync.Once
+var (
+	once    sync.Once
+	enabled syncx.AtomicBool
+)
+
+// Enabled returns if prometheus is enabled.
+func Enabled() bool {
+	return enabled.True()
+}
 
 // StartAgent starts a prometheus agent.
 func StartAgent(c Config) {
-	once.Do(func() {
-		if len(c.Host) == 0 {
-			return
-		}
+	if len(c.Host) == 0 {
+		return
+	}
 
+	once.Do(func() {
+		enabled.Set(true)
 		threading.GoSafe(func() {
 			http.Handle(c.Path, promhttp.Handler())
 			addr := fmt.Sprintf("%s:%d", c.Host, c.Port)