fix type convert error (#395)

* make sure unlock safe even if listeners panic

* fix #378

* fix #378

.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
@@ -25,10 +24,11 @@ jobs:
     - name: Get dependencies
       run: |
         go get -v -t -d ./...
-        if [ -f Gopkg.toml ]; then
-            curl https://raw.githubusercontent.com/golang/dep/master/install.sh | sh
-            dep ensure
-        fi
 
     - 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()
-	if err != nil {
-		t.Error("Miniredis could not start")
-	}
-	defer s.Close()
+	store, clean, err := redistest.CreateRedis()
+	assert.Nil(t, err)
+	defer clean()
 
-	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()
-	if err != nil {
-		t.Error("Miniredis could not start")
-	}
-	defer s.Close()
+	store, clean, err := redistest.CreateRedis()
+	assert.Nil(t, err)
+	defer clean()
 
-	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"
-)
-
-const (
-	StateClosed State = iota
-	StateOpen
+	"github.com/tal-tech/go-zero/core/timex"
 )
 
 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 returns an error instantly if the netflixBreaker rejects the request.
-		// If a panic occurs in the request, the netflixBreaker handles it as an error
+		// Do runs the given request if the Breaker accepts it.
+		// Do returns an error instantly if the Breaker rejects the request.
+		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
 		Do(req func() error) error
 
-		// DoWithAcceptable runs the given request if the netflixBreaker accepts it.
-		// Do returns an error instantly if the netflixBreaker rejects the request.
-		// If a panic occurs in the request, the netflixBreaker handles it as an error
+		// DoWithAcceptable runs the given request if the Breaker accepts it.
+		// DoWithAcceptable returns an error instantly if the Breaker rejects the request.
+		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
 		// acceptable checks if it's a successful call, even if the err is not nil.
 		DoWithAcceptable(req func() error, acceptable Acceptable) error
 
-		// DoWithFallback runs the given request if the netflixBreaker accepts it.
-		// DoWithFallback runs the fallback if the netflixBreaker rejects the request.
-		// If a panic occurs in the request, the netflixBreaker handles it as an error
+		// DoWithFallback runs the given request if the Breaker accepts it.
+		// DoWithFallback runs the fallback if the Breaker rejects the request.
+		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
 		DoWithFallback(req func() error, fallback func(err error) error) error
 
-		// DoWithFallbackAcceptable runs the given request if the netflixBreaker accepts it.
-		// DoWithFallback runs the fallback if the netflixBreaker rejects the request.
-		// If a panic occurs in the request, the netflixBreaker handles it as an error
+		// DoWithFallbackAcceptable runs the given request if the Breaker accepts it.
+		// DoWithFallbackAcceptable runs the fallback if the Breaker rejects the request.
+		// If a panic occurs in the request, the Breaker handles it as an error
 		// and causes the same panic again.
 		// acceptable checks if it's a successful call, even if the err is not nil.
 		DoWithFallbackAcceptable(req func() error, fallback func(err error) error, acceptable Acceptable) error
 	}
 
-	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-- {
-		builder.WriteString(ew.reasons[i%numHistoryReasons])
-		builder.WriteByte('\n')
+	// reverse order
+	for i := ew.index - 1; i >= ew.index-ew.count; i-- {
+		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()
-	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)
-		}
-	}
+	return execute(GetBreaker(name))
 }

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

@@ -146,15 +146,15 @@ func EcbEncryptBase64(key, src string) (string, error) {
 }
 
 func getKeyBytes(key string) ([]byte, error) {
-	if len(key) > 32 {
-		if keyBytes, err := base64.StdEncoding.DecodeString(key); err != nil {
-			return nil, err
-		} else {
-			return keyBytes, nil
-		}
+	if len(key) <= 32 {
+		return []byte(key), nil
 	}
 
-	return []byte(key), nil
+	if keyBytes, err := base64.StdEncoding.DecodeString(key); err != nil {
+		return nil, err
+	} else {
+		return keyBytes, nil
+	}
 }
 
 func pkcs5Padding(ciphertext []byte, blockSize int) []byte {

core/codec/aesecb_test.go

@@ -0,0 +1,64 @@
+package codec
+
+import (
+	"encoding/base64"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAesEcb(t *testing.T) {
+	var (
+		key     = []byte("q4t7w!z%C*F-JaNdRgUjXn2r5u8x/A?D")
+		val     = []byte("hello")
+		badKey1 = []byte("aaaaaaaaa")
+		// more than 32 chars
+		badKey2 = []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")
+	)
+	_, err := EcbEncrypt(badKey1, val)
+	assert.NotNil(t, err)
+	_, err = EcbEncrypt(badKey2, val)
+	assert.NotNil(t, err)
+	dst, err := EcbEncrypt(key, val)
+	assert.Nil(t, err)
+	_, err = EcbDecrypt(badKey1, dst)
+	assert.NotNil(t, err)
+	_, err = EcbDecrypt(badKey2, dst)
+	assert.NotNil(t, err)
+	_, err = EcbDecrypt(key, val)
+	// not enough block, just nil
+	assert.Nil(t, err)
+	src, err := EcbDecrypt(key, dst)
+	assert.Nil(t, err)
+	assert.Equal(t, val, src)
+}
+
+func TestAesEcbBase64(t *testing.T) {
+	const (
+		val     = "hello"
+		badKey1 = "aaaaaaaaa"
+		// more than 32 chars
+		badKey2 = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+	)
+	var key = []byte("q4t7w!z%C*F-JaNdRgUjXn2r5u8x/A?D")
+	b64Key := base64.StdEncoding.EncodeToString(key)
+	b64Val := base64.StdEncoding.EncodeToString([]byte(val))
+	_, err := EcbEncryptBase64(badKey1, val)
+	assert.NotNil(t, err)
+	_, err = EcbEncryptBase64(badKey2, val)
+	assert.NotNil(t, err)
+	_, err = EcbEncryptBase64(b64Key, val)
+	assert.NotNil(t, err)
+	dst, err := EcbEncryptBase64(b64Key, b64Val)
+	assert.Nil(t, err)
+	_, err = EcbDecryptBase64(badKey1, dst)
+	assert.NotNil(t, err)
+	_, err = EcbDecryptBase64(badKey2, dst)
+	assert.NotNil(t, err)
+	_, err = EcbDecryptBase64(b64Key, val)
+	assert.NotNil(t, err)
+	src, err := EcbDecryptBase64(b64Key, dst)
+	b, err := base64.StdEncoding.DecodeString(src)
+	assert.Nil(t, err)
+	assert.Equal(t, val, string(b))
+}

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 != nil {
+	if _, err = io.Copy(&c, io.LimitReader(r, unzipLimit)); 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.data[key]
-	if ok {
-		c.lruCache.add(key)
-	}
-	c.lock.Unlock()
+	value, ok := c.doGet(key)
 	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 {
-		select {
-		case <-ticker.C:
-			hit := atomic.SwapUint64(&cs.hit, 0)
-			miss := atomic.SwapUint64(&cs.miss, 0)
-			total := hit + miss
-			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)
+	for range ticker.C {
+		hit := atomic.SwapUint64(&cs.hit, 0)
+		miss := atomic.SwapUint64(&cs.miss, 0)
+		total := hit + miss
+		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)
 	}
 }

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 {
-		offset := rw.offset
-		// reset expired buckets
-		start := offset + 1
-		steps := start + span
-		var remainder int
-		if steps > rw.size {
-			remainder = steps - rw.size
-			steps = rw.size
-		}
-		for i := start; i < steps; i++ {
-			rw.win.resetBucket(i)
-			offset = i
-		}
-		for i := 0; i < remainder; i++ {
-			rw.win.resetBucket(i)
-			offset = i
-		}
-		rw.offset = offset
-		rw.lastTime = timex.Now()
+	if span <= 0 {
+		return
 	}
+
+	offset := rw.offset
+	// reset expired buckets
+	for i := 0; i < span; i++ {
+		rw.win.resetBucket((offset + i + 1) % rw.size)
+	}
+
+	rw.offset = (offset + span) % rw.size
+	now := timex.Now()
+	// align to interval time boundary
+	rw.lastTime = now - (now-rw.lastTime)%rw.interval
 }
 
+// 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 {
-		return
-	}
-
+	// s.tp can only be untyped here
 	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,73 @@
+package conf
+
+import (
+	"io/ioutil"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/fs"
+	"github.com/tal-tech/go-zero/core/hash"
+)
+
+func TestLoadConfig_notExists(t *testing.T) {
+	assert.NotNil(t, LoadConfig("not_a_file", nil))
+}
+
+func TestLoadConfig_notRecogFile(t *testing.T) {
+	filename, err := fs.TempFilenameWithText("hello")
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+	assert.NotNil(t, LoadConfig(filename, nil))
+}
+
+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
-// configuration itself along with an error that indicates if there was a problem loading the configuration.
+// Loads the properties into a properties configuration instance.
+// Returns an error that indicates if there was a problem loading the configuration.
 func LoadProperties(filename string) (Properties, error) {
 	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

@@ -24,6 +24,20 @@ func TestProperties(t *testing.T) {
 	assert.Equal(t, "test", props.GetString("app.name"))
 	assert.Equal(t, "app", props.GetString("app.program"))
 	assert.Equal(t, 5, props.GetInt("app.threads"))
+
+	val := props.ToString()
+	assert.Contains(t, val, "app.name")
+	assert.Contains(t, val, "app.program")
+	assert.Contains(t, val, "app.threads")
+}
+
+func TestLoadProperties_badContent(t *testing.T) {
+	filename, err := fs.TempFilenameWithText("hello")
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+	_, err = LoadProperties(filename)
+	assert.NotNil(t, err)
+	assert.True(t, len(err.Error()) > 0)
 }
 
 func TestSetString(t *testing.T) {
@@ -41,3 +55,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, _ = ShrinkDeadline(ctx, time.Minute)
+	ctx, cancel := context.WithDeadline(context.Background(), deadline)
+	defer cancel()
+	ctx, cancel = ShrinkDeadline(ctx, time.Minute)
+	defer cancel()
 	dl, ok := ctx.Deadline()
 	assert.True(t, ok)
 	assert.Equal(t, deadline, dl)
@@ -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, _ = ShrinkDeadline(ctx, time.Second)
+	ctx, cancel := context.WithDeadline(context.Background(), deadline)
+	defer cancel()
+	ctx, cancel = ShrinkDeadline(ctx, time.Second)
+	defer cancel()
 	dl, ok := ctx.Deadline()
 	assert.True(t, ok)
 	assert.True(t, dl.Before(deadline))

core/contextx/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) {
-	c, _ := context.WithDeadline(context.Background(), time.Now().Add(10*time.Millisecond))
+func TestContextDeadline(t *testing.T) {
+	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/facade.go

@@ -2,7 +2,7 @@ package discov
 
 import (
 	"github.com/tal-tech/go-zero/core/discov/internal"
-	"github.com/tal-tech/go-zero/core/lang"
+	"github.com/tal-tech/go-zero/core/logx"
 )
 
 type (
@@ -26,7 +26,7 @@ func NewFacade(endpoints []string) Facade {
 
 func (f Facade) Client() internal.EtcdClient {
 	conn, err := f.registry.GetConn(f.endpoints)
-	lang.Must(err)
+	logx.Must(err)
 	return conn
 }
 

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/internal/statewatcher.go

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

core/discov/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
-  namespace: discovery
+  name: etcd
+  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
-    discov_node: discov0
-  name: discov0
-  namespace: discovery
+    app: etcd
+    etcd_node: etcd0
+  name: etcd0
+  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
-      name: discov0
+        - --auto-compaction-retention=1
+      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
-  name: discov0
-  namespace: discovery
+    etcd_node: etcd0
+  name: etcd0
+  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
-    discov_node: discov1
-  name: discov1
-  namespace: discovery
+    app: etcd
+    etcd_node: etcd1
+  name: etcd1
+  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
-      name: discov1
+        - --auto-compaction-retention=1
+      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
-  name: discov1
-  namespace: discovery
+    etcd_node: etcd1
+  name: etcd1
+  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
-    discov_node: discov2
-  name: discov2
-  namespace: discovery
+    app: etcd
+    etcd_node: etcd2
+  name: etcd2
+  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
-      name: discov2
+        - --auto-compaction-retention=1
+      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
-  name: discov2
-  namespace: discovery
+    etcd_node: etcd2
+  name: etcd2
+  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
-    discov_node: discov3
-  name: discov3
-  namespace: discovery
+    app: etcd
+    etcd_node: etcd3
+  name: etcd3
+  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
-      name: discov3
+        - --auto-compaction-retention=1
+      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
-  name: discov3
-  namespace: discovery
+    etcd_node: etcd3
+  name: etcd3
+  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
-    discov_node: discov4
-  name: discov4
-  namespace: discovery
+    app: etcd
+    etcd_node: etcd4
+  name: etcd4
+  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
-      name: discov4
+        - --auto-compaction-retention=1
+      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
-  name: discov4
-  namespace: discovery
+    etcd_node: etcd4
+  name: etcd4
+  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
-	lock sync.Mutex
+	err atomic.Value // error
 }
 
 func (ae *AtomicError) Set(err error) {
-	ae.lock.Lock()
-	ae.err = err
-	ae.lock.Unlock()
+	ae.err.Store(err)
 }
 
 func (ae *AtomicError) Load() error {
-	ae.lock.Lock()
-	err := ae.err
-	ae.lock.Unlock()
-	return err
+	if v := ae.err.Load(); v != nil {
+		return v.(error)
+	}
+	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, i)
-		}
+		result = append(result, tasks...)
 	}, 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,8 +3,10 @@ package executors
 import (
 	"reflect"
 	"sync"
+	"sync/atomic"
 	"time"
 
+	"github.com/tal-tech/go-zero/core/lang"
 	"github.com/tal-tech/go-zero/core/proc"
 	"github.com/tal-tech/go-zero/core/syncx"
 	"github.com/tal-tech/go-zero/core/threading"
@@ -32,21 +34,24 @@ type (
 		container TaskContainer
 		waitGroup sync.WaitGroup
 		// avoid race condition on waitGroup when calling wg.Add/Done/Wait(...)
-		wgBarrier syncx.Barrier
-		guarded   bool
-		newTicker func(duration time.Duration) timex.Ticker
-		lock      sync.Mutex
+		wgBarrier   syncx.Barrier
+		confirmChan chan lang.PlaceholderType
+		inflight    int32
+		guarded     bool
+		newTicker   func(duration time.Duration) timex.Ticker
+		lock        sync.Mutex
 	}
 )
 
 func NewPeriodicalExecutor(interval time.Duration, container TaskContainer) *PeriodicalExecutor {
 	executor := &PeriodicalExecutor{
 		// buffer 1 to let the caller go quickly
-		commander: make(chan interface{}, 1),
-		interval:  interval,
-		container: container,
+		commander:   make(chan interface{}, 1),
+		interval:    interval,
+		container:   container,
+		confirmChan: make(chan lang.PlaceholderType),
 		newTicker: func(d time.Duration) timex.Ticker {
-			return timex.NewTicker(interval)
+			return timex.NewTicker(d)
 		},
 	}
 	proc.AddShutdownListener(func() {
@@ -59,10 +64,12 @@ func NewPeriodicalExecutor(interval time.Duration, container TaskContainer) *Per
 func (pe *PeriodicalExecutor) Add(task interface{}) {
 	if vals, ok := pe.addAndCheck(task); ok {
 		pe.commander <- vals
+		<-pe.confirmChan
 	}
 }
 
 func (pe *PeriodicalExecutor) Flush() bool {
+	pe.enterExecution()
 	return pe.executeTasks(func() interface{} {
 		pe.lock.Lock()
 		defer pe.lock.Unlock()
@@ -77,6 +84,7 @@ func (pe *PeriodicalExecutor) Sync(fn func()) {
 }
 
 func (pe *PeriodicalExecutor) Wait() {
+	pe.Flush()
 	pe.wgBarrier.Guard(func() {
 		pe.waitGroup.Wait()
 	})
@@ -85,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
 	}
 
@@ -105,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()
 
@@ -114,6 +123,9 @@ 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)
 				last = timex.Now()
 			case <-ticker.Chan():
@@ -121,13 +133,7 @@ func (pe *PeriodicalExecutor) backgroundFlush() {
 					commanded = false
 				} else if pe.Flush() {
 					last = timex.Now()
-				} else if timex.Since(last) > pe.interval*idleRound {
-					pe.lock.Lock()
-					pe.guarded = false
-					pe.lock.Unlock()
-
-					// flush again to avoid missing tasks
-					pe.Flush()
+				} else if pe.shallQuit(last) {
 					return
 				}
 			}
@@ -135,13 +141,18 @@ func (pe *PeriodicalExecutor) backgroundFlush() {
 	})
 }
 
-func (pe *PeriodicalExecutor) executeTasks(tasks interface{}) bool {
+func (pe *PeriodicalExecutor) doneExecution() {
+	pe.waitGroup.Done()
+}
+
+func (pe *PeriodicalExecutor) enterExecution() {
 	pe.wgBarrier.Guard(func() {
 		pe.waitGroup.Add(1)
 	})
-	defer pe.wgBarrier.Guard(func() {
-		pe.waitGroup.Done()
-	})
+}
+
+func (pe *PeriodicalExecutor) executeTasks(tasks interface{}) bool {
+	defer pe.doneExecution()
 
 	ok := pe.hasTasks(tasks)
 	if ok {
@@ -165,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

@@ -106,6 +106,60 @@ func TestPeriodicalExecutor_Bulk(t *testing.T) {
 	lock.Unlock()
 }
 
+func TestPeriodicalExecutor_Wait(t *testing.T) {
+	var lock sync.Mutex
+	executer := NewBulkExecutor(func(tasks []interface{}) {
+		lock.Lock()
+		defer lock.Unlock()
+		time.Sleep(10 * time.Millisecond)
+	}, WithBulkTasks(1), WithBulkInterval(time.Second))
+	for i := 0; i < 10; i++ {
+		executer.Add(1)
+	}
+	executer.Flush()
+	executer.Wait()
+}
+
+func TestPeriodicalExecutor_WaitFast(t *testing.T) {
+	const total = 3
+	var cnt int
+	var lock sync.Mutex
+	executer := NewBulkExecutor(func(tasks []interface{}) {
+		defer func() {
+			cnt++
+		}()
+		lock.Lock()
+		defer lock.Unlock()
+		time.Sleep(10 * time.Millisecond)
+	}, WithBulkTasks(1), WithBulkInterval(10*time.Millisecond))
+	for i := 0; i < total; i++ {
+		executer.Add(2)
+	}
+	executer.Flush()
+	executer.Wait()
+	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/lang/lang.go

@@ -1,16 +1,8 @@
 package lang
 
-import "log"
-
 var Placeholder PlaceholderType
 
 type (
 	GenericType     = interface{}
 	PlaceholderType = struct{}
 )
-
-func Must(err error) {
-	if err != nil {
-		log.Fatal(err)
-	}
-}

core/lang/lang_test.go

@@ -1,7 +0,0 @@
-package lang
-
-import "testing"
-
-func TestMust(t *testing.T) {
-	Must(nil)
-}

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 {
-		select {
-		case <-ticker.C:
-			if lim.store.Ping() {
-				atomic.StoreUint32(&lim.redisAlive, 1)
-				return
-			}
+	for range ticker.C {
+		if lim.store.Ping() {
+			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,9 +15,9 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/tal-tech/go-zero/core/iox"
-	"github.com/tal-tech/go-zero/core/lang"
 	"github.com/tal-tech/go-zero/core/sysx"
 	"github.com/tal-tech/go-zero/core/timex"
 )
@@ -43,9 +43,11 @@ const (
 	consoleMode = "console"
 	volumeMode  = "volume"
 
+	levelAlert  = "alert"
 	levelInfo   = "info"
 	levelError  = "error"
 	levelSevere = "severe"
+	levelFatal  = "fatal"
 	levelSlow   = "slow"
 	levelStat   = "stat"
 
@@ -96,11 +98,12 @@ type (
 		Infof(string, ...interface{})
 		Slow(...interface{})
 		Slowf(string, ...interface{})
+		WithDuration(time.Duration) Logger
 	}
 )
 
 func MustSetup(c LogConf) {
-	lang.Must(SetUp(c))
+	Must(SetUp(c))
 }
 
 // SetUp sets up the logx. If already set up, just return nil.
@@ -119,6 +122,10 @@ func SetUp(c LogConf) error {
 	}
 }
 
+func Alert(v string) {
+	output(errorLog, levelAlert, v)
+}
+
 func Close() error {
 	if writeConsole {
 		return nil
@@ -210,6 +217,15 @@ func Infof(format string, v ...interface{}) {
 	infoSync(fmt.Sprintf(format, v...))
 }
 
+func Must(err error) {
+	if err != nil {
+		msg := formatWithCaller(err.Error(), 3)
+		log.Print(msg)
+		output(severeLog, levelFatal, msg)
+		os.Exit(1)
+	}
+}
+
 func SetLevel(level uint32) {
 	atomic.StoreUint32(&logLevel, level)
 }

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 {
-	return strings.Contains(mw.builder.String(), text)
+func (mw *mockWriter) String() string {
+	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)
@@ -131,6 +194,72 @@ func TestSetLevelWithDuration(t *testing.T) {
 	assert.Equal(t, 0, writer.builder.Len())
 }
 
+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
@@ -228,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) {
@@ -248,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 {
-		defer func() {
-			if r := recover(); r != nil {
-				ErrorStack(r)
-			}
-		}()
-		compressLogFile(file)
+	if !l.compress {
+		return
 	}
+
+	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
-	logLevel = InfoLevel
+	prevLevel := atomic.LoadUint32(&logLevel)
+	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"`
-	Span  string          `json:"span,omitempty"`
-	ctx   context.Context `json:"-"`
+	Trace string `json:"trace,omitempty"`
+	Span  string `json:"span,omitempty"`
+	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 {
-		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
+	if len(options) == 0 {
+		return key, nil, 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()
-		close(output)
+		finish()
 	})
 
 	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+polyfill.go

@@ -4,12 +4,14 @@ package proc
 
 import "time"
 
+// AddShutdownListener returns fn itself on windows, lets callers call fn on their own.
 func AddShutdownListener(fn func()) func() {
-	return nil
+	return fn
 }
 
+// AddWrapUpListener returns fn itself on windows, lets callers call fn on their own.
 func AddWrapUpListener(fn func()) func() {
-	return nil
+	return fn
 }
 
 func SetTimeoutToForceQuit(duration time.Duration) {

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

@@ -0,0 +1,44 @@
+package queue
+
+import (
+	"errors"
+	"sync/atomic"
+
+	"github.com/tal-tech/go-zero/core/logx"
+)
+
+var ErrNoAvailablePusher = errors.New("no available pusher")
+
+type BalancedPusher struct {
+	name    string
+	pushers []Pusher
+	index   uint64
+}
+
+func NewBalancedPusher(pushers []Pusher) Pusher {
+	return &BalancedPusher{
+		name:    generateName(pushers),
+		pushers: pushers,
+	}
+}
+
+func (pusher *BalancedPusher) Name() string {
+	return pusher.name
+}
+
+func (pusher *BalancedPusher) Push(message string) error {
+	size := len(pusher.pushers)
+
+	for i := 0; i < size; i++ {
+		index := atomic.AddUint64(&pusher.index, 1) % uint64(size)
+		target := pusher.pushers[index]
+
+		if err := target.Push(message); err != nil {
+			logx.Error(err)
+		} else {
+			return nil
+		}
+	}
+
+	return ErrNoAvailablePusher
+}

core/queue/balancedpusher_test.go

@@ -0,0 +1,43 @@
+package queue
+
+import (
+	"fmt"
+	"strconv"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBalancedQueuePusher(t *testing.T) {
+	const numPushers = 100
+	var pushers []Pusher
+	var mockedPushers []*mockedPusher
+	for i := 0; i < numPushers; i++ {
+		p := &mockedPusher{
+			name: "pusher:" + strconv.Itoa(i),
+		}
+		pushers = append(pushers, p)
+		mockedPushers = append(mockedPushers, p)
+	}
+
+	pusher := NewBalancedPusher(pushers)
+	assert.True(t, len(pusher.Name()) > 0)
+
+	for i := 0; i < numPushers*1000; i++ {
+		assert.Nil(t, pusher.Push("item"))
+	}
+
+	var counts []int
+	for _, p := range mockedPushers {
+		counts = append(counts, p.count)
+	}
+	mean := calcMean(counts)
+	variance := calcVariance(mean, counts)
+	assert.True(t, variance < 100, fmt.Sprintf("too big variance - %.2f", variance))
+}
+
+func TestBalancedQueuePusher_NoAvailable(t *testing.T) {
+	pusher := NewBalancedPusher(nil)
+	assert.True(t, len(pusher.Name()) == 0)
+	assert.Equal(t, ErrNoAvailablePusher, pusher.Push("item"))
+}

core/queue/consumer.go

@@ -0,0 +1,10 @@
+package queue
+
+type (
+	Consumer interface {
+		Consume(string) error
+		OnEvent(event interface{})
+	}
+
+	ConsumerFactory func() (Consumer, error)
+)

core/queue/messagequeue.go

@@ -0,0 +1,6 @@
+package queue
+
+type MessageQueue interface {
+	Start()
+	Stop()
+}

core/queue/multipusher.go

@@ -0,0 +1,31 @@
+package queue
+
+import "github.com/tal-tech/go-zero/core/errorx"
+
+type MultiPusher struct {
+	name    string
+	pushers []Pusher
+}
+
+func NewMultiPusher(pushers []Pusher) Pusher {
+	return &MultiPusher{
+		name:    generateName(pushers),
+		pushers: pushers,
+	}
+}
+
+func (pusher *MultiPusher) Name() string {
+	return pusher.name
+}
+
+func (pusher *MultiPusher) Push(message string) error {
+	var batchError errorx.BatchError
+
+	for _, each := range pusher.pushers {
+		if err := each.Push(message); err != nil {
+			batchError.Add(err)
+		}
+	}
+
+	return batchError.Err()
+}

core/queue/multipusher_test.go

@@ -0,0 +1,39 @@
+package queue
+
+import (
+	"fmt"
+	"math"
+	"strconv"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMultiQueuePusher(t *testing.T) {
+	const numPushers = 100
+	var pushers []Pusher
+	var mockedPushers []*mockedPusher
+	for i := 0; i < numPushers; i++ {
+		p := &mockedPusher{
+			name: "pusher:" + strconv.Itoa(i),
+		}
+		pushers = append(pushers, p)
+		mockedPushers = append(mockedPushers, p)
+	}
+
+	pusher := NewMultiPusher(pushers)
+	assert.True(t, len(pusher.Name()) > 0)
+
+	for i := 0; i < 1000; i++ {
+		_ = pusher.Push("item")
+	}
+
+	var counts []int
+	for _, p := range mockedPushers {
+		counts = append(counts, p.count)
+	}
+	mean := calcMean(counts)
+	variance := calcVariance(mean, counts)
+	assert.True(t, math.Abs(mean-1000*(1-failProba)) < 10)
+	assert.True(t, variance < 100, fmt.Sprintf("too big variance - %.2f", variance))
+}

core/queue/producer.go

@@ -0,0 +1,15 @@
+package queue
+
+type (
+	Producer interface {
+		AddListener(listener ProduceListener)
+		Produce() (string, bool)
+	}
+
+	ProduceListener interface {
+		OnProducerPause()
+		OnProducerResume()
+	}
+
+	ProducerFactory func() (Producer, error)
+)

core/queue/queue.go

@@ -0,0 +1,239 @@
+package queue
+
+import (
+	"runtime"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/rescue"
+	"github.com/tal-tech/go-zero/core/stat"
+	"github.com/tal-tech/go-zero/core/threading"
+	"github.com/tal-tech/go-zero/core/timex"
+)
+
+const queueName = "queue"
+
+type (
+	Queue struct {
+		name                 string
+		metrics              *stat.Metrics
+		producerFactory      ProducerFactory
+		producerRoutineGroup *threading.RoutineGroup
+		consumerFactory      ConsumerFactory
+		consumerRoutineGroup *threading.RoutineGroup
+		producerCount        int
+		consumerCount        int
+		active               int32
+		channel              chan string
+		quit                 chan struct{}
+		listeners            []Listener
+		eventLock            sync.Mutex
+		eventChannels        []chan interface{}
+	}
+
+	Listener interface {
+		OnPause()
+		OnResume()
+	}
+
+	Poller interface {
+		Name() string
+		Poll() string
+	}
+
+	Pusher interface {
+		Name() string
+		Push(string) error
+	}
+)
+
+func NewQueue(producerFactory ProducerFactory, consumerFactory ConsumerFactory) *Queue {
+	queue := &Queue{
+		metrics:              stat.NewMetrics(queueName),
+		producerFactory:      producerFactory,
+		producerRoutineGroup: threading.NewRoutineGroup(),
+		consumerFactory:      consumerFactory,
+		consumerRoutineGroup: threading.NewRoutineGroup(),
+		producerCount:        runtime.NumCPU(),
+		consumerCount:        runtime.NumCPU() << 1,
+		channel:              make(chan string),
+		quit:                 make(chan struct{}),
+	}
+	queue.SetName(queueName)
+
+	return queue
+}
+
+func (queue *Queue) AddListener(listener Listener) {
+	queue.listeners = append(queue.listeners, listener)
+}
+
+func (queue *Queue) Broadcast(message interface{}) {
+	go func() {
+		queue.eventLock.Lock()
+		defer queue.eventLock.Unlock()
+
+		for _, channel := range queue.eventChannels {
+			channel <- message
+		}
+	}()
+}
+
+func (queue *Queue) SetName(name string) {
+	queue.name = name
+	queue.metrics.SetName(name)
+}
+
+func (queue *Queue) SetNumConsumer(count int) {
+	queue.consumerCount = count
+}
+
+func (queue *Queue) SetNumProducer(count int) {
+	queue.producerCount = count
+}
+
+func (queue *Queue) Start() {
+	queue.startProducers(queue.producerCount)
+	queue.startConsumers(queue.consumerCount)
+
+	queue.producerRoutineGroup.Wait()
+	close(queue.channel)
+	queue.consumerRoutineGroup.Wait()
+}
+
+func (queue *Queue) Stop() {
+	close(queue.quit)
+}
+
+func (queue *Queue) consume(eventChan chan interface{}) {
+	var consumer Consumer
+
+	for {
+		var err error
+		if consumer, err = queue.consumerFactory(); err != nil {
+			logx.Errorf("Error on creating consumer: %v", err)
+			time.Sleep(time.Second)
+		} else {
+			break
+		}
+	}
+
+	for {
+		select {
+		case message, ok := <-queue.channel:
+			if ok {
+				queue.consumeOne(consumer, message)
+			} else {
+				logx.Info("Task channel was closed, quitting consumer...")
+				return
+			}
+		case event := <-eventChan:
+			consumer.OnEvent(event)
+		}
+	}
+}
+
+func (queue *Queue) consumeOne(consumer Consumer, message string) {
+	threading.RunSafe(func() {
+		startTime := timex.Now()
+		defer func() {
+			duration := timex.Since(startTime)
+			queue.metrics.Add(stat.Task{
+				Duration: duration,
+			})
+			logx.WithDuration(duration).Infof("%s", message)
+		}()
+
+		if err := consumer.Consume(message); err != nil {
+			logx.Errorf("Error occurred while consuming %v: %v", message, err)
+		}
+	})
+}
+
+func (queue *Queue) pause() {
+	for _, listener := range queue.listeners {
+		listener.OnPause()
+	}
+}
+
+func (queue *Queue) produce() {
+	var producer Producer
+
+	for {
+		var err error
+		if producer, err = queue.producerFactory(); err != nil {
+			logx.Errorf("Error on creating producer: %v", err)
+			time.Sleep(time.Second)
+		} else {
+			break
+		}
+	}
+
+	atomic.AddInt32(&queue.active, 1)
+	producer.AddListener(routineListener{
+		queue: queue,
+	})
+
+	for {
+		select {
+		case <-queue.quit:
+			logx.Info("Quitting producer")
+			return
+		default:
+			if v, ok := queue.produceOne(producer); ok {
+				queue.channel <- v
+			}
+		}
+	}
+}
+
+func (queue *Queue) produceOne(producer Producer) (string, bool) {
+	// avoid panic quit the producer, just log it and continue
+	defer rescue.Recover()
+
+	return producer.Produce()
+}
+
+func (queue *Queue) resume() {
+	for _, listener := range queue.listeners {
+		listener.OnResume()
+	}
+}
+
+func (queue *Queue) startConsumers(number int) {
+	for i := 0; i < number; i++ {
+		eventChan := make(chan interface{})
+		queue.eventLock.Lock()
+		queue.eventChannels = append(queue.eventChannels, eventChan)
+		queue.eventLock.Unlock()
+		queue.consumerRoutineGroup.Run(func() {
+			queue.consume(eventChan)
+		})
+	}
+}
+
+func (queue *Queue) startProducers(number int) {
+	for i := 0; i < number; i++ {
+		queue.producerRoutineGroup.Run(func() {
+			queue.produce()
+		})
+	}
+}
+
+type routineListener struct {
+	queue *Queue
+}
+
+func (rl routineListener) OnProducerPause() {
+	if atomic.AddInt32(&rl.queue.active, -1) <= 0 {
+		rl.queue.pause()
+	}
+}
+
+func (rl routineListener) OnProducerResume() {
+	if atomic.AddInt32(&rl.queue.active, 1) == 1 {
+		rl.queue.resume()
+	}
+}

core/queue/queue_test.go

@@ -0,0 +1,94 @@
+package queue
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	consumers = 4
+	rounds    = 100
+)
+
+func TestQueue(t *testing.T) {
+	producer := newMockedProducer(rounds)
+	consumer := newMockedConsumer()
+	consumer.wait.Add(consumers)
+	q := NewQueue(func() (Producer, error) {
+		return producer, nil
+	}, func() (Consumer, error) {
+		return consumer, nil
+	})
+	q.AddListener(new(mockedListener))
+	q.SetName("mockqueue")
+	q.SetNumConsumer(consumers)
+	q.SetNumProducer(1)
+	q.pause()
+	q.resume()
+	go func() {
+		producer.wait.Wait()
+		q.Stop()
+	}()
+	q.Start()
+	assert.Equal(t, int32(rounds), atomic.LoadInt32(&consumer.count))
+}
+
+type mockedConsumer struct {
+	count  int32
+	events int32
+	wait   sync.WaitGroup
+}
+
+func newMockedConsumer() *mockedConsumer {
+	return new(mockedConsumer)
+}
+
+func (c *mockedConsumer) Consume(string) error {
+	atomic.AddInt32(&c.count, 1)
+	return nil
+}
+
+func (c *mockedConsumer) OnEvent(interface{}) {
+	if atomic.AddInt32(&c.events, 1) <= consumers {
+		c.wait.Done()
+	}
+}
+
+type mockedProducer struct {
+	total int32
+	count int32
+	wait  sync.WaitGroup
+}
+
+func newMockedProducer(total int32) *mockedProducer {
+	p := new(mockedProducer)
+	p.total = total
+	p.wait.Add(int(total))
+	return p
+}
+
+func (p *mockedProducer) AddListener(listener ProduceListener) {
+}
+
+func (p *mockedProducer) Produce() (string, bool) {
+	if atomic.AddInt32(&p.count, 1) <= p.total {
+		p.wait.Done()
+		return "item", true
+	} else {
+		time.Sleep(time.Second)
+		return "", false
+	}
+}
+
+type mockedListener struct {
+}
+
+func (l *mockedListener) OnPause() {
+}
+
+func (l *mockedListener) OnResume() {
+}

core/queue/util.go

@@ -0,0 +1,12 @@
+package queue
+
+import "strings"
+
+func generateName(pushers []Pusher) string {
+	names := make([]string, len(pushers))
+	for i, pusher := range pushers {
+		names[i] = pusher.Name()
+	}
+
+	return strings.Join(names, ",")
+}

core/queue/util_test.go

@@ -0,0 +1,77 @@
+package queue
+
+import (
+	"errors"
+	"math"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/mathx"
+)
+
+var (
+	proba     = mathx.NewProba()
+	failProba = 0.01
+)
+
+func init() {
+	logx.Disable()
+}
+
+func TestGenerateName(t *testing.T) {
+	pushers := []Pusher{
+		&mockedPusher{name: "first"},
+		&mockedPusher{name: "second"},
+		&mockedPusher{name: "third"},
+	}
+
+	assert.Equal(t, "first,second,third", generateName(pushers))
+}
+
+func TestGenerateNameNil(t *testing.T) {
+	var pushers []Pusher
+	assert.Equal(t, "", generateName(pushers))
+}
+
+func calcMean(vals []int) float64 {
+	if len(vals) == 0 {
+		return 0
+	}
+
+	var result float64
+	for _, val := range vals {
+		result += float64(val)
+	}
+	return result / float64(len(vals))
+}
+
+func calcVariance(mean float64, vals []int) float64 {
+	if len(vals) == 0 {
+		return 0
+	}
+
+	var result float64
+	for _, val := range vals {
+		result += math.Pow(float64(val)-mean, 2)
+	}
+	return result / float64(len(vals))
+}
+
+type mockedPusher struct {
+	name  string
+	count int
+}
+
+func (p *mockedPusher) Name() string {
+	return p.name
+}
+
+func (p *mockedPusher) Push(s string) error {
+	if proba.TrueOnProba(failProba) {
+		return errors.New("dummy")
+	}
+
+	p.count++
+	return nil
+}