avoid bigint converted into float64 when unmarshaling

* rebase upstream

* rebase

* trim no need line

* trim no need line

* trim no need line

* update doc

* remove update

* remove no need

* remove no need

* goctl add jwt support

* goctl add jwt support

* goctl add jwt support

* goctl support import

* goctl support import

* support return ()

* revert

* refactor and rename folder to group

* parser add test

Co-authored-by: kingxt <dream4kingxt@163.com>

.codecov.yml

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

.github/workflows/go.yml

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

.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

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,7 +22,6 @@ const (
 var ErrServiceUnavailable = errors.New("circuit breaker is open")
 
 type (
-	State      = int32
 	Acceptable func(err error) bool
 
 	Breaker interface {
@@ -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

@@ -27,7 +27,6 @@ func getGoogleBreaker() *googleBreaker {
 	return &googleBreaker{
 		stat:  st,
 		k:     5,
-		state: StateClosed,
 		proba: mathx.NewProba(),
 	}
 }

core/cmdline/input_test.go

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

core/codec/dh_test.go

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

core/codec/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,59 @@
+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`
+	encryptedBody = `49e7bc15640e5d927fd3f129b749536d0755baf03a0f35fc914ff1b7b8ce659e5fe3a598442eb908c5995e28bacd3d76e4420bb05b6bfc177040f66c6976f680f7123505d626ab96a9db1151f45c93bc0262db9087b9fb6801715f76f902e644a20029262858f05b0d10540842204346ac1d6d8f29cc5d47dab79af75d922ef2`
+)
+
+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

@@ -82,12 +82,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 +108,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 +145,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)

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/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_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() {

core/conf/config_test.go

@@ -0,0 +1,58 @@
+package conf
+
+import (
+	"io/ioutil"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/hash"
+)
+
+func TestConfigJson(t *testing.T) {
+	tests := []string{
+		".json",
+		".yaml",
+		".yml",
+	}
+	text := `{
+	"a": "foo",
+	"b": 1
+}`
+	for _, test := range tests {
+		test := test
+		t.Run(test, func(t *testing.T) {
+			t.Parallel()
+
+			tmpfile, err := createTempFile(test, text)
+			assert.Nil(t, err)
+			defer os.Remove(tmpfile)
+
+			var val struct {
+				A string `json:"a"`
+				B int    `json:"b"`
+			}
+			MustLoad(tmpfile, &val)
+			assert.Equal(t, "foo", val.A)
+			assert.Equal(t, 1, val.B)
+		})
+	}
+}
+
+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

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

core/contextx/valueonlycontext_test.go

@@ -34,7 +34,7 @@ func TestContextCancel(t *testing.T) {
 	assert.NotEqual(t, context.Canceled, c2.Err())
 }
 
-func TestConextDeadline(t *testing.T) {
+func TestContextDeadline(t *testing.T) {
 	c, _ := context.WithDeadline(context.Background(), time.Now().Add(10*time.Millisecond))
 	o := ValueOnlyFrom(c)
 	select {

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/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: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: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: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: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: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/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()

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

@@ -5,6 +5,7 @@ import (
 	"sync"
 	"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,19 +33,21 @@ 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
+		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)
 		},
@@ -59,10 +62,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()
@@ -114,6 +119,8 @@ func (pe *PeriodicalExecutor) backgroundFlush() {
 			select {
 			case vals := <-pe.commander:
 				commanded = true
+				pe.enterExecution()
+				pe.confirmChan <- lang.Placeholder
 				pe.executeTasks(vals)
 				last = timex.Now()
 			case <-ticker.Chan():
@@ -135,11 +142,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.waitGroup.Done()
+}
+
+func (pe *PeriodicalExecutor) executeTasks(tasks interface{}) bool {
+	defer pe.doneExecution()
 
 	ok := pe.hasTasks(tasks)
 	if ok {

core/executors/periodicalexecutor_test.go

@@ -106,6 +106,40 @@ 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)
+}
+
 // 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 {
@@ -84,6 +84,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{})
@@ -195,7 +203,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)

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{}) {

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/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"
 )
@@ -46,6 +46,7 @@ const (
 	levelInfo   = "info"
 	levelError  = "error"
 	levelSevere = "severe"
+	levelFatal  = "fatal"
 	levelSlow   = "slow"
 	levelStat   = "stat"
 
@@ -96,11 +97,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.
@@ -210,6 +212,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) {
@@ -85,6 +100,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 +186,70 @@ 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()
+	WithKeepDays(1)
+	WithGzip()
+	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 +347,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 +367,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

@@ -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,115 @@
+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
+	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
+	ctx := context.WithValue(context.Background(), tracespec.TracingKey, mock)
+	l := WithContext(ctx).(*traceLogger)
+	SetLevel(InfoLevel)
+	atomic.StoreUint32(&initialized, 1)
+	errorLog = newLogWriter(log.New(&buf, "", flags))
+	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
+	ctx := context.WithValue(context.Background(), tracespec.TracingKey, mock)
+	l := WithContext(ctx).(*traceLogger)
+	SetLevel(InfoLevel)
+	atomic.StoreUint32(&initialized, 1)
+	infoLog = newLogWriter(log.New(&buf, "", flags))
+	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
+	ctx := context.WithValue(context.Background(), tracespec.TracingKey, mock)
+	l := WithContext(ctx).(*traceLogger)
+	SetLevel(InfoLevel)
+	atomic.StoreUint32(&initialized, 1)
+	slowLog = newLogWriter(log.New(&buf, "", flags))
+	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
+	l := WithContext(context.Background()).(*traceLogger)
+	SetLevel(InfoLevel)
+	atomic.StoreUint32(&initialized, 1)
+	infoLog = newLogWriter(log.New(&buf, "", flags))
+	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/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_test.go

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

core/proc/shutdown.go

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

core/queue/balancedpusher.go

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

core/queue/balancedpusher_test.go

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

core/queue/multipusher.go

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

core/queue/multipusher_test.go

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

core/stat/alert.go

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

core/stat/alert_test.go

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

core/stat/internal/cpu_linux.go

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

core/stat/internal/cpu_linux_test.go

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

core/stat/metrics_test.go

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

core/stat/remotewriter_test.go

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

core/stores/cache/cache.go

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

core/stores/cache/cache_test.go

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

core/stores/cache/cacheconf.go

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

core/stores/cache/cachenode.go

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

core/stores/cache/cachenode_test.go

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

core/stores/cache/cacheopt.go

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

core/stores/cache/cachestat.go

@@ -1,4 +1,4 @@
-package internal
+package cache
 
 import (
 	"sync/atomic"

core/stores/cache/cleaner.go

@@ -1,11 +1,10 @@
-package internal
+package cache
 
 import (
 	"fmt"
 	"time"
 
 	"github.com/tal-tech/go-zero/core/collection"
-	"github.com/tal-tech/go-zero/core/lang"
 	"github.com/tal-tech/go-zero/core/logx"
 	"github.com/tal-tech/go-zero/core/proc"
 	"github.com/tal-tech/go-zero/core/stat"
@@ -33,7 +32,7 @@ type delayTask struct {
 func init() {
 	var err error
 	timingWheel, err = collection.NewTimingWheel(time.Second, timingWheelSlots, clean)
-	lang.Must(err)
+	logx.Must(err)
 
 	proc.AddShutdownListener(func() {
 		timingWheel.Drain(clean)

core/stores/cache/cleaner_test.go

@@ -0,0 +1,56 @@
+package cache
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNextDelay(t *testing.T) {
+	tests := []struct {
+		name   string
+		input  time.Duration
+		output time.Duration
+		ok     bool
+	}{
+		{
+			name:   "second",
+			input:  time.Second,
+			output: time.Second * 5,
+			ok:     true,
+		},
+		{
+			name:   "5 seconds",
+			input:  time.Second * 5,
+			output: time.Minute,
+			ok:     true,
+		},
+		{
+			name:   "minute",
+			input:  time.Minute,
+			output: time.Minute * 5,
+			ok:     true,
+		},
+		{
+			name:   "5 minutes",
+			input:  time.Minute * 5,
+			output: time.Hour,
+			ok:     true,
+		},
+		{
+			name:   "hour",
+			input:  time.Hour,
+			output: 0,
+			ok:     false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			next, ok := nextDelay(test.input)
+			assert.Equal(t, test.ok, ok)
+			assert.Equal(t, test.output, next)
+		})
+	}
+}

core/stores/cache/config.go

@@ -1,4 +1,4 @@
-package internal
+package cache
 
 import "github.com/tal-tech/go-zero/core/stores/redis"
 

core/stores/cache/util.go

@@ -1,4 +1,4 @@
-package internal
+package cache
 
 import "strings"
 

core/stores/cache/util_test.go

@@ -0,0 +1,26 @@
+package cache
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestFormatKeys(t *testing.T) {
+	assert.Equal(t, "a,b", formatKeys([]string{"a", "b"}))
+}
+
+func TestTotalWeights(t *testing.T) {
+	val := TotalWeights([]NodeConf{
+		{
+			Weight: -1,
+		},
+		{
+			Weight: 0,
+		},
+		{
+			Weight: 1,
+		},
+	})
+	assert.Equal(t, 1, val)
+}

core/stores/clickhouse/clickhouse.go

@@ -1,7 +1,7 @@
 package clickhouse
 
 import (
-	_ "github.com/kshvakov/clickhouse"
+	_ "github.com/ClickHouse/clickhouse-go"
 	"github.com/tal-tech/go-zero/core/stores/sqlx"
 )
 

core/stores/internal/cachenode_test.go

@@ -1,65 +0,0 @@
-package internal
-
-import (
-	"errors"
-	"math/rand"
-	"sync"
-	"testing"
-	"time"
-
-	"github.com/alicebob/miniredis"
-	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/logx"
-	"github.com/tal-tech/go-zero/core/mathx"
-	"github.com/tal-tech/go-zero/core/stat"
-	"github.com/tal-tech/go-zero/core/stores/redis"
-)
-
-func init() {
-	logx.Disable()
-	stat.SetReporter(nil)
-}
-
-func TestCacheNode_DelCache(t *testing.T) {
-	s, err := miniredis.Run()
-	assert.Nil(t, err)
-	defer s.Close()
-
-	cn := cacheNode{
-		rds:            redis.NewRedis(s.Addr(), redis.NodeType),
-		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
-		lock:           new(sync.Mutex),
-		unstableExpiry: mathx.NewUnstable(expiryDeviation),
-		stat:           NewCacheStat("any"),
-		errNotFound:    errors.New("any"),
-	}
-	assert.Nil(t, cn.DelCache())
-	assert.Nil(t, cn.DelCache([]string{}...))
-	assert.Nil(t, cn.DelCache(make([]string, 0)...))
-	cn.SetCache("first", "one")
-	assert.Nil(t, cn.DelCache("first"))
-	cn.SetCache("first", "one")
-	cn.SetCache("second", "two")
-	assert.Nil(t, cn.DelCache("first", "second"))
-}
-
-func TestCacheNode_InvalidCache(t *testing.T) {
-	s, err := miniredis.Run()
-	assert.Nil(t, err)
-	defer s.Close()
-
-	cn := cacheNode{
-		rds:            redis.NewRedis(s.Addr(), redis.NodeType),
-		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
-		lock:           new(sync.Mutex),
-		unstableExpiry: mathx.NewUnstable(expiryDeviation),
-		stat:           NewCacheStat("any"),
-		errNotFound:    errors.New("any"),
-	}
-	s.Set("any", "value")
-	var str string
-	assert.NotNil(t, cn.GetCache("any", &str))
-	assert.Equal(t, "", str)
-	_, err = s.Get("any")
-	assert.Equal(t, miniredis.ErrKeyNotFound, err)
-}

core/stores/internal/cacheopt.go

@@ -1,33 +0,0 @@
-package internal
-
-import "time"
-
-const (
-	defaultExpiry         = time.Hour * 24 * 7
-	defaultNotFoundExpiry = time.Minute
-)
-
-type (
-	Options struct {
-		Expiry         time.Duration
-		NotFoundExpiry time.Duration
-	}
-
-	Option func(o *Options)
-)
-
-func newOptions(opts ...Option) Options {
-	var o Options
-	for _, opt := range opts {
-		opt(&o)
-	}
-
-	if o.Expiry <= 0 {
-		o.Expiry = defaultExpiry
-	}
-	if o.NotFoundExpiry <= 0 {
-		o.NotFoundExpiry = defaultNotFoundExpiry
-	}
-
-	return o
-}

core/stores/kv/config.go

@@ -1,5 +1,7 @@
 package kv
 
-import "github.com/tal-tech/go-zero/core/stores/internal"
+import (
+	"github.com/tal-tech/go-zero/core/stores/cache"
+)
 
-type KvConf = internal.ClusterConf
+type KvConf = cache.ClusterConf

core/stores/kv/store.go

@@ -6,7 +6,7 @@ import (
 
 	"github.com/tal-tech/go-zero/core/errorx"
 	"github.com/tal-tech/go-zero/core/hash"
-	"github.com/tal-tech/go-zero/core/stores/internal"
+	"github.com/tal-tech/go-zero/core/stores/cache"
 	"github.com/tal-tech/go-zero/core/stores/redis"
 )
 
@@ -81,7 +81,7 @@ type (
 )
 
 func NewStore(c KvConf) Store {
-	if len(c) == 0 || internal.TotalWeights(c) <= 0 {
+	if len(c) == 0 || cache.TotalWeights(c) <= 0 {
 		log.Fatal("no cache nodes")
 	}
 

core/stores/kv/store_test.go

@@ -6,7 +6,8 @@ import (
 
 	"github.com/alicebob/miniredis"
 	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/core/stores/internal"
+	"github.com/tal-tech/go-zero/core/hash"
+	"github.com/tal-tech/go-zero/core/stores/cache"
 	"github.com/tal-tech/go-zero/core/stores/redis"
 	"github.com/tal-tech/go-zero/core/stringx"
 )
@@ -15,6 +16,10 @@ var s1, _ = miniredis.Run()
 var s2, _ = miniredis.Run()
 
 func TestRedis_Exists(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Exists("foo")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		ok, err := client.Exists("a")
 		assert.Nil(t, err)
@@ -27,6 +32,10 @@ func TestRedis_Exists(t *testing.T) {
 }
 
 func TestRedis_Eval(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Eval(`redis.call("EXISTS", KEYS[1])`, "key1")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		_, err := client.Eval(`redis.call("EXISTS", KEYS[1])`, "notexist")
 		assert.Equal(t, redis.Nil, err)
@@ -41,6 +50,12 @@ func TestRedis_Eval(t *testing.T) {
 }
 
 func TestRedis_Hgetall(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	err := store.Hset("a", "aa", "aaa")
+	assert.NotNil(t, err)
+	_, err = store.Hgetall("a")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
@@ -54,6 +69,10 @@ func TestRedis_Hgetall(t *testing.T) {
 }
 
 func TestRedis_Hvals(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Hvals("a")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
@@ -64,6 +83,10 @@ func TestRedis_Hvals(t *testing.T) {
 }
 
 func TestRedis_Hsetnx(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Hsetnx("a", "dd", "ddd")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
@@ -80,6 +103,12 @@ func TestRedis_Hsetnx(t *testing.T) {
 }
 
 func TestRedis_HdelHlen(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Hdel("a", "aa")
+	assert.NotNil(t, err)
+	_, err = store.Hlen("a")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
@@ -96,6 +125,10 @@ func TestRedis_HdelHlen(t *testing.T) {
 }
 
 func TestRedis_HIncrBy(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Hincrby("key", "field", 3)
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		val, err := client.Hincrby("key", "field", 2)
 		assert.Nil(t, err)
@@ -107,6 +140,10 @@ func TestRedis_HIncrBy(t *testing.T) {
 }
 
 func TestRedis_Hkeys(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Hkeys("a")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
@@ -117,6 +154,10 @@ func TestRedis_Hkeys(t *testing.T) {
 }
 
 func TestRedis_Hmget(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Hmget("a", "aa", "bb")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
@@ -130,6 +171,12 @@ func TestRedis_Hmget(t *testing.T) {
 }
 
 func TestRedis_Hmset(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	err := store.Hmset("a", map[string]string{
+		"aa": "aaa",
+	})
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hmset("a", map[string]string{
 			"aa": "aaa",
@@ -142,6 +189,10 @@ func TestRedis_Hmset(t *testing.T) {
 }
 
 func TestRedis_Incr(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Incr("a")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		val, err := client.Incr("a")
 		assert.Nil(t, err)
@@ -153,6 +204,10 @@ func TestRedis_Incr(t *testing.T) {
 }
 
 func TestRedis_IncrBy(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Incrby("a", 2)
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		val, err := client.Incrby("a", 2)
 		assert.Nil(t, err)
@@ -164,6 +219,20 @@ func TestRedis_IncrBy(t *testing.T) {
 }
 
 func TestRedis_List(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Lpush("key", "value1", "value2")
+	assert.NotNil(t, err)
+	_, err = store.Rpush("key", "value3", "value4")
+	assert.NotNil(t, err)
+	_, err = store.Llen("key")
+	assert.NotNil(t, err)
+	_, err = store.Lrange("key", 0, 10)
+	assert.NotNil(t, err)
+	_, err = store.Lpop("key")
+	assert.NotNil(t, err)
+	_, err = store.Lrem("key", 0, "val")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		val, err := client.Lpush("key", "value1", "value2")
 		assert.Nil(t, err)
@@ -202,6 +271,14 @@ func TestRedis_List(t *testing.T) {
 }
 
 func TestRedis_Persist(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Persist("key")
+	assert.NotNil(t, err)
+	err = store.Expire("key", 5)
+	assert.NotNil(t, err)
+	err = store.Expireat("key", time.Now().Unix()+5)
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		ok, err := client.Persist("key")
 		assert.Nil(t, err)
@@ -212,10 +289,12 @@ func TestRedis_Persist(t *testing.T) {
 		assert.Nil(t, err)
 		assert.False(t, ok)
 		err = client.Expire("key", 5)
+		assert.Nil(t, err)
 		ok, err = client.Persist("key")
 		assert.Nil(t, err)
 		assert.True(t, ok)
 		err = client.Expireat("key", time.Now().Unix()+5)
+		assert.Nil(t, err)
 		ok, err = client.Persist("key")
 		assert.Nil(t, err)
 		assert.True(t, ok)
@@ -223,8 +302,16 @@ func TestRedis_Persist(t *testing.T) {
 }
 
 func TestRedis_Sscan(t *testing.T) {
+	key := "list"
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Sadd(key, nil)
+	assert.NotNil(t, err)
+	_, _, err = store.Sscan(key, 0, "", 100)
+	assert.NotNil(t, err)
+	_, err = store.Del(key)
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
-		key := "list"
 		var list []string
 		for i := 0; i < 1550; i++ {
 			list = append(list, stringx.Randn(i))
@@ -252,6 +339,20 @@ func TestRedis_Sscan(t *testing.T) {
 }
 
 func TestRedis_Set(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Scard("key")
+	assert.NotNil(t, err)
+	_, err = store.Sismember("key", 2)
+	assert.NotNil(t, err)
+	_, err = store.Srem("key", 3, 4)
+	assert.NotNil(t, err)
+	_, err = store.Smembers("key")
+	assert.NotNil(t, err)
+	_, err = store.Srandmember("key", 1)
+	assert.NotNil(t, err)
+	_, err = store.Spop("key")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		num, err := client.Sadd("key", 1, 2, 3, 4)
 		assert.Nil(t, err)
@@ -288,6 +389,14 @@ func TestRedis_Set(t *testing.T) {
 }
 
 func TestRedis_SetGetDel(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	err := store.Set("hello", "world")
+	assert.NotNil(t, err)
+	_, err = store.Get("hello")
+	assert.NotNil(t, err)
+	_, err = store.Del("hello")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		err := client.Set("hello", "world")
 		assert.Nil(t, err)
@@ -301,6 +410,16 @@ func TestRedis_SetGetDel(t *testing.T) {
 }
 
 func TestRedis_SetExNx(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	err := store.Setex("hello", "world", 5)
+	assert.NotNil(t, err)
+	_, err = store.Setnx("newhello", "newworld")
+	assert.NotNil(t, err)
+	_, err = store.Ttl("hello")
+	assert.NotNil(t, err)
+	_, err = store.SetnxEx("newhello", "newworld", 5)
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		err := client.Setex("hello", "world", 5)
 		assert.Nil(t, err)
@@ -335,6 +454,16 @@ func TestRedis_SetExNx(t *testing.T) {
 }
 
 func TestRedis_SetGetDelHashField(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	err := store.Hset("key", "field", "value")
+	assert.NotNil(t, err)
+	_, err = store.Hget("key", "field")
+	assert.NotNil(t, err)
+	_, err = store.Hexists("key", "field")
+	assert.NotNil(t, err)
+	_, err = store.Hdel("key", "field")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		err := client.Hset("key", "field", "value")
 		assert.Nil(t, err)
@@ -354,6 +483,48 @@ func TestRedis_SetGetDelHashField(t *testing.T) {
 }
 
 func TestRedis_SortedSet(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Zadd("key", 1, "value1")
+	assert.NotNil(t, err)
+	_, err = store.Zscore("key", "value1")
+	assert.NotNil(t, err)
+	_, err = store.Zcount("key", 6, 7)
+	assert.NotNil(t, err)
+	_, err = store.Zincrby("key", 3, "value1")
+	assert.NotNil(t, err)
+	_, err = store.Zrank("key", "value2")
+	assert.NotNil(t, err)
+	_, err = store.Zrem("key", "value2", "value3")
+	assert.NotNil(t, err)
+	_, err = store.Zremrangebyscore("key", 6, 7)
+	assert.NotNil(t, err)
+	_, err = store.Zremrangebyrank("key", 1, 2)
+	assert.NotNil(t, err)
+	_, err = store.Zcard("key")
+	assert.NotNil(t, err)
+	_, err = store.Zrange("key", 0, -1)
+	assert.NotNil(t, err)
+	_, err = store.Zrevrange("key", 0, -1)
+	assert.NotNil(t, err)
+	_, err = store.ZrangeWithScores("key", 0, -1)
+	assert.NotNil(t, err)
+	_, err = store.ZrangebyscoreWithScores("key", 5, 8)
+	assert.NotNil(t, err)
+	_, err = store.ZrangebyscoreWithScoresAndLimit("key", 5, 8, 1, 1)
+	assert.NotNil(t, err)
+	_, err = store.ZrevrangebyscoreWithScores("key", 5, 8)
+	assert.NotNil(t, err)
+	_, err = store.ZrevrangebyscoreWithScoresAndLimit("key", 5, 8, 1, 1)
+	assert.NotNil(t, err)
+	_, err = store.Zadds("key", redis.Pair{
+		Key:   "value2",
+		Score: 6,
+	}, redis.Pair{
+		Key:   "value3",
+		Score: 7,
+	})
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		ok, err := client.Zadd("key", 1, "value1")
 		assert.Nil(t, err)
@@ -379,7 +550,7 @@ func TestRedis_SortedSet(t *testing.T) {
 		rank, err := client.Zrank("key", "value2")
 		assert.Nil(t, err)
 		assert.Equal(t, int64(1), rank)
-		rank, err = client.Zrank("key", "value4")
+		_, err = client.Zrank("key", "value4")
 		assert.Equal(t, redis.Nil, err)
 		num, err := client.Zrem("key", "value2", "value3")
 		assert.Nil(t, err)
@@ -469,6 +640,30 @@ func TestRedis_SortedSet(t *testing.T) {
 				Score: 5,
 			},
 		}, pairs)
+		val, err = client.Zadds("key", redis.Pair{
+			Key:   "value2",
+			Score: 6,
+		}, redis.Pair{
+			Key:   "value3",
+			Score: 7,
+		})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(2), val)
+	})
+}
+
+func TestRedis_HyperLogLog(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Pfadd("key")
+	assert.NotNil(t, err)
+	_, err = store.Pfcount("key")
+	assert.NotNil(t, err)
+
+	runOnCluster(t, func(cluster Store) {
+		_, err := cluster.Pfadd("key")
+		assert.NotNil(t, err)
+		_, err = cluster.Pfcount("key")
+		assert.NotNil(t, err)
 	})
 }
 
@@ -476,7 +671,7 @@ func runOnCluster(t *testing.T, fn func(cluster Store)) {
 	s1.FlushAll()
 	s2.FlushAll()
 
-	store := NewStore([]internal.NodeConf{
+	store := NewStore([]cache.NodeConf{
 		{
 			RedisConf: redis.RedisConf{
 				Host: s1.Addr(),

core/stores/mongo/collection.go

@@ -7,6 +7,7 @@ import (
 	"github.com/globalsign/mgo"
 	"github.com/tal-tech/go-zero/core/breaker"
 	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/stores/mongo/internal"
 	"github.com/tal-tech/go-zero/core/timex"
 )
 
@@ -29,8 +30,9 @@ type (
 	}
 
 	decoratedCollection struct {
-		*mgo.Collection
-		brk breaker.Breaker
+		name       string
+		collection internal.MgoCollection
+		brk        breaker.Breaker
 	}
 
 	keepablePromise struct {
@@ -41,7 +43,8 @@ type (
 
 func newCollection(collection *mgo.Collection) Collection {
 	return &decoratedCollection{
-		Collection: collection,
+		name:       collection.FullName,
+		collection: collection,
 		brk:        breaker.NewBreaker(),
 	}
 }
@@ -54,7 +57,7 @@ func (c *decoratedCollection) Find(query interface{}) Query {
 
 	startTime := timex.Now()
 	return promisedQuery{
-		Query: c.Collection.Find(query),
+		Query: c.collection.Find(query),
 		promise: keepablePromise{
 			promise: promise,
 			log: func(err error) {
@@ -73,7 +76,7 @@ func (c *decoratedCollection) FindId(id interface{}) Query {
 
 	startTime := timex.Now()
 	return promisedQuery{
-		Query: c.Collection.FindId(id),
+		Query: c.collection.FindId(id),
 		promise: keepablePromise{
 			promise: promise,
 			log: func(err error) {
@@ -92,7 +95,7 @@ func (c *decoratedCollection) Insert(docs ...interface{}) (err error) {
 			c.logDuration("insert", duration, err, docs...)
 		}()
 
-		return c.Collection.Insert(docs...)
+		return c.collection.Insert(docs...)
 	}, acceptable)
 }
 
@@ -104,7 +107,7 @@ func (c *decoratedCollection) Pipe(pipeline interface{}) Pipe {
 
 	startTime := timex.Now()
 	return promisedPipe{
-		Pipe: c.Collection.Pipe(pipeline),
+		Pipe: c.collection.Pipe(pipeline),
 		promise: keepablePromise{
 			promise: promise,
 			log: func(err error) {
@@ -123,7 +126,7 @@ func (c *decoratedCollection) Remove(selector interface{}) (err error) {
 			c.logDuration("remove", duration, err, selector)
 		}()
 
-		return c.Collection.Remove(selector)
+		return c.collection.Remove(selector)
 	}, acceptable)
 }
 
@@ -135,7 +138,7 @@ func (c *decoratedCollection) RemoveAll(selector interface{}) (info *mgo.ChangeI
 			c.logDuration("removeAll", duration, err, selector)
 		}()
 
-		info, err = c.Collection.RemoveAll(selector)
+		info, err = c.collection.RemoveAll(selector)
 		return err
 	}, acceptable)
 
@@ -150,7 +153,7 @@ func (c *decoratedCollection) RemoveId(id interface{}) (err error) {
 			c.logDuration("removeId", duration, err, id)
 		}()
 
-		return c.Collection.RemoveId(id)
+		return c.collection.RemoveId(id)
 	}, acceptable)
 }
 
@@ -162,7 +165,7 @@ func (c *decoratedCollection) Update(selector, update interface{}) (err error) {
 			c.logDuration("update", duration, err, selector, update)
 		}()
 
-		return c.Collection.Update(selector, update)
+		return c.collection.Update(selector, update)
 	}, acceptable)
 }
 
@@ -174,7 +177,7 @@ func (c *decoratedCollection) UpdateId(id, update interface{}) (err error) {
 			c.logDuration("updateId", duration, err, id, update)
 		}()
 
-		return c.Collection.UpdateId(id, update)
+		return c.collection.UpdateId(id, update)
 	}, acceptable)
 }
 
@@ -186,7 +189,7 @@ func (c *decoratedCollection) Upsert(selector, update interface{}) (info *mgo.Ch
 			c.logDuration("upsert", duration, err, selector, update)
 		}()
 
-		info, err = c.Collection.Upsert(selector, update)
+		info, err = c.collection.Upsert(selector, update)
 		return err
 	}, acceptable)
 
@@ -200,17 +203,17 @@ func (c *decoratedCollection) logDuration(method string, duration time.Duration,
 	} else if err != nil {
 		if duration > slowThreshold {
 			logx.WithDuration(duration).Slowf("[MONGO] mongo(%s) - slowcall - %s - fail(%s) - %s",
-				c.FullName, method, err.Error(), string(content))
+				c.name, method, err.Error(), string(content))
 		} else {
 			logx.WithDuration(duration).Infof("mongo(%s) - %s - fail(%s) - %s",
-				c.FullName, method, err.Error(), string(content))
+				c.name, method, err.Error(), string(content))
 		}
 	} else {
 		if duration > slowThreshold {
 			logx.WithDuration(duration).Slowf("[MONGO] mongo(%s) - slowcall - %s - ok - %s",
-				c.FullName, method, string(content))
+				c.name, method, string(content))
 		} else {
-			logx.WithDuration(duration).Infof("mongo(%s) - %s - ok - %s", c.FullName, method, string(content))
+			logx.WithDuration(duration).Infof("mongo(%s) - %s - ok - %s", c.name, method, string(content))
 		}
 	}
 }

core/stores/mongo/collection_test.go

@@ -5,10 +5,20 @@ import (
 	"testing"
 
 	"github.com/globalsign/mgo"
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/breaker"
+	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/stores/mongo/internal"
 	"github.com/tal-tech/go-zero/core/stringx"
 )
 
+var errDummy = errors.New("dummy")
+
+func init() {
+	logx.Disable()
+}
+
 func TestKeepPromise_accept(t *testing.T) {
 	p := new(mockPromise)
 	kp := keepablePromise{
@@ -56,6 +66,206 @@ func TestKeepPromise_keep(t *testing.T) {
 	}
 }
 
+func TestNewCollection(t *testing.T) {
+	col := newCollection(&mgo.Collection{
+		Database: nil,
+		Name:     "foo",
+		FullName: "bar",
+	})
+	assert.Equal(t, "bar", col.(*decoratedCollection).name)
+}
+
+func TestCollectionFind(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	var query mgo.Query
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().Find(gomock.Any()).Return(&query)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	actual := c.Find(nil)
+	switch v := actual.(type) {
+	case promisedQuery:
+		assert.Equal(t, &query, v.Query)
+		assert.Equal(t, errDummy, v.promise.keep(errDummy))
+	default:
+		t.Fail()
+	}
+	c.brk = new(dropBreaker)
+	actual = c.Find(nil)
+	assert.Equal(t, rejectedQuery{}, actual)
+}
+
+func TestCollectionFindId(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	var query mgo.Query
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().FindId(gomock.Any()).Return(&query)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	actual := c.FindId(nil)
+	switch v := actual.(type) {
+	case promisedQuery:
+		assert.Equal(t, &query, v.Query)
+		assert.Equal(t, errDummy, v.promise.keep(errDummy))
+	default:
+		t.Fail()
+	}
+	c.brk = new(dropBreaker)
+	actual = c.FindId(nil)
+	assert.Equal(t, rejectedQuery{}, actual)
+}
+
+func TestCollectionInsert(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().Insert(nil, nil).Return(errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	err := c.Insert(nil, nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	err = c.Insert(nil, nil)
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCollectionPipe(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	var pipe mgo.Pipe
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().Pipe(gomock.Any()).Return(&pipe)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	actual := c.Pipe(nil)
+	switch v := actual.(type) {
+	case promisedPipe:
+		assert.Equal(t, &pipe, v.Pipe)
+		assert.Equal(t, errDummy, v.promise.keep(errDummy))
+	default:
+		t.Fail()
+	}
+	c.brk = new(dropBreaker)
+	actual = c.Pipe(nil)
+	assert.Equal(t, rejectedPipe{}, actual)
+}
+
+func TestCollectionRemove(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().Remove(gomock.Any()).Return(errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	err := c.Remove(nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	err = c.Remove(nil)
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCollectionRemoveAll(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().RemoveAll(gomock.Any()).Return(nil, errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	_, err := c.RemoveAll(nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	_, err = c.RemoveAll(nil)
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCollectionRemoveId(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().RemoveId(gomock.Any()).Return(errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	err := c.RemoveId(nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	err = c.RemoveId(nil)
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCollectionUpdate(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().Update(gomock.Any(), gomock.Any()).Return(errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	err := c.Update(nil, nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	err = c.Update(nil, nil)
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCollectionUpdateId(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().UpdateId(gomock.Any(), gomock.Any()).Return(errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	err := c.UpdateId(nil, nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	err = c.UpdateId(nil, nil)
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCollectionUpsert(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().Upsert(gomock.Any(), gomock.Any()).Return(nil, errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	_, err := c.Upsert(nil, nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	_, err = c.Upsert(nil, nil)
+	assert.Equal(t, errDummy, err)
+}
+
 type mockPromise struct {
 	accepted bool
 	reason   string
@@ -68,3 +278,31 @@ func (p *mockPromise) Accept() {
 func (p *mockPromise) Reject(reason string) {
 	p.reason = reason
 }
+
+type dropBreaker struct {
+}
+
+func (d *dropBreaker) Name() string {
+	return "dummy"
+}
+
+func (d *dropBreaker) Allow() (breaker.Promise, error) {
+	return nil, errDummy
+}
+
+func (d *dropBreaker) Do(req func() error) error {
+	return nil
+}
+
+func (d *dropBreaker) DoWithAcceptable(req func() error, acceptable breaker.Acceptable) error {
+	return errDummy
+}
+
+func (d *dropBreaker) DoWithFallback(req func() error, fallback func(err error) error) error {
+	return nil
+}
+
+func (d *dropBreaker) DoWithFallbackAcceptable(req func() error, fallback func(err error) error,
+	acceptable breaker.Acceptable) error {
+	return nil
+}

core/stores/mongo/internal/collection.go

@@ -0,0 +1,17 @@
+//go:generate mockgen -package internal -destination collection_mock.go -source collection.go
+package internal
+
+import "github.com/globalsign/mgo"
+
+type MgoCollection interface {
+	Find(query interface{}) *mgo.Query
+	FindId(id interface{}) *mgo.Query
+	Insert(docs ...interface{}) error
+	Pipe(pipeline interface{}) *mgo.Pipe
+	Remove(selector interface{}) error
+	RemoveAll(selector interface{}) (*mgo.ChangeInfo, error)
+	RemoveId(id interface{}) error
+	Update(selector, update interface{}) error
+	UpdateId(id, update interface{}) error
+	Upsert(selector, update interface{}) (*mgo.ChangeInfo, error)
+}

core/stores/mongo/internal/collection_mock.go

@@ -0,0 +1,180 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: collection.go
+
+// Package internal is a generated GoMock package.
+package internal
+
+import (
+	mgo "github.com/globalsign/mgo"
+	gomock "github.com/golang/mock/gomock"
+	reflect "reflect"
+)
+
+// MockMgoCollection is a mock of MgoCollection interface
+type MockMgoCollection struct {
+	ctrl     *gomock.Controller
+	recorder *MockMgoCollectionMockRecorder
+}
+
+// MockMgoCollectionMockRecorder is the mock recorder for MockMgoCollection
+type MockMgoCollectionMockRecorder struct {
+	mock *MockMgoCollection
+}
+
+// NewMockMgoCollection creates a new mock instance
+func NewMockMgoCollection(ctrl *gomock.Controller) *MockMgoCollection {
+	mock := &MockMgoCollection{ctrl: ctrl}
+	mock.recorder = &MockMgoCollectionMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockMgoCollection) EXPECT() *MockMgoCollectionMockRecorder {
+	return m.recorder
+}
+
+// Find mocks base method
+func (m *MockMgoCollection) Find(query interface{}) *mgo.Query {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Find", query)
+	ret0, _ := ret[0].(*mgo.Query)
+	return ret0
+}
+
+// Find indicates an expected call of Find
+func (mr *MockMgoCollectionMockRecorder) Find(query interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Find", reflect.TypeOf((*MockMgoCollection)(nil).Find), query)
+}
+
+// FindId mocks base method
+func (m *MockMgoCollection) FindId(id interface{}) *mgo.Query {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "FindId", id)
+	ret0, _ := ret[0].(*mgo.Query)
+	return ret0
+}
+
+// FindId indicates an expected call of FindId
+func (mr *MockMgoCollectionMockRecorder) FindId(id interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindId", reflect.TypeOf((*MockMgoCollection)(nil).FindId), id)
+}
+
+// Insert mocks base method
+func (m *MockMgoCollection) Insert(docs ...interface{}) error {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{}
+	for _, a := range docs {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Insert", varargs...)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Insert indicates an expected call of Insert
+func (mr *MockMgoCollectionMockRecorder) Insert(docs ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Insert", reflect.TypeOf((*MockMgoCollection)(nil).Insert), docs...)
+}
+
+// Pipe mocks base method
+func (m *MockMgoCollection) Pipe(pipeline interface{}) *mgo.Pipe {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Pipe", pipeline)
+	ret0, _ := ret[0].(*mgo.Pipe)
+	return ret0
+}
+
+// Pipe indicates an expected call of Pipe
+func (mr *MockMgoCollectionMockRecorder) Pipe(pipeline interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Pipe", reflect.TypeOf((*MockMgoCollection)(nil).Pipe), pipeline)
+}
+
+// Remove mocks base method
+func (m *MockMgoCollection) Remove(selector interface{}) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Remove", selector)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Remove indicates an expected call of Remove
+func (mr *MockMgoCollectionMockRecorder) Remove(selector interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Remove", reflect.TypeOf((*MockMgoCollection)(nil).Remove), selector)
+}
+
+// RemoveAll mocks base method
+func (m *MockMgoCollection) RemoveAll(selector interface{}) (*mgo.ChangeInfo, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "RemoveAll", selector)
+	ret0, _ := ret[0].(*mgo.ChangeInfo)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// RemoveAll indicates an expected call of RemoveAll
+func (mr *MockMgoCollectionMockRecorder) RemoveAll(selector interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoveAll", reflect.TypeOf((*MockMgoCollection)(nil).RemoveAll), selector)
+}
+
+// RemoveId mocks base method
+func (m *MockMgoCollection) RemoveId(id interface{}) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "RemoveId", id)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// RemoveId indicates an expected call of RemoveId
+func (mr *MockMgoCollectionMockRecorder) RemoveId(id interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoveId", reflect.TypeOf((*MockMgoCollection)(nil).RemoveId), id)
+}
+
+// Update mocks base method
+func (m *MockMgoCollection) Update(selector, update interface{}) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Update", selector, update)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Update indicates an expected call of Update
+func (mr *MockMgoCollectionMockRecorder) Update(selector, update interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Update", reflect.TypeOf((*MockMgoCollection)(nil).Update), selector, update)
+}
+
+// UpdateId mocks base method
+func (m *MockMgoCollection) UpdateId(id, update interface{}) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateId", id, update)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// UpdateId indicates an expected call of UpdateId
+func (mr *MockMgoCollectionMockRecorder) UpdateId(id, update interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateId", reflect.TypeOf((*MockMgoCollection)(nil).UpdateId), id, update)
+}
+
+// Upsert mocks base method
+func (m *MockMgoCollection) Upsert(selector, update interface{}) (*mgo.ChangeInfo, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Upsert", selector, update)
+	ret0, _ := ret[0].(*mgo.ChangeInfo)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Upsert indicates an expected call of Upsert
+func (mr *MockMgoCollectionMockRecorder) Upsert(selector, update interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Upsert", reflect.TypeOf((*MockMgoCollection)(nil).Upsert), selector, update)
+}

core/stores/mongo/model.go

@@ -22,8 +22,8 @@ type (
 	}
 )
 
-func MustNewModel(url, database, collection string, opts ...Option) *Model {
-	model, err := NewModel(url, database, collection, opts...)
+func MustNewModel(url, collection string, opts ...Option) *Model {
+	model, err := NewModel(url, collection, opts...)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -31,15 +31,16 @@ func MustNewModel(url, database, collection string, opts ...Option) *Model {
 	return model
 }
 
-func NewModel(url, database, collection string, opts ...Option) (*Model, error) {
+func NewModel(url, collection string, opts ...Option) (*Model, error) {
 	session, err := getConcurrentSession(url)
 	if err != nil {
 		return nil, err
 	}
 
 	return &Model{
-		session:    session,
-		db:         session.DB(database),
+		session: session,
+		// If name is empty, the database name provided in the dialed URL is used instead
+		db:         session.DB(""),
 		collection: collection,
 		opts:       opts,
 	}, nil

core/stores/mongoc/cachedcollection.go

@@ -2,7 +2,7 @@ package mongoc
 
 import (
 	"github.com/globalsign/mgo"
-	"github.com/tal-tech/go-zero/core/stores/internal"
+	"github.com/tal-tech/go-zero/core/stores/cache"
 	"github.com/tal-tech/go-zero/core/stores/mongo"
 	"github.com/tal-tech/go-zero/core/syncx"
 )
@@ -12,7 +12,7 @@ var (
 
 	// can't use one SharedCalls per conn, because multiple conns may share the same cache key.
 	sharedCalls = syncx.NewSharedCalls()
-	stats       = internal.NewCacheStat("mongoc")
+	stats       = cache.NewCacheStat("mongoc")
 )
 
 type (
@@ -20,11 +20,11 @@ type (
 
 	cachedCollection struct {
 		collection mongo.Collection
-		cache      internal.Cache
+		cache      cache.Cache
 	}
 )
 
-func newCollection(collection mongo.Collection, c internal.Cache) *cachedCollection {
+func newCollection(collection mongo.Collection, c cache.Cache) *cachedCollection {
 	return &cachedCollection{
 		collection: collection,
 		cache:      c,

core/stores/mongoc/cachedcollection_test.go

@@ -16,7 +16,7 @@ import (
 	"github.com/globalsign/mgo/bson"
 	"github.com/stretchr/testify/assert"
 	"github.com/tal-tech/go-zero/core/stat"
-	"github.com/tal-tech/go-zero/core/stores/internal"
+	"github.com/tal-tech/go-zero/core/stores/cache"
 	"github.com/tal-tech/go-zero/core/stores/mongo"
 	"github.com/tal-tech/go-zero/core/stores/redis"
 )
@@ -33,7 +33,7 @@ func TestStat(t *testing.T) {
 	}
 
 	r := redis.NewRedis(s.Addr(), redis.NodeType)
-	cach := internal.NewCacheNode(r, sharedCalls, stats, mgo.ErrNotFound)
+	cach := cache.NewCacheNode(r, sharedCalls, stats, mgo.ErrNotFound)
 	c := newCollection(dummyConn{}, cach)
 
 	for i := 0; i < 10; i++ {
@@ -56,7 +56,7 @@ func TestStatCacheFails(t *testing.T) {
 	defer log.SetOutput(os.Stdout)
 
 	r := redis.NewRedis("localhost:59999", redis.NodeType)
-	cach := internal.NewCacheNode(r, sharedCalls, stats, mgo.ErrNotFound)
+	cach := cache.NewCacheNode(r, sharedCalls, stats, mgo.ErrNotFound)
 	c := newCollection(dummyConn{}, cach)
 
 	for i := 0; i < 20; i++ {
@@ -79,7 +79,7 @@ func TestStatDbFails(t *testing.T) {
 	}
 
 	r := redis.NewRedis(s.Addr(), redis.NodeType)
-	cach := internal.NewCacheNode(r, sharedCalls, stats, mgo.ErrNotFound)
+	cach := cache.NewCacheNode(r, sharedCalls, stats, mgo.ErrNotFound)
 	c := newCollection(dummyConn{}, cach)
 
 	for i := 0; i < 20; i++ {
@@ -103,7 +103,7 @@ func TestStatFromMemory(t *testing.T) {
 	}
 
 	r := redis.NewRedis(s.Addr(), redis.NodeType)
-	cach := internal.NewCacheNode(r, sharedCalls, stats, mgo.ErrNotFound)
+	cach := cache.NewCacheNode(r, sharedCalls, stats, mgo.ErrNotFound)
 	c := newCollection(dummyConn{}, cach)
 
 	var all sync.WaitGroup

core/stores/mongoc/cachedmodel.go

@@ -5,19 +5,18 @@ import (
 
 	"github.com/globalsign/mgo"
 	"github.com/tal-tech/go-zero/core/stores/cache"
-	"github.com/tal-tech/go-zero/core/stores/internal"
 	"github.com/tal-tech/go-zero/core/stores/mongo"
 	"github.com/tal-tech/go-zero/core/stores/redis"
 )
 
 type Model struct {
 	*mongo.Model
-	cache              internal.Cache
+	cache              cache.Cache
 	generateCollection func(*mgo.Session) *cachedCollection
 }
 
-func MustNewNodeModel(url, database, collection string, rds *redis.Redis, opts ...cache.Option) *Model {
-	model, err := NewNodeModel(url, database, collection, rds, opts...)
+func MustNewNodeModel(url, collection string, rds *redis.Redis, opts ...cache.Option) *Model {
+	model, err := NewNodeModel(url, collection, rds, opts...)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -25,8 +24,8 @@ func MustNewNodeModel(url, database, collection string, rds *redis.Redis, opts .
 	return model
 }
 
-func MustNewModel(url, database, collection string, c cache.CacheConf, opts ...cache.Option) *Model {
-	model, err := NewModel(url, database, collection, c, opts...)
+func MustNewModel(url, collection string, c cache.CacheConf, opts ...cache.Option) *Model {
+	model, err := NewModel(url, collection, c, opts...)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -34,16 +33,16 @@ func MustNewModel(url, database, collection string, c cache.CacheConf, opts ...c
 	return model
 }
 
-func NewNodeModel(url, database, collection string, rds *redis.Redis, opts ...cache.Option) (*Model, error) {
-	c := internal.NewCacheNode(rds, sharedCalls, stats, mgo.ErrNotFound, opts...)
-	return createModel(url, database, collection, c, func(collection mongo.Collection) *cachedCollection {
+func NewNodeModel(url, collection string, rds *redis.Redis, opts ...cache.Option) (*Model, error) {
+	c := cache.NewCacheNode(rds, sharedCalls, stats, mgo.ErrNotFound, opts...)
+	return createModel(url, collection, c, func(collection mongo.Collection) *cachedCollection {
 		return newCollection(collection, c)
 	})
 }
 
-func NewModel(url, database, collection string, conf cache.CacheConf, opts ...cache.Option) (*Model, error) {
-	c := internal.NewCache(conf, sharedCalls, stats, mgo.ErrNotFound, opts...)
-	return createModel(url, database, collection, c, func(collection mongo.Collection) *cachedCollection {
+func NewModel(url, collection string, conf cache.CacheConf, opts ...cache.Option) (*Model, error) {
+	c := cache.NewCache(conf, sharedCalls, stats, mgo.ErrNotFound, opts...)
+	return createModel(url, collection, c, func(collection mongo.Collection) *cachedCollection {
 		return newCollection(collection, c)
 	})
 }
@@ -224,9 +223,9 @@ func (mm *Model) pipe(fn func(c *cachedCollection) mongo.Pipe) (mongo.Pipe, erro
 	return fn(mm.GetCollection(session)), nil
 }
 
-func createModel(url, database, collection string, c internal.Cache,
+func createModel(url, collection string, c cache.Cache,
 	create func(mongo.Collection) *cachedCollection) (*Model, error) {
-	model, err := mongo.NewModel(url, database, collection)
+	model, err := mongo.NewModel(url, collection)
 	if err != nil {
 		return nil, err
 	}

core/stores/redis/conf_test.go

@@ -0,0 +1,110 @@
+package redis
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/stringx"
+)
+
+func TestRedisConf(t *testing.T) {
+	tests := []struct {
+		name string
+		RedisConf
+		ok bool
+	}{
+		{
+			name: "missing host",
+			RedisConf: RedisConf{
+				Host: "",
+				Type: NodeType,
+				Pass: "",
+			},
+			ok: false,
+		},
+		{
+			name: "missing type",
+			RedisConf: RedisConf{
+				Host: "localhost:6379",
+				Type: "",
+				Pass: "",
+			},
+			ok: false,
+		},
+		{
+			name: "ok",
+			RedisConf: RedisConf{
+				Host: "localhost:6379",
+				Type: NodeType,
+				Pass: "",
+			},
+			ok: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(stringx.RandId(), func(t *testing.T) {
+			if test.ok {
+				assert.Nil(t, test.RedisConf.Validate())
+				assert.NotNil(t, test.RedisConf.NewRedis())
+			} else {
+				assert.NotNil(t, test.RedisConf.Validate())
+			}
+		})
+	}
+}
+
+func TestRedisKeyConf(t *testing.T) {
+	tests := []struct {
+		name string
+		RedisKeyConf
+		ok bool
+	}{
+		{
+			name: "missing host",
+			RedisKeyConf: RedisKeyConf{
+				RedisConf: RedisConf{
+					Host: "",
+					Type: NodeType,
+					Pass: "",
+				},
+				Key: "foo",
+			},
+			ok: false,
+		},
+		{
+			name: "missing key",
+			RedisKeyConf: RedisKeyConf{
+				RedisConf: RedisConf{
+					Host: "localhost:6379",
+					Type: NodeType,
+					Pass: "",
+				},
+				Key: "",
+			},
+			ok: false,
+		},
+		{
+			name: "ok",
+			RedisKeyConf: RedisKeyConf{
+				RedisConf: RedisConf{
+					Host: "localhost:6379",
+					Type: NodeType,
+					Pass: "",
+				},
+				Key: "foo",
+			},
+			ok: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if test.ok {
+				assert.Nil(t, test.RedisKeyConf.Validate())
+			} else {
+				assert.NotNil(t, test.RedisKeyConf.Validate())
+			}
+		})
+	}
+}

core/stores/redis/redis_test.go

@@ -7,11 +7,14 @@ import (
 	"time"
 
 	"github.com/alicebob/miniredis"
+	red "github.com/go-redis/redis"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestRedis_Exists(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
+		_, err := NewRedis(client.Addr, "").Exists("a")
+		assert.NotNil(t, err)
 		ok, err := client.Exists("a")
 		assert.Nil(t, err)
 		assert.False(t, ok)
@@ -24,7 +27,9 @@ func TestRedis_Exists(t *testing.T) {
 
 func TestRedis_Eval(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
-		_, err := client.Eval(`redis.call("EXISTS", KEYS[1])`, []string{"notexist"})
+		_, err := NewRedis(client.Addr, "").Eval(`redis.call("EXISTS", KEYS[1])`, []string{"notexist"})
+		assert.NotNil(t, err)
+		_, err = client.Eval(`redis.call("EXISTS", KEYS[1])`, []string{"notexist"})
 		assert.Equal(t, Nil, err)
 		err = client.Set("key1", "value1")
 		assert.Nil(t, err)
@@ -40,6 +45,8 @@ func TestRedis_Hgetall(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
+		_, err := NewRedis(client.Addr, "").Hgetall("a")
+		assert.NotNil(t, err)
 		vals, err := client.Hgetall("a")
 		assert.Nil(t, err)
 		assert.EqualValues(t, map[string]string{
@@ -51,8 +58,11 @@ func TestRedis_Hgetall(t *testing.T) {
 
 func TestRedis_Hvals(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
+		assert.NotNil(t, NewRedis(client.Addr, "").Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
+		_, err := NewRedis(client.Addr, "").Hvals("a")
+		assert.NotNil(t, err)
 		vals, err := client.Hvals("a")
 		assert.Nil(t, err)
 		assert.ElementsMatch(t, []string{"aaa", "bbb"}, vals)
@@ -63,6 +73,8 @@ func TestRedis_Hsetnx(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
+		_, err := NewRedis(client.Addr, "").Hsetnx("a", "bb", "ccc")
+		assert.NotNil(t, err)
 		ok, err := client.Hsetnx("a", "bb", "ccc")
 		assert.Nil(t, err)
 		assert.False(t, ok)
@@ -79,6 +91,8 @@ func TestRedis_HdelHlen(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
+		_, err := NewRedis(client.Addr, "").Hlen("a")
+		assert.NotNil(t, err)
 		num, err := client.Hlen("a")
 		assert.Nil(t, err)
 		assert.Equal(t, 2, num)
@@ -93,6 +107,8 @@ func TestRedis_HdelHlen(t *testing.T) {
 
 func TestRedis_HIncrBy(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
+		_, err := NewRedis(client.Addr, "").Hincrby("key", "field", 2)
+		assert.NotNil(t, err)
 		val, err := client.Hincrby("key", "field", 2)
 		assert.Nil(t, err)
 		assert.Equal(t, 2, val)
@@ -106,6 +122,8 @@ func TestRedis_Hkeys(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
+		_, err := NewRedis(client.Addr, "").Hkeys("a")
+		assert.NotNil(t, err)
 		vals, err := client.Hkeys("a")
 		assert.Nil(t, err)
 		assert.ElementsMatch(t, []string{"aa", "bb"}, vals)
@@ -116,6 +134,8 @@ func TestRedis_Hmget(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
+		_, err := NewRedis(client.Addr, "").Hmget("a", "aa", "bb")
+		assert.NotNil(t, err)
 		vals, err := client.Hmget("a", "aa", "bb")
 		assert.Nil(t, err)
 		assert.EqualValues(t, []string{"aaa", "bbb"}, vals)
@@ -127,6 +147,7 @@ func TestRedis_Hmget(t *testing.T) {
 
 func TestRedis_Hmset(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
+		assert.NotNil(t, NewRedis(client.Addr, "").Hmset("a", nil))
 		assert.Nil(t, client.Hmset("a", map[string]string{
 			"aa": "aaa",
 			"bb": "bbb",
@@ -139,6 +160,8 @@ func TestRedis_Hmset(t *testing.T) {
 
 func TestRedis_Incr(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
+		_, err := NewRedis(client.Addr, "").Incr("a")
+		assert.NotNil(t, err)
 		val, err := client.Incr("a")
 		assert.Nil(t, err)
 		assert.Equal(t, int64(1), val)
@@ -150,6 +173,8 @@ func TestRedis_Incr(t *testing.T) {
 
 func TestRedis_IncrBy(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
+		_, err := NewRedis(client.Addr, "").Incrby("a", 2)
+		assert.NotNil(t, err)
 		val, err := client.Incrby("a", 2)
 		assert.Nil(t, err)
 		assert.Equal(t, int64(2), val)
@@ -165,26 +190,49 @@ func TestRedis_Keys(t *testing.T) {
 		assert.Nil(t, err)
 		err = client.Set("key2", "value2")
 		assert.Nil(t, err)
+		_, err = NewRedis(client.Addr, "").Keys("*")
+		assert.NotNil(t, err)
 		keys, err := client.Keys("*")
 		assert.Nil(t, err)
 		assert.ElementsMatch(t, []string{"key1", "key2"}, keys)
 	})
 }
 
+func TestRedis_HyperLogLog(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		client.Ping()
+		r := NewRedis(client.Addr, "")
+		_, err := r.Pfadd("key1")
+		assert.NotNil(t, err)
+		_, err = r.Pfcount("*")
+		assert.NotNil(t, err)
+		err = r.Pfmerge("*")
+		assert.NotNil(t, err)
+	})
+}
+
 func TestRedis_List(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
+		_, err := NewRedis(client.Addr, "").Lpush("key", "value1", "value2")
+		assert.NotNil(t, err)
 		val, err := client.Lpush("key", "value1", "value2")
 		assert.Nil(t, err)
 		assert.Equal(t, 2, val)
+		_, err = NewRedis(client.Addr, "").Rpush("key", "value3", "value4")
+		assert.NotNil(t, err)
 		val, err = client.Rpush("key", "value3", "value4")
 		assert.Nil(t, err)
 		assert.Equal(t, 4, val)
+		_, err = NewRedis(client.Addr, "").Llen("key")
+		assert.NotNil(t, err)
 		val, err = client.Llen("key")
 		assert.Nil(t, err)
 		assert.Equal(t, 4, val)
 		vals, err := client.Lrange("key", 0, 10)
 		assert.Nil(t, err)
 		assert.EqualValues(t, []string{"value2", "value1", "value3", "value4"}, vals)
+		_, err = NewRedis(client.Addr, "").Lpop("key")
+		assert.NotNil(t, err)
 		v, err := client.Lpop("key")
 		assert.Nil(t, err)
 		assert.Equal(t, "value2", v)
@@ -194,9 +242,13 @@ func TestRedis_List(t *testing.T) {
 		val, err = client.Rpush("key", "value3", "value3")
 		assert.Nil(t, err)
 		assert.Equal(t, 7, val)
+		_, err = NewRedis(client.Addr, "").Lrem("key", 2, "value1")
+		assert.NotNil(t, err)
 		n, err := client.Lrem("key", 2, "value1")
 		assert.Nil(t, err)
 		assert.Equal(t, 2, n)
+		_, err = NewRedis(client.Addr, "").Lrange("key", 0, 10)
+		assert.NotNil(t, err)
 		vals, err = client.Lrange("key", 0, 10)
 		assert.Nil(t, err)
 		assert.EqualValues(t, []string{"value2", "value3", "value4", "value3", "value3"}, vals)
@@ -215,6 +267,8 @@ func TestRedis_Mget(t *testing.T) {
 		assert.Nil(t, err)
 		err = client.Set("key2", "value2")
 		assert.Nil(t, err)
+		_, err = NewRedis(client.Addr, "").Mget("key1", "key0", "key2", "key3")
+		assert.NotNil(t, err)
 		vals, err := client.Mget("key1", "key0", "key2", "key3")
 		assert.Nil(t, err)
 		assert.EqualValues(t, []string{"value1", "", "value2", ""}, vals)
@@ -223,7 +277,9 @@ func TestRedis_Mget(t *testing.T) {
 
 func TestRedis_SetBit(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
-		err := client.SetBit("key", 1, 1)
+		err := NewRedis(client.Addr, "").SetBit("key", 1, 1)
+		assert.NotNil(t, err)
+		err = client.SetBit("key", 1, 1)
 		assert.Nil(t, err)
 	})
 }
@@ -232,6 +288,8 @@ func TestRedis_GetBit(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
 		err := client.SetBit("key", 2, 1)
 		assert.Nil(t, err)
+		_, err = NewRedis(client.Addr, "").GetBit("key", 2)
+		assert.NotNil(t, err)
 		val, err := client.GetBit("key", 2)
 		assert.Nil(t, err)
 		assert.Equal(t, 1, val)
@@ -240,6 +298,8 @@ func TestRedis_GetBit(t *testing.T) {
 
 func TestRedis_Persist(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
+		_, err := NewRedis(client.Addr, "").Persist("key")
+		assert.NotNil(t, err)
 		ok, err := client.Persist("key")
 		assert.Nil(t, err)
 		assert.False(t, ok)
@@ -248,11 +308,17 @@ func TestRedis_Persist(t *testing.T) {
 		ok, err = client.Persist("key")
 		assert.Nil(t, err)
 		assert.False(t, ok)
+		err = NewRedis(client.Addr, "").Expire("key", 5)
+		assert.NotNil(t, err)
 		err = client.Expire("key", 5)
+		assert.Nil(t, err)
 		ok, err = client.Persist("key")
 		assert.Nil(t, err)
 		assert.True(t, ok)
+		err = NewRedis(client.Addr, "").Expireat("key", time.Now().Unix()+5)
+		assert.NotNil(t, err)
 		err = client.Expireat("key", time.Now().Unix()+5)
+		assert.Nil(t, err)
 		ok, err = client.Persist("key")
 		assert.Nil(t, err)
 		assert.True(t, ok)
@@ -272,6 +338,8 @@ func TestRedis_Scan(t *testing.T) {
 		assert.Nil(t, err)
 		err = client.Set("key2", "value2")
 		assert.Nil(t, err)
+		_, _, err = NewRedis(client.Addr, "").Scan(0, "*", 100)
+		assert.NotNil(t, err)
 		keys, _, err := client.Scan(0, "*", 100)
 		assert.Nil(t, err)
 		assert.ElementsMatch(t, []string{"key1", "key2"}, keys)
@@ -292,6 +360,8 @@ func TestRedis_Sscan(t *testing.T) {
 		var cursor uint64 = 0
 		sum := 0
 		for {
+			_, _, err := NewRedis(client.Addr, "").Sscan(key, cursor, "", 100)
+			assert.NotNil(t, err)
 			keys, next, err := client.Sscan(key, cursor, "", 100)
 			assert.Nil(t, err)
 			sum += len(keys)
@@ -302,6 +372,8 @@ func TestRedis_Sscan(t *testing.T) {
 		}
 
 		assert.Equal(t, sum, 1550)
+		_, err = NewRedis(client.Addr, "").Del(key)
+		assert.NotNil(t, err)
 		_, err = client.Del(key)
 		assert.Nil(t, err)
 	})
@@ -309,46 +381,72 @@ func TestRedis_Sscan(t *testing.T) {
 
 func TestRedis_Set(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
+		_, err := NewRedis(client.Addr, "").Sadd("key", 1, 2, 3, 4)
+		assert.NotNil(t, err)
 		num, err := client.Sadd("key", 1, 2, 3, 4)
 		assert.Nil(t, err)
 		assert.Equal(t, 4, num)
+		_, err = NewRedis(client.Addr, "").Scard("key")
+		assert.NotNil(t, err)
 		val, err := client.Scard("key")
 		assert.Nil(t, err)
 		assert.Equal(t, int64(4), val)
+		_, err = NewRedis(client.Addr, "").Sismember("key", 2)
+		assert.NotNil(t, err)
 		ok, err := client.Sismember("key", 2)
 		assert.Nil(t, err)
 		assert.True(t, ok)
+		_, err = NewRedis(client.Addr, "").Srem("key", 3, 4)
+		assert.NotNil(t, err)
 		num, err = client.Srem("key", 3, 4)
 		assert.Nil(t, err)
 		assert.Equal(t, 2, num)
+		_, err = NewRedis(client.Addr, "").Smembers("key")
+		assert.NotNil(t, err)
 		vals, err := client.Smembers("key")
 		assert.Nil(t, err)
 		assert.ElementsMatch(t, []string{"1", "2"}, vals)
+		_, err = NewRedis(client.Addr, "").Srandmember("key", 1)
+		assert.NotNil(t, err)
 		members, err := client.Srandmember("key", 1)
 		assert.Nil(t, err)
 		assert.Len(t, members, 1)
 		assert.Contains(t, []string{"1", "2"}, members[0])
+		_, err = NewRedis(client.Addr, "").Spop("key")
+		assert.NotNil(t, err)
 		member, err := client.Spop("key")
 		assert.Nil(t, err)
 		assert.Contains(t, []string{"1", "2"}, member)
+		_, err = NewRedis(client.Addr, "").Smembers("key")
+		assert.NotNil(t, err)
 		vals, err = client.Smembers("key")
 		assert.Nil(t, err)
 		assert.NotContains(t, vals, member)
+		_, err = NewRedis(client.Addr, "").Sadd("key1", 1, 2, 3, 4)
+		assert.NotNil(t, err)
 		num, err = client.Sadd("key1", 1, 2, 3, 4)
 		assert.Nil(t, err)
 		assert.Equal(t, 4, num)
 		num, err = client.Sadd("key2", 2, 3, 4, 5)
 		assert.Nil(t, err)
 		assert.Equal(t, 4, num)
+		_, err = NewRedis(client.Addr, "").Sunion("key1", "key2")
+		assert.NotNil(t, err)
 		vals, err = client.Sunion("key1", "key2")
 		assert.Nil(t, err)
 		assert.ElementsMatch(t, []string{"1", "2", "3", "4", "5"}, vals)
+		_, err = NewRedis(client.Addr, "").Sunionstore("key3", "key1", "key2")
+		assert.NotNil(t, err)
 		num, err = client.Sunionstore("key3", "key1", "key2")
 		assert.Nil(t, err)
 		assert.Equal(t, 5, num)
+		_, err = NewRedis(client.Addr, "").Sdiff("key1", "key2")
+		assert.NotNil(t, err)
 		vals, err = client.Sdiff("key1", "key2")
 		assert.Nil(t, err)
 		assert.EqualValues(t, []string{"1"}, vals)
+		_, err = NewRedis(client.Addr, "").Sdiffstore("key4", "key1", "key2")
+		assert.NotNil(t, err)
 		num, err = client.Sdiffstore("key4", "key1", "key2")
 		assert.Nil(t, err)
 		assert.Equal(t, 1, num)
@@ -357,8 +455,12 @@ func TestRedis_Set(t *testing.T) {
 
 func TestRedis_SetGetDel(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
-		err := client.Set("hello", "world")
+		err := NewRedis(client.Addr, "").Set("hello", "world")
+		assert.NotNil(t, err)
+		err = client.Set("hello", "world")
 		assert.Nil(t, err)
+		_, err = NewRedis(client.Addr, "").Get("hello")
+		assert.NotNil(t, err)
 		val, err := client.Get("hello")
 		assert.Nil(t, err)
 		assert.Equal(t, "world", val)
@@ -370,8 +472,12 @@ func TestRedis_SetGetDel(t *testing.T) {
 
 func TestRedis_SetExNx(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
-		err := client.Setex("hello", "world", 5)
+		err := NewRedis(client.Addr, "").Setex("hello", "world", 5)
+		assert.NotNil(t, err)
+		err = client.Setex("hello", "world", 5)
 		assert.Nil(t, err)
+		_, err = NewRedis(client.Addr, "").Setnx("hello", "newworld")
+		assert.NotNil(t, err)
 		ok, err := client.Setnx("hello", "newworld")
 		assert.Nil(t, err)
 		assert.False(t, ok)
@@ -387,6 +493,8 @@ func TestRedis_SetExNx(t *testing.T) {
 		ttl, err := client.Ttl("hello")
 		assert.Nil(t, err)
 		assert.True(t, ttl > 0)
+		_, err = NewRedis(client.Addr, "").SetnxEx("newhello", "newworld", 5)
+		assert.NotNil(t, err)
 		ok, err = client.SetnxEx("newhello", "newworld", 5)
 		assert.Nil(t, err)
 		assert.False(t, ok)
@@ -406,12 +514,18 @@ func TestRedis_SetGetDelHashField(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
 		err := client.Hset("key", "field", "value")
 		assert.Nil(t, err)
+		_, err = NewRedis(client.Addr, "").Hget("key", "field")
+		assert.NotNil(t, err)
 		val, err := client.Hget("key", "field")
 		assert.Nil(t, err)
 		assert.Equal(t, "value", val)
+		_, err = NewRedis(client.Addr, "").Hexists("key", "field")
+		assert.NotNil(t, err)
 		ok, err := client.Hexists("key", "field")
 		assert.Nil(t, err)
 		assert.True(t, ok)
+		_, err = NewRedis(client.Addr, "").Hdel("key", "field")
+		assert.NotNil(t, err)
 		ret, err := client.Hdel("key", "field")
 		assert.Nil(t, err)
 		assert.True(t, ret)
@@ -432,23 +546,50 @@ func TestRedis_SortedSet(t *testing.T) {
 		val, err := client.Zscore("key", "value1")
 		assert.Nil(t, err)
 		assert.Equal(t, int64(2), val)
+		_, err = NewRedis(client.Addr, "").Zincrby("key", 3, "value1")
+		assert.NotNil(t, err)
 		val, err = client.Zincrby("key", 3, "value1")
 		assert.Nil(t, err)
 		assert.Equal(t, int64(5), val)
+		_, err = NewRedis(client.Addr, "").Zscore("key", "value1")
+		assert.NotNil(t, err)
 		val, err = client.Zscore("key", "value1")
 		assert.Nil(t, err)
 		assert.Equal(t, int64(5), val)
-		ok, err = client.Zadd("key", 6, "value2")
+		val, err = NewRedis(client.Addr, "").Zadds("key")
+		assert.NotNil(t, err)
+		val, err = client.Zadds("key", Pair{
+			Key:   "value2",
+			Score: 6,
+		}, Pair{
+			Key:   "value3",
+			Score: 7,
+		})
 		assert.Nil(t, err)
-		assert.True(t, ok)
-		ok, err = client.Zadd("key", 7, "value3")
+		assert.Equal(t, int64(2), val)
+		pairs, err := NewRedis(client.Addr, "").ZRevRangeWithScores("key", 1, 3)
+		assert.NotNil(t, err)
+		pairs, err = client.ZRevRangeWithScores("key", 1, 3)
 		assert.Nil(t, err)
-		assert.True(t, ok)
+		assert.EqualValues(t, []Pair{
+			{
+				Key:   "value2",
+				Score: 6,
+			},
+			{
+				Key:   "value1",
+				Score: 5,
+			},
+		}, pairs)
 		rank, err := client.Zrank("key", "value2")
 		assert.Nil(t, err)
 		assert.Equal(t, int64(1), rank)
-		rank, err = client.Zrank("key", "value4")
+		_, err = NewRedis(client.Addr, "").Zrank("key", "value4")
+		assert.NotNil(t, err)
+		_, err = client.Zrank("key", "value4")
 		assert.Equal(t, Nil, err)
+		_, err = NewRedis(client.Addr, "").Zrem("key", "value2", "value3")
+		assert.NotNil(t, err)
 		num, err := client.Zrem("key", "value2", "value3")
 		assert.Nil(t, err)
 		assert.Equal(t, 2, num)
@@ -461,31 +602,47 @@ func TestRedis_SortedSet(t *testing.T) {
 		ok, err = client.Zadd("key", 8, "value4")
 		assert.Nil(t, err)
 		assert.True(t, ok)
+		_, err = NewRedis(client.Addr, "").Zremrangebyscore("key", 6, 7)
+		assert.NotNil(t, err)
 		num, err = client.Zremrangebyscore("key", 6, 7)
 		assert.Nil(t, err)
 		assert.Equal(t, 2, num)
 		ok, err = client.Zadd("key", 6, "value2")
 		assert.Nil(t, err)
 		assert.True(t, ok)
+		_, err = NewRedis(client.Addr, "").Zadd("key", 7, "value3")
+		assert.NotNil(t, err)
 		ok, err = client.Zadd("key", 7, "value3")
 		assert.Nil(t, err)
 		assert.True(t, ok)
+		_, err = NewRedis(client.Addr, "").Zcount("key", 6, 7)
+		assert.NotNil(t, err)
 		num, err = client.Zcount("key", 6, 7)
 		assert.Nil(t, err)
 		assert.Equal(t, 2, num)
+		_, err = NewRedis(client.Addr, "").Zremrangebyrank("key", 1, 2)
+		assert.NotNil(t, err)
 		num, err = client.Zremrangebyrank("key", 1, 2)
 		assert.Nil(t, err)
 		assert.Equal(t, 2, num)
+		_, err = NewRedis(client.Addr, "").Zcard("key")
+		assert.NotNil(t, err)
 		card, err := client.Zcard("key")
 		assert.Nil(t, err)
 		assert.Equal(t, 2, card)
+		_, err = NewRedis(client.Addr, "").Zrange("key", 0, -1)
+		assert.NotNil(t, err)
 		vals, err := client.Zrange("key", 0, -1)
 		assert.Nil(t, err)
 		assert.EqualValues(t, []string{"value1", "value4"}, vals)
+		_, err = NewRedis(client.Addr, "").Zrevrange("key", 0, -1)
+		assert.NotNil(t, err)
 		vals, err = client.Zrevrange("key", 0, -1)
 		assert.Nil(t, err)
 		assert.EqualValues(t, []string{"value4", "value1"}, vals)
-		pairs, err := client.ZrangeWithScores("key", 0, -1)
+		_, err = NewRedis(client.Addr, "").ZrangeWithScores("key", 0, -1)
+		assert.NotNil(t, err)
+		pairs, err = client.ZrangeWithScores("key", 0, -1)
 		assert.Nil(t, err)
 		assert.EqualValues(t, []Pair{
 			{
@@ -497,6 +654,8 @@ func TestRedis_SortedSet(t *testing.T) {
 				Score: 8,
 			},
 		}, pairs)
+		_, err = NewRedis(client.Addr, "").ZrangebyscoreWithScores("key", 5, 8)
+		assert.NotNil(t, err)
 		pairs, err = client.ZrangebyscoreWithScores("key", 5, 8)
 		assert.Nil(t, err)
 		assert.EqualValues(t, []Pair{
@@ -509,6 +668,9 @@ func TestRedis_SortedSet(t *testing.T) {
 				Score: 8,
 			},
 		}, pairs)
+		_, err = NewRedis(client.Addr, "").ZrangebyscoreWithScoresAndLimit(
+			"key", 5, 8, 1, 1)
+		assert.NotNil(t, err)
 		pairs, err = client.ZrangebyscoreWithScoresAndLimit("key", 5, 8, 1, 1)
 		assert.Nil(t, err)
 		assert.EqualValues(t, []Pair{
@@ -517,6 +679,11 @@ func TestRedis_SortedSet(t *testing.T) {
 				Score: 8,
 			},
 		}, pairs)
+		pairs, err = client.ZrangebyscoreWithScoresAndLimit("key", 5, 8, 1, 0)
+		assert.Nil(t, err)
+		assert.Equal(t, 0, len(pairs))
+		_, err = NewRedis(client.Addr, "").ZrevrangebyscoreWithScores("key", 5, 8)
+		assert.NotNil(t, err)
 		pairs, err = client.ZrevrangebyscoreWithScores("key", 5, 8)
 		assert.Nil(t, err)
 		assert.EqualValues(t, []Pair{
@@ -529,6 +696,9 @@ func TestRedis_SortedSet(t *testing.T) {
 				Score: 5,
 			},
 		}, pairs)
+		_, err = NewRedis(client.Addr, "").ZrevrangebyscoreWithScoresAndLimit(
+			"key", 5, 8, 1, 1)
+		assert.NotNil(t, err)
 		pairs, err = client.ZrevrangebyscoreWithScoresAndLimit("key", 5, 8, 1, 1)
 		assert.Nil(t, err)
 		assert.EqualValues(t, []Pair{
@@ -537,11 +707,17 @@ func TestRedis_SortedSet(t *testing.T) {
 				Score: 5,
 			},
 		}, pairs)
+		pairs, err = client.ZrevrangebyscoreWithScoresAndLimit("key", 5, 8, 1, 0)
+		assert.Nil(t, err)
+		assert.Equal(t, 0, len(pairs))
 	})
 }
 
 func TestRedis_Pipelined(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
+		assert.NotNil(t, NewRedis(client.Addr, "").Pipelined(func(pipeliner Pipeliner) error {
+			return nil
+		}))
 		err := client.Pipelined(
 			func(pipe Pipeliner) error {
 				pipe.Incr("pipelined_counter")
@@ -551,6 +727,8 @@ func TestRedis_Pipelined(t *testing.T) {
 			},
 		)
 		assert.Nil(t, err)
+		_, err = NewRedis(client.Addr, "").Ttl("pipelined_counter")
+		assert.NotNil(t, err)
 		ttl, err := client.Ttl("pipelined_counter")
 		assert.Nil(t, err)
 		assert.Equal(t, 3600, ttl)
@@ -558,10 +736,81 @@ func TestRedis_Pipelined(t *testing.T) {
 		assert.Nil(t, err)
 		assert.Equal(t, "1", value)
 		score, err := client.Zscore("zadd", "zadd")
+		assert.Nil(t, err)
 		assert.Equal(t, int64(12), score)
 	})
 }
 
+func TestRedisString(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		client.Ping()
+		_, err := getRedis(NewRedis(client.Addr, ClusterType))
+		assert.Nil(t, err)
+		assert.Equal(t, client.Addr, client.String())
+		assert.NotNil(t, NewRedis(client.Addr, "").Ping())
+	})
+}
+
+func TestRedisScriptLoad(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		client.Ping()
+		_, err := NewRedis(client.Addr, "").scriptLoad("foo")
+		assert.NotNil(t, err)
+		_, err = client.scriptLoad("foo")
+		assert.NotNil(t, err)
+	})
+}
+
+func TestRedisToPairs(t *testing.T) {
+	pairs := toPairs([]red.Z{
+		{
+			Member: 1,
+			Score:  1,
+		},
+		{
+			Member: 2,
+			Score:  2,
+		},
+	})
+	assert.EqualValues(t, []Pair{
+		{
+			Key:   "1",
+			Score: 1,
+		},
+		{
+			Key:   "2",
+			Score: 2,
+		},
+	}, pairs)
+}
+
+func TestRedisToStrings(t *testing.T) {
+	vals := toStrings([]interface{}{1, 2})
+	assert.EqualValues(t, []string{"1", "2"}, vals)
+}
+
+func TestRedisBlpop(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		client.Ping()
+		var node mockedNode
+		_, err := client.Blpop(nil, "foo")
+		assert.NotNil(t, err)
+		_, err = client.Blpop(node, "foo")
+		assert.NotNil(t, err)
+	})
+}
+
+func TestRedisBlpopEx(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		client.Ping()
+		var node mockedNode
+		_, _, err := client.BlpopEx(nil, "foo")
+		assert.NotNil(t, err)
+		_, _, err = client.BlpopEx(node, "foo")
+		assert.NotNil(t, err)
+	})
+}
+
 func runOnRedis(t *testing.T, fn func(client *Redis)) {
 	s, err := miniredis.Run()
 	assert.Nil(t, err)
@@ -573,8 +822,18 @@ func runOnRedis(t *testing.T, fn func(client *Redis)) {
 			t.Error(err)
 		}
 
-		client.Close()
+		if client != nil {
+			client.Close()
+		}
 	}()
 
 	fn(NewRedis(s.Addr(), NodeType))
 }
+
+type mockedNode struct {
+	RedisNode
+}
+
+func (n mockedNode) BLPop(timeout time.Duration, keys ...string) *red.StringSliceCmd {
+	return red.NewStringSliceCmd("foo", "bar")
+}