make sure offset less than size even it's checked inside (#354)

* fix format

* refactor

* refactor

* optimized

* refactor

* refactor

* refactor

* add js path prefix

.codecov.yml

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

.github/workflows/codeql-analysis.yml

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

.github/workflows/go.yml

@@ -7,7 +7,6 @@ on:
     branches: [ master ]
 
 jobs:
-
   build:
     name: Build
     runs-on: ubuntu-latest
@@ -27,4 +26,9 @@ jobs:
         go get -v -t -d ./...
 
     - name: Test
-      run: go test -v -race ./...
+      run: go test -race -coverprofile=coverage.txt -covermode=atomic ./...
+
+    - name: Codecov
+      uses: codecov/codecov-action@v1.0.6
+      with:
+        token: ${{secrets.CODECOV_TOKEN}}

.gitignore

@@ -4,6 +4,7 @@
 # Unignore all with extensions
 !*.*
 !**/Dockerfile
+!**/Makefile
 
 # Unignore all dirs
 !*/
@@ -12,7 +13,6 @@
 .idea
 **/.DS_Store
 **/logs
-!Makefile
 
 # gitlab ci
 .cache

.gitlab-ci.yml

@@ -1,18 +0,0 @@
-stages:
-- analysis
-
-variables:
-  GOPATH: '/runner-cache/zero'
-  GOCACHE: '/runner-cache/zero'
-  GOPROXY: 'https://goproxy.cn,direct'
-
-analysis:
-  stage: analysis
-  image: golang
-  script:
-    - go version && go env
-    - go test -short $(go list ./...) | grep -v "no test"
-  only:
-   - merge_requests
-  tags:
-    - common

.golangci.yml

@@ -1,36 +0,0 @@
-run:
-  #  concurrency: 6
-  timeout: 5m
-  skip-dirs:
-    - core
-    - doc
-    - example
-    - rest
-    - rpcx
-    - tools
-
-
-linters:
-  disable-all: true
-  enable:
-    - bodyclose
-    - deadcode
-    - errcheck
-    - gosimple
-    - govet
-    - ineffassign
-    - staticcheck
-    - structcheck
-    - typecheck
-    - unused
-    - varcheck
-#    - dupl
-
-
-linters-settings:
-
-issues:
-  exclude-rules:
-    - linters:
-        - staticcheck
-      text: 'SA1019: (baseresponse.BoolResponse|oldresponse.FormatBadRequestResponse|oldresponse.FormatResponse)|SA5008: unknown JSON option ("optional"|"default=|"range=|"options=)'

core/bloom/bloom.go

@@ -13,15 +13,13 @@ const (
 	// maps as k in the error rate table
 	maps      = 14
 	setScript = `
-local key = KEYS[1]
 for _, offset in ipairs(ARGV) do
-	redis.call("setbit", key, offset, 1)
+	redis.call("setbit", KEYS[1], offset, 1)
 end
 `
 	testScript = `
-local key = KEYS[1]
 for _, offset in ipairs(ARGV) do
-	if tonumber(redis.call("getbit", key, offset)) == 0 then
+	if tonumber(redis.call("getbit", KEYS[1], offset)) == 0 then
 		return false
 	end
 end
@@ -39,7 +37,6 @@ type (
 
 	BloomFilter struct {
 		bits   uint
-		maps   uint
 		bitSet BitSetProvider
 	}
 )

core/bloom/bloom_test.go

@@ -3,19 +3,15 @@ package bloom
 import (
 	"testing"
 
-	"github.com/alicebob/miniredis"
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/stores/redis"
+	"github.com/tal-tech/go-zero/core/stores/redis/redistest"
 )
 
 func TestRedisBitSet_New_Set_Test(t *testing.T) {
-	s, err := miniredis.Run()
+	store, clean, err := redistest.CreateRedis()
-	if err != nil {
+	assert.Nil(t, err)
-		t.Error("Miniredis could not start")
+	defer clean()
-	}
-	defer s.Close()
 
-	store := redis.NewRedis(s.Addr(), redis.NodeType)
 	bitSet := newRedisBitSet(store, "test_key", 1024)
 	isSetBefore, err := bitSet.check([]uint{0})
 	if err != nil {
@@ -46,13 +42,10 @@ func TestRedisBitSet_New_Set_Test(t *testing.T) {
 }
 
 func TestRedisBitSet_Add(t *testing.T) {
-	s, err := miniredis.Run()
+	store, clean, err := redistest.CreateRedis()
-	if err != nil {
+	assert.Nil(t, err)
-		t.Error("Miniredis could not start")
+	defer clean()
-	}
-	defer s.Close()
 
-	store := redis.NewRedis(s.Addr(), redis.NodeType)
 	filter := New(store, "test_key", 64)
 	assert.Nil(t, filter.Add([]byte("hello")))
 	assert.Nil(t, filter.Add([]byte("world")))

core/breaker/breaker.go

@@ -5,17 +5,12 @@ import (
 	"fmt"
 	"strings"
 	"sync"
-	"time"
 
 	"github.com/tal-tech/go-zero/core/mathx"
 	"github.com/tal-tech/go-zero/core/proc"
 	"github.com/tal-tech/go-zero/core/stat"
 	"github.com/tal-tech/go-zero/core/stringx"
-)
+	"github.com/tal-tech/go-zero/core/timex"
-
-const (
-	StateClosed State = iota
-	StateOpen
 )
 
 const (
@@ -27,11 +22,10 @@ const (
 var ErrServiceUnavailable = errors.New("circuit breaker is open")
 
 type (
-	State      = int32
 	Acceptable func(err error) bool
 
 	Breaker interface {
-		// Name returns the name of the netflixBreaker.
+		// Name returns the name of the Breaker.
 		Name() string
 
 		// Allow checks if the request is allowed.
@@ -40,34 +34,34 @@ type (
 		// If not allow, ErrServiceUnavailable will be returned.
 		Allow() (Promise, error)
 
-		// Do runs the given request if the netflixBreaker accepts it.
+		// Do runs the given request if the Breaker accepts it.
-		// Do returns an error instantly if the netflixBreaker rejects the request.
+		// Do returns an error instantly if the Breaker rejects the request.
-		// If a panic occurs in the request, the netflixBreaker handles it as an error
+		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
 		Do(req func() error) error
 
-		// DoWithAcceptable runs the given request if the netflixBreaker accepts it.
+		// DoWithAcceptable runs the given request if the Breaker accepts it.
-		// Do returns an error instantly if the netflixBreaker rejects the request.
+		// DoWithAcceptable returns an error instantly if the Breaker rejects the request.
-		// If a panic occurs in the request, the netflixBreaker handles it as an error
+		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
 		// acceptable checks if it's a successful call, even if the err is not nil.
 		DoWithAcceptable(req func() error, acceptable Acceptable) error
 
-		// DoWithFallback runs the given request if the netflixBreaker accepts it.
+		// DoWithFallback runs the given request if the Breaker accepts it.
-		// DoWithFallback runs the fallback if the netflixBreaker rejects the request.
+		// DoWithFallback runs the fallback if the Breaker rejects the request.
-		// If a panic occurs in the request, the netflixBreaker handles it as an error
+		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
 		DoWithFallback(req func() error, fallback func(err error) error) error
 
-		// DoWithFallbackAcceptable runs the given request if the netflixBreaker accepts it.
+		// DoWithFallbackAcceptable runs the given request if the Breaker accepts it.
-		// DoWithFallback runs the fallback if the netflixBreaker rejects the request.
+		// DoWithFallbackAcceptable runs the fallback if the Breaker rejects the request.
-		// If a panic occurs in the request, the netflixBreaker handles it as an error
+		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
 		// acceptable checks if it's a successful call, even if the err is not nil.
 		DoWithFallbackAcceptable(req func() error, fallback func(err error) error, acceptable Acceptable) error
 	}
 
-	BreakerOption func(breaker *circuitBreaker)
+	Option func(breaker *circuitBreaker)
 
 	Promise interface {
 		Accept()
@@ -95,7 +89,7 @@ type (
 	}
 )
 
-func NewBreaker(opts ...BreakerOption) Breaker {
+func NewBreaker(opts ...Option) Breaker {
 	var b circuitBreaker
 	for _, opt := range opts {
 		opt(&b)
@@ -133,7 +127,7 @@ func (cb *circuitBreaker) Name() string {
 	return cb.name
 }
 
-func WithName(name string) BreakerOption {
+func WithName(name string) Option {
 	return func(b *circuitBreaker) {
 		b.name = name
 	}
@@ -195,23 +189,23 @@ type errorWindow struct {
 
 func (ew *errorWindow) add(reason string) {
 	ew.lock.Lock()
-	ew.reasons[ew.index] = fmt.Sprintf("%s %s", time.Now().Format(timeFormat), reason)
+	ew.reasons[ew.index] = fmt.Sprintf("%s %s", timex.Time().Format(timeFormat), reason)
 	ew.index = (ew.index + 1) % numHistoryReasons
 	ew.count = mathx.MinInt(ew.count+1, numHistoryReasons)
 	ew.lock.Unlock()
 }
 
 func (ew *errorWindow) String() string {
-	var builder strings.Builder
+	var reasons []string
 
 	ew.lock.Lock()
-	for i := ew.index + ew.count - 1; i >= ew.index; i-- {
+	// reverse order
-		builder.WriteString(ew.reasons[i%numHistoryReasons])
+	for i := ew.index - 1; i >= ew.index-ew.count; i-- {
-		builder.WriteByte('\n')
+		reasons = append(reasons, ew.reasons[(i+numHistoryReasons)%numHistoryReasons])
 	}
 	ew.lock.Unlock()
 
-	return builder.String()
+	return strings.Join(reasons, "\n")
 }
 
 type promiseWithReason struct {

core/breaker/breaker_test.go

@@ -2,7 +2,9 @@ package breaker
 
 import (
 	"errors"
+	"fmt"
 	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -33,6 +35,84 @@ func TestLogReason(t *testing.T) {
 	assert.Equal(t, numHistoryReasons, errs.count)
 }
 
+func TestErrorWindow(t *testing.T) {
+	tests := []struct {
+		name    string
+		reasons []string
+	}{
+		{
+			name: "no error",
+		},
+		{
+			name:    "one error",
+			reasons: []string{"foo"},
+		},
+		{
+			name:    "two errors",
+			reasons: []string{"foo", "bar"},
+		},
+		{
+			name:    "five errors",
+			reasons: []string{"first", "second", "third", "fourth", "fifth"},
+		},
+		{
+			name:    "six errors",
+			reasons: []string{"first", "second", "third", "fourth", "fifth", "sixth"},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var ew errorWindow
+			for _, reason := range test.reasons {
+				ew.add(reason)
+			}
+			var reasons []string
+			if len(test.reasons) > numHistoryReasons {
+				reasons = test.reasons[len(test.reasons)-numHistoryReasons:]
+			} else {
+				reasons = test.reasons
+			}
+			for _, reason := range reasons {
+				assert.True(t, strings.Contains(ew.String(), reason), fmt.Sprintf("actual: %s", ew.String()))
+			}
+		})
+	}
+}
+
+func TestPromiseWithReason(t *testing.T) {
+	tests := []struct {
+		name   string
+		reason string
+		expect string
+	}{
+		{
+			name: "success",
+		},
+		{
+			name:   "success",
+			reason: "fail",
+			expect: "fail",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			promise := promiseWithReason{
+				promise: new(mockedPromise),
+				errWin:  new(errorWindow),
+			}
+			if len(test.reason) == 0 {
+				promise.Accept()
+			} else {
+				promise.Reject(test.reason)
+			}
+
+			assert.True(t, strings.Contains(promise.errWin.String(), test.expect))
+		})
+	}
+}
+
 func BenchmarkGoogleBreaker(b *testing.B) {
 	br := NewBreaker()
 	for i := 0; i < b.N; i++ {
@@ -41,3 +121,12 @@ func BenchmarkGoogleBreaker(b *testing.B) {
 		})
 	}
 }
+
+type mockedPromise struct {
+}
+
+func (m *mockedPromise) Accept() {
+}
+
+func (m *mockedPromise) Reject() {
+}

core/breaker/breakers.go

@@ -41,10 +41,13 @@ func GetBreaker(name string) Breaker {
 	}
 
 	lock.Lock()
-	defer lock.Unlock()
+	b, ok = breakers[name]
+	if !ok {
+		b = NewBreaker(WithName(name))
+		breakers[name] = b
+	}
+	lock.Unlock()
 
-	b = NewBreaker()
-	breakers[name] = b
 	return b
 }
 
@@ -55,22 +58,5 @@ func NoBreakFor(name string) {
 }
 
 func do(name string, execute func(b Breaker) error) error {
-	lock.RLock()
+	return execute(GetBreaker(name))
-	b, ok := breakers[name]
-	lock.RUnlock()
-	if ok {
-		return execute(b)
-	} else {
-		lock.Lock()
-		b, ok = breakers[name]
-		if ok {
-			lock.Unlock()
-			return execute(b)
-		} else {
-			b = NewBreaker(WithName(name))
-			breakers[name] = b
-			lock.Unlock()
-			return execute(b)
-		}
-	}
 }

core/breaker/googlebreaker.go

@@ -2,7 +2,6 @@ package breaker
 
 import (
 	"math"
-	"sync/atomic"
 	"time"
 
 	"github.com/tal-tech/go-zero/core/collection"
@@ -21,7 +20,6 @@ const (
 // see Client-Side Throttling section in https://landing.google.com/sre/sre-book/chapters/handling-overload/
 type googleBreaker struct {
 	k     float64
-	state int32
 	stat  *collection.RollingWindow
 	proba *mathx.Proba
 }
@@ -32,7 +30,6 @@ func newGoogleBreaker() *googleBreaker {
 	return &googleBreaker{
 		stat:  st,
 		k:     k,
-		state: StateClosed,
 		proba: mathx.NewProba(),
 	}
 }
@@ -43,15 +40,9 @@ func (b *googleBreaker) accept() error {
 	// https://landing.google.com/sre/sre-book/chapters/handling-overload/#eq2101
 	dropRatio := math.Max(0, (float64(total-protection)-weightedAccepts)/float64(total+1))
 	if dropRatio <= 0 {
-		if atomic.LoadInt32(&b.state) == StateOpen {
-			atomic.CompareAndSwapInt32(&b.state, StateOpen, StateClosed)
-		}
 		return nil
 	}
 
-	if atomic.LoadInt32(&b.state) == StateClosed {
-		atomic.CompareAndSwapInt32(&b.state, StateClosed, StateOpen)
-	}
 	if b.proba.TrueOnProba(dropRatio) {
 		return ErrServiceUnavailable
 	}

core/breaker/googlebreaker_test.go

@@ -2,7 +2,6 @@ package breaker
 
 import (
 	"errors"
-	"math"
 	"math/rand"
 	"testing"
 	"time"
@@ -27,7 +26,6 @@ func getGoogleBreaker() *googleBreaker {
 	return &googleBreaker{
 		stat:  st,
 		k:     5,
-		state: StateClosed,
 		proba: mathx.NewProba(),
 	}
 }
@@ -158,7 +156,7 @@ func TestGoogleBreakerSelfProtection(t *testing.T) {
 	t.Run("total request > 100, total < 2 * success", func(t *testing.T) {
 		b := getGoogleBreaker()
 		size := rand.Intn(10000)
-		accepts := int(math.Ceil(float64(size))) + 1
+		accepts := size + 1
 		markSuccess(b, accepts)
 		markFailed(b, size-accepts)
 		assert.Nil(t, b.accept())

core/cmdline/input_test.go

@@ -0,0 +1,80 @@
+package cmdline
+
+import (
+	"fmt"
+	"os"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/iox"
+	"github.com/tal-tech/go-zero/core/lang"
+)
+
+func TestEnterToContinue(t *testing.T) {
+	restore, err := iox.RedirectInOut()
+	assert.Nil(t, err)
+	defer restore()
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		fmt.Println()
+	}()
+	go func() {
+		defer wg.Done()
+		EnterToContinue()
+	}()
+
+	wait := make(chan lang.PlaceholderType)
+	go func() {
+		wg.Wait()
+		close(wait)
+	}()
+
+	select {
+	case <-time.After(time.Second):
+		t.Error("timeout")
+	case <-wait:
+	}
+}
+
+func TestReadLine(t *testing.T) {
+	r, w, err := os.Pipe()
+	assert.Nil(t, err)
+	ow := os.Stdout
+	os.Stdout = w
+	or := os.Stdin
+	os.Stdin = r
+	defer func() {
+		os.Stdin = or
+		os.Stdout = ow
+	}()
+
+	const message = "hello"
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		fmt.Println(message)
+	}()
+	go func() {
+		defer wg.Done()
+		input := ReadLine("")
+		assert.Equal(t, message, input)
+	}()
+
+	wait := make(chan lang.PlaceholderType)
+	go func() {
+		wg.Wait()
+		close(wait)
+	}()
+
+	select {
+	case <-time.After(time.Second):
+		t.Error("timeout")
+	case <-wait:
+	}
+}

core/codec/dh_test.go

@@ -71,3 +71,12 @@ func TestDiffieHellmanMiddleManAttack(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, string(src), string(decryptedSrc))
 }
+
+func TestKeyBytes(t *testing.T) {
+	var empty DhKey
+	assert.Equal(t, 0, len(empty.Bytes()))
+
+	key, err := GenerateKey()
+	assert.Nil(t, err)
+	assert.True(t, len(key.Bytes()) > 0)
+}

core/codec/gzip.go

@@ -6,6 +6,8 @@ import (
 	"io"
 )
 
+const unzipLimit = 100 * 1024 * 1024 // 100MB
+
 func Gzip(bs []byte) []byte {
 	var b bytes.Buffer
 
@@ -24,8 +26,7 @@ func Gunzip(bs []byte) ([]byte, error) {
 	defer r.Close()
 
 	var c bytes.Buffer
-	_, err = io.Copy(&c, r)
+	if _, err = io.Copy(&c, io.LimitReader(r, unzipLimit)); err != nil {
-	if err != nil {
 		return nil, err
 	}
 

core/codec/hmac_test.go

@@ -0,0 +1,19 @@
+package codec
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHmac(t *testing.T) {
+	ret := Hmac([]byte("foo"), "bar")
+	assert.Equal(t, "f9320baf0249169e73850cd6156ded0106e2bb6ad8cab01b7bbbebe6d1065317",
+		fmt.Sprintf("%x", ret))
+}
+
+func TestHmacBase64(t *testing.T) {
+	ret := HmacBase64([]byte("foo"), "bar")
+	assert.Equal(t, "+TILrwJJFp5zhQzWFW3tAQbiu2rYyrAbe7vr5tEGUxc=", ret)
+}

core/codec/rsa_test.go

@@ -0,0 +1,58 @@
+package codec
+
+import (
+	"encoding/base64"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/fs"
+)
+
+const (
+	priKey = `-----BEGIN RSA PRIVATE KEY-----
+MIICXQIBAAKBgQC4TJk3onpqb2RYE3wwt23J9SHLFstHGSkUYFLe+nl1dEKHbD+/
+Zt95L757J3xGTrwoTc7KCTxbrgn+stn0w52BNjj/kIE2ko4lbh/v8Fl14AyVR9ms
+fKtKOnhe5FCT72mdtApr+qvzcC3q9hfXwkyQU32pv7q5UimZ205iKSBmgQIDAQAB
+AoGAM5mWqGIAXj5z3MkP01/4CDxuyrrGDVD5FHBno3CDgyQa4Gmpa4B0/ywj671B
+aTnwKmSmiiCN2qleuQYASixes2zY5fgTzt+7KNkl9JHsy7i606eH2eCKzsUa/s6u
+WD8V3w/hGCQ9zYI18ihwyXlGHIgcRz/eeRh+nWcWVJzGOPUCQQD5nr6It/1yHb1p
+C6l4fC4xXF19l4KxJjGu1xv/sOpSx0pOqBDEX3Mh//FU954392rUWDXV1/I65BPt
+TLphdsu3AkEAvQJ2Qay/lffFj9FaUrvXuftJZ/Ypn0FpaSiUh3Ak3obBT6UvSZS0
+bcYdCJCNHDtBOsWHnIN1x+BcWAPrdU7PhwJBAIQ0dUlH2S3VXnoCOTGc44I1Hzbj
+Rc65IdsuBqA3fQN2lX5vOOIog3vgaFrOArg1jBkG1wx5IMvb/EnUN2pjVqUCQCza
+KLXtCInOAlPemlCHwumfeAvznmzsWNdbieOZ+SXVVIpR6KbNYwOpv7oIk3Pfm9sW
+hNffWlPUKhW42Gc+DIECQQDmk20YgBXwXWRM5DRPbhisIV088N5Z58K9DtFWkZsd
+OBDT3dFcgZONtlmR1MqZO0pTh30lA4qovYj3Bx7A8i36
+-----END RSA PRIVATE KEY-----`
+	pubKey = `-----BEGIN PUBLIC KEY-----
+MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQC4TJk3onpqb2RYE3wwt23J9SHL
+FstHGSkUYFLe+nl1dEKHbD+/Zt95L757J3xGTrwoTc7KCTxbrgn+stn0w52BNjj/
+kIE2ko4lbh/v8Fl14AyVR9msfKtKOnhe5FCT72mdtApr+qvzcC3q9hfXwkyQU32p
+v7q5UimZ205iKSBmgQIDAQAB
+-----END PUBLIC KEY-----`
+	testBody = `this is the content`
+)
+
+func TestCryption(t *testing.T) {
+	enc, err := NewRsaEncrypter([]byte(pubKey))
+	assert.Nil(t, err)
+	ret, err := enc.Encrypt([]byte(testBody))
+	assert.Nil(t, err)
+
+	file, err := fs.TempFilenameWithText(priKey)
+	assert.Nil(t, err)
+	dec, err := NewRsaDecrypter(file)
+	assert.Nil(t, err)
+	actual, err := dec.Decrypt(ret)
+	assert.Nil(t, err)
+	assert.Equal(t, testBody, string(actual))
+
+	actual, err = dec.DecryptBase64(base64.StdEncoding.EncodeToString(ret))
+	assert.Nil(t, err)
+	assert.Equal(t, testBody, string(actual))
+}
+
+func TestBadPubKey(t *testing.T) {
+	_, err := NewRsaEncrypter([]byte("foo"))
+	assert.Equal(t, ErrPublicKey, err)
+}

core/collection/cache.go

@@ -29,7 +29,6 @@ type (
 		name           string
 		lock           sync.Mutex
 		data           map[string]interface{}
-		evicts         *list.List
 		expire         time.Duration
 		timingWheel    *TimingWheel
 		lruCache       lru
@@ -82,12 +81,7 @@ func (c *Cache) Del(key string) {
 }
 
 func (c *Cache) Get(key string) (interface{}, bool) {
-	c.lock.Lock()
+	value, ok := c.doGet(key)
-	value, ok := c.data[key]
-	if ok {
-		c.lruCache.add(key)
-	}
-	c.lock.Unlock()
 	if ok {
 		c.stats.IncrementHit()
 	} else {
@@ -113,12 +107,25 @@ func (c *Cache) Set(key string, value interface{}) {
 }
 
 func (c *Cache) Take(key string, fetch func() (interface{}, error)) (interface{}, error) {
-	val, fresh, err := c.barrier.DoEx(key, func() (interface{}, error) {
+	if val, ok := c.doGet(key); ok {
+		c.stats.IncrementHit()
+		return val, nil
+	}
+
+	var fresh bool
+	val, err := c.barrier.Do(key, func() (interface{}, error) {
+		// because O(1) on map search in memory, and fetch is an IO query
+		// so we do double check, cache might be taken by another call
+		if val, ok := c.doGet(key); ok {
+			return val, nil
+		}
+
 		v, e := fetch()
 		if e != nil {
 			return nil, e
 		}
 
+		fresh = true
 		c.Set(key, v)
 		return v, nil
 	})
@@ -137,6 +144,18 @@ func (c *Cache) Take(key string, fetch func() (interface{}, error)) (interface{}
 	return val, nil
 }
 
+func (c *Cache) doGet(key string) (interface{}, bool) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	value, ok := c.data[key]
+	if ok {
+		c.lruCache.add(key)
+	}
+
+	return value, ok
+}
+
 func (c *Cache) onEvict(key string) {
 	// already locked
 	delete(c.data, key)
@@ -258,18 +277,15 @@ func (cs *cacheStat) statLoop() {
 	ticker := time.NewTicker(statInterval)
 	defer ticker.Stop()
 
-	for {
+	for range ticker.C {
-		select {
+		hit := atomic.SwapUint64(&cs.hit, 0)
-		case <-ticker.C:
+		miss := atomic.SwapUint64(&cs.miss, 0)
-			hit := atomic.SwapUint64(&cs.hit, 0)
+		total := hit + miss
-			miss := atomic.SwapUint64(&cs.miss, 0)
+		if total == 0 {
-			total := hit + miss
+			continue
-			if total == 0 {
-				continue
-			}
-			percent := 100 * float32(hit) / float32(total)
-			logx.Statf("cache(%s) - qpm: %d, hit_ratio: %.1f%%, elements: %d, hit: %d, miss: %d",
-				cs.name, total, percent, cs.sizeCallback(), hit, miss)
 		}
+		percent := 100 * float32(hit) / float32(total)
+		logx.Statf("cache(%s) - qpm: %d, hit_ratio: %.1f%%, elements: %d, hit: %d, miss: %d",
+			cs.name, total, percent, cs.sizeCallback(), hit, miss)
 	}
 }

core/collection/cache_test.go

@@ -1,6 +1,7 @@
 package collection
 
 import (
+	"errors"
 	"strconv"
 	"sync"
 	"sync/atomic"
@@ -10,6 +11,8 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+var errDummy = errors.New("dummy")
+
 func TestCacheSet(t *testing.T) {
 	cache, err := NewCache(time.Second*2, WithName("any"))
 	assert.Nil(t, err)
@@ -63,6 +66,54 @@ func TestCacheTake(t *testing.T) {
 	assert.Equal(t, int32(1), atomic.LoadInt32(&count))
 }
 
+func TestCacheTakeExists(t *testing.T) {
+	cache, err := NewCache(time.Second * 2)
+	assert.Nil(t, err)
+
+	var count int32
+	var wg sync.WaitGroup
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			cache.Set("first", "first element")
+			cache.Take("first", func() (interface{}, error) {
+				atomic.AddInt32(&count, 1)
+				time.Sleep(time.Millisecond * 100)
+				return "first element", nil
+			})
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+
+	assert.Equal(t, 1, cache.size())
+	assert.Equal(t, int32(0), atomic.LoadInt32(&count))
+}
+
+func TestCacheTakeError(t *testing.T) {
+	cache, err := NewCache(time.Second * 2)
+	assert.Nil(t, err)
+
+	var count int32
+	var wg sync.WaitGroup
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			_, err := cache.Take("first", func() (interface{}, error) {
+				atomic.AddInt32(&count, 1)
+				time.Sleep(time.Millisecond * 100)
+				return "", errDummy
+			})
+			assert.Equal(t, errDummy, err)
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+
+	assert.Equal(t, 0, cache.size())
+	assert.Equal(t, int32(1), atomic.LoadInt32(&count))
+}
+
 func TestCacheWithLruEvicts(t *testing.T) {
 	cache, err := NewCache(time.Minute, WithLimit(3))
 	assert.Nil(t, err)
@@ -72,9 +123,9 @@ func TestCacheWithLruEvicts(t *testing.T) {
 	cache.Set("third", "third element")
 	cache.Set("fourth", "fourth element")
 
-	value, ok := cache.Get("first")
+	_, ok := cache.Get("first")
 	assert.False(t, ok)
-	value, ok = cache.Get("second")
+	value, ok := cache.Get("second")
 	assert.True(t, ok)
 	assert.Equal(t, "second element", value)
 	value, ok = cache.Get("third")
@@ -94,9 +145,9 @@ func TestCacheWithLruEvicted(t *testing.T) {
 	cache.Set("third", "third element")
 	cache.Set("fourth", "fourth element")
 
-	value, ok := cache.Get("first")
+	_, ok := cache.Get("first")
 	assert.False(t, ok)
-	value, ok = cache.Get("second")
+	value, ok := cache.Get("second")
 	assert.True(t, ok)
 	assert.Equal(t, "second element", value)
 	cache.Set("fifth", "fifth element")

core/collection/ring.go

@@ -6,6 +6,10 @@ type Ring struct {
 }
 
 func NewRing(n int) *Ring {
+	if n < 1 {
+		panic("n should be greater than 0")
+	}
+
 	return &Ring{
 		elements: make([]interface{}, n),
 	}

core/collection/ring_test.go

@@ -6,6 +6,12 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func TestNewRing(t *testing.T) {
+	assert.Panics(t, func() {
+		NewRing(0)
+	})
+}
+
 func TestRingLess(t *testing.T) {
 	ring := NewRing(5)
 	for i := 0; i < 3; i++ {

core/collection/rollingwindow.go

@@ -8,8 +8,10 @@ import (
 )
 
 type (
+	// RollingWindowOption let callers customize the RollingWindow.
 	RollingWindowOption func(rollingWindow *RollingWindow)
 
+	// RollingWindow defines a rolling window to calculate the events in buckets with time interval.
 	RollingWindow struct {
 		lock          sync.RWMutex
 		size          int
@@ -17,11 +19,17 @@ type (
 		interval      time.Duration
 		offset        int
 		ignoreCurrent bool
-		lastTime      time.Duration
+		lastTime      time.Duration // start time of the last bucket
 	}
 )
 
+// NewRollingWindow returns a RollingWindow that with size buckets and time interval,
+// use opts to customize the RollingWindow.
 func NewRollingWindow(size int, interval time.Duration, opts ...RollingWindowOption) *RollingWindow {
+	if size < 1 {
+		panic("size must be greater than 0")
+	}
+
 	w := &RollingWindow{
 		size:     size,
 		win:      newWindow(size),
@@ -34,6 +42,7 @@ func NewRollingWindow(size int, interval time.Duration, opts ...RollingWindowOpt
 	return w
 }
 
+// Add adds value to current bucket.
 func (rw *RollingWindow) Add(v float64) {
 	rw.lock.Lock()
 	defer rw.lock.Unlock()
@@ -41,6 +50,7 @@ func (rw *RollingWindow) Add(v float64) {
 	rw.win.add(rw.offset, v)
 }
 
+// Reduce runs fn on all buckets, ignore current bucket if ignoreCurrent was set.
 func (rw *RollingWindow) Reduce(fn func(b *Bucket)) {
 	rw.lock.RLock()
 	defer rw.lock.RUnlock()
@@ -70,29 +80,23 @@ func (rw *RollingWindow) span() int {
 
 func (rw *RollingWindow) updateOffset() {
 	span := rw.span()
-	if span > 0 {
+	if span <= 0 {
-		offset := rw.offset
+		return
-		// reset expired buckets
-		start := offset + 1
-		steps := start + span
-		var remainder int
-		if steps > rw.size {
-			remainder = steps - rw.size
-			steps = rw.size
-		}
-		for i := start; i < steps; i++ {
-			rw.win.resetBucket(i)
-			offset = i
-		}
-		for i := 0; i < remainder; i++ {
-			rw.win.resetBucket(i)
-			offset = i
-		}
-		rw.offset = offset
-		rw.lastTime = timex.Now()
 	}
+
+	offset := rw.offset
+	// reset expired buckets
+	for i := 0; i < span; i++ {
+		rw.win.resetBucket((offset + i + 1) % rw.size)
+	}
+
+	rw.offset = (offset + span) % rw.size
+	now := timex.Now()
+	// align to interval time boundary
+	rw.lastTime = now - (now-rw.lastTime)%rw.interval
 }
 
+// Bucket defines the bucket that holds sum and num of additions.
 type Bucket struct {
 	Sum   float64
 	Count int64
@@ -114,9 +118,9 @@ type window struct {
 }
 
 func newWindow(size int) *window {
-	var buckets []*Bucket
+	buckets := make([]*Bucket, size)
 	for i := 0; i < size; i++ {
-		buckets = append(buckets, new(Bucket))
+		buckets[i] = new(Bucket)
 	}
 	return &window{
 		buckets: buckets,
@@ -130,14 +134,15 @@ func (w *window) add(offset int, v float64) {
 
 func (w *window) reduce(start, count int, fn func(b *Bucket)) {
 	for i := 0; i < count; i++ {
-		fn(w.buckets[(start+i)%len(w.buckets)])
+		fn(w.buckets[(start+i)%w.size])
 	}
 }
 
 func (w *window) resetBucket(offset int) {
-	w.buckets[offset].reset()
+	w.buckets[offset%w.size].reset()
 }
 
+// IgnoreCurrentBucket lets the Reduce call ignore current bucket.
 func IgnoreCurrentBucket() RollingWindowOption {
 	return func(w *RollingWindow) {
 		w.ignoreCurrent = true

core/collection/rollingwindow_test.go

@@ -11,6 +11,13 @@ import (
 
 const duration = time.Millisecond * 50
 
+func TestNewRollingWindow(t *testing.T) {
+	assert.NotNil(t, NewRollingWindow(10, time.Second))
+	assert.Panics(t, func() {
+		NewRollingWindow(0, time.Second)
+	})
+}
+
 func TestRollingWindowAdd(t *testing.T) {
 	const size = 3
 	r := NewRollingWindow(size, duration)
@@ -81,7 +88,7 @@ func TestRollingWindowReduce(t *testing.T) {
 	for _, test := range tests {
 		t.Run(stringx.Rand(), func(t *testing.T) {
 			r := test.win
-			for x := 0; x < size; x = x + 1 {
+			for x := 0; x < size; x++ {
 				for i := 0; i <= x; i++ {
 					r.Add(float64(i))
 				}
@@ -98,6 +105,37 @@ func TestRollingWindowReduce(t *testing.T) {
 	}
 }
 
+func TestRollingWindowBucketTimeBoundary(t *testing.T) {
+	const size = 3
+	interval := time.Millisecond * 30
+	r := NewRollingWindow(size, interval)
+	listBuckets := func() []float64 {
+		var buckets []float64
+		r.Reduce(func(b *Bucket) {
+			buckets = append(buckets, b.Sum)
+		})
+		return buckets
+	}
+	assert.Equal(t, []float64{0, 0, 0}, listBuckets())
+	r.Add(1)
+	assert.Equal(t, []float64{0, 0, 1}, listBuckets())
+	time.Sleep(time.Millisecond * 45)
+	r.Add(2)
+	r.Add(3)
+	assert.Equal(t, []float64{0, 1, 5}, listBuckets())
+	// sleep time should be less than interval, and make the bucket change happen
+	time.Sleep(time.Millisecond * 20)
+	r.Add(4)
+	r.Add(5)
+	r.Add(6)
+	assert.Equal(t, []float64{1, 5, 15}, listBuckets())
+	time.Sleep(time.Millisecond * 100)
+	r.Add(7)
+	r.Add(8)
+	r.Add(9)
+	assert.Equal(t, []float64{0, 0, 24}, listBuckets())
+}
+
 func TestRollingWindowDataRace(t *testing.T) {
 	const size = 3
 	r := NewRollingWindow(size, duration)

core/collection/set.go

@@ -15,6 +15,7 @@ const (
 	stringType
 )
 
+// Set is not thread-safe, for concurrent use, make sure to use it with synchronization.
 type Set struct {
 	data map[interface{}]lang.PlaceholderType
 	tp   int
@@ -182,10 +183,7 @@ func (s *Set) add(i interface{}) {
 }
 
 func (s *Set) setType(i interface{}) {
-	if s.tp != untyped {
+	// s.tp can only be untyped here
-		return
-	}
-
 	switch i.(type) {
 	case int:
 		s.tp = intType

core/collection/set_test.go

@@ -5,8 +5,13 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/logx"
 )
 
+func init() {
+	logx.Disable()
+}
+
 func BenchmarkRawSet(b *testing.B) {
 	m := make(map[interface{}]struct{})
 	for i := 0; i < b.N; i++ {
@@ -147,3 +152,51 @@ func TestCount(t *testing.T) {
 	// then
 	assert.Equal(t, set.Count(), 3)
 }
+
+func TestKeysIntMismatch(t *testing.T) {
+	set := NewSet()
+	set.add(int64(1))
+	set.add(2)
+	vals := set.KeysInt()
+	assert.EqualValues(t, []int{2}, vals)
+}
+
+func TestKeysInt64Mismatch(t *testing.T) {
+	set := NewSet()
+	set.add(1)
+	set.add(int64(2))
+	vals := set.KeysInt64()
+	assert.EqualValues(t, []int64{2}, vals)
+}
+
+func TestKeysUintMismatch(t *testing.T) {
+	set := NewSet()
+	set.add(1)
+	set.add(uint(2))
+	vals := set.KeysUint()
+	assert.EqualValues(t, []uint{2}, vals)
+}
+
+func TestKeysUint64Mismatch(t *testing.T) {
+	set := NewSet()
+	set.add(1)
+	set.add(uint64(2))
+	vals := set.KeysUint64()
+	assert.EqualValues(t, []uint64{2}, vals)
+}
+
+func TestKeysStrMismatch(t *testing.T) {
+	set := NewSet()
+	set.add(1)
+	set.add("2")
+	vals := set.KeysStr()
+	assert.EqualValues(t, []string{"2"}, vals)
+}
+
+func TestSetType(t *testing.T) {
+	set := NewUnmanagedSet()
+	set.add(1)
+	set.add("2")
+	vals := set.Keys()
+	assert.ElementsMatch(t, []interface{}{1, "2"}, vals)
+}

core/collection/timingwheel.go

@@ -204,6 +204,7 @@ func (tw *TimingWheel) removeTask(key interface{}) {
 
 	timer := val.(*positionEntry)
 	timer.item.removed = true
+	tw.timers.Del(key)
 }
 
 func (tw *TimingWheel) run() {
@@ -248,7 +249,6 @@ func (tw *TimingWheel) scanAndRunTasks(l *list.List) {
 		if task.removed {
 			next := e.Next()
 			l.Remove(e)
-			tw.timers.Del(task.key)
 			e = next
 			continue
 		} else if task.circle > 0 {
@@ -301,6 +301,7 @@ func (tw *TimingWheel) setTask(task *timingEntry) {
 func (tw *TimingWheel) setTimerPosition(pos int, task *timingEntry) {
 	if val, ok := tw.timers.Get(task.key); ok {
 		timer := val.(*positionEntry)
+		timer.item = task
 		timer.pos = pos
 	} else {
 		tw.timers.Set(task.key, &positionEntry{

core/collection/timingwheel_test.go

@@ -213,7 +213,10 @@ func TestTimingWheel_SetTimer(t *testing.T) {
 	}
 
 	for _, test := range tests {
+		test := test
 		t.Run(stringx.RandId(), func(t *testing.T) {
+			t.Parallel()
+
 			var count int32
 			ticker := timex.NewFakeTicker()
 			tick := func() {
@@ -291,7 +294,10 @@ func TestTimingWheel_SetAndMoveThenStart(t *testing.T) {
 	}
 
 	for _, test := range tests {
+		test := test
 		t.Run(stringx.RandId(), func(t *testing.T) {
+			t.Parallel()
+
 			var count int32
 			ticker := timex.NewFakeTicker()
 			tick := func() {
@@ -376,7 +382,10 @@ func TestTimingWheel_SetAndMoveTwice(t *testing.T) {
 	}
 
 	for _, test := range tests {
+		test := test
 		t.Run(stringx.RandId(), func(t *testing.T) {
+			t.Parallel()
+
 			var count int32
 			ticker := timex.NewFakeTicker()
 			tick := func() {
@@ -454,7 +463,10 @@ func TestTimingWheel_ElapsedAndSet(t *testing.T) {
 	}
 
 	for _, test := range tests {
+		test := test
 		t.Run(stringx.RandId(), func(t *testing.T) {
+			t.Parallel()
+
 			var count int32
 			ticker := timex.NewFakeTicker()
 			tick := func() {
@@ -542,7 +554,10 @@ func TestTimingWheel_ElapsedAndSetThenMove(t *testing.T) {
 	}
 
 	for _, test := range tests {
+		test := test
 		t.Run(stringx.RandId(), func(t *testing.T) {
+			t.Parallel()
+
 			var count int32
 			ticker := timex.NewFakeTicker()
 			tick := func() {
@@ -579,6 +594,31 @@ func TestTimingWheel_ElapsedAndSetThenMove(t *testing.T) {
 	}
 }
 
+func TestMoveAndRemoveTask(t *testing.T) {
+	ticker := timex.NewFakeTicker()
+	tick := func(v int) {
+		for i := 0; i < v; i++ {
+			ticker.Tick()
+		}
+	}
+	var keys []int
+	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
+		assert.Equal(t, "any", k)
+		assert.Equal(t, 3, v.(int))
+		keys = append(keys, v.(int))
+		ticker.Done()
+	}, ticker)
+	defer tw.Stop()
+	tw.SetTimer("any", 3, testStep*8)
+	tick(6)
+	tw.MoveTimer("any", testStep*7)
+	tick(3)
+	tw.RemoveTimer("any")
+	tick(30)
+	time.Sleep(time.Millisecond)
+	assert.Equal(t, 0, len(keys))
+}
+
 func BenchmarkTimingWheel(b *testing.B) {
 	b.ReportAllocs()
 

core/conf/config.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"log"
+	"os"
 	"path"
 
 	"github.com/tal-tech/go-zero/core/mapping"
@@ -19,7 +20,7 @@ func LoadConfig(file string, v interface{}) error {
 	if content, err := ioutil.ReadFile(file); err != nil {
 		return err
 	} else if loader, ok := loaders[path.Ext(file)]; ok {
-		return loader(content, v)
+		return loader([]byte(os.ExpandEnv(string(content))), v)
 	} else {
 		return fmt.Errorf("unrecoginized file type: %s", file)
 	}

core/conf/config_test.go

@@ -0,0 +1,61 @@
+package conf
+
+import (
+	"io/ioutil"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/hash"
+)
+
+func TestConfigJson(t *testing.T) {
+	tests := []string{
+		".json",
+		".yaml",
+		".yml",
+	}
+	text := `{
+	"a": "foo",
+	"b": 1,
+	"c": "${FOO}"
+}`
+	for _, test := range tests {
+		test := test
+		t.Run(test, func(t *testing.T) {
+			os.Setenv("FOO", "2")
+			defer os.Unsetenv("FOO")
+			tmpfile, err := createTempFile(test, text)
+			assert.Nil(t, err)
+			defer os.Remove(tmpfile)
+
+			var val struct {
+				A string `json:"a"`
+				B int    `json:"b"`
+				C string `json:"c"`
+			}
+			MustLoad(tmpfile, &val)
+			assert.Equal(t, "foo", val.A)
+			assert.Equal(t, 1, val.B)
+			assert.Equal(t, "2", val.C)
+		})
+	}
+}
+
+func createTempFile(ext, text string) (string, error) {
+	tmpfile, err := ioutil.TempFile(os.TempDir(), hash.Md5Hex([]byte(text))+"*"+ext)
+	if err != nil {
+		return "", err
+	}
+
+	if err := ioutil.WriteFile(tmpfile.Name(), []byte(text), os.ModeTemporary); err != nil {
+		return "", err
+	}
+
+	filename := tmpfile.Name()
+	if err = tmpfile.Close(); err != nil {
+		return "", err
+	}
+
+	return filename, nil
+}

core/conf/properties.go

@@ -30,12 +30,12 @@ type mapBasedProperties struct {
 	lock       sync.RWMutex
 }
 
-// Loads the properties into a properties configuration instance. May return the
+// Loads the properties into a properties configuration instance.
-// configuration itself along with an error that indicates if there was a problem loading the configuration.
+// Returns an error that indicates if there was a problem loading the configuration.
 func LoadProperties(filename string) (Properties, error) {
 	lines, err := iox.ReadTextLines(filename, iox.WithoutBlank(), iox.OmitWithPrefix("#"))
 	if err != nil {
-		return nil, nil
+		return nil, err
 	}
 
 	raw := make(map[string]string)

core/conf/properties_test.go

@@ -41,3 +41,8 @@ func TestSetInt(t *testing.T) {
 	props.SetInt(key, value)
 	assert.Equal(t, value, props.GetInt(key))
 }
+
+func TestLoadBadFile(t *testing.T) {
+	_, err := LoadProperties("nosuchfile")
+	assert.NotNil(t, err)
+}

core/contextx/deadline_test.go

@@ -10,8 +10,10 @@ import (
 
 func TestShrinkDeadlineLess(t *testing.T) {
 	deadline := time.Now().Add(time.Second)
-	ctx, _ := context.WithDeadline(context.Background(), deadline)
+	ctx, cancel := context.WithDeadline(context.Background(), deadline)
-	ctx, _ = ShrinkDeadline(ctx, time.Minute)
+	defer cancel()
+	ctx, cancel = ShrinkDeadline(ctx, time.Minute)
+	defer cancel()
 	dl, ok := ctx.Deadline()
 	assert.True(t, ok)
 	assert.Equal(t, deadline, dl)
@@ -19,8 +21,10 @@ func TestShrinkDeadlineLess(t *testing.T) {
 
 func TestShrinkDeadlineMore(t *testing.T) {
 	deadline := time.Now().Add(time.Minute)
-	ctx, _ := context.WithDeadline(context.Background(), deadline)
+	ctx, cancel := context.WithDeadline(context.Background(), deadline)
-	ctx, _ = ShrinkDeadline(ctx, time.Second)
+	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/valueonlycontext_test.go

@@ -12,7 +12,8 @@ func TestContextCancel(t *testing.T) {
 	c := context.WithValue(context.Background(), "key", "value")
 	c1, cancel := context.WithCancel(c)
 	o := ValueOnlyFrom(c1)
-	c2, _ := context.WithCancel(o)
+	c2, cancel2 := context.WithCancel(o)
+	defer cancel2()
 	contexts := []context.Context{c1, c2}
 
 	for _, c := range contexts {
@@ -34,8 +35,9 @@ func TestContextCancel(t *testing.T) {
 	assert.NotEqual(t, context.Canceled, c2.Err())
 }
 
-func TestConextDeadline(t *testing.T) {
+func TestContextDeadline(t *testing.T) {
-	c, _ := context.WithDeadline(context.Background(), time.Now().Add(10*time.Millisecond))
+	c, cancel := context.WithDeadline(context.Background(), time.Now().Add(10*time.Millisecond))
+	cancel()
 	o := ValueOnlyFrom(c)
 	select {
 	case <-time.After(100 * time.Millisecond):
@@ -43,9 +45,11 @@ func TestConextDeadline(t *testing.T) {
 		t.Fatal("ValueOnlyContext: context should not have timed out")
 	}
 
-	c, _ = context.WithDeadline(context.Background(), time.Now().Add(10*time.Millisecond))
+	c, cancel = context.WithDeadline(context.Background(), time.Now().Add(10*time.Millisecond))
+	cancel()
 	o = ValueOnlyFrom(c)
-	c, _ = context.WithDeadline(o, time.Now().Add(20*time.Millisecond))
+	c, cancel = context.WithDeadline(o, time.Now().Add(20*time.Millisecond))
+	defer cancel()
 	select {
 	case <-time.After(100 * time.Millisecond):
 		t.Fatal("ValueOnlyContext+Deadline: context should have timed out")

core/discov/clients.go

@@ -8,7 +8,7 @@ import (
 )
 
 const (
-	indexOfKey = iota
+	_ = iota
 	indexOfId
 )
 

core/discov/internal/listener.go

@@ -1,4 +1,3 @@
-//go:generate mockgen -package internal -destination listener_mock.go -source listener.go Listener
 package internal
 
 type Listener interface {

core/discov/internal/listener_mock.go

@@ -1,45 +0,0 @@
-// Code generated by MockGen. DO NOT EDIT.
-// Source: listener.go
-
-// Package internal is a generated GoMock package.
-package internal
-
-import (
-	gomock "github.com/golang/mock/gomock"
-	reflect "reflect"
-)
-
-// MockListener is a mock of Listener interface
-type MockListener struct {
-	ctrl     *gomock.Controller
-	recorder *MockListenerMockRecorder
-}
-
-// MockListenerMockRecorder is the mock recorder for MockListener
-type MockListenerMockRecorder struct {
-	mock *MockListener
-}
-
-// NewMockListener creates a new mock instance
-func NewMockListener(ctrl *gomock.Controller) *MockListener {
-	mock := &MockListener{ctrl: ctrl}
-	mock.recorder = &MockListenerMockRecorder{mock}
-	return mock
-}
-
-// EXPECT returns an object that allows the caller to indicate expected use
-func (m *MockListener) EXPECT() *MockListenerMockRecorder {
-	return m.recorder
-}
-
-// OnUpdate mocks base method
-func (m *MockListener) OnUpdate(keys, values []string, newKey string) {
-	m.ctrl.T.Helper()
-	m.ctrl.Call(m, "OnUpdate", keys, values, newKey)
-}
-
-// OnUpdate indicates an expected call of OnUpdate
-func (mr *MockListenerMockRecorder) OnUpdate(keys, values, newKey interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnUpdate", reflect.TypeOf((*MockListener)(nil).OnUpdate), keys, values, newKey)
-}

core/discov/kubernetes/discov-namespace.yaml

@@ -1,4 +1,4 @@
 apiVersion: v1
 kind: Namespace
 metadata:
-  name: discovery
+  name: discov

core/discov/kubernetes/etcd.yaml

@@ -1,16 +1,16 @@
 apiVersion: v1
 kind: Service
 metadata:
-  name: discov
+  name: etcd
-  namespace: discovery
+  namespace: discov
 spec:
   ports:
-    - name: discov-port
+    - name: etcd-port
       port: 2379
       protocol: TCP
       targetPort: 2379
   selector:
-    app: discov
+    app: etcd
 
 ---
 
@@ -18,30 +18,31 @@ apiVersion: v1
 kind: Pod
 metadata:
   labels:
-    app: discov
+    app: etcd
-    discov_node: discov0
+    etcd_node: etcd0
-  name: discov0
+  name: etcd0
-  namespace: discovery
+  namespace: discov
 spec:
   containers:
     - command:
         - /usr/local/bin/etcd
         - --name
-        - discov0
+        - etcd0
         - --initial-advertise-peer-urls
-        - http://discov0:2380
+        - http://etcd0:2380
         - --listen-peer-urls
         - http://0.0.0.0:2380
         - --listen-client-urls
         - http://0.0.0.0:2379
         - --advertise-client-urls
-        - http://discov0:2379
+        - http://etcd0.discov:2379
         - --initial-cluster
-        - discov0=http://discov0:2380,discov1=http://discov1:2380,discov2=http://discov2:2380,discov3=http://discov3:2380,discov4=http://discov4:2380
+        - etcd0=http://etcd0:2380,etcd1=http://etcd1:2380,etcd2=http://etcd2:2380,etcd3=http://etcd3:2380,etcd4=http://etcd4:2380
         - --initial-cluster-state
         - new
-      image: registry-vpc.cn-hangzhou.aliyuncs.com/xapp/etcd:latest
+        - --auto-compaction-retention=1
-      name: discov0
+      image: quay.io/coreos/etcd:latest
+      name: etcd0
       ports:
         - containerPort: 2379
           name: client
@@ -49,8 +50,6 @@ spec:
         - containerPort: 2380
           name: server
           protocol: TCP
-  imagePullSecrets:
-    - name: aliyun
   affinity:
     podAntiAffinity:
       requiredDuringSchedulingIgnoredDuringExecution:
@@ -59,7 +58,7 @@ spec:
               - key: app
                 operator: In
                 values:
-                  - discov
+                  - etcd
           topologyKey: "kubernetes.io/hostname"
   restartPolicy: Always
 
@@ -69,9 +68,9 @@ apiVersion: v1
 kind: Service
 metadata:
   labels:
-    discov_node: discov0
+    etcd_node: etcd0
-  name: discov0
+  name: etcd0
-  namespace: discovery
+  namespace: discov
 spec:
   ports:
     - name: client
@@ -83,7 +82,7 @@ spec:
       protocol: TCP
       targetPort: 2380
   selector:
-    discov_node: discov0
+    etcd_node: etcd0
 
 ---
 
@@ -91,30 +90,31 @@ apiVersion: v1
 kind: Pod
 metadata:
   labels:
-    app: discov
+    app: etcd
-    discov_node: discov1
+    etcd_node: etcd1
-  name: discov1
+  name: etcd1
-  namespace: discovery
+  namespace: discov
 spec:
   containers:
     - command:
         - /usr/local/bin/etcd
         - --name
-        - discov1
+        - etcd1
         - --initial-advertise-peer-urls
-        - http://discov1:2380
+        - http://etcd1:2380
         - --listen-peer-urls
         - http://0.0.0.0:2380
         - --listen-client-urls
         - http://0.0.0.0:2379
         - --advertise-client-urls
-        - http://discov1:2379
+        - http://etcd1.discov:2379
         - --initial-cluster
-        - discov0=http://discov0:2380,discov1=http://discov1:2380,discov2=http://discov2:2380,discov3=http://discov3:2380,discov4=http://discov4:2380
+        - etcd0=http://etcd0:2380,etcd1=http://etcd1:2380,etcd2=http://etcd2:2380,etcd3=http://etcd3:2380,etcd4=http://etcd4:2380
         - --initial-cluster-state
         - new
-      image: registry-vpc.cn-hangzhou.aliyuncs.com/xapp/etcd:latest
+        - --auto-compaction-retention=1
-      name: discov1
+      image: quay.io/coreos/etcd:latest
+      name: etcd1
       ports:
         - containerPort: 2379
           name: client
@@ -122,8 +122,6 @@ spec:
         - containerPort: 2380
           name: server
           protocol: TCP
-  imagePullSecrets:
-    - name: aliyun
   affinity:
     podAntiAffinity:
       requiredDuringSchedulingIgnoredDuringExecution:
@@ -132,7 +130,7 @@ spec:
               - key: app
                 operator: In
                 values:
-                  - discov
+                  - etcd
           topologyKey: "kubernetes.io/hostname"
   restartPolicy: Always
 
@@ -142,9 +140,9 @@ apiVersion: v1
 kind: Service
 metadata:
   labels:
-    discov_node: discov1
+    etcd_node: etcd1
-  name: discov1
+  name: etcd1
-  namespace: discovery
+  namespace: discov
 spec:
   ports:
     - name: client
@@ -156,7 +154,7 @@ spec:
       protocol: TCP
       targetPort: 2380
   selector:
-    discov_node: discov1
+    etcd_node: etcd1
 
 ---
 
@@ -164,30 +162,31 @@ apiVersion: v1
 kind: Pod
 metadata:
   labels:
-    app: discov
+    app: etcd
-    discov_node: discov2
+    etcd_node: etcd2
-  name: discov2
+  name: etcd2
-  namespace: discovery
+  namespace: discov
 spec:
   containers:
     - command:
         - /usr/local/bin/etcd
         - --name
-        - discov2
+        - etcd2
         - --initial-advertise-peer-urls
-        - http://discov2:2380
+        - http://etcd2:2380
         - --listen-peer-urls
         - http://0.0.0.0:2380
         - --listen-client-urls
         - http://0.0.0.0:2379
         - --advertise-client-urls
-        - http://discov2:2379
+        - http://etcd2.discov:2379
         - --initial-cluster
-        - discov0=http://discov0:2380,discov1=http://discov1:2380,discov2=http://discov2:2380,discov3=http://discov3:2380,discov4=http://discov4:2380
+        - etcd0=http://etcd0:2380,etcd1=http://etcd1:2380,etcd2=http://etcd2:2380,etcd3=http://etcd3:2380,etcd4=http://etcd4:2380
         - --initial-cluster-state
         - new
-      image: registry-vpc.cn-hangzhou.aliyuncs.com/xapp/etcd:latest
+        - --auto-compaction-retention=1
-      name: discov2
+      image: quay.io/coreos/etcd:latest
+      name: etcd2
       ports:
         - containerPort: 2379
           name: client
@@ -195,8 +194,6 @@ spec:
         - containerPort: 2380
           name: server
           protocol: TCP
-  imagePullSecrets:
-    - name: aliyun
   affinity:
     podAntiAffinity:
       requiredDuringSchedulingIgnoredDuringExecution:
@@ -205,7 +202,7 @@ spec:
               - key: app
                 operator: In
                 values:
-                  - discov
+                  - etcd
           topologyKey: "kubernetes.io/hostname"
   restartPolicy: Always
 
@@ -215,9 +212,9 @@ apiVersion: v1
 kind: Service
 metadata:
   labels:
-    discov_node: discov2
+    etcd_node: etcd2
-  name: discov2
+  name: etcd2
-  namespace: discovery
+  namespace: discov
 spec:
   ports:
     - name: client
@@ -229,7 +226,7 @@ spec:
       protocol: TCP
       targetPort: 2380
   selector:
-    discov_node: discov2
+    etcd_node: etcd2
 
 ---
 
@@ -237,30 +234,31 @@ apiVersion: v1
 kind: Pod
 metadata:
   labels:
-    app: discov
+    app: etcd
-    discov_node: discov3
+    etcd_node: etcd3
-  name: discov3
+  name: etcd3
-  namespace: discovery
+  namespace: discov
 spec:
   containers:
     - command:
         - /usr/local/bin/etcd
         - --name
-        - discov3
+        - etcd3
         - --initial-advertise-peer-urls
-        - http://discov3:2380
+        - http://etcd3:2380
         - --listen-peer-urls
         - http://0.0.0.0:2380
         - --listen-client-urls
         - http://0.0.0.0:2379
         - --advertise-client-urls
-        - http://discov3:2379
+        - http://etcd3.discov:2379
         - --initial-cluster
-        - discov0=http://discov0:2380,discov1=http://discov1:2380,discov2=http://discov2:2380,discov3=http://discov3:2380,discov4=http://discov4:2380
+        - etcd0=http://etcd0:2380,etcd1=http://etcd1:2380,etcd2=http://etcd2:2380,etcd3=http://etcd3:2380,etcd4=http://etcd4:2380
         - --initial-cluster-state
         - new
-      image: registry-vpc.cn-hangzhou.aliyuncs.com/xapp/etcd:latest
+        - --auto-compaction-retention=1
-      name: discov3
+      image: quay.io/coreos/etcd:latest
+      name: etcd3
       ports:
         - containerPort: 2379
           name: client
@@ -268,8 +266,6 @@ spec:
         - containerPort: 2380
           name: server
           protocol: TCP
-  imagePullSecrets:
-    - name: aliyun
   affinity:
     podAntiAffinity:
       requiredDuringSchedulingIgnoredDuringExecution:
@@ -278,7 +274,7 @@ spec:
               - key: app
                 operator: In
                 values:
-                  - discov
+                  - etcd
           topologyKey: "kubernetes.io/hostname"
   restartPolicy: Always
 
@@ -288,9 +284,9 @@ apiVersion: v1
 kind: Service
 metadata:
   labels:
-    discov_node: discov3
+    etcd_node: etcd3
-  name: discov3
+  name: etcd3
-  namespace: discovery
+  namespace: discov
 spec:
   ports:
     - name: client
@@ -302,7 +298,7 @@ spec:
       protocol: TCP
       targetPort: 2380
   selector:
-    discov_node: discov3
+    etcd_node: etcd3
 
 ---
 
@@ -310,30 +306,31 @@ apiVersion: v1
 kind: Pod
 metadata:
   labels:
-    app: discov
+    app: etcd
-    discov_node: discov4
+    etcd_node: etcd4
-  name: discov4
+  name: etcd4
-  namespace: discovery
+  namespace: discov
 spec:
   containers:
     - command:
         - /usr/local/bin/etcd
         - --name
-        - discov4
+        - etcd4
         - --initial-advertise-peer-urls
-        - http://discov4:2380
+        - http://etcd4:2380
         - --listen-peer-urls
         - http://0.0.0.0:2380
         - --listen-client-urls
         - http://0.0.0.0:2379
         - --advertise-client-urls
-        - http://discov4:2379
+        - http://etcd4.discov:2379
         - --initial-cluster
-        - discov0=http://discov0:2380,discov1=http://discov1:2380,discov2=http://discov2:2380,discov3=http://discov3:2380,discov4=http://discov4:2380
+        - etcd0=http://etcd0:2380,etcd1=http://etcd1:2380,etcd2=http://etcd2:2380,etcd3=http://etcd3:2380,etcd4=http://etcd4:2380
         - --initial-cluster-state
         - new
-      image: registry-vpc.cn-hangzhou.aliyuncs.com/xapp/etcd:latest
+        - --auto-compaction-retention=1
-      name: discov4
+      image: quay.io/coreos/etcd:latest
+      name: etcd4
       ports:
         - containerPort: 2379
           name: client
@@ -341,8 +338,6 @@ spec:
         - containerPort: 2380
           name: server
           protocol: TCP
-  imagePullSecrets:
-    - name: aliyun
   affinity:
     podAntiAffinity:
       requiredDuringSchedulingIgnoredDuringExecution:
@@ -351,7 +346,7 @@ spec:
               - key: app
                 operator: In
                 values:
-                  - discov
+                  - etcd
           topologyKey: "kubernetes.io/hostname"
   restartPolicy: Always
 
@@ -361,9 +356,9 @@ apiVersion: v1
 kind: Service
 metadata:
   labels:
-    discov_node: discov4
+    etcd_node: etcd4
-  name: discov4
+  name: etcd4
-  namespace: discovery
+  namespace: discov
 spec:
   ports:
     - name: client
@@ -375,4 +370,4 @@ spec:
       protocol: TCP
       targetPort: 2380
   selector:
-    discov_node: discov4
+    etcd_node: etcd4

core/discov/publisher_test.go

@@ -111,6 +111,10 @@ func TestPublisher_keepAliveAsyncQuit(t *testing.T) {
 	defer ctrl.Finish()
 	const id clientv3.LeaseID = 1
 	cli := internal.NewMockEtcdClient(ctrl)
+	cli.EXPECT().ActiveConnection()
+	cli.EXPECT().Close()
+	defer cli.Close()
+	cli.ActiveConnection()
 	restore := setMockClient(cli)
 	defer restore()
 	cli.EXPECT().Ctx().AnyTimes()

core/errorx/atomicerror.go

@@ -1,21 +1,18 @@
 package errorx
 
-import "sync"
+import "sync/atomic"
 
 type AtomicError struct {
-	err  error
+	err atomic.Value // error
-	lock sync.Mutex
 }
 
 func (ae *AtomicError) Set(err error) {
-	ae.lock.Lock()
+	ae.err.Store(err)
-	ae.err = err
-	ae.lock.Unlock()
 }
 
 func (ae *AtomicError) Load() error {
-	ae.lock.Lock()
+	if v := ae.err.Load(); v != nil {
-	err := ae.err
+		return v.(error)
-	ae.lock.Unlock()
+	}
-	return err
+	return nil
 }

core/errorx/atomicerror_test.go

@@ -2,6 +2,8 @@ package errorx
 
 import (
 	"errors"
+	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -19,3 +21,53 @@ func TestAtomicErrorNil(t *testing.T) {
 	var err AtomicError
 	assert.Nil(t, err.Load())
 }
+
+func BenchmarkAtomicError(b *testing.B) {
+	var aerr AtomicError
+	wg := sync.WaitGroup{}
+
+	b.Run("Load", func(b *testing.B) {
+		var done uint32
+		go func() {
+			for {
+				if atomic.LoadUint32(&done) != 0 {
+					break
+				}
+				wg.Add(1)
+				go func() {
+					aerr.Set(errDummy)
+					wg.Done()
+				}()
+			}
+		}()
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_ = aerr.Load()
+		}
+		b.StopTimer()
+		atomic.StoreUint32(&done, 1)
+		wg.Wait()
+	})
+	b.Run("Set", func(b *testing.B) {
+		var done uint32
+		go func() {
+			for {
+				if atomic.LoadUint32(&done) != 0 {
+					break
+				}
+				wg.Add(1)
+				go func() {
+					_ = aerr.Load()
+					wg.Done()
+				}()
+			}
+		}()
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			aerr.Set(errDummy)
+		}
+		b.StopTimer()
+		atomic.StoreUint32(&done, 1)
+		wg.Wait()
+	})
+}

core/errorx/callchain.go

@@ -0,0 +1,11 @@
+package errorx
+
+func Chain(fns ...func() error) error {
+	for _, fn := range fns {
+		if err := fn(); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

core/errorx/callchain_test.go

@@ -0,0 +1,27 @@
+package errorx
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestChain(t *testing.T) {
+	var errDummy = errors.New("dummy")
+	assert.Nil(t, Chain(func() error {
+		return nil
+	}, func() error {
+		return nil
+	}))
+	assert.Equal(t, errDummy, Chain(func() error {
+		return errDummy
+	}, func() error {
+		return nil
+	}))
+	assert.Equal(t, errDummy, Chain(func() error {
+		return nil
+	}, func() error {
+		return errDummy
+	}))
+}

core/executors/bulkexecutor_test.go

@@ -12,14 +12,14 @@ func TestBulkExecutor(t *testing.T) {
 	var values []int
 	var lock sync.Mutex
 
-	exeutor := NewBulkExecutor(func(items []interface{}) {
+	executor := NewBulkExecutor(func(items []interface{}) {
 		lock.Lock()
 		values = append(values, len(items))
 		lock.Unlock()
 	}, WithBulkTasks(10), WithBulkInterval(time.Minute))
 
 	for i := 0; i < 50; i++ {
-		exeutor.Add(1)
+		executor.Add(1)
 		time.Sleep(time.Millisecond)
 	}
 
@@ -40,13 +40,13 @@ func TestBulkExecutorFlushInterval(t *testing.T) {
 	var wait sync.WaitGroup
 
 	wait.Add(1)
-	exeutor := NewBulkExecutor(func(items []interface{}) {
+	executor := NewBulkExecutor(func(items []interface{}) {
 		assert.Equal(t, size, len(items))
 		wait.Done()
 	}, WithBulkTasks(caches), WithBulkInterval(time.Millisecond*100))
 
 	for i := 0; i < size; i++ {
-		exeutor.Add(1)
+		executor.Add(1)
 	}
 
 	wait.Wait()
@@ -86,9 +86,7 @@ func TestBuldExecutorFlushSlowTasks(t *testing.T) {
 		time.Sleep(time.Millisecond * 100)
 		lock.Lock()
 		defer lock.Unlock()
-		for _, i := range tasks {
+		result = append(result, tasks...)
-			result = append(result, i)
-		}
 	}, WithBulkTasks(1000))
 	for i := 0; i < total; i++ {
 		assert.Nil(t, exec.Add(i))

core/executors/chunkexecutor_test.go

@@ -12,14 +12,14 @@ func TestChunkExecutor(t *testing.T) {
 	var values []int
 	var lock sync.Mutex
 
-	exeutor := NewChunkExecutor(func(items []interface{}) {
+	executor := NewChunkExecutor(func(items []interface{}) {
 		lock.Lock()
 		values = append(values, len(items))
 		lock.Unlock()
 	}, WithChunkBytes(10), WithFlushInterval(time.Minute))
 
 	for i := 0; i < 50; i++ {
-		exeutor.Add(1, 1)
+		executor.Add(1, 1)
 		time.Sleep(time.Millisecond)
 	}
 
@@ -40,13 +40,13 @@ func TestChunkExecutorFlushInterval(t *testing.T) {
 	var wait sync.WaitGroup
 
 	wait.Add(1)
-	exeutor := NewChunkExecutor(func(items []interface{}) {
+	executor := NewChunkExecutor(func(items []interface{}) {
 		assert.Equal(t, size, len(items))
 		wait.Done()
 	}, WithChunkBytes(caches), WithFlushInterval(time.Millisecond*100))
 
 	for i := 0; i < size; i++ {
-		exeutor.Add(1, 1)
+		executor.Add(1, 1)
 	}
 
 	wait.Wait()

core/executors/periodicalexecutor.go

@@ -3,6 +3,7 @@ package executors
 import (
 	"reflect"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/tal-tech/go-zero/core/lang"
@@ -35,6 +36,7 @@ type (
 		// avoid race condition on waitGroup when calling wg.Add/Done/Wait(...)
 		wgBarrier   syncx.Barrier
 		confirmChan chan lang.PlaceholderType
+		inflight    int32
 		guarded     bool
 		newTicker   func(duration time.Duration) timex.Ticker
 		lock        sync.Mutex
@@ -82,6 +84,7 @@ func (pe *PeriodicalExecutor) Sync(fn func()) {
 }
 
 func (pe *PeriodicalExecutor) Wait() {
+	pe.Flush()
 	pe.wgBarrier.Guard(func() {
 		pe.waitGroup.Wait()
 	})
@@ -90,18 +93,16 @@ func (pe *PeriodicalExecutor) Wait() {
 func (pe *PeriodicalExecutor) addAndCheck(task interface{}) (interface{}, bool) {
 	pe.lock.Lock()
 	defer func() {
-		var start bool
 		if !pe.guarded {
 			pe.guarded = true
-			start = true
+			// defer to unlock quickly
+			defer pe.backgroundFlush()
 		}
 		pe.lock.Unlock()
-		if start {
-			pe.backgroundFlush()
-		}
 	}()
 
 	if pe.container.AddTask(task) {
+		atomic.AddInt32(&pe.inflight, 1)
 		return pe.container.RemoveAll(), true
 	}
 
@@ -110,6 +111,9 @@ func (pe *PeriodicalExecutor) addAndCheck(task interface{}) (interface{}, bool)
 
 func (pe *PeriodicalExecutor) backgroundFlush() {
 	threading.GoSafe(func() {
+		// flush before quit goroutine to avoid missing tasks
+		defer pe.Flush()
+
 		ticker := pe.newTicker(pe.interval)
 		defer ticker.Stop()
 
@@ -119,6 +123,7 @@ func (pe *PeriodicalExecutor) backgroundFlush() {
 			select {
 			case vals := <-pe.commander:
 				commanded = true
+				atomic.AddInt32(&pe.inflight, -1)
 				pe.enterExecution()
 				pe.confirmChan <- lang.Placeholder
 				pe.executeTasks(vals)
@@ -128,13 +133,7 @@ func (pe *PeriodicalExecutor) backgroundFlush() {
 					commanded = false
 				} else if pe.Flush() {
 					last = timex.Now()
-				} else if timex.Since(last) > pe.interval*idleRound {
+				} else if pe.shallQuit(last) {
-					pe.lock.Lock()
-					pe.guarded = false
-					pe.lock.Unlock()
-
-					// flush again to avoid missing tasks
-					pe.Flush()
 					return
 				}
 			}
@@ -177,3 +176,19 @@ func (pe *PeriodicalExecutor) hasTasks(tasks interface{}) bool {
 		return true
 	}
 }
+
+func (pe *PeriodicalExecutor) shallQuit(last time.Duration) (stop bool) {
+	if timex.Since(last) <= pe.interval*idleRound {
+		return
+	}
+
+	// checking pe.inflight and setting pe.guarded should be locked together
+	pe.lock.Lock()
+	if atomic.LoadInt32(&pe.inflight) == 0 {
+		pe.guarded = false
+		stop = true
+	}
+	pe.lock.Unlock()
+
+	return
+}

core/executors/periodicalexecutor_test.go

@@ -140,6 +140,26 @@ func TestPeriodicalExecutor_WaitFast(t *testing.T) {
 	assert.Equal(t, total, cnt)
 }
 
+func TestPeriodicalExecutor_Deadlock(t *testing.T) {
+	executor := NewBulkExecutor(func(tasks []interface{}) {
+	}, WithBulkTasks(1), WithBulkInterval(time.Millisecond))
+	for i := 0; i < 1e5; i++ {
+		executor.Add(1)
+	}
+}
+
+func TestPeriodicalExecutor_hasTasks(t *testing.T) {
+	ticker := timex.NewFakeTicker()
+	defer ticker.Stop()
+
+	exec := NewPeriodicalExecutor(time.Millisecond, newContainer(time.Millisecond, nil))
+	exec.newTicker = func(d time.Duration) timex.Ticker {
+		return ticker
+	}
+	assert.False(t, exec.hasTasks(nil))
+	assert.True(t, exec.hasTasks(1))
+}
+
 // go test -benchtime 10s -bench .
 func BenchmarkExecutor(b *testing.B) {
 	b.ReportAllocs()

core/filex/file_test.go

@@ -4,9 +4,8 @@ import (
 	"os"
 	"testing"
 
-	"github.com/tal-tech/go-zero/core/fs"
-
 	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/fs"
 )
 
 const (
@@ -15,34 +14,34 @@ const (
 	text          = `first line
 Cum sociis natoque penatibus et magnis dis parturient. Phasellus laoreet lorem vel dolor tempus vehicula. Vivamus sagittis lacus vel augue laoreet rutrum faucibus. Integer legentibus erat a ante historiarum dapibus.
 Quisque ut dolor gravida, placerat libero vel, euismod. Quam temere in vitiis, legem sancimus haerentia. Qui ipsorum lingua Celtae, nostra Galli appellantur. Quis aute iure reprehenderit in voluptate velit esse. Fabio vel iudice vincam, sunt in culpa qui officia. Cras mattis iudicium purus sit amet fermentum.
-Quo usque tandem abutere, Catilina, patientia nostra? Gallia est omnis divisa in partes tres, quarum. Quam diu etiam furor iste tuus nos eludet? Quid securi etiam tamquam eu fugiat nulla pariatur. Curabitur blandit tempus ardua ridiculus sed magna.
+Quo usque tandem abutere, Catilina, patientia nostra? Gallia est omnis divisa in partes tres, quarum. Quam diu etiam furor iste tuus nos eludet? Quid securi etiam tamquam eu fugiat nulla pariatur. Curabitur blandit tempus ardua ridiculous sed magna.
 Magna pars studiorum, prodita quaerimus. Cum ceteris in veneratione tui montes, nascetur mus. Morbi odio eros, volutpat ut pharetra vitae, lobortis sed nibh. Plura mihi bona sunt, inclinet, amari petere vellent. Idque Caesaris facere voluntate liceret: sese habere. Tu quoque, Brute, fili mi, nihil timor populi, nihil!
 Tityre, tu patulae recubans sub tegmine fagi dolor. Inmensae subtilitatis, obscuris et malesuada fames. Quae vero auctorem tractata ab fiducia dicuntur.
 Cum sociis natoque penatibus et magnis dis parturient. Phasellus laoreet lorem vel dolor tempus vehicula. Vivamus sagittis lacus vel augue laoreet rutrum faucibus. Integer legentibus erat a ante historiarum dapibus.
 Quisque ut dolor gravida, placerat libero vel, euismod. Quam temere in vitiis, legem sancimus haerentia. Qui ipsorum lingua Celtae, nostra Galli appellantur. Quis aute iure reprehenderit in voluptate velit esse. Fabio vel iudice vincam, sunt in culpa qui officia. Cras mattis iudicium purus sit amet fermentum.
-Quo usque tandem abutere, Catilina, patientia nostra? Gallia est omnis divisa in partes tres, quarum. Quam diu etiam furor iste tuus nos eludet? Quid securi etiam tamquam eu fugiat nulla pariatur. Curabitur blandit tempus ardua ridiculus sed magna.
+Quo usque tandem abutere, Catilina, patientia nostra? Gallia est omnis divisa in partes tres, quarum. Quam diu etiam furor iste tuus nos eludet? Quid securi etiam tamquam eu fugiat nulla pariatur. Curabitur blandit tempus ardua ridiculous sed magna.
 Magna pars studiorum, prodita quaerimus. Cum ceteris in veneratione tui montes, nascetur mus. Morbi odio eros, volutpat ut pharetra vitae, lobortis sed nibh. Plura mihi bona sunt, inclinet, amari petere vellent. Idque Caesaris facere voluntate liceret: sese habere. Tu quoque, Brute, fili mi, nihil timor populi, nihil!
 Tityre, tu patulae recubans sub tegmine fagi dolor. Inmensae subtilitatis, obscuris et malesuada fames. Quae vero auctorem tractata ab fiducia dicuntur.
 Cum sociis natoque penatibus et magnis dis parturient. Phasellus laoreet lorem vel dolor tempus vehicula. Vivamus sagittis lacus vel augue laoreet rutrum faucibus. Integer legentibus erat a ante historiarum dapibus.
 Quisque ut dolor gravida, placerat libero vel, euismod. Quam temere in vitiis, legem sancimus haerentia. Qui ipsorum lingua Celtae, nostra Galli appellantur. Quis aute iure reprehenderit in voluptate velit esse. Fabio vel iudice vincam, sunt in culpa qui officia. Cras mattis iudicium purus sit amet fermentum.
-Quo usque tandem abutere, Catilina, patientia nostra? Gallia est omnis divisa in partes tres, quarum. Quam diu etiam furor iste tuus nos eludet? Quid securi etiam tamquam eu fugiat nulla pariatur. Curabitur blandit tempus ardua ridiculus sed magna.
+Quo usque tandem abutere, Catilina, patientia nostra? Gallia est omnis divisa in partes tres, quarum. Quam diu etiam furor iste tuus nos eludet? Quid securi etiam tamquam eu fugiat nulla pariatur. Curabitur blandit tempus ardua ridiculous sed magna.
 Magna pars studiorum, prodita quaerimus. Cum ceteris in veneratione tui montes, nascetur mus. Morbi odio eros, volutpat ut pharetra vitae, lobortis sed nibh. Plura mihi bona sunt, inclinet, amari petere vellent. Idque Caesaris facere voluntate liceret: sese habere. Tu quoque, Brute, fili mi, nihil timor populi, nihil!
 Tityre, tu patulae recubans sub tegmine fagi dolor. Inmensae subtilitatis, obscuris et malesuada fames. Quae vero auctorem tractata ab fiducia dicuntur.
 ` + longLine
 	textWithLastNewline = `first line
 Cum sociis natoque penatibus et magnis dis parturient. Phasellus laoreet lorem vel dolor tempus vehicula. Vivamus sagittis lacus vel augue laoreet rutrum faucibus. Integer legentibus erat a ante historiarum dapibus.
 Quisque ut dolor gravida, placerat libero vel, euismod. Quam temere in vitiis, legem sancimus haerentia. Qui ipsorum lingua Celtae, nostra Galli appellantur. Quis aute iure reprehenderit in voluptate velit esse. Fabio vel iudice vincam, sunt in culpa qui officia. Cras mattis iudicium purus sit amet fermentum.
-Quo usque tandem abutere, Catilina, patientia nostra? Gallia est omnis divisa in partes tres, quarum. Quam diu etiam furor iste tuus nos eludet? Quid securi etiam tamquam eu fugiat nulla pariatur. Curabitur blandit tempus ardua ridiculus sed magna.
+Quo usque tandem abutere, Catilina, patientia nostra? Gallia est omnis divisa in partes tres, quarum. Quam diu etiam furor iste tuus nos eludet? Quid securi etiam tamquam eu fugiat nulla pariatur. Curabitur blandit tempus ardua ridiculous sed magna.
 Magna pars studiorum, prodita quaerimus. Cum ceteris in veneratione tui montes, nascetur mus. Morbi odio eros, volutpat ut pharetra vitae, lobortis sed nibh. Plura mihi bona sunt, inclinet, amari petere vellent. Idque Caesaris facere voluntate liceret: sese habere. Tu quoque, Brute, fili mi, nihil timor populi, nihil!
 Tityre, tu patulae recubans sub tegmine fagi dolor. Inmensae subtilitatis, obscuris et malesuada fames. Quae vero auctorem tractata ab fiducia dicuntur.
 Cum sociis natoque penatibus et magnis dis parturient. Phasellus laoreet lorem vel dolor tempus vehicula. Vivamus sagittis lacus vel augue laoreet rutrum faucibus. Integer legentibus erat a ante historiarum dapibus.
 Quisque ut dolor gravida, placerat libero vel, euismod. Quam temere in vitiis, legem sancimus haerentia. Qui ipsorum lingua Celtae, nostra Galli appellantur. Quis aute iure reprehenderit in voluptate velit esse. Fabio vel iudice vincam, sunt in culpa qui officia. Cras mattis iudicium purus sit amet fermentum.
-Quo usque tandem abutere, Catilina, patientia nostra? Gallia est omnis divisa in partes tres, quarum. Quam diu etiam furor iste tuus nos eludet? Quid securi etiam tamquam eu fugiat nulla pariatur. Curabitur blandit tempus ardua ridiculus sed magna.
+Quo usque tandem abutere, Catilina, patientia nostra? Gallia est omnis divisa in partes tres, quarum. Quam diu etiam furor iste tuus nos eludet? Quid securi etiam tamquam eu fugiat nulla pariatur. Curabitur blandit tempus ardua ridiculous sed magna.
 Magna pars studiorum, prodita quaerimus. Cum ceteris in veneratione tui montes, nascetur mus. Morbi odio eros, volutpat ut pharetra vitae, lobortis sed nibh. Plura mihi bona sunt, inclinet, amari petere vellent. Idque Caesaris facere voluntate liceret: sese habere. Tu quoque, Brute, fili mi, nihil timor populi, nihil!
 Tityre, tu patulae recubans sub tegmine fagi dolor. Inmensae subtilitatis, obscuris et malesuada fames. Quae vero auctorem tractata ab fiducia dicuntur.
 Cum sociis natoque penatibus et magnis dis parturient. Phasellus laoreet lorem vel dolor tempus vehicula. Vivamus sagittis lacus vel augue laoreet rutrum faucibus. Integer legentibus erat a ante historiarum dapibus.
 Quisque ut dolor gravida, placerat libero vel, euismod. Quam temere in vitiis, legem sancimus haerentia. Qui ipsorum lingua Celtae, nostra Galli appellantur. Quis aute iure reprehenderit in voluptate velit esse. Fabio vel iudice vincam, sunt in culpa qui officia. Cras mattis iudicium purus sit amet fermentum.
-Quo usque tandem abutere, Catilina, patientia nostra? Gallia est omnis divisa in partes tres, quarum. Quam diu etiam furor iste tuus nos eludet? Quid securi etiam tamquam eu fugiat nulla pariatur. Curabitur blandit tempus ardua ridiculus sed magna.
+Quo usque tandem abutere, Catilina, patientia nostra? Gallia est omnis divisa in partes tres, quarum. Quam diu etiam furor iste tuus nos eludet? Quid securi etiam tamquam eu fugiat nulla pariatur. Curabitur blandit tempus ardua ridiculous sed magna.
 Magna pars studiorum, prodita quaerimus. Cum ceteris in veneratione tui montes, nascetur mus. Morbi odio eros, volutpat ut pharetra vitae, lobortis sed nibh. Plura mihi bona sunt, inclinet, amari petere vellent. Idque Caesaris facere voluntate liceret: sese habere. Tu quoque, Brute, fili mi, nihil timor populi, nihil!
 Tityre, tu patulae recubans sub tegmine fagi dolor. Inmensae subtilitatis, obscuris et malesuada fames. Quae vero auctorem tractata ab fiducia dicuntur.
 ` + longLine + "\n"

core/fx/fn.go

@@ -49,7 +49,7 @@ func From(generate GenerateFunc) Stream {
 	return Range(source)
 }
 
-// Just converts the given arbitary items to a Stream.
+// Just converts the given arbitrary items to a Stream.
 func Just(items ...interface{}) Stream {
 	source := make(chan interface{}, len(items))
 	for _, item := range items {
@@ -68,6 +68,7 @@ func Range(source <-chan interface{}) Stream {
 }
 
 // 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 {
 	if n < 0 {
 		n = 0
@@ -84,6 +85,14 @@ func (p Stream) Buffer(n int) Stream {
 	return Range(source)
 }
 
+// Count counts the number of elements in the result.
+func (p Stream) Count() (count int) {
+	for range p.source {
+		count++
+	}
+	return
+}
+
 // Distinct removes the duplicated items base on the given KeyFunc.
 func (p Stream) Distinct(fn KeyFunc) Stream {
 	source := make(chan interface{})
@@ -151,6 +160,10 @@ func (p Stream) Group(fn KeyFunc) Stream {
 }
 
 func (p Stream) Head(n int64) Stream {
+	if n < 1 {
+		panic("n must be greater than 0")
+	}
+
 	source := make(chan interface{})
 
 	go func() {
@@ -195,7 +208,7 @@ func (p Stream) Merge() Stream {
 	return Range(source)
 }
 
-// Parallel applies the given ParallenFunc to each item concurrently with given number of workers.
+// 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{}) {
 		fn(item)
@@ -235,7 +248,37 @@ func (p Stream) Sort(less LessFunc) Stream {
 	return Just(items...)
 }
 
+// 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 {
+	if n < 1 {
+		panic("n should be greater than 0")
+	}
+
+	source := make(chan interface{})
+	go func() {
+		var chunk []interface{}
+		for item := range p.source {
+			chunk = append(chunk, item)
+			if len(chunk) == n {
+				source <- chunk
+				chunk = nil
+			}
+		}
+		if chunk != nil {
+			source <- chunk
+		}
+		close(source)
+	}()
+
+	return Range(source)
+}
+
 func (p Stream) Tail(n int64) Stream {
+	if n < 1 {
+		panic("n should be greater than 0")
+	}
+
 	source := make(chan interface{})
 
 	go func() {

core/fx/fn_test.go

@@ -49,6 +49,36 @@ func TestBufferNegative(t *testing.T) {
 	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{}) {
@@ -139,6 +169,14 @@ func TestHead(t *testing.T) {
 	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) {
@@ -245,6 +283,22 @@ func TestSort(t *testing.T) {
 	})
 }
 
+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) {
@@ -256,6 +310,14 @@ func TestTail(t *testing.T) {
 	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{}) {

core/iox/pipe.go

@@ -0,0 +1,23 @@
+package iox
+
+import "os"
+
+// RedirectInOut redirects stdin to r, stdout to w, and callers need to call restore afterwards.
+func RedirectInOut() (restore func(), err error) {
+	var r, w *os.File
+	r, w, err = os.Pipe()
+	if err != nil {
+		return
+	}
+
+	ow := os.Stdout
+	os.Stdout = w
+	or := os.Stdin
+	os.Stdin = r
+	restore = func() {
+		os.Stdin = or
+		os.Stdout = ow
+	}
+
+	return
+}

core/iox/pipe_test.go

@@ -0,0 +1,13 @@
+package iox
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRedirectInOut(t *testing.T) {
+	restore, err := RedirectInOut()
+	assert.Nil(t, err)
+	defer restore()
+}

core/limit/periodlimit_test.go

@@ -3,9 +3,10 @@ package limit
 import (
 	"testing"
 
-	"github.com/alicebob/miniredis"
+	"github.com/alicebob/miniredis/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/tal-tech/go-zero/core/stores/redis"
+	"github.com/tal-tech/go-zero/core/stores/redis/redistest"
 )
 
 func TestPeriodLimit_Take(t *testing.T) {
@@ -33,16 +34,16 @@ func TestPeriodLimit_RedisUnavailable(t *testing.T) {
 }
 
 func testPeriodLimit(t *testing.T, opts ...LimitOption) {
-	s, err := miniredis.Run()
+	store, clean, err := redistest.CreateRedis()
 	assert.Nil(t, err)
-	defer s.Close()
+	defer clean()
 
 	const (
 		seconds = 1
 		total   = 100
 		quota   = 5
 	)
-	l := NewPeriodLimit(seconds, quota, redis.NewRedis(s.Addr(), redis.NodeType), "periodlimit", opts...)
+	l := NewPeriodLimit(seconds, quota, store, "periodlimit", opts...)
 	var allowed, hitQuota, overQuota int
 	for i := 0; i < total; i++ {
 		val, err := l.Take("first")

core/limit/tokenlimit.go

@@ -153,13 +153,10 @@ func (lim *TokenLimiter) waitForRedis() {
 		lim.rescueLock.Unlock()
 	}()
 
-	for {
+	for range ticker.C {
-		select {
+		if lim.store.Ping() {
-		case <-ticker.C:
+			atomic.StoreUint32(&lim.redisAlive, 1)
-			if lim.store.Ping() {
+			return
-				atomic.StoreUint32(&lim.redisAlive, 1)
-				return
-			}
 		}
 	}
 }

core/limit/tokenlimit_test.go

@@ -4,10 +4,11 @@ import (
 	"testing"
 	"time"
 
-	"github.com/alicebob/miniredis"
+	"github.com/alicebob/miniredis/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/tal-tech/go-zero/core/logx"
 	"github.com/tal-tech/go-zero/core/stores/redis"
+	"github.com/tal-tech/go-zero/core/stores/redis/redistest"
 )
 
 func init() {
@@ -44,16 +45,16 @@ func TestTokenLimit_Rescue(t *testing.T) {
 }
 
 func TestTokenLimit_Take(t *testing.T) {
-	s, err := miniredis.Run()
+	store, clean, err := redistest.CreateRedis()
 	assert.Nil(t, err)
-	defer s.Close()
+	defer clean()
 
 	const (
 		total = 100
 		rate  = 5
 		burst = 10
 	)
-	l := NewTokenLimiter(rate, burst, redis.NewRedis(s.Addr(), redis.NodeType), "tokenlimit")
+	l := NewTokenLimiter(rate, burst, store, "tokenlimit")
 	var allowed int
 	for i := 0; i < total; i++ {
 		time.Sleep(time.Second / time.Duration(total))
@@ -66,16 +67,16 @@ func TestTokenLimit_Take(t *testing.T) {
 }
 
 func TestTokenLimit_TakeBurst(t *testing.T) {
-	s, err := miniredis.Run()
+	store, clean, err := redistest.CreateRedis()
 	assert.Nil(t, err)
-	defer s.Close()
+	defer clean()
 
 	const (
 		total = 100
 		rate  = 5
 		burst = 10
 	)
-	l := NewTokenLimiter(rate, burst, redis.NewRedis(s.Addr(), redis.NodeType), "tokenlimit")
+	l := NewTokenLimiter(rate, burst, store, "tokenlimit")
 	var allowed int
 	for i := 0; i < total; i++ {
 		if l.Allow() {

core/load/adaptiveshedder_test.go

@@ -135,6 +135,7 @@ func TestAdaptiveShedderShouldDrop(t *testing.T) {
 		passCounter:     passCounter,
 		rtCounter:       rtCounter,
 		windows:         buckets,
+		dropTime:        syncx.NewAtomicDuration(),
 		droppedRecently: syncx.NewAtomicBool(),
 	}
 	// cpu >=  800, inflight < maxPass
@@ -160,6 +161,40 @@ func TestAdaptiveShedderShouldDrop(t *testing.T) {
 	}
 	shedder.avgFlying = 80
 	assert.False(t, shedder.shouldDrop())
+
+	// cpu >=  800, inflight < maxPass
+	systemOverloadChecker = func(int64) bool {
+		return true
+	}
+	shedder.avgFlying = 80
+	shedder.flying = 80
+	_, err := shedder.Allow()
+	assert.NotNil(t, err)
+}
+
+func TestAdaptiveShedderStillHot(t *testing.T) {
+	logx.Disable()
+	passCounter := newRollingWindow()
+	rtCounter := newRollingWindow()
+	for i := 0; i < 10; i++ {
+		if i > 0 {
+			time.Sleep(bucketDuration)
+		}
+		passCounter.Add(float64((i + 1) * 100))
+		for j := i*10 + 1; j <= i*10+10; j++ {
+			rtCounter.Add(float64(j))
+		}
+	}
+	shedder := &adaptiveShedder{
+		passCounter:     passCounter,
+		rtCounter:       rtCounter,
+		windows:         buckets,
+		dropTime:        syncx.NewAtomicDuration(),
+		droppedRecently: syncx.ForAtomicBool(true),
+	}
+	assert.False(t, shedder.stillHot())
+	shedder.dropTime.Set(-coolOffDuration * 2)
+	assert.False(t, shedder.stillHot())
 }
 
 func BenchmarkAdaptiveShedder_Allow(b *testing.B) {

core/load/sheddergroup_test.go

@@ -13,3 +13,8 @@ func TestGroup(t *testing.T) {
 		assert.NotNil(t, limiter)
 	})
 }
+
+func TestShedderClose(t *testing.T) {
+	var nop nopCloser
+	assert.Nil(t, nop.Close())
+}

core/logx/durationlogger.go

@@ -8,55 +8,60 @@ import (
 	"github.com/tal-tech/go-zero/core/timex"
 )
 
-const customCallerDepth = 3
+const durationCallerDepth = 3
 
-type customLog logEntry
+type durationLogger logEntry
 
 func WithDuration(d time.Duration) Logger {
-	return customLog{
+	return &durationLogger{
 		Duration: timex.ReprOfDuration(d),
 	}
 }
 
-func (l customLog) Error(v ...interface{}) {
+func (l *durationLogger) Error(v ...interface{}) {
 	if shouldLog(ErrorLevel) {
-		l.write(errorLog, levelError, formatWithCaller(fmt.Sprint(v...), customCallerDepth))
+		l.write(errorLog, levelError, formatWithCaller(fmt.Sprint(v...), durationCallerDepth))
 	}
 }
 
-func (l customLog) Errorf(format string, v ...interface{}) {
+func (l *durationLogger) Errorf(format string, v ...interface{}) {
 	if shouldLog(ErrorLevel) {
-		l.write(errorLog, levelError, formatWithCaller(fmt.Sprintf(format, v...), customCallerDepth))
+		l.write(errorLog, levelError, formatWithCaller(fmt.Sprintf(format, v...), durationCallerDepth))
 	}
 }
 
-func (l customLog) Info(v ...interface{}) {
+func (l *durationLogger) Info(v ...interface{}) {
 	if shouldLog(InfoLevel) {
 		l.write(infoLog, levelInfo, fmt.Sprint(v...))
 	}
 }
 
-func (l customLog) Infof(format string, v ...interface{}) {
+func (l *durationLogger) Infof(format string, v ...interface{}) {
 	if shouldLog(InfoLevel) {
 		l.write(infoLog, levelInfo, fmt.Sprintf(format, v...))
 	}
 }
 
-func (l customLog) Slow(v ...interface{}) {
+func (l *durationLogger) Slow(v ...interface{}) {
 	if shouldLog(ErrorLevel) {
 		l.write(slowLog, levelSlow, fmt.Sprint(v...))
 	}
 }
 
-func (l customLog) Slowf(format string, v ...interface{}) {
+func (l *durationLogger) Slowf(format string, v ...interface{}) {
 	if shouldLog(ErrorLevel) {
 		l.write(slowLog, levelSlow, fmt.Sprintf(format, v...))
 	}
 }
 
-func (l customLog) write(writer io.Writer, level, content string) {
+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))
+	outputJson(writer, logEntry(*l))
 }

core/logx/durationlogger_test.go

@@ -0,0 +1,52 @@
+package logx
+
+import (
+	"log"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWithDurationError(t *testing.T) {
+	var builder strings.Builder
+	log.SetOutput(&builder)
+	WithDuration(time.Second).Error("foo")
+	assert.True(t, strings.Contains(builder.String(), "duration"), builder.String())
+}
+
+func TestWithDurationErrorf(t *testing.T) {
+	var builder strings.Builder
+	log.SetOutput(&builder)
+	WithDuration(time.Second).Errorf("foo")
+	assert.True(t, strings.Contains(builder.String(), "duration"), builder.String())
+}
+
+func TestWithDurationInfo(t *testing.T) {
+	var builder strings.Builder
+	log.SetOutput(&builder)
+	WithDuration(time.Second).Info("foo")
+	assert.True(t, strings.Contains(builder.String(), "duration"), builder.String())
+}
+
+func TestWithDurationInfof(t *testing.T) {
+	var builder strings.Builder
+	log.SetOutput(&builder)
+	WithDuration(time.Second).Infof("foo")
+	assert.True(t, strings.Contains(builder.String(), "duration"), builder.String())
+}
+
+func TestWithDurationSlow(t *testing.T) {
+	var builder strings.Builder
+	log.SetOutput(&builder)
+	WithDuration(time.Second).Slow("foo")
+	assert.True(t, strings.Contains(builder.String(), "duration"), builder.String())
+}
+
+func TestWithDurationSlowf(t *testing.T) {
+	var builder strings.Builder
+	log.SetOutput(&builder)
+	WithDuration(time.Second).WithDuration(time.Hour).Slowf("foo")
+	assert.True(t, strings.Contains(builder.String(), "duration"), builder.String())
+}

core/logx/logs.go

@@ -15,6 +15,7 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/tal-tech/go-zero/core/iox"
 	"github.com/tal-tech/go-zero/core/sysx"
@@ -42,6 +43,7 @@ const (
 	consoleMode = "console"
 	volumeMode  = "volume"
 
+	levelAlert  = "alert"
 	levelInfo   = "info"
 	levelError  = "error"
 	levelSevere = "severe"
@@ -96,6 +98,7 @@ type (
 		Infof(string, ...interface{})
 		Slow(...interface{})
 		Slowf(string, ...interface{})
+		WithDuration(time.Duration) Logger
 	}
 )
 
@@ -119,6 +122,10 @@ func SetUp(c LogConf) error {
 	}
 }
 
+func Alert(v string) {
+	output(errorLog, levelAlert, v)
+}
+
 func Close() error {
 	if writeConsole {
 		return nil
@@ -213,6 +220,7 @@ func Infof(format string, v ...interface{}) {
 func Must(err error) {
 	if err != nil {
 		msg := formatWithCaller(err.Error(), 3)
+		log.Print(msg)
 		output(severeLog, levelFatal, msg)
 		os.Exit(1)
 	}

core/logx/logs_test.go

@@ -6,8 +6,10 @@ import (
 	"io"
 	"io/ioutil"
 	"log"
+	"os"
 	"runtime"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -21,10 +23,13 @@ var (
 )
 
 type mockWriter struct {
+	lock    sync.Mutex
 	builder strings.Builder
 }
 
 func (mw *mockWriter) Write(data []byte) (int, error) {
+	mw.lock.Lock()
+	defer mw.lock.Unlock()
 	return mw.builder.Write(data)
 }
 
@@ -32,12 +37,22 @@ func (mw *mockWriter) Close() error {
 	return nil
 }
 
+func (mw *mockWriter) Contains(text string) bool {
+	mw.lock.Lock()
+	defer mw.lock.Unlock()
+	return strings.Contains(mw.builder.String(), text)
+}
+
 func (mw *mockWriter) Reset() {
+	mw.lock.Lock()
+	defer mw.lock.Unlock()
 	mw.builder.Reset()
 }
 
-func (mw *mockWriter) Contains(text string) bool {
+func (mw *mockWriter) String() string {
-	return strings.Contains(mw.builder.String(), text)
+	mw.lock.Lock()
+	defer mw.lock.Unlock()
+	return mw.builder.String()
 }
 
 func TestFileLineFileMode(t *testing.T) {
@@ -69,6 +84,14 @@ func TestFileLineConsoleMode(t *testing.T) {
 	assert.True(t, writer.Contains(fmt.Sprintf("%s:%d", file, line+1)))
 }
 
+func TestStructedLogAlert(t *testing.T) {
+	doTestStructedLog(t, levelAlert, func(writer io.WriteCloser) {
+		errorLog = writer
+	}, func(v ...interface{}) {
+		Alert(fmt.Sprint(v...))
+	})
+}
+
 func TestStructedLogInfo(t *testing.T) {
 	doTestStructedLog(t, levelInfo, func(writer io.WriteCloser) {
 		infoLog = writer
@@ -85,6 +108,46 @@ func TestStructedLogSlow(t *testing.T) {
 	})
 }
 
+func TestStructedLogSlowf(t *testing.T) {
+	doTestStructedLog(t, levelSlow, func(writer io.WriteCloser) {
+		slowLog = writer
+	}, func(v ...interface{}) {
+		Slowf(fmt.Sprint(v...))
+	})
+}
+
+func TestStructedLogStat(t *testing.T) {
+	doTestStructedLog(t, levelStat, func(writer io.WriteCloser) {
+		statLog = writer
+	}, func(v ...interface{}) {
+		Stat(v...)
+	})
+}
+
+func TestStructedLogStatf(t *testing.T) {
+	doTestStructedLog(t, levelStat, func(writer io.WriteCloser) {
+		statLog = writer
+	}, func(v ...interface{}) {
+		Statf(fmt.Sprint(v...))
+	})
+}
+
+func TestStructedLogSevere(t *testing.T) {
+	doTestStructedLog(t, levelSevere, func(writer io.WriteCloser) {
+		severeLog = writer
+	}, func(v ...interface{}) {
+		Severe(v...)
+	})
+}
+
+func TestStructedLogSeveref(t *testing.T) {
+	doTestStructedLog(t, levelSevere, func(writer io.WriteCloser) {
+		severeLog = writer
+	}, func(v ...interface{}) {
+		Severef(fmt.Sprint(v...))
+	})
+}
+
 func TestStructedLogWithDuration(t *testing.T) {
 	const message = "hello there"
 	writer := new(mockWriter)
@@ -135,6 +198,68 @@ func TestMustNil(t *testing.T) {
 	Must(nil)
 }
 
+func TestSetup(t *testing.T) {
+	MustSetup(LogConf{
+		ServiceName: "any",
+		Mode:        "console",
+	})
+	MustSetup(LogConf{
+		ServiceName: "any",
+		Mode:        "file",
+		Path:        os.TempDir(),
+	})
+	MustSetup(LogConf{
+		ServiceName: "any",
+		Mode:        "volume",
+		Path:        os.TempDir(),
+	})
+	assert.NotNil(t, setupWithVolume(LogConf{}))
+	assert.NotNil(t, setupWithFiles(LogConf{}))
+	assert.Nil(t, setupWithFiles(LogConf{
+		ServiceName: "any",
+		Path:        os.TempDir(),
+		Compress:    true,
+		KeepDays:    1,
+	}))
+	setupLogLevel(LogConf{
+		Level: levelInfo,
+	})
+	setupLogLevel(LogConf{
+		Level: levelError,
+	})
+	setupLogLevel(LogConf{
+		Level: levelSevere,
+	})
+	_, err := createOutput("")
+	assert.NotNil(t, err)
+	Disable()
+}
+
+func TestDisable(t *testing.T) {
+	Disable()
+
+	var opt logOptions
+	WithKeepDays(1)(&opt)
+	WithGzip()(&opt)
+	assert.Nil(t, Close())
+	writeConsole = false
+	assert.Nil(t, Close())
+}
+
+func TestWithGzip(t *testing.T) {
+	fn := WithGzip()
+	var opt logOptions
+	fn(&opt)
+	assert.True(t, opt.gzipEnabled)
+}
+
+func TestWithKeepDays(t *testing.T) {
+	fn := WithKeepDays(1)
+	var opt logOptions
+	fn(&opt)
+	assert.Equal(t, 1, opt.keepDays)
+}
+
 func BenchmarkCopyByteSliceAppend(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		var buf []byte
@@ -232,7 +357,7 @@ func doTestStructedLog(t *testing.T, level string, setup func(writer io.WriteClo
 		t.Error(err)
 	}
 	assert.Equal(t, level, entry.Level)
-	assert.Equal(t, message, entry.Content)
+	assert.True(t, strings.Contains(entry.Content, message))
 }
 
 func testSetLevelTwiceWithMode(t *testing.T, mode string) {
@@ -252,4 +377,10 @@ func testSetLevelTwiceWithMode(t *testing.T, mode string) {
 	atomic.StoreUint32(&initialized, 1)
 	Info(message)
 	assert.Equal(t, 0, writer.builder.Len())
+	Infof(message)
+	assert.Equal(t, 0, writer.builder.Len())
+	ErrorStack(message)
+	assert.Equal(t, 0, writer.builder.Len())
+	ErrorStackf(message)
+	assert.Equal(t, 0, writer.builder.Len())
 }

core/logx/rotatelogger.go

@@ -192,14 +192,16 @@ func (l *RotateLogger) init() error {
 }
 
 func (l *RotateLogger) maybeCompressFile(file string) {
-	if l.compress {
+	if !l.compress {
-		defer func() {
+		return
-			if r := recover(); r != nil {
-				ErrorStack(r)
-			}
-		}()
-		compressLogFile(file)
 	}
+
+	defer func() {
+		if r := recover(); r != nil {
+			ErrorStack(r)
+		}
+	}()
+	compressLogFile(file)
 }
 
 func (l *RotateLogger) maybeDeleteOutdatedFiles() {

core/logx/rotatelogger_test.go

@@ -0,0 +1,119 @@
+package logx
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/fs"
+)
+
+func TestDailyRotateRuleMarkRotated(t *testing.T) {
+	var rule DailyRotateRule
+	rule.MarkRotated()
+	assert.Equal(t, getNowDate(), rule.rotatedTime)
+}
+
+func TestDailyRotateRuleOutdatedFiles(t *testing.T) {
+	var rule DailyRotateRule
+	assert.Empty(t, rule.OutdatedFiles())
+	rule.days = 1
+	assert.Empty(t, rule.OutdatedFiles())
+	rule.gzip = true
+	assert.Empty(t, rule.OutdatedFiles())
+}
+
+func TestDailyRotateRuleShallRotate(t *testing.T) {
+	var rule DailyRotateRule
+	rule.rotatedTime = time.Now().Add(time.Hour * 24).Format(dateFormat)
+	assert.True(t, rule.ShallRotate())
+}
+
+func TestRotateLoggerClose(t *testing.T) {
+	filename, err := fs.TempFilenameWithText("foo")
+	assert.Nil(t, err)
+	if len(filename) > 0 {
+		defer os.Remove(filename)
+	}
+	logger, err := NewLogger(filename, new(DailyRotateRule), false)
+	assert.Nil(t, err)
+	assert.Nil(t, logger.Close())
+}
+
+func TestRotateLoggerGetBackupFilename(t *testing.T) {
+	filename, err := fs.TempFilenameWithText("foo")
+	assert.Nil(t, err)
+	if len(filename) > 0 {
+		defer os.Remove(filename)
+	}
+	logger, err := NewLogger(filename, new(DailyRotateRule), false)
+	assert.Nil(t, err)
+	assert.True(t, len(logger.getBackupFilename()) > 0)
+	logger.backup = ""
+	assert.True(t, len(logger.getBackupFilename()) > 0)
+}
+
+func TestRotateLoggerMayCompressFile(t *testing.T) {
+	filename, err := fs.TempFilenameWithText("foo")
+	assert.Nil(t, err)
+	if len(filename) > 0 {
+		defer os.Remove(filename)
+	}
+	logger, err := NewLogger(filename, new(DailyRotateRule), false)
+	assert.Nil(t, err)
+	logger.maybeCompressFile(filename)
+	_, err = os.Stat(filename)
+	assert.Nil(t, err)
+}
+
+func TestRotateLoggerMayCompressFileTrue(t *testing.T) {
+	filename, err := fs.TempFilenameWithText("foo")
+	assert.Nil(t, err)
+	logger, err := NewLogger(filename, new(DailyRotateRule), true)
+	assert.Nil(t, err)
+	if len(filename) > 0 {
+		defer func() {
+			os.Remove(filename)
+			os.Remove(filepath.Base(logger.getBackupFilename()) + ".gz")
+		}()
+	}
+	logger.maybeCompressFile(filename)
+	_, err = os.Stat(filename)
+	assert.NotNil(t, err)
+}
+
+func TestRotateLoggerRotate(t *testing.T) {
+	filename, err := fs.TempFilenameWithText("foo")
+	assert.Nil(t, err)
+	logger, err := NewLogger(filename, new(DailyRotateRule), true)
+	assert.Nil(t, err)
+	if len(filename) > 0 {
+		defer func() {
+			os.Remove(filename)
+			os.Remove(logger.getBackupFilename())
+			os.Remove(filepath.Base(logger.getBackupFilename()) + ".gz")
+		}()
+	}
+	err = logger.rotate()
+	assert.Nil(t, err)
+}
+
+func TestRotateLoggerWrite(t *testing.T) {
+	filename, err := fs.TempFilenameWithText("foo")
+	assert.Nil(t, err)
+	rule := new(DailyRotateRule)
+	logger, err := NewLogger(filename, rule, true)
+	assert.Nil(t, err)
+	if len(filename) > 0 {
+		defer func() {
+			os.Remove(filename)
+			os.Remove(logger.getBackupFilename())
+			os.Remove(filepath.Base(logger.getBackupFilename()) + ".gz")
+		}()
+	}
+	logger.write([]byte(`foo`))
+	rule.rotatedTime = time.Now().Add(-time.Hour * 24).Format(dateFormat)
+	logger.write([]byte(`bar`))
+}

core/logx/syslog_test.go

@@ -15,7 +15,7 @@ const testlog = "Stay hungry, stay foolish."
 func TestCollectSysLog(t *testing.T) {
 	CollectSysLog()
 	content := getContent(captureOutput(func() {
-		log.Printf(testlog)
+		log.Print(testlog)
 	}))
 	assert.True(t, strings.Contains(content, testlog))
 }
@@ -33,10 +33,10 @@ func captureOutput(f func()) string {
 	writer := new(mockWriter)
 	infoLog = writer
 
-	prevLevel := logLevel
+	prevLevel := atomic.LoadUint32(&logLevel)
-	logLevel = InfoLevel
+	SetLevel(InfoLevel)
 	f()
-	logLevel = prevLevel
+	SetLevel(prevLevel)
 
 	return writer.builder.String()
 }

core/logx/tracelog_test.go

@@ -1,49 +0,0 @@
-package logx
-
-import (
-	"context"
-	"strings"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/trace/tracespec"
-)
-
-const (
-	mockTraceId = "mock-trace-id"
-	mockSpanId  = "mock-span-id"
-)
-
-var mock tracespec.Trace = new(mockTrace)
-
-func TestTraceLog(t *testing.T) {
-	var buf strings.Builder
-	ctx := context.WithValue(context.Background(), tracespec.TracingKey, mock)
-	WithContext(ctx).(tracingEntry).write(&buf, levelInfo, testlog)
-	assert.True(t, strings.Contains(buf.String(), mockTraceId))
-	assert.True(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) {
-}

core/logx/tracelogger.go

@@ -4,54 +4,61 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"time"
 
+	"github.com/tal-tech/go-zero/core/timex"
 	"github.com/tal-tech/go-zero/core/trace/tracespec"
 )
 
-type tracingEntry struct {
+type traceLogger struct {
 	logEntry
-	Trace string          `json:"trace,omitempty"`
+	Trace string `json:"trace,omitempty"`
-	Span  string          `json:"span,omitempty"`
+	Span  string `json:"span,omitempty"`
-	ctx   context.Context `json:"-"`
+	ctx   context.Context
 }
 
-func (l tracingEntry) Error(v ...interface{}) {
+func (l *traceLogger) Error(v ...interface{}) {
 	if shouldLog(ErrorLevel) {
-		l.write(errorLog, levelError, formatWithCaller(fmt.Sprint(v...), customCallerDepth))
+		l.write(errorLog, levelError, formatWithCaller(fmt.Sprint(v...), durationCallerDepth))
 	}
 }
 
-func (l tracingEntry) Errorf(format string, v ...interface{}) {
+func (l *traceLogger) Errorf(format string, v ...interface{}) {
 	if shouldLog(ErrorLevel) {
-		l.write(errorLog, levelError, formatWithCaller(fmt.Sprintf(format, v...), customCallerDepth))
+		l.write(errorLog, levelError, formatWithCaller(fmt.Sprintf(format, v...), durationCallerDepth))
 	}
 }
 
-func (l tracingEntry) Info(v ...interface{}) {
+func (l *traceLogger) Info(v ...interface{}) {
 	if shouldLog(InfoLevel) {
 		l.write(infoLog, levelInfo, fmt.Sprint(v...))
 	}
 }
 
-func (l tracingEntry) Infof(format string, v ...interface{}) {
+func (l *traceLogger) Infof(format string, v ...interface{}) {
 	if shouldLog(InfoLevel) {
 		l.write(infoLog, levelInfo, fmt.Sprintf(format, v...))
 	}
 }
 
-func (l tracingEntry) Slow(v ...interface{}) {
+func (l *traceLogger) Slow(v ...interface{}) {
 	if shouldLog(ErrorLevel) {
 		l.write(slowLog, levelSlow, fmt.Sprint(v...))
 	}
 }
 
-func (l tracingEntry) Slowf(format string, v ...interface{}) {
+func (l *traceLogger) Slowf(format string, v ...interface{}) {
 	if shouldLog(ErrorLevel) {
 		l.write(slowLog, levelSlow, fmt.Sprintf(format, v...))
 	}
 }
 
-func (l tracingEntry) write(writer io.Writer, level, content string) {
+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
@@ -61,7 +68,7 @@ func (l tracingEntry) write(writer io.Writer, level, content string) {
 }
 
 func WithContext(ctx context.Context) Logger {
-	return tracingEntry{
+	return &traceLogger{
 		ctx: ctx,
 	}
 }

core/logx/tracelogger_test.go

@@ -0,0 +1,116 @@
+package logx
+
+import (
+	"context"
+	"log"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/trace/tracespec"
+)
+
+const (
+	mockTraceId = "mock-trace-id"
+	mockSpanId  = "mock-span-id"
+)
+
+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)
+	WithContext(ctx).(*traceLogger).write(&buf, levelInfo, testlog)
+	assert.True(t, strings.Contains(buf.String(), mockTraceId))
+	assert.True(t, strings.Contains(buf.String(), mockSpanId))
+}
+
+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)
+	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))
+	buf.Reset()
+	l.WithDuration(time.Second).Errorf(testlog)
+	assert.True(t, strings.Contains(buf.String(), mockTraceId))
+	assert.True(t, strings.Contains(buf.String(), mockSpanId))
+}
+
+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)
+	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))
+	buf.Reset()
+	l.WithDuration(time.Second).Infof(testlog)
+	assert.True(t, strings.Contains(buf.String(), mockTraceId))
+	assert.True(t, strings.Contains(buf.String(), mockSpanId))
+}
+
+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)
+	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))
+	buf.Reset()
+	l.WithDuration(time.Second).Slowf(testlog)
+	assert.True(t, strings.Contains(buf.String(), mockTraceId))
+	assert.True(t, strings.Contains(buf.String(), mockSpanId))
+}
+
+func TestTraceWithoutContext(t *testing.T) {
+	var buf mockWriter
+	atomic.StoreUint32(&initialized, 1)
+	infoLog = newLogWriter(log.New(&buf, "", flags))
+	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))
+	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) {
+}

core/mapping/fieldoptions_test.go

@@ -0,0 +1,31 @@
+package mapping
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type Bar struct {
+	Val string `json:"val"`
+}
+
+func TestFieldOptionOptionalDep(t *testing.T) {
+	var bar Bar
+	rt := reflect.TypeOf(bar)
+	for i := 0; i < rt.NumField(); i++ {
+		field := rt.Field(i)
+		val, opt, err := parseKeyAndOptions(jsonTagKey, field)
+		assert.Equal(t, "val", val)
+		assert.Nil(t, opt)
+		assert.Nil(t, err)
+	}
+
+	// check nil working
+	var o *fieldOptions
+	check := func(o *fieldOptions) {
+		assert.Equal(t, 0, len(o.optionalDep()))
+	}
+	check(o)
+}

core/mapping/unmarshaler.go

@@ -23,6 +23,7 @@ const (
 var (
 	errTypeMismatch     = errors.New("type mismatch")
 	errValueNotSettable = errors.New("value is not settable")
+	errValueNotStruct   = errors.New("value type is not struct")
 	keyUnmarshaler      = NewUnmarshaler(defaultKeyName)
 	cacheKeys           atomic.Value
 	cacheKeysLock       sync.Mutex
@@ -80,6 +81,10 @@ func (u *Unmarshaler) unmarshalWithFullName(m Valuer, v interface{}, fullName st
 	}
 
 	rte := reflect.TypeOf(v).Elem()
+	if rte.Kind() != reflect.Struct {
+		return errValueNotStruct
+	}
+
 	rve := rv.Elem()
 	numFields := rte.NumField()
 	for i := 0; i < numFields; i++ {
@@ -345,7 +350,7 @@ func (u *Unmarshaler) processNamedFieldWithValue(field reflect.StructField, valu
 			options := opts.options()
 			if len(options) > 0 {
 				if !stringx.Contains(options, mapValue.(string)) {
-					return fmt.Errorf(`error: value "%s" for field "%s" is not defined in opts "%v"`,
+					return fmt.Errorf(`error: value "%s" for field "%s" is not defined in options "%v"`,
 						mapValue, key, options)
 				}
 			}

core/mapping/unmarshaler_test.go

@@ -14,6 +14,13 @@ import (
 // so we only can test to 62 bits.
 const maxUintBitsToTest = 62
 
+func TestUnmarshalWithFullNameNotStruct(t *testing.T) {
+	var s map[string]interface{}
+	content := []byte(`{"name":"xiaoming"}`)
+	err := UnmarshalJsonBytes(content, &s)
+	assert.Equal(t, errValueNotStruct, err)
+}
+
 func TestUnmarshalWithoutTagName(t *testing.T) {
 	type inner struct {
 		Optional bool `key:",optional"`
@@ -2380,6 +2387,13 @@ func TestUnmarshalNestedMapSimpleTypeMatch(t *testing.T) {
 	assert.Equal(t, "1", c.Anything["id"])
 }
 
+func TestUnmarshalValuer(t *testing.T) {
+	unmarshaler := NewUnmarshaler(jsonTagKey)
+	var foo string
+	err := unmarshaler.UnmarshalValuer(nil, foo)
+	assert.NotNil(t, err)
+}
+
 func BenchmarkUnmarshalString(b *testing.B) {
 	type inner struct {
 		Value string `key:"value"`

core/mapping/utils.go

@@ -153,58 +153,57 @@ func doParseKeyAndOptions(field reflect.StructField, value string) (string, *fie
 	key := strings.TrimSpace(segments[0])
 	options := segments[1:]
 
-	if len(options) > 0 {
+	if len(options) == 0 {
-		var fieldOpts fieldOptions
+		return key, nil, nil
-
-		for _, segment := range options {
-			option := strings.TrimSpace(segment)
-			switch {
-			case option == stringOption:
-				fieldOpts.FromString = true
-			case strings.HasPrefix(option, optionalOption):
-				segs := strings.Split(option, equalToken)
-				switch len(segs) {
-				case 1:
-					fieldOpts.Optional = true
-				case 2:
-					fieldOpts.Optional = true
-					fieldOpts.OptionalDep = segs[1]
-				default:
-					return "", nil, fmt.Errorf("field %s has wrong optional", field.Name)
-				}
-			case option == optionalOption:
-				fieldOpts.Optional = true
-			case strings.HasPrefix(option, optionsOption):
-				segs := strings.Split(option, equalToken)
-				if len(segs) != 2 {
-					return "", nil, fmt.Errorf("field %s has wrong options", field.Name)
-				} else {
-					fieldOpts.Options = strings.Split(segs[1], optionSeparator)
-				}
-			case strings.HasPrefix(option, defaultOption):
-				segs := strings.Split(option, equalToken)
-				if len(segs) != 2 {
-					return "", nil, fmt.Errorf("field %s has wrong default option", field.Name)
-				} else {
-					fieldOpts.Default = strings.TrimSpace(segs[1])
-				}
-			case strings.HasPrefix(option, rangeOption):
-				segs := strings.Split(option, equalToken)
-				if len(segs) != 2 {
-					return "", nil, fmt.Errorf("field %s has wrong range", field.Name)
-				}
-				if nr, err := parseNumberRange(segs[1]); err != nil {
-					return "", nil, err
-				} else {
-					fieldOpts.Range = nr
-				}
-			}
-		}
-
-		return key, &fieldOpts, nil
 	}
 
-	return key, nil, nil
+	var fieldOpts fieldOptions
+	for _, segment := range options {
+		option := strings.TrimSpace(segment)
+		switch {
+		case option == stringOption:
+			fieldOpts.FromString = true
+		case strings.HasPrefix(option, optionalOption):
+			segs := strings.Split(option, equalToken)
+			switch len(segs) {
+			case 1:
+				fieldOpts.Optional = true
+			case 2:
+				fieldOpts.Optional = true
+				fieldOpts.OptionalDep = segs[1]
+			default:
+				return "", nil, fmt.Errorf("field %s has wrong optional", field.Name)
+			}
+		case option == optionalOption:
+			fieldOpts.Optional = true
+		case strings.HasPrefix(option, optionsOption):
+			segs := strings.Split(option, equalToken)
+			if len(segs) != 2 {
+				return "", nil, fmt.Errorf("field %s has wrong options", field.Name)
+			} else {
+				fieldOpts.Options = strings.Split(segs[1], optionSeparator)
+			}
+		case strings.HasPrefix(option, defaultOption):
+			segs := strings.Split(option, equalToken)
+			if len(segs) != 2 {
+				return "", nil, fmt.Errorf("field %s has wrong default option", field.Name)
+			} else {
+				fieldOpts.Default = strings.TrimSpace(segs[1])
+			}
+		case strings.HasPrefix(option, rangeOption):
+			segs := strings.Split(option, equalToken)
+			if len(segs) != 2 {
+				return "", nil, fmt.Errorf("field %s has wrong range", field.Name)
+			}
+			if nr, err := parseNumberRange(segs[1]); err != nil {
+				return "", nil, err
+			} else {
+				fieldOpts.Range = nr
+			}
+		}
+	}
+
+	return key, &fieldOpts, nil
 }
 
 func implicitValueRequiredStruct(tag string, tp reflect.Type) (bool, error) {

core/mathx/int_test.go

@@ -31,6 +31,7 @@ func TestMaxInt(t *testing.T) {
 	}
 
 	for _, each := range cases {
+		each := each
 		t.Run(stringx.Rand(), func(t *testing.T) {
 			actual := MaxInt(each.a, each.b)
 			assert.Equal(t, each.expect, actual)

core/mr/mapreduce.go

@@ -16,7 +16,10 @@ const (
 	minWorkers     = 1
 )
 
-var ErrCancelWithNil = errors.New("mapreduce cancelled with nil")
+var (
+	ErrCancelWithNil  = errors.New("mapreduce cancelled with nil")
+	ErrReduceNoOutput = errors.New("reduce not writing value")
+)
 
 type (
 	GenerateFunc    func(source chan<- interface{})
@@ -76,7 +79,7 @@ func Map(generate GenerateFunc, mapper MapFunc, opts ...Option) chan interface{}
 	collector := make(chan interface{}, options.workers)
 	done := syncx.NewDoneChan()
 
-	go mapDispatcher(mapper, source, collector, done.Done(), options.workers)
+	go executeMappers(mapper, source, collector, done.Done(), options.workers)
 
 	return collector
 }
@@ -93,7 +96,14 @@ func MapReduceWithSource(source <-chan interface{}, mapper MapperFunc, reducer R
 	collector := make(chan interface{}, options.workers)
 	done := syncx.NewDoneChan()
 	writer := newGuardedWriter(output, done.Done())
+	var closeOnce sync.Once
 	var retErr errorx.AtomicError
+	finish := func() {
+		closeOnce.Do(func() {
+			done.Close()
+			close(output)
+		})
+	}
 	cancel := once(func(err error) {
 		if err != nil {
 			retErr.Set(err)
@@ -102,19 +112,24 @@ func MapReduceWithSource(source <-chan interface{}, mapper MapperFunc, reducer R
 		}
 
 		drain(source)
-		done.Close()
+		finish()
-		close(output)
 	})
 
 	go func() {
 		defer func() {
 			if r := recover(); r != nil {
 				cancel(fmt.Errorf("%v", r))
+			} else {
+				finish()
 			}
 		}()
 		reducer(collector, writer, cancel)
+		drain(collector)
 	}()
-	go mapperDispatcher(mapper, source, collector, done.Done(), cancel, options.workers)
+
+	go executeMappers(func(item interface{}, w Writer) {
+		mapper(item, w, cancel)
+	}, source, collector, done.Done(), options.workers)
 
 	value, ok := <-output
 	if err := retErr.Load(); err != nil {
@@ -122,13 +137,14 @@ func MapReduceWithSource(source <-chan interface{}, mapper MapperFunc, reducer R
 	} else if ok {
 		return value, nil
 	} else {
-		return nil, nil
+		return nil, ErrReduceNoOutput
 	}
 }
 
 func MapReduceVoid(generator GenerateFunc, mapper MapperFunc, reducer VoidReducerFunc, opts ...Option) error {
 	_, err := MapReduce(generator, 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)
@@ -213,20 +229,6 @@ func executeMappers(mapper MapFunc, input <-chan interface{}, collector chan<- i
 	}
 }
 
-func mapDispatcher(mapper MapFunc, input <-chan interface{}, collector chan<- interface{},
-	done <-chan lang.PlaceholderType, workers int) {
-	executeMappers(func(item interface{}, writer Writer) {
-		mapper(item, writer)
-	}, input, collector, done, workers)
-}
-
-func mapperDispatcher(mapper MapperFunc, input <-chan interface{}, collector chan<- interface{},
-	done <-chan lang.PlaceholderType, cancel func(error), workers int) {
-	executeMappers(func(item interface{}, writer Writer) {
-		mapper(item, writer, cancel)
-	}, input, collector, done, workers)
-}
-
 func newOptions() *mapReduceOptions {
 	return &mapReduceOptions{
 		workers: defaultWorkers,

core/mr/mapreduce_test.go

@@ -378,6 +378,22 @@ func TestMapReduceVoidCancelWithRemains(t *testing.T) {
 	assert.True(t, done.True())
 }
 
+func TestMapReduceWithoutReducerWrite(t *testing.T) {
+	uids := []int{1, 2, 3}
+	res, err := MapReduce(func(source chan<- interface{}) {
+		for _, uid := range uids {
+			source <- uid
+		}
+	}, func(item interface{}, writer Writer, cancel func(error)) {
+		writer.Write(item)
+	}, func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
+		drain(pipe)
+		// not calling writer.Write(...), should not panic
+	})
+	assert.Equal(t, ErrReduceNoOutput, err)
+	assert.Nil(t, res)
+}
+
 func BenchmarkMapReduce(b *testing.B) {
 	b.ReportAllocs()
 

core/proc/goroutines_test.go

@@ -0,0 +1,16 @@
+package proc
+
+import (
+	"log"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDumpGoroutines(t *testing.T) {
+	var buf strings.Builder
+	log.SetOutput(&buf)
+	dumpGoroutines()
+	assert.True(t, strings.Contains(buf.String(), ".dump"))
+}

core/proc/profile.go

@@ -26,10 +26,6 @@ var started uint32
 
 // Profile represents an active profiling session.
 type Profile struct {
-	// path holds the base path where various profiling files are  written.
-	// If blank, the base path will be generated by ioutil.TempDir.
-	path string
-
 	// closers holds cleanup functions that run after each profile
 	closers []func()
 

core/proc/profile_test.go

@@ -0,0 +1,21 @@
+package proc
+
+import (
+	"log"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestProfile(t *testing.T) {
+	var buf strings.Builder
+	log.SetOutput(&buf)
+	profiler := StartProfile()
+	// start again should not work
+	assert.NotNil(t, StartProfile())
+	profiler.Stop()
+	// stop twice
+	profiler.Stop()
+	assert.True(t, strings.Contains(buf.String(), ".pprof"))
+}

core/proc/shutdown.go

@@ -32,7 +32,7 @@ func AddWrapUpListener(fn func()) (waitForCalled func()) {
 	return wrapUpListeners.addListener(fn)
 }
 
-func SetTimeoutToForceQuit(duration time.Duration) {
+func SetTimeToForceQuit(duration time.Duration) {
 	delayTimeBeforeForceQuit = duration
 }
 

core/prof/profilecenter_test.go

@@ -0,0 +1,16 @@
+package prof
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestReport(t *testing.T) {
+	once.Do(func() {})
+	assert.NotContains(t, generateReport(), "foo")
+	report("foo", time.Second)
+	assert.Contains(t, generateReport(), "foo")
+	report("foo", time.Second)
+}

core/prof/profiler_test.go

@@ -0,0 +1,23 @@
+package prof
+
+import (
+	"testing"
+
+	"github.com/tal-tech/go-zero/core/utils"
+)
+
+func TestProfiler(t *testing.T) {
+	EnableProfiling()
+	Start()
+	Report("foo", ProfilePoint{
+		ElapsedTimer: utils.NewElapsedTimer(),
+	})
+}
+
+func TestNullProfiler(t *testing.T) {
+	p := newNullProfiler()
+	p.Start()
+	p.Report("foo", ProfilePoint{
+		ElapsedTimer: utils.NewElapsedTimer(),
+	})
+}

core/queue/balancedpusher.go

@@ -9,24 +9,24 @@ import (
 
 var ErrNoAvailablePusher = errors.New("no available pusher")
 
-type BalancedQueuePusher struct {
+type BalancedPusher struct {
 	name    string
 	pushers []Pusher
 	index   uint64
 }
 
-func NewBalancedQueuePusher(pushers []Pusher) Pusher {
+func NewBalancedPusher(pushers []Pusher) Pusher {
-	return &BalancedQueuePusher{
+	return &BalancedPusher{
 		name:    generateName(pushers),
 		pushers: pushers,
 	}
 }
 
-func (pusher *BalancedQueuePusher) Name() string {
+func (pusher *BalancedPusher) Name() string {
 	return pusher.name
 }
 
-func (pusher *BalancedQueuePusher) Push(message string) error {
+func (pusher *BalancedPusher) Push(message string) error {
 	size := len(pusher.pushers)
 
 	for i := 0; i < size; i++ {

core/queue/balancedpusher_test.go

@@ -20,7 +20,7 @@ func TestBalancedQueuePusher(t *testing.T) {
 		mockedPushers = append(mockedPushers, p)
 	}
 
-	pusher := NewBalancedQueuePusher(pushers)
+	pusher := NewBalancedPusher(pushers)
 	assert.True(t, len(pusher.Name()) > 0)
 
 	for i := 0; i < numPushers*1000; i++ {
@@ -37,7 +37,7 @@ func TestBalancedQueuePusher(t *testing.T) {
 }
 
 func TestBalancedQueuePusher_NoAvailable(t *testing.T) {
-	pusher := NewBalancedQueuePusher(nil)
+	pusher := NewBalancedPusher(nil)
 	assert.True(t, len(pusher.Name()) == 0)
 	assert.Equal(t, ErrNoAvailablePusher, pusher.Push("item"))
 }

core/queue/multipusher.go

@@ -2,23 +2,23 @@ package queue
 
 import "github.com/tal-tech/go-zero/core/errorx"
 
-type MultiQueuePusher struct {
+type MultiPusher struct {
 	name    string
 	pushers []Pusher
 }
 
-func NewMultiQueuePusher(pushers []Pusher) Pusher {
+func NewMultiPusher(pushers []Pusher) Pusher {
-	return &MultiQueuePusher{
+	return &MultiPusher{
 		name:    generateName(pushers),
 		pushers: pushers,
 	}
 }
 
-func (pusher *MultiQueuePusher) Name() string {
+func (pusher *MultiPusher) Name() string {
 	return pusher.name
 }
 
-func (pusher *MultiQueuePusher) Push(message string) error {
+func (pusher *MultiPusher) Push(message string) error {
 	var batchError errorx.BatchError
 
 	for _, each := range pusher.pushers {

core/queue/multipusher_test.go

@@ -21,7 +21,7 @@ func TestMultiQueuePusher(t *testing.T) {
 		mockedPushers = append(mockedPushers, p)
 	}
 
-	pusher := NewMultiQueuePusher(pushers)
+	pusher := NewMultiPusher(pushers)
 	assert.True(t, len(pusher.Name()) > 0)
 
 	for i := 0; i < 1000; i++ {

core/service/servicegroup_test.go

@@ -27,7 +27,7 @@ func newMockedService(multiplier int) *mockedService {
 
 func (s *mockedService) Start() {
 	mutex.Lock()
-	number = number * s.multiplier
+	number *= s.multiplier
 	mutex.Unlock()
 	done <- struct{}{}
 	<-s.quit

core/stat/alert.go

@@ -11,10 +11,10 @@ import (
 	"time"
 
 	"github.com/tal-tech/go-zero/core/executors"
+	"github.com/tal-tech/go-zero/core/logx"
 	"github.com/tal-tech/go-zero/core/proc"
 	"github.com/tal-tech/go-zero/core/sysx"
 	"github.com/tal-tech/go-zero/core/timex"
-	"github.com/tal-tech/go-zero/core/utils"
 )
 
 const (
@@ -24,7 +24,7 @@ const (
 )
 
 var (
-	reporter     = utils.Report
+	reporter     = logx.Alert
 	lock         sync.RWMutex
 	lessExecutor = executors.NewLessExecutor(time.Minute * 5)
 	dropped      int32

core/stat/alert_test.go

@@ -0,0 +1,22 @@
+// +build linux
+
+package stat
+
+import (
+	"strconv"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestReport(t *testing.T) {
+	var count int32
+	SetReporter(func(s string) {
+		atomic.AddInt32(&count, 1)
+	})
+	for i := 0; i < 10; i++ {
+		Report(strconv.Itoa(i))
+	}
+	assert.Equal(t, int32(1), count)
+}

core/stat/internal/cpu_linux.go

@@ -22,19 +22,30 @@ var (
 	cores     uint64
 )
 
+// if /proc not present, ignore the cpu calcuation, like wsl linux
 func init() {
 	cpus, err := perCpuUsage()
-	logx.Must(err)
+	if err != nil {
-	cores = uint64(len(cpus))
+		logx.Error(err)
+		return
+	}
 
+	cores = uint64(len(cpus))
 	sets, err := cpuSets()
-	logx.Must(err)
+	if err != nil {
+		logx.Error(err)
+		return
+	}
+
 	quota = float64(len(sets))
 	cq, err := cpuQuota()
 	if err == nil {
 		if cq != -1 {
 			period, err := cpuPeriod()
-			logx.Must(err)
+			if err != nil {
+				logx.Error(err)
+				return
+			}
 
 			limit := float64(cq) / float64(period)
 			if limit < quota {
@@ -44,10 +55,16 @@ func init() {
 	}
 
 	preSystem, err = systemCpuUsage()
-	logx.Must(err)
+	if err != nil {
+		logx.Error(err)
+		return
+	}
 
 	preTotal, err = totalCpuUsage()
-	logx.Must(err)
+	if err != nil {
+		logx.Error(err)
+		return
+	}
 }
 
 func RefreshCpu() uint64 {

core/stat/internal/cpu_linux_test.go

@@ -1,6 +1,14 @@
 package internal
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRefreshCpu(t *testing.T) {
+	assert.True(t, RefreshCpu() >= 0)
+}
 
 func BenchmarkRefreshCpu(b *testing.B) {
 	for i := 0; i < b.N; i++ {

core/stat/metrics_test.go

@@ -0,0 +1,37 @@
+package stat
+
+import (
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMetrics(t *testing.T) {
+	counts := []int{1, 5, 10, 100, 1000, 1000}
+	for _, count := range counts {
+		m := NewMetrics("foo")
+		m.SetName("bar")
+		for i := 0; i < count; i++ {
+			m.Add(Task{
+				Duration:    time.Millisecond * time.Duration(i),
+				Description: strconv.Itoa(i),
+			})
+		}
+		m.AddDrop()
+		var writer mockedWriter
+		SetReportWriter(&writer)
+		m.executor.Flush()
+		assert.Equal(t, "bar", writer.report.Name)
+	}
+}
+
+type mockedWriter struct {
+	report *StatReport
+}
+
+func (m *mockedWriter) Write(report *StatReport) error {
+	m.report = report
+	return nil
+}

core/stat/remotewriter_test.go

@@ -0,0 +1,30 @@
+package stat
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"gopkg.in/h2non/gock.v1"
+)
+
+func TestRemoteWriter(t *testing.T) {
+	defer gock.Off()
+
+	gock.New("http://foo.com").Reply(200).BodyString("foo")
+	writer := NewRemoteWriter("http://foo.com")
+	err := writer.Write(&StatReport{
+		Name: "bar",
+	})
+	assert.Nil(t, err)
+}
+
+func TestRemoteWriterFail(t *testing.T) {
+	defer gock.Off()
+
+	gock.New("http://foo.com").Reply(503).BodyString("foo")
+	writer := NewRemoteWriter("http://foo.com")
+	err := writer.Write(&StatReport{
+		Name: "bar",
+	})
+	assert.NotNil(t, err)
+}

core/stores/cache/cache.go

@@ -1,4 +1,4 @@
-package internal
+package cache
 
 import (
 	"fmt"

core/stores/cache/cache_test.go

@@ -1,4 +1,4 @@
-package internal
+package cache
 
 import (
 	"encoding/json"
@@ -8,11 +8,11 @@ import (
 	"testing"
 	"time"
 
-	"github.com/alicebob/miniredis"
 	"github.com/stretchr/testify/assert"
 	"github.com/tal-tech/go-zero/core/errorx"
 	"github.com/tal-tech/go-zero/core/hash"
 	"github.com/tal-tech/go-zero/core/stores/redis"
+	"github.com/tal-tech/go-zero/core/stores/redis/redistest"
 	"github.com/tal-tech/go-zero/core/syncx"
 )
 
@@ -76,23 +76,23 @@ func (mc *mockedNode) TakeWithExpire(v interface{}, key string, query func(v int
 
 func TestCache_SetDel(t *testing.T) {
 	const total = 1000
-	r1 := miniredis.NewMiniRedis()
+	r1, clean1, err := redistest.CreateRedis()
-	assert.Nil(t, r1.Start())
+	assert.Nil(t, err)
-	defer r1.Close()
+	defer clean1()
-	r2 := miniredis.NewMiniRedis()
+	r2, clean2, err := redistest.CreateRedis()
-	assert.Nil(t, r2.Start())
+	assert.Nil(t, err)
-	defer r2.Close()
+	defer clean2()
 	conf := ClusterConf{
 		{
 			RedisConf: redis.RedisConf{
-				Host: r1.Addr(),
+				Host: r1.Addr,
 				Type: redis.NodeType,
 			},
 			Weight: 100,
 		},
 		{
 			RedisConf: redis.RedisConf{
-				Host: r2.Addr(),
+				Host: r2.Addr,
 				Type: redis.NodeType,
 			},
 			Weight: 100,
@@ -111,6 +111,45 @@ func TestCache_SetDel(t *testing.T) {
 		assert.Nil(t, c.GetCache(fmt.Sprintf("key/%d", i), &v))
 		assert.Equal(t, i, v)
 	}
+	assert.Nil(t, c.DelCache())
+	for i := 0; i < total; i++ {
+		assert.Nil(t, c.DelCache(fmt.Sprintf("key/%d", i)))
+	}
+	for i := 0; i < total; i++ {
+		var v int
+		assert.Equal(t, errPlaceholder, c.GetCache(fmt.Sprintf("key/%d", i), &v))
+		assert.Equal(t, 0, v)
+	}
+}
+
+func TestCache_OneNode(t *testing.T) {
+	const total = 1000
+	r, clean, err := redistest.CreateRedis()
+	assert.Nil(t, err)
+	defer clean()
+	conf := ClusterConf{
+		{
+			RedisConf: redis.RedisConf{
+				Host: r.Addr,
+				Type: redis.NodeType,
+			},
+			Weight: 100,
+		},
+	}
+	c := NewCache(conf, syncx.NewSharedCalls(), NewCacheStat("mock"), errPlaceholder)
+	for i := 0; i < total; i++ {
+		if i%2 == 0 {
+			assert.Nil(t, c.SetCache(fmt.Sprintf("key/%d", i), i))
+		} else {
+			assert.Nil(t, c.SetCacheWithExpire(fmt.Sprintf("key/%d", i), i, 0))
+		}
+	}
+	for i := 0; i < total; i++ {
+		var v int
+		assert.Nil(t, c.GetCache(fmt.Sprintf("key/%d", i), &v))
+		assert.Equal(t, i, v)
+	}
+	assert.Nil(t, c.DelCache())
 	for i := 0; i < total; i++ {
 		assert.Nil(t, c.DelCache(fmt.Sprintf("key/%d", i)))
 	}
@@ -188,6 +227,25 @@ func TestCache_Balance(t *testing.T) {
 	assert.Equal(t, total/10, count)
 }
 
+func TestCacheNoNode(t *testing.T) {
+	dispatcher := hash.NewConsistentHash()
+	c := cacheCluster{
+		dispatcher:  dispatcher,
+		errNotFound: errPlaceholder,
+	}
+	assert.NotNil(t, c.DelCache("foo"))
+	assert.NotNil(t, c.DelCache("foo", "bar", "any"))
+	assert.NotNil(t, c.GetCache("foo", nil))
+	assert.NotNil(t, c.SetCache("foo", nil))
+	assert.NotNil(t, c.SetCacheWithExpire("foo", nil, time.Second))
+	assert.NotNil(t, c.Take(nil, "foo", func(v interface{}) error {
+		return nil
+	}))
+	assert.NotNil(t, c.TakeWithExpire(nil, "foo", func(v interface{}, duration time.Duration) error {
+		return nil
+	}))
+}
+
 func calcEntropy(m map[int]int, total int) float64 {
 	var entropy float64
 

core/stores/cache/cacheconf.go

@@ -1,5 +1,3 @@
 package cache
 
-import "github.com/tal-tech/go-zero/core/stores/internal"
+type CacheConf = ClusterConf
-
-type CacheConf = internal.ClusterConf

core/stores/cache/cachenode.go

@@ -1,13 +1,13 @@
-package internal
+package cache
 
 import (
-	"encoding/json"
 	"errors"
 	"fmt"
 	"math/rand"
 	"sync"
 	"time"
 
+	"github.com/tal-tech/go-zero/core/jsonx"
 	"github.com/tal-tech/go-zero/core/logx"
 	"github.com/tal-tech/go-zero/core/mathx"
 	"github.com/tal-tech/go-zero/core/stat"
@@ -79,7 +79,7 @@ func (c cacheNode) SetCache(key string, v interface{}) error {
 }
 
 func (c cacheNode) SetCacheWithExpire(key string, v interface{}, expire time.Duration) error {
-	data, err := json.Marshal(v)
+	data, err := jsonx.Marshal(v)
 	if err != nil {
 		return err
 	}
@@ -168,24 +168,24 @@ func (c cacheNode) doTake(v interface{}, key string, query func(v interface{}) e
 			}
 		}
 
-		return json.Marshal(v)
+		return jsonx.Marshal(v)
 	})
 	if err != nil {
 		return err
 	}
 	if fresh {
 		return nil
-	} else {
-		// got the result from previous ongoing query
-		c.stat.IncrementTotal()
-		c.stat.IncrementHit()
 	}
 
-	return json.Unmarshal(val.([]byte), v)
+	// got the result from previous ongoing query
+	c.stat.IncrementTotal()
+	c.stat.IncrementHit()
+
+	return jsonx.Unmarshal(val.([]byte), v)
 }
 
 func (c cacheNode) processCache(key string, data string, v interface{}) error {
-	err := json.Unmarshal([]byte(data), v)
+	err := jsonx.Unmarshal([]byte(data), v)
 	if err == nil {
 		return nil
 	}

core/stores/cache/cachenode_test.go

@@ -0,0 +1,207 @@
+package cache
+
+import (
+	"errors"
+	"fmt"
+	"math/rand"
+	"strconv"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/alicebob/miniredis/v2"
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/mathx"
+	"github.com/tal-tech/go-zero/core/stat"
+	"github.com/tal-tech/go-zero/core/stores/redis"
+	"github.com/tal-tech/go-zero/core/stores/redis/redistest"
+	"github.com/tal-tech/go-zero/core/syncx"
+)
+
+var errTestNotFound = errors.New("not found")
+
+func init() {
+	logx.Disable()
+	stat.SetReporter(nil)
+}
+
+func TestCacheNode_DelCache(t *testing.T) {
+	store, clean, err := redistest.CreateRedis()
+	assert.Nil(t, err)
+	defer clean()
+
+	cn := cacheNode{
+		rds:            store,
+		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
+		lock:           new(sync.Mutex),
+		unstableExpiry: mathx.NewUnstable(expiryDeviation),
+		stat:           NewCacheStat("any"),
+		errNotFound:    errTestNotFound,
+	}
+	assert.Nil(t, cn.DelCache())
+	assert.Nil(t, cn.DelCache([]string{}...))
+	assert.Nil(t, cn.DelCache(make([]string, 0)...))
+	cn.SetCache("first", "one")
+	assert.Nil(t, cn.DelCache("first"))
+	cn.SetCache("first", "one")
+	cn.SetCache("second", "two")
+	assert.Nil(t, cn.DelCache("first", "second"))
+}
+
+func TestCacheNode_InvalidCache(t *testing.T) {
+	s, err := miniredis.Run()
+	assert.Nil(t, err)
+	defer s.Close()
+
+	cn := cacheNode{
+		rds:            redis.NewRedis(s.Addr(), redis.NodeType),
+		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
+		lock:           new(sync.Mutex),
+		unstableExpiry: mathx.NewUnstable(expiryDeviation),
+		stat:           NewCacheStat("any"),
+		errNotFound:    errTestNotFound,
+	}
+	s.Set("any", "value")
+	var str string
+	assert.NotNil(t, cn.GetCache("any", &str))
+	assert.Equal(t, "", str)
+	_, err = s.Get("any")
+	assert.Equal(t, miniredis.ErrKeyNotFound, err)
+}
+
+func TestCacheNode_Take(t *testing.T) {
+	store, clean, err := redistest.CreateRedis()
+	assert.Nil(t, err)
+	defer clean()
+
+	cn := cacheNode{
+		rds:            store,
+		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
+		barrier:        syncx.NewSharedCalls(),
+		lock:           new(sync.Mutex),
+		unstableExpiry: mathx.NewUnstable(expiryDeviation),
+		stat:           NewCacheStat("any"),
+		errNotFound:    errTestNotFound,
+	}
+	var str string
+	err = cn.Take(&str, "any", func(v interface{}) error {
+		*v.(*string) = "value"
+		return nil
+	})
+	assert.Nil(t, err)
+	assert.Equal(t, "value", str)
+	assert.Nil(t, cn.GetCache("any", &str))
+	val, err := store.Get("any")
+	assert.Nil(t, err)
+	assert.Equal(t, `"value"`, val)
+}
+
+func TestCacheNode_TakeNotFound(t *testing.T) {
+	store, clean, err := redistest.CreateRedis()
+	assert.Nil(t, err)
+	defer clean()
+
+	cn := cacheNode{
+		rds:            store,
+		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
+		barrier:        syncx.NewSharedCalls(),
+		lock:           new(sync.Mutex),
+		unstableExpiry: mathx.NewUnstable(expiryDeviation),
+		stat:           NewCacheStat("any"),
+		errNotFound:    errTestNotFound,
+	}
+	var str string
+	err = cn.Take(&str, "any", func(v interface{}) error {
+		return errTestNotFound
+	})
+	assert.Equal(t, errTestNotFound, err)
+	assert.Equal(t, errTestNotFound, cn.GetCache("any", &str))
+	val, err := store.Get("any")
+	assert.Nil(t, err)
+	assert.Equal(t, `*`, val)
+
+	store.Set("any", "*")
+	err = cn.Take(&str, "any", func(v interface{}) error {
+		return nil
+	})
+	assert.Equal(t, errTestNotFound, err)
+	assert.Equal(t, errTestNotFound, cn.GetCache("any", &str))
+
+	store.Del("any")
+	var errDummy = errors.New("dummy")
+	err = cn.Take(&str, "any", func(v interface{}) error {
+		return errDummy
+	})
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCacheNode_TakeWithExpire(t *testing.T) {
+	store, clean, err := redistest.CreateRedis()
+	assert.Nil(t, err)
+	defer clean()
+
+	cn := cacheNode{
+		rds:            store,
+		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
+		barrier:        syncx.NewSharedCalls(),
+		lock:           new(sync.Mutex),
+		unstableExpiry: mathx.NewUnstable(expiryDeviation),
+		stat:           NewCacheStat("any"),
+		errNotFound:    errors.New("any"),
+	}
+	var str string
+	err = cn.TakeWithExpire(&str, "any", func(v interface{}, expire time.Duration) error {
+		*v.(*string) = "value"
+		return nil
+	})
+	assert.Nil(t, err)
+	assert.Equal(t, "value", str)
+	assert.Nil(t, cn.GetCache("any", &str))
+	val, err := store.Get("any")
+	assert.Nil(t, err)
+	assert.Equal(t, `"value"`, val)
+}
+
+func TestCacheNode_String(t *testing.T) {
+	store, clean, err := redistest.CreateRedis()
+	assert.Nil(t, err)
+	defer clean()
+
+	cn := cacheNode{
+		rds:            store,
+		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
+		barrier:        syncx.NewSharedCalls(),
+		lock:           new(sync.Mutex),
+		unstableExpiry: mathx.NewUnstable(expiryDeviation),
+		stat:           NewCacheStat("any"),
+		errNotFound:    errors.New("any"),
+	}
+	assert.Equal(t, store.Addr, cn.String())
+}
+
+func TestCacheValueWithBigInt(t *testing.T) {
+	store, clean, err := redistest.CreateRedis()
+	assert.Nil(t, err)
+	defer clean()
+
+	cn := cacheNode{
+		rds:            store,
+		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
+		barrier:        syncx.NewSharedCalls(),
+		lock:           new(sync.Mutex),
+		unstableExpiry: mathx.NewUnstable(expiryDeviation),
+		stat:           NewCacheStat("any"),
+		errNotFound:    errors.New("any"),
+	}
+
+	const (
+		key         = "key"
+		value int64 = 323427211229009810
+	)
+
+	assert.Nil(t, cn.SetCache(key, value))
+	var val interface{}
+	assert.Nil(t, cn.GetCache(key, &val))
+	assert.Equal(t, strconv.FormatInt(value, 10), fmt.Sprintf("%v", val))
+}

core/stores/cache/cacheopt.go

@@ -1,21 +1,45 @@
 package cache
 
-import (
+import "time"
-	"time"
 
-	"github.com/tal-tech/go-zero/core/stores/internal"
+const (
+	defaultExpiry         = time.Hour * 24 * 7
+	defaultNotFoundExpiry = time.Minute
 )
 
-type Option = internal.Option
+type (
+	Options struct {
+		Expiry         time.Duration
+		NotFoundExpiry time.Duration
+	}
+
+	Option func(o *Options)
+)
+
+func newOptions(opts ...Option) Options {
+	var o Options
+	for _, opt := range opts {
+		opt(&o)
+	}
+
+	if o.Expiry <= 0 {
+		o.Expiry = defaultExpiry
+	}
+	if o.NotFoundExpiry <= 0 {
+		o.NotFoundExpiry = defaultNotFoundExpiry
+	}
+
+	return o
+}
 
 func WithExpiry(expiry time.Duration) Option {
-	return func(o *internal.Options) {
+	return func(o *Options) {
 		o.Expiry = expiry
 	}
 }
 
 func WithNotFoundExpiry(expiry time.Duration) Option {
-	return func(o *internal.Options) {
+	return func(o *Options) {
 		o.NotFoundExpiry = expiry
 	}
 }

core/stores/cache/cachestat.go

@@ -1,4 +1,4 @@
-package internal
+package cache
 
 import (
 	"sync/atomic"
@@ -48,20 +48,17 @@ func (cs *CacheStat) statLoop() {
 	ticker := time.NewTicker(statInterval)
 	defer ticker.Stop()
 
-	for {
+	for range ticker.C {
-		select {
+		total := atomic.SwapUint64(&cs.Total, 0)
-		case <-ticker.C:
+		if total == 0 {
-			total := atomic.SwapUint64(&cs.Total, 0)
+			continue
-			if total == 0 {
-				continue
-			}
-
-			hit := atomic.SwapUint64(&cs.Hit, 0)
-			percent := 100 * float32(hit) / float32(total)
-			miss := atomic.SwapUint64(&cs.Miss, 0)
-			dbf := atomic.SwapUint64(&cs.DbFails, 0)
-			logx.Statf("dbcache(%s) - qpm: %d, hit_ratio: %.1f%%, hit: %d, miss: %d, db_fails: %d",
-				cs.name, total, percent, hit, miss, dbf)
 		}
+
+		hit := atomic.SwapUint64(&cs.Hit, 0)
+		percent := 100 * float32(hit) / float32(total)
+		miss := atomic.SwapUint64(&cs.Miss, 0)
+		dbf := atomic.SwapUint64(&cs.DbFails, 0)
+		logx.Statf("dbcache(%s) - qpm: %d, hit_ratio: %.1f%%, hit: %d, miss: %d, db_fails: %d",
+			cs.name, total, percent, hit, miss, dbf)
 	}
 }

core/stores/cache/cleaner.go

@@ -1,4 +1,4 @@
-package internal
+package cache
 
 import (
 	"fmt"